YOLOv5 — Детектор лиц
Привет Хабр, недавно у меня возникла необходимость сделать детектор лиц в реальном времени для дальнейшей сегментации и распознавания. В качестве детектора я использовал очень популярную сеть YOLOv5.
Сразу скажу что в Android используя GPU скорость детекции составляет 0.009 сек, около 110 fps, размер изображения на входе 160х160, для детекции в режиме селфи этого хватает за глаза, при этом точьность детекции очень высока.
Как обычно я буду тренировать сеть в бесплатном Google Colaboratory.
Подключаем Google Диск на котором находятся размеченные данные для тренировки сети:
from google.colab import drive
drive.mount('/content/drive', force_remount=True)
Устанавливаем YOLOv5:
!git clone https://github.com/ultralytics/yolov5 !pip install -r /content/yolov5/requirements.txt
Я вручную переразметил чуть более 15 000 изображений из датасета Wider добавив несколько тысяч картинок людей в масках.
Сам датасет можно скачать здесь, собственно он и использовался для тренировки.
Сам процесс тренировки нейронной сети предельно прост и подробно описан здесь.
Распаковываем тренировочные данные в папку /content/yolov5/datasets/faces:
from zipfile import ZipFile
with ZipFile('/content/drive/MyDrive/faces.zip', 'r') as zipObj:
zipObj.extractall('/content/yolov5/datasets/faces')
Запускаем процесс тренировки:
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
%cd /content/yolov5
!python "/content/yolov5/train.py" --img 640 --batch 64 --epochs 900 --data "/content/drive/MyDrive/faces.yaml" \
--project '/content/drive/MyDrive/face_detect/' --name 'weights' --weights "/content/drive/MyDrive/last.pt"
--project
указывает где будут сохраняться веса;
--data
https://drive.google.com/file/d/1DeN5E24Janwsc5dAnxh5Oejyuh1-OKmf/view?usp=sharing;
--img
размер изображений для тренировки;
--batch
размер батча;
--weights
вес с которого продолжить или начать новую тренировку.
Натренированный вес можно скачать здесь.
Для старта тренировки я использовал https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5n-7.pt
Выбрать подходящую вам по скорости и качеству модель можно здесь.
Для достижения максимальной точности детекции авторы YOLOv5 рекомендуют тренировать сетку 600 эпох и более, это занимает примерно 5 дней на Tesla с 16 GB памяти
В своих проектах под Android и Windows я использую Delphi и TensorFlow Lite, сам проект, библиотеки для Android и Windows и исходный код для Delphi можно скачать здесь.
Можно потестить в Windows x64, внутри скомпилированный .exe и dll, так же понадобиться Visual C++ Redistributable.