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

AI-визуальный поиск товаров по фотографии

Пользователь фотографирует вещь, которую хочет купить — система находит похожие товары в каталоге. Это reverse image search, адаптированный для ecommerce: нужно находить визуально похожие товары, а не идентичные изображения.

Архитектура: embedding + vector search

import torch
import torch.nn.functional as F
from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import numpy as np
import qdrant_client
from qdrant_client.models import Distance, VectorParams, PointStruct

class VisualSearchEngine:
    """
    CLIP-embedding + Qdrant vector DB для поиска по фото.
    CLIP умеет text→image и image→image поиск из коробки.
    """
    def __init__(
        self,
        qdrant_url: str = 'http://localhost:6333',
        collection_name: str = 'products',
        embedding_dim: int = 768    # CLIP ViT-L/14
    ):
        self.clip_model = CLIPModel.from_pretrained(
            'openai/clip-vit-large-patch14'
        ).eval().cuda()
        self.clip_processor = CLIPProcessor.from_pretrained(
            'openai/clip-vit-large-patch14'
        )

        self.client = qdrant_client.QdrantClient(url=qdrant_url)
        self.collection_name = collection_name
        self._ensure_collection(embedding_dim)

    def _ensure_collection(self, dim: int) -> None:
        if not self.client.collection_exists(self.collection_name):
            self.client.create_collection(
                collection_name=self.collection_name,
                vectors_config=VectorParams(
                    size=dim,
                    distance=Distance.COSINE
                )
            )

    @torch.no_grad()
    def embed_image(self, image: Image.Image) -> np.ndarray:
        inputs = self.clip_processor(
            images=image, return_tensors='pt'
        ).to('cuda')
        emb = self.clip_model.get_image_features(**inputs)
        return F.normalize(emb, dim=-1).cpu().numpy().squeeze()

    @torch.no_grad()
    def embed_text(self, text: str) -> np.ndarray:
        inputs = self.clip_processor(
            text=[text], return_tensors='pt', padding=True
        ).to('cuda')
        emb = self.clip_model.get_text_features(**inputs)
        return F.normalize(emb, dim=-1).cpu().numpy().squeeze()

    def index_product(
        self,
        product_id: str,
        product_image: Image.Image,
        metadata: dict
    ) -> None:
        embedding = self.embed_image(product_image)
        self.client.upsert(
            collection_name=self.collection_name,
            points=[PointStruct(
                id=hash(product_id) % (2**63),
                vector=embedding.tolist(),
                payload={
                    'product_id': product_id,
                    'category': metadata.get('category'),
                    'price': metadata.get('price'),
                    'brand': metadata.get('brand'),
                    **metadata
                }
            )]
        )

    def search_by_image(
        self,
        query_image: Image.Image,
        top_k: int = 20,
        filters: dict = None,       # {'category': 'shoes', 'max_price': 5000}
        score_threshold: float = 0.65
    ) -> list[dict]:
        query_embedding = self.embed_image(query_image)

        # Строим фильтры Qdrant
        qdrant_filter = None
        if filters:
            from qdrant_client.models import Filter, FieldCondition, MatchValue, Range
            conditions = []
            for key, value in filters.items():
                if key == 'max_price':
                    conditions.append(
                        FieldCondition(key='price', range=Range(lte=value))
                    )
                elif key == 'min_price':
                    conditions.append(
                        FieldCondition(key='price', range=Range(gte=value))
                    )
                else:
                    conditions.append(
                        FieldCondition(key=key, match=MatchValue(value=value))
                    )
            qdrant_filter = Filter(must=conditions)

        results = self.client.search(
            collection_name=self.collection_name,
            query_vector=query_embedding.tolist(),
            limit=top_k,
            query_filter=qdrant_filter,
            score_threshold=score_threshold,
            with_payload=True
        )

        return [
            {
                'product_id': r.payload['product_id'],
                'score': round(r.score, 4),
                'metadata': {k: v for k, v in r.payload.items()
                             if k != 'product_id'}
            }
            for r in results
        ]

    def search_by_text(
        self, query_text: str, top_k: int = 20
    ) -> list[dict]:
        """Text-to-image поиск: 'красные кроссовки Nike' → результаты"""
        query_embedding = self.embed_text(query_text)
        results = self.client.search(
            collection_name=self.collection_name,
            query_vector=query_embedding.tolist(),
            limit=top_k,
            with_payload=True
        )
        return [{'product_id': r.payload['product_id'],
                 'score': r.score} for r in results]

Кропирование объекта перед поиском

Если пользователь загрузил фото со сложным фоном — CLIP-embedding будет «загрязнён» фоном. Предварительное удаление фона улучшает точность поиска на 10–18%.

from rembg import remove as rembg_remove

def prepare_search_query(
    user_image: Image.Image,
    remove_background: bool = True,
    crop_to_object: bool = True
) -> Image.Image:
    if remove_background:
        # rembg → RGBA
        rgba = rembg_remove(user_image)
        if crop_to_object:
            # Авто-кроп по bounding box непрозрачных пикселей
            bbox = rgba.getbbox()   # (left, upper, right, lower)
            if bbox:
                rgba = rgba.crop(bbox)
        # Белый фон под объектом
        background = Image.new('RGB', rgba.size, (255, 255, 255))
        background.paste(rgba, mask=rgba.split()[3])
        return background
    return user_image

Fine-tuning CLIP на fashion-домене

CLIP обучен на общих данных, для специфических доменов (мода, мебель, электроника) есть смысл fine-tuning:

from transformers import CLIPModel, CLIPProcessor
import torch
from torch.optim import AdamW

def finetune_clip_for_domain(
    model: CLIPModel,
    train_loader,          # (image_tensor, text_tensor) пары
    num_epochs: int = 10,
    learning_rate: float = 1e-6   # очень малый LR — CLIP уже хорошо обучен
) -> CLIPModel:
    """
    Fine-tuning только visual encoder.
    Text encoder замораживаем — он нам нужен для text→image поиска.
    """
    for param in model.text_model.parameters():
        param.requires_grad = False

    optimizer = AdamW(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=learning_rate, weight_decay=0.01
    )

    for epoch in range(num_epochs):
        model.train()
        for batch_images, batch_texts in train_loader:
            outputs = model(
                input_ids=batch_texts['input_ids'].cuda(),
                attention_mask=batch_texts['attention_mask'].cuda(),
                pixel_values=batch_images.cuda()
            )
            # InfoNCE loss
            logits_per_image = outputs.logits_per_image
            labels = torch.arange(
                logits_per_image.shape[0], device='cuda'
            )
            loss = (F.cross_entropy(logits_per_image, labels) +
                    F.cross_entropy(logits_per_image.T, labels)) / 2

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    return model

Производительность

Сроки