SearchIntroSerials

Автоматическое обнаружение заставок в эпизодах сериалов с использованием трансформерной модели и CRF.

Описание

Проект SearchIntroSerials решает задачу автоматического обнаружения временных интервалов заставок в эпизодах сериалов. Это полезно для стриминговых платформ, автоматизации обработки видео или улучшения пользовательского опыта (например, пропуск заставок). Модель использует эмбеддинги кадров, извлечённые с помощью CLIP, и трансформер с CRF-слоем для предсказания последовательности меток (заставка/не заставка). Проект включает обучение, инференс и оценку с метриками на уровне кадров (Precision, Recall, F1) и интервалов (IoU).

Содержание

• Установка
• Использование
• Подход
• Архитектура модели
• Обучение
• Инференс
• Результаты
• Примеры предсказаний
• Зависимости
• Лицензия
• Будущие улучшения

Установка

Для работы с проектом требуется Python 3.8+. Клонируйте репозиторий и установите зависимости:

git clone https://github.com/serttyzar/SearchIntroSerials.git
cd SearchIntroSerials
pip install -r requirements.txt

Убедитесь, что у вас есть доступ к видеофайлам в формате data_train_short/video_id/video_id.mp4 и меткам в data_train_short/labels.json.

Использование

1. Подготовка данных:

   • Извлеките кадры из видео с частотой 1 FPS (до 180 секунд).
   • Сгенерируйте эмбеддинги кадров с помощью CLIP (openai/clip-vit-large-patch14).
   • Используйте src/generate_dataset.py для создания dataset.pt:

   python src/generate_dataset.py

2. Обучение модели:

   • Запустите Jupyter Notebook training_notebook.ipynb для обучения модели:

   jupyter notebook training_notebook.ipynb

3. Инференс:

   • Используйте src/inference.py для предсказания заставок в новом видео:

   python src/inference.py --video_path videos/test_video.mp4 --model_path processed_dataset/model.pt

   Пример вывода:

   Заставка обнаружена: 20.00с - 35.00с
   

Подход

Проект использует следующий подход:

1. Извлечение кадров: Видео разбиваются на кадры с частотой 1 FPS (до 180 секунд).
2. Эмбеддинги: Кадры преобразуются в эмбеддинги с помощью CLIP для кодирования визуальной информации.
3. Моделирование последовательности: Трансформер обрабатывает последовательности эмбеддингов, учитывая временные зависимости.
4. Последовательная разметка: CRF-слой предсказывает метки (0 — не заставка, 1 — заставка) с учётом контекста.
5. Постобработка: Применяется binary_closing для сглаживания предсказаний и удаления коротких шумов.

Исследование выбора инструментов

• CLIP (CLIP Paper):

  • Выбран за способность создавать семантически богатые эмбеддинги, которые эффективно кодируют визуальные особенности заставок (анимация, логотипы, текст).
  • Альтернативы (ResNet, VGG) менее подходят, так как требуют дообучения для специфических задач.
  • Преимущество: Предобученная модель, не требующая дообучения, что экономит ресурсы.

• Трансформер (Attention Is All You Need):

  • Эффективен для обработки последовательностей благодаря механизму внимания, который улавливает зависимости между кадрами на большом временном диапазоне.
  • Альтернативы (LSTM, GRU) менее эффективны для длинных последовательностей из-за проблем с градиентами.
  • Выбраны параметры: n_heads=12, n_layers=6 для баланса между производительностью и вычислительными затратами.

• CRF (CRF Paper):

  • Обеспечивает согласованность предсказанных меток, что критично для выделения непрерывных интервалов заставок.
  • Кастомная реализация CRF выбрана для избежания зависимости от внешних библиотек (например, torchcrf) и гибкости настройки переходов.
  • Альтернативы (Softmax) не учитывают зависимости между метками, что снижает качество интервалов.

• Веса классов:

  • Веса классов (class_weights=[1.0, 2.0]) использованы для борьбы с дисбалансом (4.5175% меток 1).
  • Веса увеличивают значимость положительного класса (заставка), улучшая recall.

• Постобработка (binary_closing):

  • Сглаживает предсказания, устраняя короткие шумы, что повышает precision и IoU.

Почему эти инструменты?

• Эффективность: CLIP и трансформеры обеспечивают высокое качество без необходимости дообучения на больших датасетах.
• Гибкость: Кастомный CRF позволяет настраивать переходы (например, увеличивать вероятность "заставка → заставка").
• Робастность: Постобработка повышает качество интервалов, что критично для реальных приложений.

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

Модель IntroDetectionTransformer состоит из:

• Позиционное кодирование: Добавляет временную информацию к эмбеддингам [batch_size, 60, 768].
• Трансформерный энкодер: 6 слоёв, 12 голов внимания, обрабатывает последовательность эмбеддингов.
• Классификатор: Линейный слой, преобразующий выход трансформера в логиты для 2 классов (0, 1).
• CRF-слой: Моделирует зависимости между метками, вычисляя оптимальную последовательность.

