Решение для хакатона "Цифровой прорыв" от команды "Русская красавица 3.0"

Описание результата

1. Поиск по картинке

Реализован путём генерации эмбеддингов (семантических векторов) через Open Source модель M-CLIP и вставки в БД ClickHouse. Затем, при получении пользовательского изображения высчитывается его эмбеддинг, и выполняется поиск в ClickHouse для нахождения наиболее близкого вектора. Топ самых близких по семантике изображений возвращается пользователю. Данный подход позволяет получить наибольшую масштабируемость благодаря свойствам ClickHouse и алгоритму, который позволяет провести предобработку изображений в датасете один раз и больше не производить тяжелых вычислительных операций.

Один INSERT изображения в БД и его обработка занимают менее одной секунды на обычном ноутбуке.

А SELECT (поиск изображения пользователем, требуемый в кейсе) выполняется мгновенно для базы с десятками тысяч записей.

2. Классификация категории

Реализована с применением той же модели M-CLIP, что уменьшает затраты по RAM и времени. Алгоритм таков: предгенерируем эмбеддинг для каждой из категорий, существующих в датасете, генерируем эмбеддинг для входного изображения и находим самую подходящую категорию, основываясь на схожести эмбеддинга изображения и этой категории. Так как здесь мы тоже используем семантический вектор входного изображения, его можно переиспользовать (посчитать только 1 раз), что увеличивает производительность.

Мы много экспериментировали с анализом категорий, улучшением промптов (префиксов), изменением названий категорий и т.д. Это позволило добиться прироста в 10% - с 53% до 63% точности.

На предоставленном датасете модель получает тестовую метрику accuracy = 63%

2.1 Улучшенная классификация категории !!

За пару часов до конца хакатона нам пришла в голову идея реализовать классификацию категории невероятно простым способом: найдя наиболее похожую картинку на введённую среди тысяч строк в ClickHouse, просто взять сохранённую категорию этого объекта и использовать в качестве результата классификации.

Мы еле успели реализовать попробовать его, и это простое решение дало шокирующе хороший результат - 86.1% accuracy на заданном датасете на неизвестных модели экспонатах.

В итоге мы решили полностью перейти на этот метод и предоставить его в качестве решения второй подзадачи. У него есть много плюсов: он простой, даёт большую точность, работает быстрее и представляет наибольшую масштабируемость, ведь никак не зависит от изменения категорий и структуры датасета.

3. Генерация описания

Генерация описания реализована с помощью модели GIT. Были перепробованы многие LLM и другие модели (BLIP, vit-gpt2, nougat-base, trocr и др.), остановились на этой модели, потому что она не потребляет слишком много ресурсов и предоставляет достойные ответы. Ответ от git-base занимает несколько секунд, ответ от git-large - больше (около 10 сек), но результат качественнее. Реализован выбор подходящей модели.

4. Фронтенд

Фронтенд в качестве демонстрации реализован на Python-фреймворке Gradio специально для демо ИИ-моделей.

5. Поиск по тексту (дополнительно)

Дополнительно реализовали поиск по тексту. Генерируем семантический вектор текста и ищем в БД.

6. Экспорт JSON (дополнительно)

Трэкер предложил нам классную идею: добавить на сайт кнопку экспорта полученных данных в JSON, чтобы впоследствии можно было интегрировать нашу ИИ-систему в существующие информационные системы музеев.

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

Схема JSON

{
    "similar_object_id": string,
    "category": string,
    "description": string
}

Локальное разворачивание

1. Зайти в папку solution
2. Выполнить pip install -r requirements.txt
3. Запустить ClickHouse с помощью docker compose up -d
4. Затем нужно загрузить данные из .csv в БД. Выполняем python db.py. Переменная DATASET_PATH должна быть установлена в корень распакованного mincult-train.zip.
   Это может занять некоторое время, поэтому, если приемлемо загрузить только часть датасета в БД, можно написать, например python db.py 1200, чтобы добавить в БД 1200 изображений.
5. Теперь запускаем python main.py и можем использовать сервис по адресу localhost:8042.