На днях мне наконец-то удалось поиграть в Cyberpunk 2077, и я заметил, что в игре есть одна интересная особенность: Когда персонаж говорит на иностранном языке, текст сначала появляется над ним в оригинале, а затем как бы вживую переводится на английский.

Тогда я задался вопросом: сколько работы потребуется, чтобы создать нечто подобное с помощью современного DL-стека? Можно ли сделать это за выходные?

В результаты работы хотелось иметь систему, которая могла бы:

• Обрабатывать короткие видеоклипы (например, одну сцену)

• Работать с несколькими персонажами / дикторами

• Распознавать речь на английском и польском языках

• Переводить речь на любой язык

• Определять кто говорит какую фразу

• Показывать того, кто говорит на экране

• Добавлять субтитры к оригинальному видео, подражая примеру Cyberpunk.

• Запускаться удаленно в облаке

TL;DR

С экосистемой ML, которую мы имеем сегодня, возможно построить PoC с такими требованиями за пару вечеров (как мне казалось).

Готовые инструменты достаточно просты в интеграции. Более того, обилие предварительно обученных моделей означает, что я мог создать все приложение, не обучая модели совсем.

Что касается сроков - это определенно заняло больше времени, чем я ожидал, но на самом деле большая часть времени была потрачена на вопросы, не связанные с ML (например, выяснение того, как добавить символы Unicode в кадр видео).

Вот клип 60-х годов из интервью, проведенного на польском языке и переведенного на английский.

А вот часть интервью Киану Ривза (который играет главного героя в Cyberpunk 2077) со Стивеном Колбертом, переведенная на польский язык.

Для того чтобы все летало будем использовать следующие технологии:

• ffmpeg-python для обработки видео

• Whisper для speech recognition

• NVIDIA NeMo для задачи speaker diarization

• DeepL для перевода

• RetinaFace для детекции лиц

• DeepFace для построения эмбеддингов по лицам

• scikit-learn для кластеризации лиц в кадре

• Gradio для demo фронта

• Modal для serverless deployment

Также использовался PIL и OpenCV для аннотации видеокадров и yt-dlp для загрузки роликов из YT.

Вот набросок того, как все это работает вместе для получения конечного результата:

Детали реализации

1. Обработка звука

1.1 Извлекаем звук

Извлечь звук из webm / mp4 очень просто с помощью ffmpeg:

def extract_audio(path: str, path_out: Optional[str] = None):
    """Extract audio from a video file using ffmpeg"""

    audio = ffmpeg.input(path).audio
    output = ffmpeg.output(audio, path_out)
    output = ffmpeg.overwrite_output(output)
    ffmpeg.run(output, quiet=True)

    return path_out

1.2 Speech to text

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

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

Я решил подавать весь аудиопоток в Whisper как один вход, но если вы хотите улучшить эту часть кода, вы можете поэкспериментировать с разделением аудио для каждого диктора, но я уверен, что это не даст лучших результатов.

1.3 Перевод

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

По моему опыту, DeepL работает лучше, чем Google Translate, а их API дает вам 500 тыс. символов в месяц бесплатно. Они также предоставляют удобный интерфейс на Python.

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

1.4 Определяем спикеров

Первоначально я использовал для этой цели PyAnnote, поскольку он доступен на HuggingFace и очень прост для интеграции.

pipeline = Pipeline.from_pretrained(
    "pyannote/speaker-diarization@2.1.1",
    use_auth_token=auth_token,
    cache_dir=cache_dir,
)

dia = pipeline(path_audio)

К сожалению, качество было так себе, и ошибки в этой части пайплайна вредили всем последующим шагам.

Поэтому далее я обратился к NeMo от NVIDIA. По их словам: "NVIDIA NeMo - это набор инструментов разговорного ИИ, созданный для исследователей, работающих над автоматическим распознаванием речи (ASR), синтезом текста в речь (TTS), большими языковыми моделями (LLM) и обработкой естественного языка (NLP)". NeMO отлично работает, особенно для английского языка. Есть проблемы с короткими сегментами перекрывающейся речи, но он определенно достаточно хорош для POC.

Самым большим недостатком является то, что NeMo - это исследовательский набор инструментов. Поэтому простые задачи типа "дай мне уникальные идентификаторы для этого аудиофайла" приводят к достаточно большим кускам кода.

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

1.5 Метчим спикеров и их фразы

