Разработка AI-системы для распознавания товаров на полке

AI-система распознавания товаров на полках

Задача распознавания товаров на полках ритейла — это комбинация детекции (где товар) и идентификации (что за товар, какой именно SKU). Сложность в масштабе: крупный ритейлер имеет 10 000–50 000 уникальных SKU, причём упаковки регулярно меняются.

Детекция товаров: fine-tuning на полочных снимках

from ultralytics import YOLO
import yaml
from pathlib import Path

def prepare_retail_dataset_config(
    data_dir: str,
    class_names: list[str]
) -> str:
    """
    Конфиг датасета для YOLOv8.
    Для ритейл-полок рекомендуем imgsz=1280 — детали упаковок важны.
    """
    config = {
        'path': data_dir,
        'train': 'images/train',
        'val':   'images/val',
        'test':  'images/test',
        'nc':    len(class_names),
        'names': class_names
    }
    config_path = Path(data_dir) / 'dataset.yaml'
    with open(config_path, 'w') as f:
        yaml.dump(config, f, allow_unicode=True)
    return str(config_path)

# Обучение детектора товаров
model = YOLO('yolov8l.pt')
model.train(
    data='retail_dataset.yaml',
    imgsz=1280,        # важно: мелкие ценники и надписи требуют разрешения
    batch=8,           # при 1280 батч меньше
    epochs=200,
    device='0',
    augment=True,
    mosaic=0.5,        # снижаем mosaic — не хотим менять масштаб товаров
    copy_paste=0.3,    # полезно для retail
    rect=False         # прямоугольные батчи ухудшают детекцию мелких объектов
)

Идентификация SKU: embedding + kNN

При 10 000+ SKU softmax-классификатор не масштабируется: при добавлении нового товара нужно переобучать всю модель. Embedding-подход (metric learning) решает это: новый SKU = добавляем его embedding в индекс без переобучения.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import timm
import faiss
import numpy as np

class SKUEmbeddingModel(nn.Module):
    """
    ArcFace-like metric learning для идентификации товаров.
    Обучаем на кропах товаров → embedding 512-dim.
    """
    def __init__(self, num_skus: int, embedding_dim: int = 512):
        super().__init__()
        self.backbone = timm.create_model(
            'efficientnet_b4',
            pretrained=True,
            num_classes=0
        )
        self.embedding = nn.Sequential(
            nn.Linear(self.backbone.num_features, embedding_dim),
            nn.BatchNorm1d(embedding_dim)
        )
        # ArcFace head для обучения
        self.arcface = ArcFaceHead(embedding_dim, num_skus)

    def forward(self, x: torch.Tensor, labels: torch.Tensor = None):
        feat = self.backbone(x)
        emb  = F.normalize(self.embedding(feat), dim=1)
        if labels is not None:
            return self.arcface(emb, labels)
        return emb

class ArcFaceHead(nn.Module):
    def __init__(self, dim: int, num_classes: int,
                 margin: float = 0.3, scale: float = 32.0):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(num_classes, dim))
        self.margin = margin
        self.scale  = scale

    def forward(self, emb: torch.Tensor, labels: torch.Tensor):
        import math
        W = F.normalize(self.weight, dim=1)
        cosine = F.linear(emb, W)
        # Применяем margin только к правильному классу
        one_hot = torch.zeros_like(cosine)
        one_hot.scatter_(1, labels.unsqueeze(1), 1)
        phi = cosine - self.margin
        output = (one_hot * phi + (1 - one_hot) * cosine) * self.scale
        return F.cross_entropy(output, labels)

class SKUFAISSIndex:
    """FAISS-индекс для быстрого поиска похожих SKU"""
    def __init__(self, embedding_dim: int = 512):
        self.index = faiss.IndexFlatIP(embedding_dim)  # inner product = cosine при нормализации
        self.sku_ids = []

    def add_sku(self, sku_id: str, embedding: np.ndarray) -> None:
        emb_norm = embedding / (np.linalg.norm(embedding) + 1e-8)
        self.index.add(emb_norm.reshape(1, -1).astype(np.float32))
        self.sku_ids.append(sku_id)

    def search(
        self, query_embedding: np.ndarray, top_k: int = 5
    ) -> list[dict]:
        q = (query_embedding / (np.linalg.norm(query_embedding) + 1e-8)
             ).reshape(1, -1).astype(np.float32)
        scores, indices = self.index.search(q, top_k)
        return [
            {'sku_id': self.sku_ids[idx], 'score': float(scores[0][i])}
            for i, idx in enumerate(indices[0])
            if idx < len(self.sku_ids)
        ]

Обработка смены упаковки

Главная операционная сложность ритейла — packaging refresh. Раз в год бренды меняют дизайн, и модель начинает ошибаться на новых упаковках.

Решение: онлайн-обновление индекса. Для embedding-подхода достаточно сфотографировать новую упаковку и добавить embedding в FAISS-индекс. Старый embedding можно удалить или оставить (модель автоматически предпочтёт более близкий).

def update_sku_appearance(
    sku_index: SKUFAISSIndex,
    model: SKUEmbeddingModel,
    sku_id: str,
    new_product_images: list,
    keep_old: bool = False      # False = заменяем, True = добавляем вариант
) -> None:
    model.eval()
    embeddings = []

    with torch.no_grad():
        for img in new_product_images:
            emb = model(img.unsqueeze(0).cuda()).cpu().numpy()
            embeddings.append(emb.squeeze())

    # Усредняем по нескольким ракурсам
    mean_emb = np.mean(embeddings, axis=0)

    if not keep_old:
        # Удаляем старые записи (FAISS IDMap для удаления)
        pass  # требует IndexIDMap

    sku_index.add_sku(sku_id, mean_emb)
    print(f'Updated SKU {sku_id} with {len(new_product_images)} images')

Точность на реальных данных ритейла

Сроки