Читая очередную статью по OpenCV, я обнаружил, что в этой библиотеке есть собственная нейросеть для распознавания лиц с высокой точностью.
Я решил опробовать OpenCV и придумал проект. Задача в том, чтобы определять лица на изображениях и сохранять их в отдельные файлы.
Весь проект доступен в репозитории GitHub. Я решил использовать свою фотографию из инстаграм.
Цель
Мы будем работать с двумя основными элементами:
- Белая рамка вокруг распознанного лица. Нам понадобится файл data_generator.py
- Вырезанные лица из изображений. Каждое из них будет сохранятся отдельно в папку. Файл face_extractor.py
Для начала выполним основные шаги: импортируем библиотеки, загрузим модель распознавания, создадим папки для входных изображений и вырезанных лиц.
Проект
Импорт библиотек
Импортируем os
для получения доступа к файлам. Далее cv2
— для работы с изображениями. И наконец numpy
, который поможет работать с многомерными массивами.
import os
import cv2
import numpy as np
Определяем пути и загружаем модель
Модели prototxt
и caffemodel
можно взять в репозитории OpenCV. Я разместил их в папке model_data
. Файл prototxt
содержит текстовое описание сети, а caffemodel
— вес. Читаем оба файла и загружаем модель.
base_dir = os.path.dirname(__file__)
prototxt_path = os.path.join(base_dir + 'model_data/deploy.prototxt')
caffemodel_path = os.path.join(base_dir + 'model_data/weights.caffemodel')
model = cv2.dnn.readNetFromCaffe(prototxt_path, caffemodel_path)
Создаём папку
Реализуем проверку. Если папка не существует, то она будет создана. Папка для выходных файлов updated_images.
if not os.path.exists('updated_images'):
print("New directory created")
os.makedirs('updated_images')
Так же поступаем и с папкой для вырезанных лиц faces.
if not os.path.exists('faces'):
print("New directory created")
os.makedirs('faces')
Чтение изображений
Перебираем все изображения в папке images. Далее делаем проверку на соответствие расширений файлов, чтобы работать только с файлами .png
и .jpg
.
for file in os.listdir(base_dir + 'images'):
file_name, file_extension = os.path.splitext(file)
if (file_extension in ['.png','.jpg']):
print("Image path: {}".format(base_dir + 'images/' + file))
Обнаружение лиц
Используя cv2.imread
, читаем изображение и создаем blob с помощью cv2.dnn.blobFromImage
. Далее отправляем этот blob в модель и получаем обнаружения, используя model.forward()
.
image = cv2.imread(base_dir + 'images/' + file)
(h, w) = image.shape[:2]
blob = cv2.dnn.blobFromImage(cv2.resize(image, (300, 300)), 1.0, (300, 300), (104.0, 177.0, 123.0))
model.setInput(blob)
detections = model.forward()
Мы выполнили основные шаги. Следующим шагом будет построение белых рамок вокруг лиц и сохранение их в каталоге updated_images. Далее сохраняем извлечённые лица в каталоге faces.
1. Создаём боксы вокруг лиц
Далее перебираем все лица, обнаруженные на изображении, и извлекаем начальную и конечную точки. Затем извлекаем значение достоверности. Если алгоритм более чем на 50% уверен, что обнаружение является лицом, то вокруг него показываем прямоугольник.
# Create frame around face
for i in range(0, detections.shape[2]):
box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
(startX, startY, endX, endY) = box.astype("int")
confidence = detections[0, 0, i, 2]
# If confidence > 0.5, show box around face
if (confidence > 0.5):
cv2.rectangle(image, (startX, startY), (endX, endY), (255, 255, 255), 2)
cv2.imwrite(base_dir + 'updated_images/' + file, image)
print("Image " + file + " converted successfully")
Далее, используя cv2.imwrite
, сохраняем изображение в папку updated_images
с тем же именем файла.
2. Извлекаем лицо
Как говорилось выше, перебираем все лица, вычисляем достоверность обнаружения и, если она больше 50%, извлекаем лицо. Обратите внимание на строку frame = image[startY:endY, startX:endX]
. Здесь происходит извлечение.
# Identify each face
for i in range(0, detections.shape[2]):
box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
(startX, startY, endX, endY) = box.astype("int")
confidence = detections[0, 0, i, 2]
# If confidence > 0.5, save it as a separate file
if (confidence > 0.5):
count += 1
frame = image[startY:endY, startX:endX]
cv2.imwrite(base_dir + 'faces/' + str(i) + '_' + file, frame)
Далее сохраняем это новое изображение в папку faces. Имена файлов присваиваются следующим образом: если имя исходного изображения было sampleImage.png, имя файла после обработки будет 0_sampleImage.png С каждым лицом увеличиваем счётчик и после полного выполнения выводим его на консоль.
Теперь всё готово. Вы можете обрабатывать столько изображений сколько захотите, создавая наборы данных для дальнейших проектов.
Читайте также:
Перевод статьи Karan Bhanot: Extracting faces using OpenCV Face Detection Neural Network