Поиск изображений

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

Используемый датасет

Perceptual hash

[Colab]

Детальное описание работы phash

Из изображений мы создаем хеши заданной длины. Чем меньше расстояние Хэмминга между двумя хешами, тем больше схожесть изображений.

Наивный способ поиска - линейный поиск. Сравниваем хеш целевого изображения с хешами всех изображений и возвращаем N изображений с наименьшим расстоянием Хэмминга (либо отсекаем по какому-нибудь threshold'у).

Можно ли быстрее? Можно! Используя структуру данных Vantage-point tree, можно построить дерево за O(n log n) и осуществлять поиск за O(log n).

Немного о производительности

Внимательные читатели, которые просмотрели ноутбук, могли заметить, что скорость vantage-point tree либо наравне, либо медленнее простого цикла for. В ноутбуке есть закомментированный блок кода, который позволяет сгенерировать массив длиной 100к строк. Этот массив можно подставить вместо массива хешей и убедится, что с увеличением количества строк... ничего не меняется, vptree все так же проигрывает. В чем же дело? А дело в том, что если вы начнете искать реализацию vantage point tree в PyPI, то найдете лишь 1 пакет - vptree. Он работает довольно плохо, из-за чего прироста производительности нет. Я нашел реализацию vp-tree на javascript и написал небольшой тест производительности. for-loop перестает быть быстрее, чем vptree после 10к строк и дальше отрыв только увеличивается. Кто-то может сказать, что для получения top N элементов, весь массив сортировать необязательно. Согласен, но vp-tree быстрее, даже если мы не проводим сортировку массива вовсе. Ссылка на gist

Интересный факт - я не смог найти реализацию vp-tree, где есть функция добавления новых данных в дерево. Перестраивать дерево каждый раз, когда у вас появляется новый хеш может быть очень затратно. Если вы сможете найти/написать быструю реализацию vp-tree c возможностью добавления/удаления строк, напишите мне и я обновлю ноутбук/статью.

Плюсы

• Хеш имеет малый размер

• Быстро вычисляется

• Быстро ищется

• Устойчив к ресайзу

Минусы

• Не устойчив к кропу

• Не устойчив к поворотам

• Не усточив к {transformation_name}

Одна из самых частых трансформаций изображения - это кроп. Решение - создавать несколько хешей на одно изображение, хешируя "интересные" места.

https://habr.com/ru/post/205398/
https://habr.com/ru/post/211773/

Вывод: phash отлично подходит для дедупликации, поиска оригиналов изображений по их preview/thumbnail. Не устойчив к кропу.

RGB Histogram

[Colab]

Генерируем RGB гистограммы, сравниваем гистограммы, если гистограммы похожи, то изображения схожи по цвету.

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

После применение операции flatten, можно заметить, что гистограмма превращается в вектор из 4096 значений.

Попробуем применить k-nearest neighbor search, в качестве метрики выберем евклидово расстояние.

Результат неплохой. Попробуем применить approximate nearest neighbor search. Используем библиотеку hnswlib, которая реализует структуру данных под названием Hierarchical Navigable Small World. Теперь поиск занимает не 50-70ms, а всего лишь 0.4ms.

Новые данные можно добавлять в индекс без его полной перестройки.
Больше про approximate nearest neighbor search - https://habr.com/ru/company/mailru/blog/338360/

Плюсы

• Устойчив к трансформациям, не сильно меняющим гистограмму изображения

• Более устойчив к кропу, чем phash

Минусы

• Большой вес (Для 16 бинов RGB гистограммы вектор имеет размер 4096)

• Учитывает только цвета, не учитывает геометрию

SIFT/SURF/ORB

[SIFT Colab]

Подробнее про SIFT.

Можно использовать более быстрые алгоритмы, например SURF или ORB. Мы будем использовать модификацию SIFT - Root SIFT.

Суть модификации:

descs /= (descs.sum(axis=1, keepdims=True) + eps)
descs = np.sqrt(descs)

Вот таким нехитрым способом точность SIFT повышается на ~5 процентов.
План действий: генерируем SIFT features, применяем Brute-Force Matcher(cv2.BFMatcher), сортируем изображения по среднему расстоянию хороших matches.

Плюсы:

• SIFT инвариантен пропорциональному масштабированию, ориентации, изменению освещённости и частично инвариантен аффинным искажениям

Минусы:

• Занимает много места

• Медленно вычисляется

• Медленно ищется (Есть методы значительно ускоряющие поиск, но я не нашел понятного примера на python)

NN features

[Colab ResNet50] [Colab CLIP]

В процессе своего обучения нейросети решающие задачу классификации изображений становятся способны сжимать изображения до достаточно маленьких векторов. Вектора похожих изображений находятся близко друг к другу. Длина вектора в ResNet50 - 2048. "Открутим" полносвязный классифицирующий слой от ResNet50, сгенерируем векторы всех наших изображений и с помощью knn найдем наиболее близкие к целевому изображению. Используем евклидово расстояние в качестве метрики.

model = ResNet50(weights='imagenet', include_top=False,input_shape=(224, 224, 3),pooling='max')

Попробуем использовать другую нейросеть - CLIP. В ней нет классифицирующего слоя, просто используем метод encode_image и получаем вектор длинной 512.

CLIP справляется c поиском немного хуже, возможно из-за того, что его функция препроцессинга сжимает изображение до 224 с сохранением aspect ratio, а затем берет Center Crop, т.е не все детали могут попасть в кадр.

Визуализируем полученные вектора. Для этого используем алгоритм t-SNE.

t-SNE ResNet50 (10100x10100 7.91MB)

t-SNE CLIP (10100x10100 7.04MB)

Features которые генерирует CLIP более точно группируют похожие изображения. На визуализации векторов можно увидеть, что CLIP группирует вместе изображения, которые схожи не только визуально, но и по смыслу.

Плюсы:

• Хорошо справляется с большинством трансформаций

• Быстро ищется благодаря approximate nearest neighbor search

Минусы:

• Для быстрого вычисления features нужен GPU

CLIP text search

[Colab CLIP]

Подробнее про CLIP:

• https://habr.com/ru/post/539312/

• https://habr.com/ru/post/540312/

CLIP способен генерировать вектор из текстового запроса, причем этот вектор будет близок к векторам изображений, которые описываются в запросе. В наш knn search мы можем передать не features изображения, а features текстового запроса. Магия.

text_tokenized = clip.tokenize(["a picture of a windows xp wallpaper"]).to(device)
with torch.no_grad():
        text_features = model.encode_text(text_tokenized)
        text_features /= text_features.norm(dim=-1, keepdim=True)

Github со всеми ноутбуками