Третья жизнь пет-проекта по распознаванию рукописных цифр

В этом блогпосте я поделюсь историей о том, как я обновлял свой старенький пет-проект по распознаванию цифр, как делал разметку для него, и почему модель предсказывает 12 классов, хотя цифр всего 10.

Вступление

Пять лет назад, когда я получил свою первую работу в DS, я хотел как можно быстрее набрать побольше опыта. Среди прочего, я работал над пет-проектом: приложением на Flask, которое позволяло пользователям рисовать цифры и распознавать их с помощью ML-модели. На его разработку у меня ушло несколько месяцев, но оно того стоило как с точки зрения прокачивания навыков, так и в плане развития карьеры. Я даже писал о нём статью ^[1] на хабре.

Два года спустя я опубликовал новую версию с различными улучшениями; например, я использовал OpenCV для распознавания отдельных цифр, а модель была расширена до 11 классов, предсказывающих не-цифры.

Эти два приложения были развернуты на Heroku с использованием бесплатного плана, но некоторое время назад эти планы были прекращены. Я не хотел, чтобы проект канул в лету, поэтому решил сделать новую версию. Делать просто передеплой проекта на новой платформе было бы неинтересно, поэтому я обучил модель YOLOv3 с нуля на 12 классах. Несмотря на то, что это всего лишь пет-проект, в нём было много проблем, которые встречаются и в реальных проектах. В этой статье я хочу поделиться своим опытом работы над этим проектом, начиная со сбора данных и заканчивая деплоем.

Вот ссылка ^[2] на само приложение.

Сбор и подготовка данных

Сбор и разметка данных - важная часть любого проекта. Благодаря предыдущим версиям этого приложения у меня был датасет из примерно 19 тыс. изображений, которые хранились на Amazon S3. Лейблы для этих изображений были изначально сгенерированы моими моделями, и я знал, что часть из них ошибочны, ибо никакие модели не могут быть идеальными. По моим оценкам, уровень ошибок составлял около 10%, что означало, что около 2 тысяч изображений имели неправильные метки.

Помимо ошибок в самих лейблах, было много кейсов, когда мне самому было непонятно, что же показано на картинке. Например, люди иногда рисовали цифры так, что было трудно определить, что изображено (2 или 8, 1 или 7), или рисовали несколько цифр на одном изображении, что добавляло дополнительную головную боль. Кроме того, в моей предыдущей модели был реализован класс "other" для распознавания объектов, не являющихся цифрами, но мне все равно нужно было проверить все метки.

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

Image classification

Когда я начал работать над своим обновленным проектом по распознаванию цифр, я начал с обучения модели CNN на Pytorch на моем MacBook. Ради интереса также обучил модель ViT, используя этот гайд ^[3]. Обе модели были обучены с помощью MPS ^[4] Pytorch, что намного быстрее тренировки на CPU (пусть и уступает полнцоценным GPU).

Ранее я уже разработал пайплайн для тренировки моделей на PyTorch-lightning и Hydra, и я смог его легко допилить для этого проекта. Код можно посмотреть здесь ^[5].

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

Стоит отметить, что на данный момент у меня было 12 классов, которые модели должны были распознать: 10 для цифр, один для "other" и последний класс, который я назвал "censored". Думаю, что вы легко сможете найти примеры на картинке ниже ;) У меня собралось немало примеров, и модели смогли распознать этот класс весьма хорошо.

Хотя гонять эти эксперименты была весело, они были лишь промежуточным шагом на пути к моей цели - обучению object detection.

Разметка данных для object detection

Как я уже упоминал, моей целью в этом проекте было обучение модели object detection, для чего требовались bounding box для каждого объекта на каждом изображении. Для начала я использовал cv2.findContours и cv2.boundingRect из OpenCV для построения bounding box вокруг объектов на изображениях. Чтобы упростить первый шаг, я сначала работал только с изображениями, содержащими один объект.

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

