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

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

Планограмма — план расстановки товаров на полке. Контроль соответствия реальной полки планограмме традиционно делается вручную мерчандайзером раз в неделю. Автоматизация через CV-камеры или фото от торговых представителей снижает время реакции с дней до минут.

Задача: от фотографии полки к отчёту о нарушениях

Пайплайн состоит из трёх шагов: детекция продуктов на полке → идентификация каждого продукта → сравнение с эталонной планограммой.

from ultralytics import YOLO
import numpy as np
from PIL import Image
import torch
import torch.nn.functional as F

class PlanogramComplianceChecker:
    """
    Шаг 1: YOLOv8 детектирует все товары на полке (bbox + класс)
    Шаг 2: CLIP/ViT идентифицирует конкретный SKU по cropу
    Шаг 3: Сравнение с планограммой
    """
    def __init__(
        self,
        detector_path: str,       # fine-tuned YOLO на полках
        sku_embeddings_path: str, # CLIP embeddings всех SKU
        planogram: dict           # {position: sku_id}
    ):
        self.detector = YOLO(detector_path)
        sku_data = np.load(sku_embeddings_path)
        self.sku_embeddings = torch.from_numpy(
            sku_data['embeddings']
        ).float()                 # (N_SKU, embedding_dim)
        self.sku_ids = sku_data['sku_ids'].tolist()
        self.planogram = planogram

        # CLIP для идентификации SKU
        from transformers import CLIPProcessor, CLIPModel
        self.clip_model = CLIPModel.from_pretrained(
            'openai/clip-vit-large-patch14'
        ).eval().cuda()
        self.clip_processor = CLIPProcessor.from_pretrained(
            'openai/clip-vit-large-patch14'
        )

    def analyze_shelf(
        self,
        shelf_image: Image.Image,
        confidence_threshold: float = 0.5
    ) -> dict:
        img_array = np.array(shelf_image)

        # Шаг 1: детекция
        detections = self.detector.predict(
            img_array, conf=confidence_threshold, verbose=False
        )[0]

        shelf_products = []
        for box in detections.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            crop = shelf_image.crop((x1, y1, x2, y2))

            # Шаг 2: идентификация SKU через CLIP
            sku_id, similarity = self._identify_sku(crop)

            shelf_products.append({
                'bbox': [x1, y1, x2, y2],
                'sku_id': sku_id,
                'confidence': float(box.conf),
                'sku_similarity': float(similarity),
                'position': self._get_shelf_position(
                    [x1, y1, x2, y2], img_array.shape
                )
            })

        # Шаг 3: сравнение с планограммой
        compliance = self._check_compliance(shelf_products)
        return compliance

    @torch.no_grad()
    def _identify_sku(
        self, crop: Image.Image
    ) -> tuple[str, float]:
        inputs = self.clip_processor(
            images=crop, return_tensors='pt'
        ).to('cuda')
        features = self.clip_model.get_image_features(**inputs)
        features = F.normalize(features, dim=-1).cpu()

        # Cosine similarity со всеми SKU embeddings
        similarities = (features @ self.sku_embeddings.T).squeeze()
        best_idx = similarities.argmax().item()
        return self.sku_ids[best_idx], float(similarities[best_idx])

    def _get_shelf_position(
        self, bbox: list, img_shape: tuple
    ) -> dict:
        """Горизонтальная позиция + ряд полки"""
        h, w = img_shape[:2]
        cx = (bbox[0] + bbox[2]) / 2
        cy = (bbox[1] + bbox[3]) / 2
        return {
            'col': int(cx / w * 10),   # 0-9 — десять колонок
            'row': int(cy / h * 5)     # 0-4 — пять рядов
        }

    def _check_compliance(self, shelf_products: list) -> dict:
        violations = []
        actual_positions = {
            f"{p['position']['row']}_{p['position']['col']}": p['sku_id']
            for p in shelf_products
        }

        for position_key, expected_sku in self.planogram.items():
            actual_sku = actual_positions.get(position_key)
            if actual_sku is None:
                violations.append({
                    'type': 'out_of_stock',
                    'position': position_key,
                    'expected_sku': expected_sku
                })
            elif actual_sku != expected_sku:
                violations.append({
                    'type': 'wrong_product',
                    'position': position_key,
                    'expected_sku': expected_sku,
                    'actual_sku': actual_sku
                })

        compliance_score = 1.0 - len(violations) / max(len(self.planogram), 1)
        return {
            'compliance_score': round(compliance_score, 3),
            'violations': violations,
            'total_positions': len(self.planogram),
            'violations_count': len(violations),
            'detected_products': len(shelf_products)
        }

Индексация SKU через CLIP embeddings

from transformers import CLIPProcessor, CLIPModel
import torch
import numpy as np
from pathlib import Path

def build_sku_index(
    product_images_dir: str,   # директория: {sku_id}/{image1.jpg, ...}
    output_path: str,
    model_name: str = 'openai/clip-vit-large-patch14',
    images_per_sku: int = 5    # усредняем embedding по нескольким фото
) -> None:
    """
    Строим CLIP-индекс всех SKU.
    Несколько фото на товар → усреднённый embedding стабильнее.
    """
    model = CLIPModel.from_pretrained(model_name).eval().cuda()
    processor = CLIPProcessor.from_pretrained(model_name)

    sku_embeddings = []
    sku_ids = []

    for sku_dir in sorted(Path(product_images_dir).iterdir()):
        if not sku_dir.is_dir():
            continue
        sku_id = sku_dir.name
        image_files = list(sku_dir.glob('*.{jpg,jpeg,png}'))[:images_per_sku]

        if not image_files:
            continue

        batch_embeddings = []
        for img_path in image_files:
            image = Image.open(img_path).convert('RGB')
            inputs = processor(images=image, return_tensors='pt').to('cuda')
            with torch.no_grad():
                emb = model.get_image_features(**inputs)
                emb = F.normalize(emb, dim=-1).cpu().numpy()
            batch_embeddings.append(emb)

        mean_emb = np.mean(batch_embeddings, axis=0)
        mean_emb = mean_emb / np.linalg.norm(mean_emb)
        sku_embeddings.append(mean_emb.squeeze())
        sku_ids.append(sku_id)

    np.savez(
        output_path,
        embeddings=np.array(sku_embeddings),
        sku_ids=np.array(sku_ids)
    )
    print(f'Indexed {len(sku_ids)} SKUs')

Сроки