class IntroDetectionTransformer(nn.Module):
    def __init__(self, d_model=768, n_heads=8, n_layers=4, num_labels=2, class_weights=None):
        super().__init__()
        self.pos_encoding = nn.Parameter(torch.zeros(1, 60, d_model))
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=n_heads, batch_first=True, dropout=0.1)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
        self.classifier = nn.Linear(d_model, num_labels)
        self.crf = CRF(num_labels, batch_first=True)
        if class_weights is not None:
            self.register_buffer('class_weights', class_weights)
        else:
            self.class_weights = None

    def forward(self, embeddings, labels=None, mask=None):
        x = embeddings + self.pos_encoding
        x = self.transformer(x) 
        logits = self.classifier(x)
        if labels is not None:
            if self.class_weights is not None:
                logits[:, :, 1] = logits[:, :, 1] * self.class_weights[1]
            return self.crf(logits, labels.long(), mask=mask)
        else:
            return self.crf.decode(logits, mask=mask)

CRF

Кастомная реализация CRF ввиду проблем с импортом библиотеки на Kaggle, где обучалась модель:

• Переходы: Матрица [num_tags, num_tags] (2x2) для вероятностей переходов между метками.
• Алгоритм Витерби: Используется для декодирования оптимальной последовательности меток.
• Лосс: Отрицательное логарифмическое правдоподобие, учитывающее истинные метки и маску.

Пример изменения весов CRF для повышения recall:

model.crf.transitions.data[1, 1] += 1.5  # Увеличиваем вероятность "заставка → заставка"
model.crf.transitions.data[1, 0] -= 1.5  # Уменьшаем вероятность "заставка → не заставка"

Обучение

Модель обучалась в Jupyter Notebook training_notebook.ipynb:

• Датасет: train_dataset.pt с эмбеддингами CLIP [60, 768] и метками [60]. Сформирован из эмбеддингов сериалов, обрезанных первыми 3 минутами(зачастую заставка в этом интервале)
• Предобработка:
  • Многократная аугментация (2 копии) для кадров с меткой 1.
  • Oversampling с весами [20.0, 1.0] для окон с метками 1.
• Лосс: CRF-лосс с весами классов [1.0, 2.0].
• Оптимизатор: Adam (lr=3e-5, weight_decay=1e-5).
• Эпохи: 30.
• Устройство: CUDA (GPU).

Код обучения:

for epoch in range(30):
    total_loss = 0
    for embeddings, labels in tqdm(dataloader, desc=f"Эпоха {epoch+1}/30"):
        embeddings, labels = embeddings.to(device), labels.to(device)
        mask = torch.ones(embeddings.shape[0], 60, device=device).bool()
        optimizer.zero_grad()
        loss = model(embeddings, labels, mask)
        if loss.dim() > 0:
            loss = loss.mean()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Эпоха {epoch+1}, Средняя потеря: {total_loss / len(dataloader):.4f}")

Инференс

Инференс выполняется с помощью inference.py или в training_notebook.ipynb:

1. Загружается модель и данные (test_dataset2.pt).
2. Модель предсказывает последовательности меток с помощью CRF-декодирования.
3. Применяется постобработка (binary_closing с окном 3) для сглаживания.
4. Вычисляются интервалы заставок с помощью scipy.ndimage.label.

Пример кода инференса:

model.eval()
with torch.no_grad():
    for embeddings, labels, video_ids, start_indices in test_dataloader:
        embeddings = embeddings.to(device)
        mask = torch.ones(embeddings.shape[0], 60, device=device).bool()
        predicted_sequences = model(embeddings, mask=mask)
        predicted_labels = np.array(predicted_sequences)
        closed_labels = [binary_closing(seq, structure=np.ones(3)).astype(int) for seq in predicted_labels]

Результаты

После обучения и постобработки достигнуты следующие метрики:

Метрика	Значение
Precision	0.6586
Recall	0.7726
F1-мера	0.7111
Средний IoU	0.6295

• Прогресс: F1-мера (0.7111) и Recall (0.7726) превышают целевые 0.7, но IoU (0.6295) ниже цели 0.7.
• Анализ: Постобработка улучшает precision и F1, и делаетrecall стабильным. Низкий IoU указывает на неточности в границах интервалов.

Примеры предсказаний

Пример 1: Видео -220020068_456255414.mp4.

• Предсказание: .
• Значение в labels: {"url": "https://vkvideo.ru/video-220020068_456249358", "name": "24 \u0447\u0430\u0441\u0430. 8 \u0441\u0435\u0437\u043e\u043d, 3 \u0441\u0435\u0440\u0438\u044f", "start": "0:00:05", "end": "0:00:15"}

Пример 2: Видео 220020068_456255389.mp4.

• Предсказание: .
• Значение в labels: {"url": "https://vkvideo.ru/video-220020068_456255389", "name": "\u0411\u0430\u0441\u043a\u0435\u0442\u0441. 1 \u0441\u0435\u0437\u043e\u043d, 9 \u0441\u0435\u0440\u0438\u044f", "start": "0:00:00", "end": "0:00:08"}.

Зависимости

• Python 3.8+
• torch
• torchvision
• transformers
• numpy
• scipy
• scikit-learn
• tqdm
• opencv-python

Установите зависимости:

pip install -r requirements.txt

Лицензия

Проект распространяется под лицензией MIT. См. LICENSE.

Будущие улучшения

• Дообучение CLIP для улучшения эмбеддингов.
• Добавление звуковых эмбеддингов может очень сильно повысить качество, к примеру MFCC или спектрограммы, так как заставка часто имеет яркое звуковое сопровождение
• Эксперименты с другими архитектурами (например, LSTM+CRF).
• Увеличение датасета с добавлением новых видео. Очень важный шаг + более сбалансированный датасет в плане классов.
• Настройка гиперпараметров CRF для повышения IoU.