Введение
Недавно был опубликован анонс новой 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 etreeimport cv2from glob import globimport reCAR_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/yolov5pip install -U -r yolov5/requirements.txt
Информация о данных будет храниться в файле data.yaml. Создадим этот файл со следующим содержанием:
train: ../train/imagesval: ../valid/imagesnc: 10names: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
где train путь до директории, содержащей изображения для трейна, valid валидационные изображения, nc количество классов, names имена классов. Также необходимо положить архитектуру в директорию yolov5/models/. Возьмем за основу архитектуру yolov5s.yaml и поместим в custom_yolov5s.yaml следующую информацию:
# parametersnc: 10 # number of classesdepth_multiple: 0.33 # model depth multiplewidth_multiple: 0.50 # layer channel multiple# anchorsanchors: - [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 backbonebackbone: # [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 headhead: [[-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). Так что следим за новостями новых архитектур нейронных сетей и пробуем на собственных данных! Надеюсь, данная статья была вам полезна.
