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

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

В мире есть много вещей, которые интуитивно понятны и очевидны для нас. Например, если перед нами два похожих цветка, мы можем определить их принадлежность одному виду, даже не зная названий этих растений. Этот навык позволяет нам распознавать объекты и определять их в группы. Разумеется, подобные алгоритмы уже давно существуют в современных поисковиках Google, Яндекс и прочих. Но что, если вы проектируете обособленную систему с собственной базой изображений одной или нескольких конкретных тематик и вам необходим функционал поиска похожих изображений?

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

Изображение для компьютера

Изображение на самом низком уровне (grayscale, цвет варьируется от 0 до 255, где 0 — черный цвет, а 255 — белый цвет) для компьютера представляет собой матрицу интенсивности цвета.

Казалось бы, чтобы оценить, насколько похожи 2 картинки, достаточно сравнить интенсивность в каждом пикселе. 

На самом деле описанная выше гипотеза сравнения изображений примитивна. 

Рассмотрим 2 изображения:

.

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

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

Векторизация изображения

Сети Vision Transformer (ViT) показали выдающуюся эффективность в векторизации изображений, превосходя традиционные методы. Их способность улавливать глобальные и локальные зависимости в данных делает их востребованными в различных областях компьютерного зрения.

Опишем процесс векторизации изображения с помощью базовой сети-трансформера.

1. Разбивка на патчи

Изображение делится на небольшие фрагменты, называемые патчами. Каждый патч представляется в виде одномерного вектора, который содержит информацию о содержимом данного фрагмента. 

Предположим, у нас есть изображение размером 48 x 48. Мы будем разбивать его на патчи фиксированного размера 16 x 16, так что получим (48/16) x (48/16) = 9 патчей. 

Стоит отметить, что на текущий момент изображения все еще находятся в 2D-формате (патчи имеют размер 16×16×3, предполагая изображение RGB), поэтому нужно их развернуть в 1D-формат. Для этого участки изображения линейно проецируются в вектор с использованием матрицы вложений, которая чаще всего обозначается как E.

Теперь, чтобы сохранить позиционную информацию (расположение объектов на исходном изображении), в трансформер кодировщик передается порядковый номер патча изображения вместе с этим самым патчем. Это просто обучаемый вектор. Каждый из векторов параметризован,  вместе они формируют обучающуюся таблицу вложений.

2. Векторизация с помощью трансформера

Для обучения трансформера кодировщика в каждую входную последовательность добавляется токен CLS. Он является своеобразным маркером, который выступает в роли агрегированного представления всей последовательности для решения конкретных задач (в контексте задачи векторизации изображения — это соотношение изображение-вектор). 

После получения последовательности эмбеддингов патчей они передаются в трансформер-кодировщик. Он состоит из L идентичных слоев, и в каждом слое он имеет два основных блока: многоголовое самовнимание (MSA) и многослойный персептрона (MLP), между которыми применяется слой нормализации. Также после каждого блока идёт остаточное соединение (Residual Connection). 

MLP состоит из двух полносвязных слоев (FC) с нелинейной функцией активации GeLU между ними.

Self-attention (SA) — это механизм внимания, который позволяет учитывать взаимные зависимости между различными элементами в последовательности данных. Обычно SA имеет последовательность, Q, ключ, K и значение V в качестве входных данных. Результатом является взвешенная сумма векторов значений, а вес вычисляется путем скалярного произведения Q и K, деленного на корень из размерности ключевого вектора. Затем для вычисления веса применяется функция Softmax. 

Слой множественного внимания (MSA) разбивает входные данные на несколько частей и параллельно вычисляет точечное произведение по отдельности. Это связано с тем, что многократное проецирование Q, K и V может принести пользу системе в целом. Таким образом, MSA может обрабатывать информацию из разных подпространств в разном положении.

3. Агрегация векторов

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

4. Расстояние между векторами

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

Для двух векторов A и B косинусное расстояние вычисляется по формуле:

Принято считать*, что косинусное расстояние принимает значения от 0 до 1:

• 1 означает, что векторы ортогональны (нет сходства),

• 0 означает, что векторы сонаправлены и полностью совпадают.

  * бывают случаи, когда косинусное расстояние превышает 1, но рассматривать их природу в данной статье мы не будем.

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

Реализация алгоритма

1. Подготовка функций

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

В свободном доступе существует огромное количество моделей-трансформеров, способных производить векторизацию изображений. Мы будем использовать модель clip-ViT-B-16 из библиотеки SentenceTransformer. Данная модель обучается с использованием контрастивного подхода, что позволяет ей улавливать и геометрические, и семантические закономерности объектов на изображениях. На самом деле платформа Hugging Face предоставляет доступ и к другим моделям векторизации изображений, например clip-ViT-B-32 и clip-ViT-L-14. Мы остановили свой выбор на clip-ViT-B-16, так как это компромисс в плане размера модели и её качества. Считывать изображения будем с помощью библиотеки PIL. 

