Как мы с помощью ИИ выбираем обложки для сериалов в KION: кейс MTS AI

Привет, Хабр! На связи вновь Андрей Дугин, руководитель группы видеоаналитики компании MTS AI. Сегодня я закончу рассказ о том, как мы с помощью ИИ выбираем обложки для сериалов в KION. Первую часть можно прочитать здесь.

b2068af7fa74a0d136172f0307139f54.jpg

В прошлой части мы закончили на скоринге по лицам. Сейчас расскажу, что происходит после этого этапа.

Оцениваем композицию и обучаем нейронную сеть

Проверка оригинальных гипотез — наше всё, и для оценки композиции у нас тоже нашлась одна. Мы предположили, что главные объекты в кадре всегда будут в фокусе и находятся в определённых типичных местах. Нельзя ли в таком случае обойтись вообще без Object Detection? Оказалось — можно! Сейчас расскажу, как именно.

Мы умеем считать Лапласиан и его вариацию в качестве меры резкости. Попробуем разбить ранее посчитанный Лапласиан (изображение) на сетку — например, 16 × 9 — и посчитаем резкость в каждой ячейке. Если изображение не бьётся нацело, сделаем «резиновую» сетку, где ячейки могут слегка накладываться друг на друга или отстоять на небольшое расстояние. Погрешность в 1–2 пикселя нам погоду не сделает.

Пример на кадре из мультфильма «Энканто»

Пример на кадре из мультфильма «Энканто»

Если посчитать вариацию Лапласиана по каждой ячейке, то мы получим грубую карту глубины кадра, а точнее, карту ГРИП (ведь расфокусированные объекты могут быть как на переднем, так и на заднем плане). Полученную матрицу размерности (9, 16) вытянем в вектор (1, 144).

Далее нужно кластеризовать полученные векторы и найти центры кластеров. Но даже для одного фильма датасет размерности будет порядка (~165 000, 144), и кластеризация по алгоритмам scikit-learn может занять от получаса до бесконечности. Так что просто обучим элементарную однослойную нейросеть (это даже не Deep Learning).

Базово сеть состоит из Dense-слоя со 144 юнитами и kernel_constraint=«unit_norm» (все векторы-столбцы матрицы весов будут единичной длины), инициализировать можно случайными неотрицательными числами. Для выбора максимального значения ставим слой типа GlobalMaxPooling. Если на вход подавать нормированные векторы единичной длины, то нейросеть будет умножать входной вектор на матрицу весов и выдавать оценку одной (из 144) наиболее подходящей композиции в диапазоне [0.0, 1.0]. Если подавать ненормированные векторы, то в оценку композиции будет входить и оценка резкости изображения, но значение результата может стать значительно больше единицы.

Обучать можно только на позитивных примерах (мы так и делали), подавая на вход все кадры из набора фильмов, а на выход — всегда единицу. Если хочется поэкспериментировать с негативными примерами, то «добудем» их, просто перемешав компоненты каждого вектора, а на выход в качестве значения будем подавать ноль. В качестве loss-функции подойдёт как bce, так и mae или mse — в столбцах матрицы весов в любом случае будут формироваться векторы-центроиды кластеров самых распространённых композиций, при этом обучение нейронки занимает не более нескольких минут даже на CPU.

Где взять качественные данные? Сначала мы насобирали с десяток 3D-мультфильмов от студий Disney, Pixar, Sony Pictures и т. п. В них почти нет шумов и смазанности, хорошие карты глубины и отличная композиция кадров. Если в обычном фильме кадр может быть неудачным, то в мультфильме практически всё идеально. Первые эксперименты показали жизнеспособность такого подхода, однако одним из недостатков стали «пластиковые» лица мультяшных героев, практически лишённые текстур и, соответственно, резкости. Стали искать качественный материал среди сериалов и нашли турецкий ремейк «Доктора Хауса» под названием «Доктор Хаос» (73 серии есть на KION).

Там всё прекрасно — узкая ГРИП, объекты хорошо отделены от фона, хорошее боке и полное отсутствие шумов. Вот пример кадра:

adc176cf3ed46cc35f0ddbfa642662ef.png

Посмотрим на векторы-столбцы матрицы весов обученной нейросети, преобразовав каждый вектор из (144, 1) в мини-картинку (9, 16) и разместив 144 такие картинки в сетку 9 × 16:

546fee4dc37bb7d57ce0db4dd016d126.png