Далее мне нужно было получить bbox для изображений, содержащих несколько объектов. Сначала я пытался извлечь их автоматически, но обнаружил слишком много ошибок, поскольку люди часто рисовали цифры несколькими несвязанными штрихами.

После небольшого ресерча я нашёл (точнее мне подсказали), что https://www.makesense.ai/ ^[6] является полезным инструментом для разметки bbox. На просмотр и разметку всех изображений ушло несколько часов, но в итоге я получил разметку bbox для 16,5 тыс. изображений.

Одна из проблем, с которой я столкнулся при маркировке данных, заключалась в том, чтобы решить, что размечать, а что нет. Например, на скриншоте ниже не очень понятно, что является "0", а что является каким-то мусором.

Тренировка YOLOv3

Для начала я использовал для обучения только изображения с одним объектом, чтобы убедиться, что все работает нормально. Я использовал этот шикарный туториал ^[7] и натренировал модельку с хорошим качеством.

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

Я долго дебажил, чтобы понять в чём же были ошибки.

Одна из проблем, с которой я столкнулся, была связана с аугментациями: некоторые ауги из библиотеки Albumentations приводили к искажению bbox. Вот старое issue ^[8] на GitHub об этой проблеме. В результате я начал использовать imgaug для аугментаций и использовал albumentations только на последнем этапе нормализации и ресайза.

import imgaug.augmenters as iaa
from imgaug.augmentables.bbs import BoundingBox, BoundingBoxesOnImage
import numpy as np
import albumentations

# an example of bbox
original_bboxes = [[14, 17, 28, 75, 1], [63, 74, 63, 69, 2], [140, 102, 39, 78, 3]]
# an example of image
max_x = 200
max_y = 200
original_image = np.random.randint(0, 255, size=(max_x, max_y, 3))

# creating bounding boxes in imgaug format
bbi = []
for b in [[b[0], b[1], b[0] + b[2], b[1] + b[3], b[4]] for b in original_bboxes]:
    bbi.append(BoundingBox(x1=b[0], y1=b[1], x2=b[2], y2=b[3], label=b[4]))

bbs = BoundingBoxesOnImage(bbi, shape=original_image.shape)

# defining the transformation
seq = iaa.Sequential([
    iaa.Affine(rotate=(-45, 45), scale=(0.9, 1.0),
               translate_px={"x": (-20, 20), "y": (-20, 20)}, cval=255
               )
])
# applying the transformation
im, bbs_aug = seq(image=im, bounding_boxes=bbs)

# convert the bounding boxes into yolo format
bboxes = [[b_.center_x / max_x, b_.center_y / max_y,
                       b_.width / max_x, b_.height / max_y, b_.label] for b_ in bbs_aug.bounding_boxes]

# define albumentations transforms
albumentations.Compose([albumentations.Resize(always_apply=False, p=1, height=192, width=192, interpolation=1),
                        albumentations.Normalize(always_apply=False, p=1.0, mean=(0, 0, 0), std=(1, 1, 1), max_pixel_value=255.0),
                        albumentations.pytorch.transforms.ToTensorV2(always_apply=True, p=1.0, transpose_mask=False)],
                       p=1.0,
                       bbox_params={'format': 'yolo', 'label_fields': None, 'min_area': 0.0, 'min_visibility': 0.0, 'check_each_transform': True},
                       keypoint_params=None, additional_targets={})

Еще одна проблема, которую я обнаружил, связана с форматом bbox: я ошибся и использовал формат coco, в то время как код модели ожидал, что они будут в формате yolo. Исправление этого косяка помогло, но метрики модели все равно были недостаточно высокими.

Чтобы эту ошибку не повторили другие, давайте рассмотрим эти форматы подробнее:

# an example of converting the bounding boxes between different formats.

import numpy as np
import albumentations

