Как стать автором
Обновить

YOLOv5 — Детектор лиц

Время на прочтение2 мин
Количество просмотров4.5K

Привет Хабр, недавно у меня возникла необходимость сделать детектор лиц в реальном времени для дальнейшей сегментации и распознавания. В качестве детектора я использовал очень популярную сеть 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.

Теги:
Хабы:
Всего голосов 4: ↑4 и ↓0+4
Комментарии3

Публикации