Система рекомендаций для изображений: пример на Python и CLIP / Хабр

Всем привет! В этой статье хочу поделиться с вами моим кейсом создания быстрой, а главное — рабочей системы рекомендаций для картин.

Первые шаги

Все началось с того, что мне написал мой друг с просьбой проконсультировать его в области создания веб-приложений. Мы долго обсуждали его идею, и в какой-то момент я сказал:
«Давай я займусь сайтом. Я всё-таки программист».

Мы договорились начать.

Идея сайта была не новой, но довольно интересной — что-то вроде Pinterest, но для картин и предметов искусства. Главная фича нашего продукта должна была быть в:

подборе похожих картин
персонализированных рекомендациях для пользователя

Сначала мы занимались привычными вещами:

проектировали UI
продумывали UX
делали базовый каркас приложения

Но когда дело дошло до системы рекомендаций — я внезапно понял, что не знаю, как её реализовать.

Я начал искать статьи, читать документацию, спрашивать у ChatGPT. Самым очевидным решением оказалась система рекомендаций по тегам.

Например:

картина:  теги: импрессионизм, природа, вода
картина:  теги: портрет, природа, вода

Алгоритм был простой:

Берем картину
Смотрим ее теги
Находим картины с такими же тегами
Сортируем по количеству совпадений

query = """
SELECT
    p.id,
    p.title,
    COUNT(*) AS score
FROM painting_tags pt1
JOIN painting_tags pt2
    ON pt1.tag_id = pt2.tag_id
JOIN paintings p
    ON p.id = pt2.painting_id
WHERE pt1.painting_id = :painting_id
  AND pt2.painting_id != :painting_id
GROUP BY p.id
ORDER BY score DESC
LIMIT 10
"""

result = session.execute(query, {"painting_id": 1})

Но очень быстро всплыла проблема.

Каждую картину приходилось размечать вручную:

писать теги
проверять их
следить за качеством
подбор был не точный

Когда картин стало несколько сотен — это стало медленно, дорого и неудобно.

Переход к эмбеддингам

Я решил использовать эмбеддинги изображений.

Идея проста:

изображение → нейросеть → вектор признаков

Дальше можно искать похожие вектора, а значит — похожие картины.

Модель для эмбеддингов CLIP / OpenCLIP

Для генерации эмбеддингов существует много моделей, но я остановился на CLIP.

CLIP — это модель, обученная связывать изображения и текст в одном пространстве признаков.

Это дает несколько важных преимуществ:

1. Понимание содержимого изображения

CLIP обучен на огромном количестве изображений и подписей, поэтому он может улавливать:

стиль
композицию
объекты
настроение изображения

Например, две картины с похожей атмосферой будут иметь близкие вектора, даже если у них нет одинаковых тегов.

2. Единое пространство для текста и изображений

CLIP умеет кодировать:

изображение → embedding
текст → embedding

Это открывает возмо��ность:

искать картины по тексту
делать мультимодальные рекомендации

Например:

"impressionist landscape" → похожие картины

Сначала создадим сервис, который будет отвечать за генерацию эмбеддингов.

from transformers import CLIPModel, CLIPProcessor
import torch

Загрузка модели:

class CLIPEmbeddingService:

    def __init__(self, model_name="openai/clip-vit-large-patch14-336"):

        self.device = self._setup_device()

        self.model = CLIPModel.from_pretrained(model_name).to(self.device)
        self.processor = CLIPProcessor.from_pretrained(model_name)

        self.model.eval()

Модель загружается один раз и используется для inference.

Для ускорения вычислений сервис автоматически выбирает устройство.

def _setup_device(self):

    if torch.cuda.is_available():
        return torch.device("cuda")
    elif torch.backends.mps.is_available():
        return torch.device("mps")
    else:
        return torch.device("cpu")

Это позволяет использовать:

GPU
Apple Silicon
CPU

в зависимости от доступного оборудования.

Перед обработкой нужно привести изображение к формату PIL.

def _load_image(self, image_input):

    if image_input.startswith(("http://", "https://")):
        with urllib.request.urlopen(image_input) as response:
            image_bytes = response.read()

    elif os.path.isfile(image_input):
        with open(image_input, "rb") as file:
            image_bytes = file.read()

    image = Image.open(BytesIO(image_bytes)).convert("RGB")

    return image

Сервис умеет работать с:

URL изображений
файлами
base64 строками

Генерация embedding

Теперь можно получить вектор признаков изображения.

def _get_image_embedding(self, image_input):

    image = self._load_image(image_input)

    inputs = self.processor(
        images=image,
        return_tensors="pt"
    ).to(self.device)

    with torch.no_grad():
        embedding = self.model.get_image_features(**inputs)

    embedding = embedding / embedding.norm(p=2, dim=-1, keepdim=True)

    return embedding

Мы также нормализуем вектор, чтобы корректно использовать cosine similarity.

Получение embedding для картины

Теперь можно создать метод, который возвращает embedding для конкретного артворка.