Здесь я использовал простую эвристику, где для каждого участка речи (вывод NeMo) мы находим фразу, обнаруженную Whisper с наибольшим перекрытием.

Эта часть кода определенно может быть улучшена с помощью более сложного подхода.

2. Обработка видео

Тут все тривиально: cv2 и ffmpeg. Главный совет заключается в том, что для обработки видео лучше всего использовать генераторы - вероятно, вы не захотите загружать 1-минутное видео в массив numpy (1920 * 1080 * 3 * 24 * 60 записей займут ~35 ГБ оперативной памяти).

2.1 Face detection

К счастью, это уже давно решенная задача, решается с помощью RetinaFace или MTCNN. На этом этапе мы запускаем предварительно обученную модель для обнаружения всех лиц, видимых в каждом кадре. Затем мы обрезаем, выравниваем и изменяем их размер в соответствии с требованиями последующей модели.

На современном GPU обработка ~60 секунд видео занимает несколько минут.

К счастью, в Modal мы можем распараллеливать код, поэтому время выполнения сокращается, даже если обработка происходит на машинах с одним процессором.

2.2 Извлечение эмбеддингов из лиц

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

Для этого я взял модель FaceNet512 из библиотеки DeepFace.

После извлечения эмбеддингов нам нужно присвоить им уникальные идентификаторы. Для этого я использовал простой алгоритм иерархической кластеризации (а точнее, агломеративную кластеризацию из scikit-learn).

Агломеративная кластеризация будет рекурсивно объединять кластеры до тех пор, пока расстояние между ними будет меньше определенного порога. Этот порог зависит от модели и метрики. Здесь я использовал то же значение, которое используется DeepFace при выполнении "верификации лица".

2.3 Соединяем лица и звуки

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

Затем мы делаем то же самое для дикторов - создаем набор кадров, где слышен данный диктор.

Теперь для каждого ID лица мы находим ID диктора, для которого Jaccard index между двумя наборами минимален.

2.4 Объединяем все вместе

После того, как мы аннотировали каждый кадр с id диктора, id лица, фразой на оригинальном языке и фразой на переведенном языке, мы можем добавить субтитры.

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

Теперь осталось выяснить, как разместить субтитры так, чтобы они поместились на экране и т.д.

Эта часть кода может быть улучшена для работы с большим количеством языков. Чтобы разместить на экране символы UTF-8, мне потребовалось явно передать PIL путь к файлу шрифта. Проблема в том, что для разных языков требуются разные шрифты, поэтому текущее решение не будет работать, например, для корейского языка.

Деплоим

Подготовка фронтенда может быть выполнена с помощью Gradio за пару часов.

Serverless backend

Мы могли бы попробовать развернуть модель с помощью Huggingface Spaces, но я хотел попробовать что-то более "продакшн-готовное".

Я выбрал Modal - serverless платформу, созданную Эриком Бернхардссоном и его командой. Вы можете прочитать больше об этом в его блоге.

Пример локального кода:

def run_asr(audio_path: str):
    return whisper.transcribe(audio_path)


def process_single_frame(frame: np.ndarray, text: str):
    frame = add_subtitles(frame, text)
    return frame

Пример кода для Modal:

@stub.function(image=gpu_image, gpu=True)
def run_asr(audio_path: str):
    return whisper.transcribe(audio_path)


@stub.function(image=cpu_image)
def process_single_frame(frame: np.ndarray, text: str):
    frame = add_subtitles(frame, text)
    return frame

Очевидно, что еще есть некоторые шероховатости (Modal все еще находится в бета-версии), и мне пришлось решить одну последнюю проблему: при запуске приложения FastAPI существует ограничение в 45 секунд на каждый запрос. А поскольку обработка видео занимает больше времени, я воспользовался костылем: при первом нажатии кнопки Submit вы получаете id, и вы можете использовать этот id для получения конечного результата:

В конце хотел бы привести некоторые ограничения этого POC:

• Обработка 30-секундного видео занимает несколько минут

• Сопоставление лиц и голосов основывается на простой эвристике не будет работать в определенных сценариях (например, если весь разговор между двумя людьми записан с одного угла)

• Многие этапы пайплайна опираются на упрощенные эвристики (например, поиск уникальных лиц с помощью агломеративной кластеризации).

• Пайплайн был протестирован только на нескольких примерах

Github репозиторий с кодом

Еще больше примеров использования ML в современных сервисах можно посмотреть в моем телеграм канале. Я пишу про ML, стартапы и релокацию в UK для IT специалистов.