Привет, я Дима Абакумов, разработчик в диджитал-агентстве ДАЛЕЕ. Расскажу, как я написал бота на Python, который находит дубли мемов в нашем мем-чате, и какие методы сравнения изображений для этого использовал.

Как появился кейс?

Есть у нас классный чат с мемами в компании, который используется для поддержания командного духа. Это самая популярная активность в компании: у нас работает 300 человек, а в мем-чате сидит 179 из них. Это уже отдельная субкультура в рамках агентства. Чат появился стихийно в прошлом году, но за короткое время завоевал сердца сотрудников.

Наш HR-менеджер Аня Евсеева объясняет его популярность так: «Секрет один — все любят мемы! Тут каждый может найти единомышленника и скидывать картинки на любимую тему. Любишь мемы с котиками? — круто! Сделал самодельный мем со смешной фотографией с корпоратива? — ещё круче!». Возможно, вы и кликнули на эту статью, потому что в названии было слово «мем».

При этом в чате есть несколько правил:

1. Кидать можно мемы из интернета

2. Можно присылать самодельные, про компанию и команду

3. Но главное — нельзя повторяться, а то ждут репрессии от бан-полиции

Раньше повторы отслеживались вручную: кто-то скинул мем, ему отвечают, что такое уже было. Но это ещё надо доказать, поэтому ребятам приходилось листать бесконечную историю сообщений, чтобы найти пруфы. «Непорядок», — подумал я, и решил этот процесс автоматизировать. Вот что у меня получилось.

Первая версия — EfficientNet

Писать свою модель и обучать я не стал, это слишком затратно для задачи «напиши бота как можно быстрее», поэтому в первом варианте алгоритма использовал EfficientNet-Lite. Эта свёрточная нейросеть анализирует признаки изображения и задаёт для каждого из них вектор. Для сопоставления изображений используется принцип косинусного сходства.

Как это работает?

1. Я делаю ресайз изображения до 512×512 пикселей

2. Перевожу изображение в оттенки серого

3. Каждый пиксель нормализуется до 1

4. Модель вытягивает около двух тысяч признаков изображения и сохраняет их как векторы

5. Затем для каждого вектора исходного изображения и изображения из базы вычисляется косинус

def _get_fingerprint(self, filename):
  file = Image.open(filename).convert('L').resize(self.IMAGE_SHAPE) # self.IMAGE_SHAPE = (512, 512)
  file = np.stack((file,) * 3, axis=-1)
  file = np.array(file) / 255.0

  embedding = self.model.predict(file[np.newaxis, ...], verbose=2)
  embedding_np = np.array(embedding)
  flattened_feature = embedding_np.flatten()

  return flattened_feature

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

Бот определяет сходство новой картинки и мемов из собственной базы и выносит свой вердикт. Оптимальный порог идентичности  — 86%. При >86% бот не видел сходства там, где они были, а ниже — считал похожими абсолютно разные картинки.

У первой версии бота было много ложных срабатываний, которые сами стали источником мемов. Когда я только выкатил бота в чат, он вообще все картинки прогонял по базе и выдавал «Такого мема ещё не было», и так на каждую уникальную картинку, сообщений по 40 в день. В чате даже устраивали голосование, стоит ли удалить бота. Я всё починил за полчаса, и над ним смилостивились. Но мемы в истории остались.

Где бот лажал

После того, как бот несколько месяцев проработал в мем-чате, появились ложноположительные и ложноотрицательные срабатывания. 

Во-первых, бот плохо распознавал текстовые мемы. На разные скриншоты твитов бот реагировал как на одинаковые картинки, даже если текст не совпадал совершенно. За это доставалось от пользователей — мемами, подколами и даже посылами куда больше.

Во-вторых, бот мог определить идентичные изображения с разным размером или объёмом фона как непохожие. Некоторые участники пытались пользоваться этим лайфхаком: обрезать или уменьшать изображение, чтобы обойти бота. Но в большинстве случаев бан-комитет мем-чата все равно их разоблачал. А ещё упрекал бота в том, что он халатно относится к своим обязанностям.

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

Вторая версия — CLIP

Короче, все эти проблемы побудили меня начать работать над второй версией бота. Я провёл небольшое исследование, чтобы определить лучший метод сравнения изображений. В русскоязычном сегменте интернета, как водится, мало свежего контента по теме. Даже на Хабре последние релевантные статьи, если верить поиску, выходили в 2014 году. В тот момент мне очень помогла вот эта статья с описанием подходов.

Вот какие методы я рассматривал.

Цветовые гистограммы

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

Индекс структурной схожести (SSIM)

SSIM — метрика, оценивающая структурную схожесть между двумя изображениями. Она учитывает яркость, контраст и структуру, проставляя оценку между -1 (не схожи) и 1 (идентичны). 

Подход на основе глубокого обучения (deep learning)

Для второй версии бота я отказался от свёрточной нейронной сети и выбрал CLIP, нейросеть, обученную на парах «картинка-текст». К слову, алгоритм CLIP встроен в генеративные нейронные сети, такие как MidJourney или DALLE-3, для связи векторных представлений текста и изображений, что позволяет создавать новые изображения на их основе. Этот подход также помогает оптимизировать хранение медиаконтента, предотвращая загрузку дубликатов, а также реализовать рекомендации в интернет-магазинах на основе сходства модели или цвета. В общем, штука крайне разносторонняя и полезная. 

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