def get_embedding(self, art_id, image_input):

    embedding = self._get_image_embedding(image_input)

    return {
        "id": art_id,
        "embedding": embedding
    }

На выходе получаем структуру:

{
  "id": 123,
  "embedding": [0.12, -0.44, ...]
}

Почему мы выбрали Elasticsearch

После генерации эмбеддингов возникает следующая задача:

где хранить вектора и как быстро искать похожие?

Для этого существует несколько специализированных векторных баз данных, но мы решили использовать Elasticsearch.

Причина довольно простая.

Нашему сервису нужно было не только искать похожие изображения, но и поддерживать обычный текстовый поиск.

Например:

поиск по названию картины
поиск по автору
поиск по описанию

Elasticsearch уже отлично решает такие задачи.

Гибридный поиск

Использование Elasticsearch позволило объединить два типа поиска:

Текстовый поиск

title: "water lilies"

Векторный поиск

embedding → похожие картины

В результате система рекомендаций стала выглядеть так:

изображение
   ↓
CLIP
   ↓
embedding
   ↓
Elasticsearch
   ↓
поиск похожих картин

Такой подход дал несколько преимуществ:

система не требует ручной разметки тегов
алгоритм анализирует само изображение
можно делать текстовый и векторный поиск в одной системе

Персонализированные рекомендации

Поиск похожих картин — это полезная функция, но нам хотелось сделать следующий шаг — персонализированные рекомендации.

Идея проста: если мы знаем, какие картины нравятся пользователю, можно попытаться понять его вкус и рекомендовать новые работы, которые могут ему понравиться.

Какие данные мы использовали

На первом этапе мы решили не строить сложную систему аналитики, а использовать простые пользовательские сигналы.

Мы учитывали:

лайки картин
посещения страниц картин

Эти события уже фиксировались в базе данных, поэтому их можно было использовать для построения профиля.

Где мы храним события

Пока что все события хранятся в PostgreSQL.

Причина довольно простая — нагрузки на систему пока небольшие:

пользователей не так много
поток событий относительно маленький
сложная аналитическая инфраструктура пока не требуется

Поэтому добавлять отдельную систему для аналитики было бы преждевременной оптимизацией.

Построение профиля пользователя

Каждая картина в системе уже имеет embedding, который хранится в Elasticsearch.

профиль пользователя = средний embedding картин, с которыми он взаимодействовал

Схематично это выглядит так:

картины пользователя
        ↓
получаем embeddings
        ↓
усредняем
        ↓
получаем embedding пользователя

Пример кода построения профиля

Сначала получаем embeddings картин, с которыми взаимодействовал пользователь.

like
view

Теперь получаем embeddings этих картин из Elasticsearch.

def get_embeddings(art_ids, es):

    embeddings = []

    for art_id in art_ids:

        doc = es.get(
            index="painting_embeddings",
            id=art_id
        )

        embeddings.append(doc["_source"]["embedding"])

    return embeddings

После этого можно построить embedding пользователя.

import numpy as np

def build_user_profile(embeddings):

    if not embeddings:
        return None

    return np.mean(embeddings, axis=0).tolist()

Получается вектор, который отражает предпочтения пользователя.

Теперь можно искать картины, похожие на профиль пользователя.

def recommend_for_user(user_embedding, es):

    query = {
        "knn": {
            "field": "embedding",
            "query_vector": user_embedding,
            "k": 10,
            "num_candidates": 100
        }
    }

    res = es.search(
        index="painting_embeddings",
        knn=query
    )

    return res["hits"]["hits"]

Таким образом пользователь получает персональную ленту картин.

Почему пока хватает PostgreSQL

Пока система работает на PostgreSQL, потому что объем данных относительно небольшой.

Мы не храним миллионы событий и не строим сложные аналитические модели, поэтому обычная реляционная база данных полностью справляется с задачей.

Но по мере роста продукта ситуация может измениться.

Планируемый переход на ClickHouse

Если количество пользователей и событий вырастет, мы планируем перенести хранение событий в ClickHouse.

ClickHouse хорошо подходит для:

хранения больших объемов событий
аналитических запросов
построения пользовательских профилей

Это позволит:

хранить значительно больше событий
учитывать временные факторы
строить более точные профили пользователей

Например, можно будет учитывать:

вес лайков
частоту просмотров
недавние интересы пользователя

Итог

В результате система рекомендаций стала состоять из нескольких частей:

изображение
    ↓
CLIP
    ↓
embedding
    ↓
Elasticsearch
    ↓
поиск похожих картин

+

события пользователя
    ↓
профиль пользователя
    ↓
персональные рекомендации

Даже относительно простая модель профиля позволяет значительно улучшить качество рекомендаций.

А по мере роста системы можно постепенно усложнять алгоритмы и инфраструктуру.

Надеюсь, что этот опыт и примеры кода будут полезны кому-то, кто планирует построить свою собственную систему рекомендаций или просто хочет понять, как реализовать такой кейс на практике.