В предыдущей статье я рассказывал как собираю данные для детектирования текста на изображениях.
Теперь используя собранный набор данных попробуем натренировать одну из самых популярных сетей для сегментации и детектирования объектов YOLOv5.
Для этого будем использовать бесплатный Google Colab.
Подключаем диск из аккаунта Google
from google.colab import drive
drive.mount('/content/drive', force_remount=True)
В корневой папке должен находиться архив с тренировочными данными который можно найти здесь GitHub.
Также скачайте и поместите в корневую папку файлы last.pt, text_segment.yaml, hyp.scratch-low.yaml, labels_yolo.zip
Клонируем из GitHub YOLOv5 и устанавливаем
!git clone https://github.com/ultralytics/yolov5
!pip install -r /content/yolov5/requirements.txt
Я выбрал самую лёгкую и быструю модель yolov5n-seg и изменил файл конфигурации модели yolov5n-seg.yaml что бы уменьшить размер модели и увеличить скорость вывода.
Распаковываем изображения и файлы разметки в папку /content/yolov5/datasets/text_detection
from zipfile import ZipFile
with ZipFile('/content/drive/MyDrive/sd_text_detection_train.zip', 'r') as zipObj:
zipObj.extractall('/content/yolov5/datasets/text_detection')
with ZipFile('/content/drive/MyDrive/labels_yolo.zip', 'r') as zipObj:
zipObj.extractall('/content/yolov5/datasets/text_detection/labels')
Всё готово для тренировки сети, приступаем
%cd /content/yolov5
!python "/content/yolov5/segment/train.py" --img 800 --batch 12 --epochs 1200 --data "/content/drive/MyDrive/text_segment.yaml" \
--hyp "/content/drive/MyDrive/hyp.scratch-low.yaml" --project '/content/drive/MyDrive/text_detect/' --name 'weights' \
--weights "/content/drive/MyDrive/last.pt"
Авторы YOLOv5 рекомендуют тренировать сеть 300 эпох и более, тренировка одной эпохи в бесплатной версии Google Colab занимает примерно 1 час.
На моём домашнем компьютере в Windows 10 тренировка 1 эпохи на видеокарте RTX 2060 с 12 ГБ занимает примерно 6 минут.
Проблема бесплатной версии Google Colab отсутствие многопоточности и SSD диска.
Спустя примерно дофига времени в Google Colab, протестируем нашу модель
!python "segment/predict.py" --imgsz 800 --iou-thres 0.25 --conf-thres 0.5 --hide-labels --hide-conf --line-thickness 2 --device "cpu" --weights "/content/drive/MyDrive/text_detect/weights/weights/last.pt" --source "/content/yolov5/test"
На мой взгляд результат просто отличный, учитывая что это самая лёгкая модель и к тому же обрезанная в двое.
После конвертации в TensorFlow Lite вес модели составляет примерно 5 МБ, скорость вывода на андроид телефоне Redmi Note 10 используя GPU составляет 300 мс.
Комментарии (15)
ZlodeiBaal
00.00.0000 00:00Не проще ли взять https://github.com/open-mmlab/mmocr или https://github.com/PaddlePaddle/PaddleOCR или Easy OCR?
По таким задачам как чеки там детекторы будут почти идеально работать даже из коробки Возможно вам потом надо будет переобучить только детектор, но часто даже это не надо.
Их плюс в том, что там часто детектор и распознавание текста совмещены, что позволяет более стабильно работать с искривленными или погнутыми текстами.
SmallDonkey Автор
00.00.0000 00:00Их плюс в том, что там часто детектор и распознавание текста совмещены
Это не так, вы ошибаетесь
ZlodeiBaal
00.00.0000 00:00+1https://github.com/open-mmlab/mmocr/blob/1.x/projects/ABCNet/README.md
https://github.com/open-mmlab/mmocr/blob/1.x/projects/ABCNet/README_V2.md
https://github.com/open-mmlab/mmocr/blob/1.x/configs/kie/sdmgr/README.md
https://github.com/open-mmlab/mmocr/blob/1.x/projects/SPTS/README.mdНо в целом, я скорее про сети которые можно разделить на детектор-классификатор, но которые стабильны к изгибу текста:
https://github.com/open-mmlab/mmocr/blob/1.x/configs/textdet/drrg/README.md
https://github.com/open-mmlab/mmocr/blob/1.x/configs/textdet/fcenet/README.md
и.т.д.
А целом, если у вас смысл чеки распознавать, то даже такое есть, но это уже немного не в тему - https://github.com/clovaai/donut
SmallDonkey Автор
00.00.0000 00:00Не имеет значение какую сеть вы возьмете, это решать вам исходя из ваших потребностей, смысл статьи не в этом.
ZlodeiBaal
00.00.0000 00:00Собственно в этом мой вопрос и есть. Зачем вам понадобилось тренировать свою сеть, если есть много предтренированных. Я не вижу какого-то нового домена данных, я не вижу каких-то повышенных требований на точность/скорость. Кажется, что mmocr тут все должен решить.
SmallDonkey Автор
00.00.0000 00:00Можно было взять любую сеть, тот же Easy OCR дообучить, смысл проверить качество датасета, мне кажется порог вхождения с yolo самый низкий, там разберется даже далекий от этого всего человек.
ZlodeiBaal
00.00.0000 00:00Цимес в том, что по большинстве задач mmocr не придется обучать. И порог ниже.
dimnsk
00.00.0000 00:00+1А зачем для OCR - Yolo ? Из статьи не понятно извините...
можно взять уже обученные OCR на своих датасетах и дообучить...
PS у вас в Colab GPU включен ? не может одна эпоха на GPU так долго учится ...SmallDonkey Автор
00.00.0000 00:00Да GPU включен был, там Tesla видюха, 11.7 GB жрала пока тренировал.
mairon2300
00.00.0000 00:00А почему именно yolov5 для детекции текста используется, а не тот же DBNet, но облегченный? Проводилось ли какое-то сравнение между ними?
SmallDonkey Автор
00.00.0000 00:00Я не проводил ни какого сравнения между этими сетями.
Можно было взять любую сеть, смысл проверить качество датасета
mairon2300
00.00.0000 00:00-1Не увидел у вас в статье упоминание валидационной выборки. По файлу с путями к выборкам кажется, что валидация у вас проходит на той же самой выборе, на которой проходит обучение - то есть неизвестно, как поведет себя модель на тех данных которые она не видела (безотносительно качества обучающей выборки)
UtrobinMV
В данном случае Yolov5 детектирует только сами прямоугольники с буквами? Она же не распознает сама, что там написано?! Верно?
SmallDonkey Автор
Текст не распознает, только детекция
UtrobinMV
Ждем часть 3!