Что такое YOLO?

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

Почему я буду использовать эту архитектуру?

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

YOLO предоставляет же другой концепт. Эта архитектура будет просматривать изображение единожды, во время этого просмотра она будет сжимать изображение до квадратной матрицы размером 13x13. Внутри каждой из клеток этой матрицы будет записана информация о содержании искомого объекта и его классе

Что нужно?

Работать будем с пятой stable версией YOLO , что нужно для обучения:

  • установленная библиотека для python, можно установить из pypi при помощи менеджера пакетов pip или напрямую склонировать с github

  • необходимо будет скачать преобученные веса из репозитория yolo на github с которых можно будет начать тренировку или продолжить, я буду использовать YOLOv5x

Предобученные веса, ссылка - https://github.com/ultralytics/yolov5/releases
Предобученные веса, ссылка - https://github.com/ultralytics/yolov5/releases
  • датасет, я нашел дасатет на kaggle с нужной иерархией расположения изображений и и их labels (файлов, содержащих координаты найденного номера) для обучения с российскими и зарубежными номерами, очень важно правильно размеченный датасет, в отличие от ANN датасетов, файл с координатами искомого объекта на изображении здесь в формате .txt, а не .json

Датасет на kaggle, ссылка - https://www.kaggle.com/datasets/kirillpribludenko/number-plates-50-russain-50-others

Структура:

.
├── dataset
│   ├── classes.txt
│   ├── test
│   │   ├── images
│   │   └── labels
│   ├── train
│   │   ├── images
│   │   └── labels
│   └── val
│       ├── images
│       └── labels
├── dataset.yaml
└── yolov5x.pt
  • dataset.yaml файл, который должен выглядеть следующим образом

train: /content/drive/MyDrive/yolo_training/dataset/train/images 

val: /content/drive/MyDrive/yolo_training/dataset/val/images 

test:  /content/drive/MyDrive/yolo_training/dataset/test/images



# Classes

nc: 3  

names: [ 'car', 'plate', 'truck']  

где train, test и val пути до этих папок с изображениями в датасете из kaggle, nc - количество классов искомых объектов, в моем случае их 3, names - список из названий этих классов.

Процесс обучения (google colab)

После того как папку, которая содержит датасет, предобученные веса и yaml файл, необходимо будет выгрузить в свой google drive, бесплатный объем составляет 15гб, что вполне хватит для наших целей.

Обязательно перед всеми манипуляциями необходимо в новом notebook поменять runtime на GPU. Для этого Runtime -> Change Runtime Type и в появившемся окне ставим GPU (он нужен для корректной работы CUDA).

Далее переходим в colab (https://colab.research.google.com/) и создаем в нем новый notebook.

  • обязательно делаем mount и удостоверяемся, что в нашем .yaml файле прописан путь /content/drive/MyDrive/ перед нашей папкой

from google.colab import drive
drive.mount('/content/drive')
  • клонируем репозиторий, он появится рядом с директорией drive

!git clone https://github.com/ultralytics/yolov5.git
  • переходим в нашу папку и устанавливем зависимости из requirements.txt

%cd yolov5/
!pip install -r requirements.txt
  • запускаем команду на обучение по предобученным весам


!python train.py --img 896 --batch 8 --epochs 50 --data '/content/drive/MyDrive/yolo_training/dataset.yaml' --weights '/content/drive/MyDrive/yolo_training/yolov5x.pt' --nosave --cache

Параметры для train.py:

  • img - размер изображения, должен быть кратным 32

  • batch - батч(колво изображений обрабатываемых одновременных)

  • epochs - количество эпох

  • data - путь к датасету

  • weights - путь до весов с которых начать или продолжить тренировку

Результат:

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

50 epochs completed in 0.736 hours.
Optimizer stripped from runs/train/exp/weights/last.pt, 173.3MB
Results saved to runs/train/exp

Проверка:

Для того чтобы протестить наши новые веса, необходимо будет запустить следующую команду:


!python detect.py --source '/content/drive/MyDrive/yolo_training/dataset/test/images' --weights '/content/yolov5/runs/train/exp/weights/last.pt' --img 896 --save-txt --save-crop

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

Speed: 0.7ms pre-process, 82.2ms inference, 3.0ms NMS per image at shape (1, 3, 896, 896)
Results saved to runs/detect/exp
10 labels saved to runs/detect/exp/labels

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

Что дальше?

Планируется еще несколько статей, в которых при помощи openCV и Tesseract будет реализовано распознавание символов на автомобильных номерах.

Комментарии (0)