# an example of bbox in coco format: x_min, y_min, width, height
coco_bboxes = [[14, 17, 28, 75], [63, 74, 63, 69], [140, 102, 39, 78]]

# convert coco format to pascal_voc format
pascal_voc_bboxes = [[box[0], box[1], box[0] + box[2], box[1] + box[3]]] for box in coco_bboxes]

max_x = 200
max_y = 200
# convert coco format to yolo format
yolo_bboxes = [[(box[0] + box[2] / 2) / max_x,
                (box[1] + box[3] / 2) / max_y,
                box[2] / max_x, box[3] / max_y] for box in coco_bboxes]

После дальнейших экспериментов я обнаружил последнюю проблему: при обучении моделей классификации изображений я извлекал нарисованные цифры с помощью OpenCV и делал ресайз до 32x32 или 64x64. Это означало, что цифры занимали все пространство на изображении. Однако когда я начал обучать модели object detection, я брал весь canvas и ресайзил его до 64x64. В результате многие объекты становились слишком маленькими и кривыми для эффективного распознавания. Увеличение размера изображений до 192x192 помогло улучшить работу модели.

Если интересно, вот ссылка ^[9] на репорт Weight and Biases.

Ранее я упоминал, что обучал модели классификации изображений с помощью Pytorch MPS на MacBook. Однако, когда я попытался обучить модель object detection таким же образом, я столкнулся с некоторыми проблемами в Pytorch MPS. Одна внутренняя операция падала, поэтому мне пришлось перейти на CPU. На GitHub есть специальное issue ^[10], где люди могут поделиться подобными проблемами.

При обучении на изображениях размером 64x64 это работало достаточно быстро (хотя и занимало 15 минут), но увеличение размера изображения до 192x192 делало обучение непомерно медленным. В результате я решил использовать Google Colab. К сожалению, бесплатной версии оказалось недостаточно, поэтому мне пришлось приобрести 100 кредитов. Одна эпоха на Colab заняла всего 3 минуты. Запуска нескольких экспериментов было достаточно, чтобы получить хорошие метрики.

Разработка и деплой

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

Streamlit

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

Я использовал этот инструмент canvas ^[11], чтобы позволить пользователям рисовать цифры, которые модель будет распознавать. Процесс разработки приложения был относительно быстрым и занял всего пару часов. Как только приложение было готово, я смог перейти к этапу деплоя.

Весь код приложения можно увидеть тут ^[12].

В прошлых версиях проекта я хранил веса на Amazon S3, но эта моделька была намного тяжелее, и каждый раз грузить веса оттуда - это дорого и затратно. Так что я просто использовал Git LFS.

Deployment

Изначально я планировал захостить приложение на облаке streamlit ^[13], поскольку это отличная платформа для быстрого развертывания и шаринга небольших приложений. Я успешно развернул приложение на streamlit cloud, но когда я поделился им в одном чате, оно быстро уперлось в лимиты. Это означало, что мне нужно было найти альтернативное решение.

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

Тогда я вспомнил о Hugging Face Spaces ^[14], платформе, специально разработанной для деплоя ML-приложений. Я смог легко развернуть свое приложение на этой платформе (ушло меньше часа), и оно заработало без каких-либо проблем.

CI/CD

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

Я установил pre-commit hooks с black, mypy, flake8 и прочим.

Запретил прямой пуш в мастер.

При создании PR триггерятся проверки deepsource и пайплайн на Github Actions.

После успешного мерджа PR, триггерится ещё один пайплайн - для синхронизации кода с репо на Hugging Face Spaces.

Фейл со Style Transfer

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

Итоги

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

Этот проект был ценным и приятным опытом обучения для меня, и я надеюсь, что вы также нашли его интересным :)

Дополнительные ссылки:

• Ссылка на проект на моём сайте ^[15]

• Датасет на каггле с картинками и bbox ^[16]

• Код для тренировки ^[5]

• Код проекта на Github ^[12]