Интро

Для всех, кто знаком со свертками, задача мэтчинга персонажейне кажется сверхсложной. На Kaggle есть даже соревнования с подобной задачей и размеченный датасет с персонажами мультсериала Симпсоны. Но здесь ключевое слово — «размеченный».

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

Какие данные

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

Итого, у нас в датасете оказалось около 25 тысяч гравюр. Да-да, это только гравюры, о количестве персонажей пока речи не идет. А учитывая любовь граверов 18-19 веков к изображению сцен с массовыми скоплениями людей, можем утверждать сразу, что персонажей будет намного больше.

Как мы добывали персонажей, или кровь, пот, слезы и… lang SAM

Задача сегментации персонажей с подобного рода изображений нетривиальная: гравюры — совсем не фотографии, пропорции человечков на гравюрах сильно отличаются от пропорций обычных людей, на которых обучались алгоритмы сегментации. А многие персонажи вообще не похожи на людей, ведь в Европе 18-19 веков гравюры могли быть карикатурами, как, например, эта:

Обучать и размечать — сами мы не хотели этим заниматься, поэтому сравнивали несколько готовых решений (SAM и HQ-SAM в комбинации с детектором YOLOv8, lang SAM).

Пара слов о модели детекции: без разметки тут не обошлось, к сожалению. Мы дообучили модель детекции YOLOv8 large на размеченных нами данных. Разметка осуществлялась при помощи платформы Roboflow, датасет состоял из 500 гравюр, или 2666 персонажей (на обучение пошли 2130 персонажей, на валидацию — 370, остальное — на тест). Мы сравнили несколько обученных моделей и пришли к выводу: с увеличением эпох особого улучшения результата не было, поэтому взяли модель, учившуюся 70 эпох (ниже представлены результаты обучения и примеры работы алгоритма на тестовых данных).

Алгоритм работает неидеально: где-то границы bounding box’ов обрезают конечности, где-то, наоборот, — выделяют слишком много, а где-то и вовсе не выделяют ничего. Наверное, если бы мы разметили больше данных, то результаты детекции были бы лучше...

На вход SAM и HQ-SAM подавались предсказанные YOLO bounding box’ы, и в границах этих bounding box’ов алгоритм сегментации обводил персонажей. С lang SAM другая история — там на вход подавался промпт (путем экспериментов мы выяснили, что лучше всего для нашей задачи работает "human characters"), на основе которого встроенный в lang SAM алгоритм детекции GroundingDINO делал свои предсказания, а модель осуществляла сегментацию. Ниже представлены гравюры с результатами разных алгоритмов.

В итоге мы остановились на lang SAM — она разделяет персонажей лучше, чем остальные, что заметно по рисунку выше.

Вырезав всех сегментированных персонажей, мы получили датасет из 41421 персонажа. Стоит сделать оговорку, что мы брали не всех персонажей, а только тех, в которых модель была уверена (порог уверенности conf = 0.4 подобрали эмпирически). Ниже представлены примеры сегментированных персонажей.

Классификация всему голова

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

Итого: у нас задача — научиться классифицировать 41 тысячу с чем-то классов персонажей (если быть точными — 41421), где в каждом классе всего одно изображение, а этого недостаточно. На помощь приходит аугментация, с ней мы расширили исходный датасет. Единственное изображение в классе мы видоизменили при помощи библиотеки Albumentation. Мы использовали два вида аугментаций: слабую (или же base_aug — так она фигурирует на графиках и в описании экспериментов) и расширенную (она же extended_aug).

Базовая заключалась в несильном изменении персонажа. Предполагалось, что при таком виде аугментации алгоритму классификации не составит труда найти пары (персонаж; его аугментированная версия). Итак, в нее входили следующие манипуляции с исходным изображением:

1. CLAHE или перевод в ч/б — трансформации цвета изображения. Вероятность применения — 0.3. CLAHE — адаптивное выравнивание гистограммы с ограничением контраста, которое помогает улучшить контраст изображения.

2. RandomCrop — случайная обрезка изображения (высота = 300, ширина = 200). Вероятность применения также 0.3.

3. Изменение размера изображения на (400, 220) необходимо для дальнейшего расчета эмбеддингов при помощи бэкбона ResNet18, применялось ко всем изображениям.

4. Нормализации (mean = [0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225]) — рекомендация для исполизования ResNet18, применялась как и решейп ко всем изображениям (ссылка на документацию ResNet18).

Расширенная аугментация (она же extended aug) задумывалась для усложнения задачи нахождения пар (персонаж; его аугментированная версия), так как в этом случае аугментированная версия уже сильнее отличается от оригинала. В нее входила базовая аугментация и дополнительно к ней следующие манипуляции:

1. ShiftScaleRotate — поворот исходного изображения с ограничением угла вращения от -15° до 15°. Вероятность применения — 0.3.

2. HorizontalFlip — зеркальное отображение оригинала по горизонтали. Вероятность применения — 0.5.

3. Гаусовский шум или кастомная трансформация с рандомным расположением полос на изрображении для создания эффекта разрезанного изображения. Вероятность применения — 0.3.

Далее поговорим про self-supervised learning в виде контрастивного обучения. Суть алгоритма заключается в акценте на извлечении значимых представлений данных путем сопоставления схожих/положительных (в нашем случае это аугментированные персонажи) и не схожих/отрицательных (просто других персонажей). Метод применяется из предположения, что схожие персонажи лежат рядом в пространстве эмбеддингов, а остальные расположены на бОльшем расстоянии.

Мы работали с PyTorch реализацией модели SimCLR (вдохновлялись кодом).

Сейчас будет немного скучное перечисление технической составляющей экспериментов: как упоминалось ранее, в качестве бэкбона мы взяли ResNet18 и дообучили его, дополнительно учили projection head (с тремя линейными слоями и выходной размеренностью 128). В роли оптимизатора выступал AdamW (lr=0.25, trust_coef=1e03). В качестве лосса брали классический для подобного рода задач NT-Xent Loss, он же Normalized Temperature-scaled Cross Entropy Loss.

И вишенка на торте — валидация

Валидация состояла из двух больших частей.

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

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

Когда датасеты готовы, приступаем ко второй части. Ее проще расписать пошаговым алгоритмом:

1. Случайным образом из всех аугментированных изображений выбираем 2000 персонажей — они составят валидационный датасет. Каждый элемент валидационного датасета представляет собой пару (номер класса —аугментированное изображение). Валидационных датасетов в итоге было два (с фоном и без).

Далее для каждого элемента валидационного датасета повторяем (всего 2000 раз):

• Берем аугментированное изображение и считаем его эмбеддинг при помощи модели SimCLR.

• Измеряем косинусное расстояние между эмбеддингом из пп. 1. и эмбеддингами из соответствующего датасета эмбеддингов из первого шага. Получаем массив косинусных расстояний.

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

• Ищем позицию номера класса аугментированного изображения в массиве классов и на основе этой позиции рассчитываем метрики Top-1, Top-3, Top-10.

Всего здесь задокументируем 3 эксперимента:

Изменение функции потерь во время обучения SimCLR для каждого из экспериментов представлено на графике ниже.

Ниже представим графики со сравнением аутпутов моделей из всех трех экспериментов:

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

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

Послесловие и рефлексия

Разумеется, в ходе работы над классификатором мы экспериментировали еще и с замороженным бэкбоном ResNet18 и трансформером VIT (замороженным и размороженным), но результаты оказались сильно хуже описанных выше экспериментов, поэтому мы не включили эти части работы в статью.

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

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