Соревнование, проводимое в рамках олимпиады НТО
6 место на паблике и 7 на привате
Вам нужно разработать алгоритм, который способен распознать рукописный текст в школьных тетрадях. В качестве входных данных вам будут предоставлены фотографии целых листов. Предсказание модели — список распознанных строк с координатами полигонов и получившимся текстом.
- Найти области каждого слова на изображении
- Определить на каком языке писали на этой фотографии (англиский/русский)
- Каждую найденную область прогнать через OCR модель и распознать что там написано
- Instance Segmentation = crop + maskrcnn_resnet50_fpn
- Language classification = resnet_34
- Ocr = microsoft/trocr-small-handwritten + hkr + generating dataset + HandWrittenBlot
В качестве модели распознавания мы попробовали две модели: detectoRS - одна из моделей зоопарка mmdetection и предобученный на COCO maskrcnn_resnet50_fpn (реализовали на торче).
Использовали дефолтную версию из mmdetection. При обучении ресайзили картинку в (1333, 800) и предсказывали сразу все слова на изображение
Чтобы размеры картинок не были помехой (были и вертикальные и горизонтальные сэмплы), модель обучили на кропах размерами (256, 256), а предсказание полигонов происходило в несколько этапов:
- Сжать размер картинки в 4 раза
- Пройтись окном (256, 256) и сделать предсказания
- Склеить слова на границах окон
Для того, чтобы склеивание происходило быстрее алгоритм сопоставления цветов был переписна на numba.jit
Так как на одной фотографии мог быть либо только русский, либо только англиский, то мы брали несколько слов из тех, что нашли при помощи сегментации и распознавали на каком языке они были написаны. Потом смотрели каких слов больше (русских или англиских) и в зависимости от этого определяли язык.
Было два способа распознавать язык
- Обучить trOcr на русском и англиском языке и просто получать с них предикт. Хоть CER у такой мультиязычной модели был и больше чем у тех, что обучены по отдельности, язык она практически не путала, однако довольно долго работала
- Дообучить resnet_34, чтобы он определял какой это язык. Данная модель косячила заметно чаще, но зато была шустренькой
В итоге остановились на втором варианте, так как первый чуть-чуть не залезал в тайм лимит
При обучение модели распознавания текста мы во многом вдохновились статьей ребят, которы заняли первое место на Digital Petr
В качестве модели распознавания текста была использована state-of-the-art модель trOcr: Transformer-based Optical Character Recognition. Мы дообучали microsoft/trocr-small-handwritten
на нашем датасете и смогли выбить CER=0.04
для русского языка и CER=0.06
для англиского.
Для обучения русскоязычной модели мы собрали соединили несколько датасетов:
- Тот, что дали огранизаторы соревнования
- HKR dataset - датасет русских и казахских рукописных слов
- Сгенерированные нами данные. Так как из модели, которая используют CTC Loss можно достать границы букв в словах, то мы смогли разбить данные нам слова на буквы и потом склеивать собственные слова (подробнее про генерацию датасета туть).
Для англиского языка было взято всего два датасета
- Тот, что дали организаторы соревнования
- IAM dataset - стандартный датасет для OCR на англиском
Чтобы добится лучшего качества распознавния мы добавили аугов
transforms = A.Compose([
A.Resize(384, 384),
AlbuHandWrittenBlot(blots, p=0.3),
A.Rotate(limit=[-7, 7]),
A.OneOf([
A.ToGray(always_apply=True),
A.CLAHE(always_apply=True, clip_limit=15),
], 0.3)
])
Из нестандартных аугментаций были использованы HandWrittenBlot - имитация рукописных почеркушек.