YoloV5 для распознавания марок автомобилей

Введение

Недавно был опубликован анонс новой YOLOv5, которая идейно дает гораздо лучший процент распознавания на датасете COCO, чем предыдущие версии. Автор решил испробовать новую модель на задаче распознавания марок автомобилей.

Данные

Перво-наперво нам необходимы данные. Датасет был собран вручную, путем фотографирования стоящих на стоянке автомобилей с видимым значком марки (для этого пришлось выйти из дому в 5 утра, дабы не пугать удивленных прохожих). Разметка данных производилась с помощью инструмента labelImg. Всего было размечено 118 фотографий автомобилей следующих марок: Lada, Kia, Nissan, Volkswagen, Chevrolet, Ford, Mitsubishi, Renault, Hyundai, Opel. Довольно оптимистичная цель — пытаться обучить нейросеть на таком небольшом наборе данных, — однако, только практика покажет!

Предобработка

Поскольку labelImg использует формат Pascal VOC и пишет разметку в XML-файлы, а YOLOv5 использует TXT-файлы со строками следующего формата: «номер-класса x-центра y-центра width height», все нормированное на единицу, — мы должны преобразовать разметку к этому формату. Помещаем фотографии в директорию CAR_BRANDS/ и пишем код для преобразования:

from lxml import etree
import cv2
from glob import glob
import re

CAR_BRANDS_PATH = 'CAR_BRANDS/'

classes = {'lada': 0, 'kia': 1, 'nissan': 2, 'volkswagen': 3, 'chevrolet': 4,
           'ford': 5, 'mitsubishi': 6, 'renault': 7, 'hyundai': 8, 'opel': 9}

def extract_txt(img_path, xml_path, txt_path):
    img = cv2.imread(img_path)
    W = img.shape[1]
    H = img.shape[0]
    tree = etree.parse(xml_path)

    root = tree.getroot()
    with open(txt_path, 'w') as f:
        for child in root:
            if child.tag == 'object':
                for object_child in child:
                    coords = []
                    if object_child.tag == 'name':
                        f.write(str(classes.get(object_child.text))+' ')
                    if object_child.tag == 'bndbox':
                        for bbox_child in object_child:
                            coords.append(int(bbox_child.text))
                        xmin = coords[0]
                        xmax = coords[2]
                        ymin = coords[1]
                        ymax = coords[3]

                        w = (xmax - xmin) / W
                        h = (ymax - ymin) / H
                        xc = (xmin + (xmax - xmin) / 2) / W
                        yc = (ymin + (ymax - ymin) / 2) / H

                        f.write(str(xc)+' '+str(yc)+' '+str(w)+' '+str(h)+'\n')


xmls = sorted(glob(CAR_BRANDS_PATH + '*.xml'))

for xml_path in xmls:
    img_path = re.findall('[A-Za-z0-9_/]+', xml_path)[0]+'.jpg'
    txt_path = re.findall('[A-Za-z0-9_/]+', img_path)[0]+'.txt'
    extract_txt(img_path, xml_path, txt_path)

Мы будем загонять в нейросеть фото размера (416, 416), поэтому следующая функция поможет нам их отресайзить:

def resize(folder):
    files = glob(folder + '*.jpg')
    for file in files:
        img = cv2.imread(file)
        img = cv2.resize(img, (416, 416))
        cv2.imwrite(file, img)

Аугментация

Для аугментации был использован ресурс roboflow, однако вы можете пользоваться любым доступным вам методом аугментации. Фото были размножены до количества 4130, типы аугментации — Crop, Brightness, Exposure, Noise. Итоговые файлы были разделены на трейн, валидацию и тест в соотношении 7:2:1.

Обучение

Скачаем YOLOv5 с ее репозитория на гитхабе и установим requirements:

git clone https://github.com/roboflow-ai/yolov5
pip install -U -r yolov5/requirements.txt

Информация о данных будет храниться в файле data.yaml. Создадим этот файл со следующим содержанием:

train: ../train/images
val: ../valid/images

nc: 10
names: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

где train — путь до директории, содержащей изображения для трейна, valid — валидационные изображения, nc — количество классов, names — имена классов. Также необходимо положить архитектуру в директорию yolov5/models/. Возьмем за основу архитектуру yolov5s.yaml и поместим в custom_yolov5s.yaml следующую информацию:

# parameters
nc: 10  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple

# anchors
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# yolov5 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Focus, [64, 3]],  # 1-P1/2
   [-1, 1, Conv, [128, 3, 2]],  # 2-P2/4
   [-1, 3, Bottleneck, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 4-P3/8
   [-1, 9, BottleneckCSP, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 6-P4/16
   [-1, 9, BottleneckCSP, [512]],
   [-1, 1, Conv, [1024, 3, 2]], # 8-P5/32
   [-1, 1, SPP, [1024, [5, 9, 13]]],
   [-1, 6, BottleneckCSP, [1024]],  # 10
  ]

# yolov5 head
head:
  [[-1, 3, BottleneckCSP, [1024, False]],  # 11
   [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]],  # 12 (P5/32-large)

   [-2, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
   [-1, 1, Conv, [512, 1, 1]],
   [-1, 3, BottleneckCSP, [512, False]],
   [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]],  # 17 (P4/16-medium)

   [-2, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
   [-1, 1, Conv, [256, 1, 1]],
   [-1, 3, BottleneckCSP, [256, False]],
   [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]],  # 22 (P3/8-small)

   [[], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

Теперь можно приступить к обучению. Перейдем в yolov5/:

cd yolov5/

и обучим нашу модель:

python train.py --img 416 --batch 4 --epochs 50 --data '../data.yaml' --cfg ./models/custom_yolov5s.yaml --weights '' --name yolov5s_results --nosave --cache

Результаты

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

python detect.py --weights weights/last_yolov5s_results.pt --img 416 --conf 0.4 --save-txt --source ../test/images

При желании можно увеличить или уменьшить значение параметра conf — это позволит отбирать предсказания с вероятностью, большей, чем указано в conf.

Результаты сохраняются в директорию yolov5/inference/output/ (как изображения, так и TXT-файлы (при условии указания флага --save-txt). Из 12 фото, отобранных для теста, марки распознались в 10 случаях, при этом все распознанные оказались верными. Довольно неплохой результат для исходной выборки в 118 фотографий! Примеры вывода нейронной сети (пользуемся словарем, записанным ранее, для идентификации марок):



И, напоследок, посмотрим повнимательнее на процесс обучения. Во время обучения YOLOv5 делает с фотографиями нечто подобное:



Это т.н. «мозаичная аугментация» — к фотографиям применяются отражения и нарезка фото на кусочки. Убрать ее можно, если залезть в код train.py и сменить флаг augmet на значение False. Когда ее использование может навредить? Например, если бы мы задались целью идентифицировать номера автомобилей — отражение и нарезка отдельных букв и цифр привели бы не к улучшению, а ухудшению качества распознавания.

Что дальше?

Во-первых, одним из просторов для творчества может послужить перебирание разных вариаций архитектур YOLOv5 — все они лежат в yolov5/models/. Напомню, что мы с вами использовали yolov5s — она считается наиболее быстрой, но не самой лучшей по качеству распознавания. Также классическая тема для YOLO — поиграться со значениями якорей (anchors). Так что следим за новостями новых архитектур нейронных сетей и пробуем на собственных данных! Надеюсь, данная статья была вам полезна.