from sentence_transformers import SentenceTransformer
from PIL import Image

model_name = 'clip-ViT-B-16'
st_model = SentenceTransformer(model_name)

def vectorize_img(img_path, model=st_model):
    img = Image.open(img_path)
    return st_model.encode(img)

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

import pandas as pd
import os

def create_images_db(images_folder, model=st_model):
    data_dict = dict()
    for file_name in os.listdir(images_folder):
        if os.path.isfile(images_folder + file_name):
            image_path = images_folder + file_name
            emb = vectorize_img(image_path)
            data_dict[file_name] = emb
    return pd.DataFrame(data_dict.items(), columns=['Image', 'Embedding'])

Так как обычно база данных создаётся не при каждом запуске алгоритма поиска, то нужно предусмотреть сценарий использования без регулярного запуска функции create_images_db(). Для этого достаточно, во-первых, сохранить в любой удобный формат датафрейм (мы используем JSON) и, во-вторых, написать функцию по корректному импорту этого файла. Если просто прочитать JSON, то значения в столбце с эмбеддингами будут храниться как строки, что нам не подходит.

import numpy as np

def import_images_db(file_path):
    data_df = pd.read_json(file_path)
    data_df['Embedding'] = data_df['Embedding'].apply(lambda x: np.array(x))
    return data_df

Также подготовим функцию для вычисления косинусного расстояния между двумя векторами. Будем использовать модуль spatial из библиотеки scipy.

from scipy import spatial

def calculate_cos_dist(emb_a, emb_b):
    result_distance = spatial.distance.cosine(emb_a, emb_b)
    return result_distance

Осталось оформить финальную функцию для поиска по базе. Она будет возвращать n-самых близких изображений из базы по отношению ко входному изображению.

import copy

def found_similar_images(input_img_path, images_db, n=1):
    input_vec = vectorize_img(input_img_path)
    result_df = copy.deepcopy(images_db)
    result_df['Distance_with_input'] = result_df.apply(lambda x: calculate_cos_dist(input_vec, x['Embedding']), axis=1)
    result_df_sorted = result_df.sort_values('Distance_with_input').reset_index()
    return result_df_sorted['Image'].head(n)

В итоге наш проект выглядит следующим образом:

from sentence_transformers import SentenceTransformer
from scipy import spatial
from PIL import Image
import pandas as pd
import numpy as np
import copy
import os

model_name = 'clip-ViT-B-16'
st_model = SentenceTransformer(model_name)

def vectorize_img(img_path: str, model: SentenceTransformer=st_model) -> np.array:
    img = Image.open(img_path)
    return st_model.encode(img)

def create_images_db(images_folder: str, model: SentenceTransformer=st_model) -> pd.DataFrame:
    data_dict = dict()
    for file_name in os.listdir(images_folder):
        if os.path.isfile(images_folder + file_name):
            image_path = images_folder + file_name
            emb = vectorize_img(image_path)
            data_dict[file_name] = emb
    return pd.DataFrame(data_dict.items(), columns=['Image', 'Embedding'])

def get_df(df_path: str) -> pd.DataFrame:
    data_df = pd.read_json(df_path)
    data_df['Embedding'] = data_df['Embedding'].apply(lambda x: np.array(x))
    return data_df

def calculate_cos_dist(emb_a: np.array, emb_b: np.array) -> float:
    result_distance = spatial.distance.cosine(emb_a, emb_b)
    return result_distance

def found_similar_images(input_img_path: str, images_db: pd.DataFrame, n: int=1) -> pd.DataFrame:
    input_vec = vectorize_img(input_img_path)
    result_df = copy.deepcopy(images_db)
    result_df['Distance_with_input'] = result_df.apply(lambda x: calculate_cos_dist(input_vec, x['Embedding']), axis=1)
    result_df_sorted =    result_df.sort_values('Distance_with_input').reset_index()
    result_df_sorted = result_df_sorted[['Image', 'Distance_with_input']]
    return result_df_sorted.head(n)

2. Демонстрация

Добавим в папку data изображения для создания нашей базы:

В качестве тестового изображения будем использовать следующее:

Создадим базу (для демонстрации мы пропустим этап экспорта базы в JSON) и найдём самое близкое изображение для bunny_check.jpeg.

images_folder = 'data/'
images_db = create_images_db(images_folder)
input_img_path = 'bunny_check.jpeg'
result_df = found_similar_images(input_img_path, images_db)
print(result_df.iloc[0])
# bunny_is.jpg

Заключение

Сегодня мы познакомили вас с задачей векторизации и сравнения изображений, а также рассказали, как можно спроектировать алгоритм поиска изображений на Python. Разумеется, всегда остаются варианты возможных улучшений и доработок алгоритма. Например, вы можете написать функцию добавления нового изображения в уже имеющуюся базу или поэкспериментировать с другими моделями построения эмбеддингов изображений. Будем рады обратной связи, делитесь своими мыслями по тематике статьи в комментариях.