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

Вода, камни, небо

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

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 возможностью добавления/удаления строк, напишите мне и я обновлю ноутбук/статью.

В нашем датасете оказалось 2 дубликата первого изображения

Плюсы

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

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

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

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

Минусы

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

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

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

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

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

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

RGB Histogram

[Colab]

RGB гистограмма

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

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

Линейный поиск. Сравниваем гистограммы методом cv2.HISTCMP_INTERSECT (53ms)

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

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

Bruteforce knn (73ms) и hnsw(0.4ms) выдают одинаковые изображения

Результат неплохой. Попробуем применить 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.

Поиск кропов (30s)

Плюсы:

  • 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')
Изображение по которому ищем
Водопады (ResNet50)
Поиск сильно пикселизированного изображения (ResNet50)
Поиск по кропу (ResNet50)

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

Водопады (CLIP)
Поиск сильно пикселизированного изображения (CLIP)
Поиск по кропу (CLIP)

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:

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)
"a picture of a sunset near the sea"
"a picture of a fog near the mountains"
"a picture of a windows xp wallpaper"

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