Прежде чем перейти к статье, хочу вам представить, экономическую онлайн игру Brave Knights, в которой вы можете играть и зарабатывать. Регистируйтесь, играйте и зарабатывайте!
Привет, Хабр! Около полугода назад я наткнулся на статью с практически аналогичным названием . Сама идея мне показалась интересной и захотелось сделать нечто похожее только с использованием компьютерного зрения, и вот что у меня получилось:
Общая архитектура проекта
В основе проекта лежит три сервиса, где первый отвечает за детекцию субтитров и преобразования изображения текста в машиночитаемый текстовый формат, второй сервис предназначен для перевода текста( на данный момент только с английского на русский язык) и последний сервис отвечает за визуализацию и озвучку текста. Общение между сервисами реализовано с помощью библиотеки ZeroMQ.
Сбор данных
Решение практически любой ML задачи начинается со сбора данных — эта задача не исключение, поэтому необходимо было собрать датасет, который включает в себя скриншоты видеороликов с YouTube, на которых присутствуют субтитры и выделить ограничительные рамки самих субтитров. Вот основные требования к датасету:
1) Субтитры должны быть на разных языках.
2) Картинка с самим видеороликом должна быть разного размера, то есть видеоролик может быть на весь экран, так и на части экрана.
3) Субтитры находятся в разных областях экрана.
4) Требования к самим субтитрам:
могут быть разного размера;
семейство шрифтов: обычный и пропорциональный без засечек;
цвет шрифта: белый;
цвет фона: чёрный;
прозрачность фона: 75%;
прозрачность окна: 0%.
Датасет был успешно собран сотрудниками компании LabelMe , за что им огромное спасибо! Исходный датасет вы можете найти на Kaggle .
Детекция субтитров
После того как датасет был собран, необходимо обучить модельку находить субтитры, да ещё чтоб и в реальном времени детектировало. Выбор пал на yolov5 . И после 50 эпох обучения на предобученной модельке получились довольно впечатляющие метрики:
И вот как детекция работает в real-time:
Подготовка изображения к оптическому распознаванию текста
Большинству библиотек для распознавания текста, о которых поговорим чуть ниже, лучше передавать grayscale картинку, поэтому первым делом я именно это и делаю:
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 1))
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
_, thresh_img = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
transformation_image = cv2.morphologyEx(thresh_img, cv2.MORPH_OPEN,
kernel, iterations=1)На выходе получается нечто подобное:
Теперь-то изображение уже готово, казалось бы, бери да и передавай изображение в какую-нибудь библиотеку для распознавания текста, но не тут то было. Дело в том, что yolo выдаёт у меня примерно 10 кадров в секунду и соответственно будет получаться очень много изображений с одинаковым текстом, а зачем нам переводить и озвучивать один и тот же текст много раз?
В OpenCV есть такая прекрасная функция как matchTemplate(). Данную функцию можно рассматривать как очень простую форму обнаружения объектов. Используя сопоставление шаблонов, мы можем детектировать объекты на входном изображении, используя «шаблон», содержащий объект, который мы хотим обнаружить:
Чтобы найти шаблон на исходном изображении, мы перемещаем шаблон слева направо и сверху вниз по исходному изображению:
Данная функция одним из параметров возвращает условно вероятность от 0 до 1 насколько два изображения похожи между собой. Используя это число, мы можем предположить, что если оно больше, чем 0.75, то на этих двух изображениях расположен один и тот же текст. Более подробно про данную функцию можно почитать тут .
Вторая проблема это то, что текст на следующем кадре дополняется к предыдущему по 1-2 слова, то есть:
Как можно увидеть текст первой картинки отличается от текста второй картинки одним лишь «is». Думаю, это не вариант вычленять по одному слову и передавать это одно слово на следующий шаг, поэтому предлагаю определить каким-то образом, что строка закончилась и только лишь изображение на котором закончилась строка передавать на следующий этап. О том, как мы будем это определять поговорим чуть ниже.
С данной трудностью нам поможет справиться функция bitwise_and() . Эта побитовая операция объединяет два изображения таким образом, что только одинаковые части обеих изображений остаются на выходном изображении. В идеальном случае, если бы мы применили эту функцию к двум изображениям выше у нас бы должно было получилось что-то наподобие такого:
Но на деле получается следующее:
Всё дело в том, что детекция отрабатывает не всегда одинаково, то есть часто бывает, что изображение с таким же текстом, но на другом кадре, будет отличаться на 2-3 пикселя по ширине и/или высоте. В данном случае можно найти ограничительные рамки каждой буквы и обрезать изображение по этим координатам и таким образом мы сможем убрать чёрный фон вокруг текста. Про алгоритмы нахождения контуров и ограничительных рамок я писал в одной из прошлых своих статей. И по итогу, у нас должно получиться изображение подобное идеальному случаю. А дальше всё просто: мы передаём полученное изображение вместе с первым ( у которого отсутствует слово is в конце) изображением, в функцию matchTemplate() и действуем точно таким же образом как было описано несколькими абзацами выше. Вот и изображение готово к распознаванию текста:)
OCR
После успешной, на мой взгляд, детекции, хотелось бы так же успешно преобразовывать изображения текста в машиночитаемый текстовый формат, но не тут то было. В качестве библиотек для распознавания текста присмотрел следующие : Tesseract, EasyOCR и PaddleOCR. Теперь необходимо выбрать одну из предложенных. Что ж, было решено провести небольшое исследование и проверить насколько хороша каждая из библиотек работает на моих данных по трём критериям: CER, WER и время работы алгоритма, где про первые две метрики можно почитать тут . Для этого было взято 50 изображений из датасета и субтитры данных картинок я переписал в отдельный JSON файл . Прогнал данные через инференс каждого решения и получил следующие результаты:
По правде сказать, думал, что результаты будут лучше, ведь у меня была подготовлена довольно хорошо grayscale картинка: минимум шума на фоне, текст ровный, без поворотов и искажений.
Теперь нужно определиться с библиотекой, которую стоит использовать. Кажется практически очевидным, использовать PaddleOCR, но проблема в том, что у меня доступно всего лишь 2Гб видеопамяти, из которых ~1.5Гб потребляет моделька детекции, а на CPU данное решение работает достаточно долго, что для real-time не очень-то подходит, такая же ситуация и с EasyOCR, за исключением того, что на CPU мне вообще не удалось запустить данное решение, но на будущее, когда у меня появится достойная видюха, буду иметь в виду:) Остался только Tesseract — его я и буду использовать.
Обработка и перевод текста
Так часто бывает, что Tesseract в слове может неправильно распознать 1-2 буквы, из-за чего в дальнейшем слово может быть переведено неправильно. Поэтому я решил использовать библиотеку pyenchant, которая проверяет орфографию слова и если слово написано неправильно, библиотека советует слово, которое похоже на текущее слово, и уже получившийся текст передаём на следующий этап - перевод.
С переводом всё просто - Google Translate API. Переводит местами не совсем корректно, но быстро и никаких ограничений на количество запросов я не заметил.
Озвучка и визуализация текста
Для озвучки я использовал не навороченную либу pyttsx3. Для русского языка звучит так себе, но со временем привыкаешь. А визуализация текста была реализована при помощи библиотеки PySimpleGUI
Заключение
Данный проект вряд ли будет использоваться людьми для решения поставленной задачи, по крайне мере в текущей реализации, так как проблем тут ещё хватает, например со скоростью(отстаёт от реальности на секунд 5 ), но идеи по улучшению проекта у меня тоже имеются, например, улучшить перевод текста как описано здесь, ну и обучить модель распознавания текста на своих данных. На этом у меня всё:) Исходный код как обычно доступен на github . Всем хорошего дня!