Тестирование методов и выбор лучшего

Я тестировал методы на трёх разных типах пар картинок.

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

   

2. Изображения одинаковые, разница только в размере или в объёме фона.  Идеальный результат, который должен выдать метод — 1

   

3. Изображения почти идентичны — отличаются только небольшие участки. Здесь чем ниже процент схожести, тем лучше

   

Что получилось

Для каждой пары картинок я вычислил процент идентичности по трём методам сравнения изображений.

import cv2
import torch
import open_clip

from skimage import metrics
from image_similarity_measures.evaluate import evaluation
from sentence_transformers import util
from PIL import Image


def histogram_based(image1, image2):
    hist_img1 = cv2.calcHist([image1], [0, 1, 2], None, [256, 256, 256], [0, 256, 0, 256, 0, 256])
    hist_img1[255, 255, 255] = 0  # ignore all white pixels
    cv2.normalize(hist_img1, hist_img1, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX)

    hist_img2 = cv2.calcHist([image2], [0, 1, 2], None, [256, 256, 256], [0, 256, 0, 256, 0, 256])
    hist_img2[255, 255, 255] = 0  # ignore all white pixels
    cv2.normalize(hist_img2, hist_img2, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX)

    # Find the metric value
    metric_val = cv2.compareHist(hist_img1, hist_img2, cv2.HISTCMP_CORREL)

    return round(metric_val, 2)


def structural_similarity_index(image1, image2):
    image2 = cv2.resize(image2, (image1.shape[1], image1.shape[0]), interpolation=cv2.INTER_AREA)

    # Convert images to grayscale
    image1_gray = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)
    image2_gray = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)

    # Calculate SSIM
    ssim_score = metrics.structural_similarity(image1_gray, image2_gray, full=True)

    return round(ssim_score[0], 2)


def clip_cnn(image1, image2):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-16-plus-240', pretrained="laion400m_e32")
    model.to(device)

    def image_encoder(img):
        img1 = Image.fromarray(img).convert('RGB')
        img1 = preprocess(img1).unsqueeze(0).to(device)
        img1 = model.encode_image(img1)
        return img1

    def generate_score(test_img, data_img):
        img1 = image_encoder(test_img)
        img2 = image_encoder(data_img)
        cos_scores = util.pytorch_cos_sim(img1, img2)
        score = round(float(cos_scores[0][0]) * 100, 2)
        return score

    return round(generate_score(image1, image2), 2)


COMPARATORS = [
    histogram_based,
    structural_similarity_index,
    clip_cnn
]

print("Different tweet. Lower is better")
image_path_1 = 'data/3/photo_2023-09-22_07-46-25.jpg'
image_path_2 = 'data/3/photo_2023-09-22_13-49-42.jpg'
image1 = cv2.imread(image_path_1)
image2 = cv2.imread(image_path_2)
for comparator in COMPARATORS:
    print(comparator(image1, image2))

print("Similar images but with padding. 1 is for ideal")
image1 = cv2.imread('data/2/photo_2023-09-26_22-50-15.jpg')
image2 = cv2.imread('data/2/photo_2023-09-26_22-54-42.jpg')
for comparator in COMPARATORS:
    print(comparator(image1, image2))

print("A different photo, but with slightly different text. Lower is better")
image1 = cv2.imread('data/1/photo_2023-10-11_10-02-11.jpg')
image2 = cv2.imread('data/1/photo_2023-10-11_10-02-14.jpg')
for comparator in COMPARATORS:
    print(comparator(image1, image2))

Для каждой пары я пытался найти оптимальный порог, для референса брал 86%. Несмотря на то, что в некоторых случаях метод гистограмм тоже давал приемлемые значения, в среднем, CLIP лучше справлялся с задачей. А вот SSIM подкачал по результатам тестирования. 

Кроме алгоритма бот сменил и библиотеку для предварительной обработки изображений: я отказался от PIL и выбрал OpenCV. Большая часть методов сравнения использовали именно её, к тому же у меня раньше не было опыта работы с ней. Хотелось это исправить.

Чтобы сравнивать новые картинки с уже присланными, нужно было придумать, как обойти техническую особенность ботов в Telegram. Получить доступ к истории чата можно через бота, который туда добавлен, — тогда будут видны только новые сообщения. Есть еще вариант через аккаунт пользователя. Но тут опасность в том, что ключ доступа от аккаунта (не пароль) хранится в коде, и если кто-то обратится к исходникам, то он сможет действовать от моего имени: это и переписки, и сообщения, и звонки. Поэтому в первой версии я просто выгрузил историю чата и скормил её боту, чтобы сформировать библиотеку в формате «картинка — дата отправки». Бот, когда видел повторяющийся мем, присылал не ссылку на конкретное сообщение, а изображение из базы.

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

Что впереди

Мемы в чате существуют не просто так: есть целая схема конкурсов за самые смешные картинки. Раз в неделю пост с самым большим количеством реакций становится мемом недели; из них выбирается мем месяца, а потом и мем года. Победители номинаций получают славу и почёт мем-чата, сертификат на онлайн-покупки или лимитированный мерч ДАЛЕЕ и диплом с уникальным дизайном. 

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

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