Поиск коротких заставок в сериалах

Разные подходы

Сначала подумаем, как в принципе можно решать эту задачу.

1. Самым простым решением на сегодняшний день кажется использование мультимодальной LLM c последующим извлечением таргета.

   • Системный промпт:

   Ты видеоаналитик, который обрабатывает сериалы. 
   Твоя задача — по содержимому видео найти заставку, которая обычно длится несколько секунд. 
   Заставка — это повторяющийся от серии к серии фрагмент, который связан с представлением сериала, часто с музыкой, логотипом, вступительными титрами и т. п.
   
   Если такая заставка есть, верни ВСЕ наиболее вероятные временные интервалы, в формате JSON. Формат вывода:
   
   {
       "<video_id>": {
           "url": "<video_url>",
           "name": "<название серии>",
           "start": "<чч:мм:сс>",
           "end": "<чч:мм:сс>"
       }
   }
   
   Важно: выдай **только JSON**, без пояснений. Не пиши ничего вне JSON. Если заставка не обнаружена — не возвращай объект.
   
   

   • Ещё можно испоьзовать few-shot: покажем модели на нескольких примерах из трейна, как нужно находить заставку и возврещать ответ.

2. Но мы будем обучать свою модель. Можно обрабатывать кадры независимо с помощью CNN (или трансформера, например) + линейный слой для классификации, но мы теряем свойство "связанности" кадров между собой, а это уже весомое допущение.

3. Можно разбить каждое видео на небольшие клипы с пересечениями. Длину клипа делаем гиперпараметром. Далее пропускаем клипы через видеоэнкодер. Получившиеся эмбеддинги кластеризуем, например, DBSCAN'ом. В классах, где много похожих клипов, находятся заставки.

4. Прошлый метод можно назвать "self-supervised", но у нас же есть трейн-выборка. Тогда можно поверх эмбеддинов поставить MLP Head для бинарной классификации: заставка/не заставка. Давайте рассмотрим этот метод подробнее.

Архитектура модели

В качестве видеоэнкодера можно использовать VideoMAE (Hugging Face (base model), arXiv). Выход модели пропускаем через MLP слой с выходной размерностью 1 (после применения сигмоиды - просто вероятность быть заставкой).

Общая схема похожа на архитектуру ViT (тем более, используемая модель основана на ViT):

Класс модели можно реализовать так:

class IntroModel(nn.Module):
    def __init__(self, pretrained_name="MCG-NJU/videomae-base", hidden_dim=256, freeze_backbone=True):
        super().__init__()
        
        # Загружаем модели с hugging face
        self.processor = AutoImageProcessor.from_pretrained(pretrained_name)
        self.backbone = AutoModelForPreTraining.from_pretrained(pretrained_name)

        # Замораживаем или нет слои трансформера (если нет, то полный fine turning модели)
        if freeze_backbone:
            for param in self.backbone.parameters():
                param.requires_grad = False
        
        # Классификация
        self.mlp = nn.Sequential(
            nn.Linear(self.backbone.config.hidden_size, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, 1)
        )
    
    def forward(self, video_frames):
        inputs = self.processor(video_frames, return_tensors="pt")
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        
        # Достаём CLS-токен
        outputs = self.backbone(**inputs)
        cls_emb = outputs.last_hidden_state.mean(dim=1)
        logits = self.mlp(cls_emb)
        probs = torch.sigmoid(logits)
        return probs

Подготовка данных

Нам нужно нарезать видео на клипы с пересечениями. Делать это будем непосредственно в классе датасета (его наследуем от torch.utils.data.Dataset).

Читать видео можно с помощью decord.VideoReader. Этот модуль позволяет вычислять fps и прочие нужные нам характеристики видео.

Перед тем, как вернуть элемент датасета, надо обработать данные через transformers.AutoImageProcessor и убрать батч-размерность. Далее преобразовать в "понятный" формат для hugging face трансформеров:

inputs = self.processor(frames, return_tensors="pt")
inputs = {k: v.squeeze(0) for k, v in inputs.items()}
to_return = {
    **inputs,
    "labels": torch.tensor(label, dtype=torch.float)
}

Обучение

Для обучения удобнее всего использовать transformers.Trainer c логированием в системы сбора метрик (Wandb, например).

Формально, мы выполняем задачу классификации с минимизацией бинарной кросс-энтропии: $$BCE = -[y\ log(p) + (1-y)\ log(1-p)]$$ где $y$ - метка таргета, $p$ - предсказанная вероятность.

В качестве метрики качества лучше использовать F1-score, так как классы заставка/не заставка сильно не сбалансированы (заставок сильно меньше, конечно).

Инференс

Наверное, тут уже ничего не сказать нового про ML, скорее интересен один алгоритмический момент:

1. Делим видео из теста на пересекающиеся клипы.
2. Пропускаем через модель, предсказываем класс.
3. Если предсказали заставку, то сохраним таймкоды.

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

Когда ответ для видео готов, формируем итоговый json.

Улучшения решения

• Главная проблема подхода в том, что мы режем видео сеткой, то есть у нас могут быть фрагменты, границы которых не совпадают с сеткой. Например, такие ситуации:

  

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

• Надо продумать, как оценивать точность модели не на промежуточной задачи классификации, а уже на итоговой - для поиска таймкодов. При этом нужно учитывать то, что мы режем видео сеткой.

• Можно использовать мультимодальные модели (например, CLIP) с указанимем промпта "заставка" или "интро". При этом использовать тот же принцип решения, что написан здесь.