Как я делаю OCR — Часть 2

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

59af0c6098e99e0652fe7148b33c2bac.jpg

Теперь используя собранный набор данных попробуем натренировать одну из самых популярных сетей для сегментации и детектирования объектов 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')   

944a59091ad7743f53f1d568ea29b32d.jpg

Всё готово для тренировки сети, приступаем

%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 час.

adada03a185d8573aa07a5cb279cfe02.jpg

На моём домашнем компьютере в 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" 

21f763688b58dbcffeaff04bf19222e8.jpg

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

После конвертации в TensorFlow Lite вес модели составляет примерно 5 МБ, скорость вывода на андроид телефоне Redmi Note 10 ипользуя GPU составляет 300 мс.

© Habrahabr.ru