Для ещё большей наглядности отсортируем по среднеквадратичному отклонению:

a11832aabdf30b3b61336861a8bc5bf6.png

Красиво получилось, правда? Видно, что примерно половина выученных карт ГРИП — это хорошо узнаваемые «голова + плечи» (это в принципе особенность сериалов и является некоторым недостатком), также встречаются стоящие рядом люди и другие композиции.

Теперь, обучив нейросеть, мы можем оценивать каждый кадр любого фильма и считать как оценку композиции (опционально с мерой резкости кадра), так и номер композиции. Это позволяет разложить кадры фильма по отдельным папкам. Вот, например, композиция номер 10 (заметьте — никакого AI для поиска лиц!):

0f7f9f6fe9d96a06f0c2f57acf30bcb5.png

А вот композиция номер 4 — силуэты с наклоном головы:

c4b624447be0685ecf685c603dd9e018.png

При этом названия файлов соответствуют значениям оценки композиции с учётом резкости и отсортированы по этому признаку. Слева вверху — «худший» кадр, справа внизу — «лучший». Вот они:

549a54a9678f238beb4adc413ffc2a02.jpg

Любопытно, что следующие два кадра принадлежат к одному кластеру и имеют одинаковую композицию, поскольку в каждом из них есть зоны резкости слева и справа от центра:

35a4801fceb425e2ec8c9dc0bfdb09e0.jpg

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

Поиск текста в кадре

Для простого детектирования текста (без распознавания) мы использовали готовую модель из OpenVino, она превосходно справляется с поиском текста, а при низких порогах confidence находит даже единичные буквы. Это позволяет нам избавиться от кадров с титрами/субтитрами и другими нежелательными надписями.

Модель для детекции текста

Модель для детекции текста

Оценка фотографического качества

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

Гистограммы качества изображений

Гистограммы качества изображений

Для того чтобы оценить фотографическое качество кадра, стоит использовать два цветовых пространства:

  • LAB, где L — яркость, А — насыщенность и баланс белого (красный/зелёный), B — насыщенность и баланс белого (синий/жёлтый)

  • HSL/HSV, где H — оттенки цвета (HUE), S — насыщенность (saturation) и L/V — яркость и контраст

Нас интересуют каналы L, A/B и H. Довольно просто смоделировать некие «целевые» гистограммы, которые будут давать хороший результат. Распределение яркости пикселей в канале L должно быть похоже на широкое гауссовское с максимумом в средних тонах. Каналы A/B также похожи на центрированное гауссовское распределение, но значительно более узкое. Если мы хотим видеть разнообразие и богатство (не путать с насыщенностью) цветов на картинке, то канал H должен быть близок к равномерному распределению.

Такие гистограммы довольно хороши:

«Идеальные» канальные гистограммы

«Идеальные» канальные гистограммы

Для оценки того, насколько гистограмма изображения близка к идеальной, воспользуемся «расстоянием» Кульбака-Лейблера, которое позволяет оценить степень сходства двух распределений. Воспользуемся функцией kl_div из библиотеки scipy:

47347a7e4a085c3d2b0250e0f1e7ee6a.png

Ниже приведён пример класса, который одновременно оценивает все три гистограммы:

class CinematicQuality:
    def __init__(self):
        self.metrics = {
            'many_colors': \
                HistogramQuality(sigmas=1e-32, cvtcolor=cv2.COLOR_BGR2HSV_FULL, channels=[0]),
            'vibe_colors': \
                HistogramQuality(sigmas=1.0, cvtcolor=cv2.COLOR_BGR2LAB, channels=[1, 2]),
            'hist_spread': \
                HistogramQuality(sigmas=3.0, cvtcolor=cv2.COLOR_BGR2LAB, channels=[0]),
        }

    def score(self, image):
        scores = [metric.score(image) for metric in self.metrics.values()]
        score = np.linalg.norm(scores) / np.sqrt(len(scores))
        return score

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

Сортировка изображений по насыщенности цветов (слева вверху — низкая, справа внизу — высокая)

Сортировка изображений по насыщенности цветов (слева вверху — низкая, справа внизу — высокая)

Сортировка изображений по количеству цветов (слева вверху — мало цветов, справа внизу — много)

Сортировка изображений по количеству цветов (слева вверху — мало цветов, справа внизу — много)

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

Гистограмма титров

Гистограмма титров

Собираем результаты вместе

Для сведения всех параметров воедино используется класс CandidatesSelection:

class CandidatesSelection:
    def __init__(self,
                 path_to_video: Union[str, Path],
                 labeling: Optional[ActorsLabeling] = None,
                 num_of_posters: int = 150,
                 require_faces_presence: Literal["all", "arbitrary", "none"] = "arbitrary",
                 min_scene_duration_seconds: float = 3.0,
                 text_confidence_threshold: Optional[float] = 0.25,
                 face_confidence_threshold: Optional[float] = 0.0,
                 min_face_area: float = 0.0,
                 max_face_area: float = 1.0,
                 face_reduction_step: float = 0.05,
                 qual_reduction_step: float = 0.01,
                 autocrop: bool = True,
                 image_digits: int = 8,
                 image_suffix: str = ".png",
                 output_video_suffix: str = ".posters.mp4",
                 output_archive_suffix: str = ".posters.zip",
                 temp_folder_name: str = "frames",
                 cleanup: bool = False,
                 visualize: bool = False):

Значение передаваемых параметров очевидно из их названий, необходимо только пояснить face_reduction_step и qual_reduction_step: эти параметры обозначают долю отбрасываемых «худших» кадров, когда мы итеративно усекаем количество оставшихся кадров-кандидатов, проверяя то скоринг по лицам, то скоринг по фотографическому качеству — до тех пор, пока не останется заданное количество кадров num_of_posters.

Наконец, полученные кадры собираем в MP4 (короткий lossless-видеоролик, который будет передан legacy-коду для выбора единственного постера) или ZIP (архив для дизайнеров, которые смогут использовать подборку кадров для отрисовки альтернативного постера к фильму, помогающего бороться с «баннерной слепотой»).

def extract_short_video(self,
                        # Кадры можно собрать в видео (для генератора постеров) или в zip-архив (для дизайнеров)
                        container: Literal["video", "archive"] = "video",
                        # Для дизайнеров можно пропустить выборку лучших кадров и упаковать просто "сырую" нарезку
                        unfiltered: bool = False) -> Path:
    # Рядом с видео создадим папку для временного сохранения постеров
    self._prepare_temp_folder()
    # Извлекаем избранные кадры во временную папку
    self._extract_candidate_frames_to_temp_folder()
    # Фильтруем и удаляем недостаточно качественные кадры
    if not unfiltered:
        self._filter_extracted_candidate_frames()
    # Собираем оставшиеся кадры в выбранный контейнер
    if container == "video":
        path_to_container = self._create_video_from_frames()
    elif container == "archive":
        path_to_container = self._create_archive_from_frames()
    else:
        raise ValueError
    # Подчищаем за собой
    if self.cleanup:
        self._prepare_temp_folder(rmdir=True)
    # Возвращаем путь к созданному видео или архиву
    return path_to_container

Запуск команды ffmpeg для сборки короткого ролика выглядит так:

command = [
    "ffmpeg",
    "-y", "-hide_banner",
    "-v", "error",
    "-framerate", "1",
    "-i", path_to_hardlink.as_posix(),
    "-c:v", "libx264",
    "-tune", "stillimage",
    "-profile:v", "high444",
    "-preset", "ultrafast",
    "-crf", "0",
    "-pix_fmt", "yuv444p",
    path_to_output_video.as_posix()
]
sp.check_call(command)

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

Подборка исходников постеров для фильма «Мира»

Подборка исходников постеров для фильма «Мира»

Ура, у нас есть результат!

Выше я писал о том, что мы готовим TOP-150 постеров. И ниже — примеры заготовок для них. Как видите, кадры выглядят весьма впечатляюще, наша система отлично справилась с задачей:

64a78d110e9d05ebaec3178869270d34.pngОбработка мультфильма LadyBug and Cat Noir с подготовкой исходников постеров

Обработка мультфильма LadyBug and Cat Noir с подготовкой исходников постеров

Что дальше?

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

  • доработать алгоритм определения оптимальной композиции с добавлением его в наш пайплайн, в том числе с опциями выбирать кадры с людьми/без людей или как придётся

  • реализовать случайное чередование различных типов кадров, описанных выше, в пределах одного сезона сериала

  • заменить весь legacy-код на мультимодальную LLM-модель, которая позволяет проанализировать каждый постер-кандидат и ответить на вопросы о нём, чтобы выбрать одно лучшее изображение — такие модели получили сильное развитие за последний год

На этом всё. Если у вас есть вопросы или замечания по теме статьи, пишите в комментариях — обсудим.

© Habrahabr.ru