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

Вступление

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

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

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

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

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

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

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

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

Image classification

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

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

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

Стоит отметить, что на данный момент у меня было 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/ является полезным инструментом для разметки bbox. На просмотр и разметку всех изображений ушло несколько часов, но в итоге я получил разметку bbox для 16,5 тыс. изображений.

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

Тренировка YOLOv3

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

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

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

Одна из проблем, с которой я столкнулся, была связана с аугментациями: некоторые ауги из библиотеки Albumentations приводили к искажению bbox. Вот старое issue на 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 помогло улучшить работу модели.

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

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

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

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

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

Streamlit

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

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

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

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

Deployment

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

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

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

CI/CD

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

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

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

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

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

Фейл со Style Transfer

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

Итоги

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

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

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

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

• Датасет на каггле с картинками и bbox

• Код для тренировки

• Код проекта на Github