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