Распознавание лиц с InsightFace или как CatBoost имена угадывал

Цель статьи - рассказать про простой и в тоже время рабочий вариант создания системы распознавания лиц, используя только модели из коробки, а именно ,библиотеку InsightFace для предобработки изображений и Catboost для их классификации.

Прежде чем приступить разделим задачу на этапы.

Этапы

1)Нахождение  лиц на изображениях в обучающем наборе данных

2) Перевод изъятых изображений с лицами в векторный формат

3) Обучение модели классификации на получившихся данных

4) Классификация лица на неизвестных для модели изображениях (тестовой выборке)

Подготовка

Для начала установим необходимые библиотеки.

pip install -U insightface
pip install onnx 
pip install onnxruntime

Для запуска моделей InsightFace на GPU необходимо установить onnx и onnxruntime. Onnx (Open Neural Network Exchange) это открытая библиотеке для разработки моделей
и быстрого обмена ими между разработчиками Про совместимость версий onnx и CUDA можно посмотреть в таблице здесь.

Вот полный список использованных мной библиотек

import os
import pickle
from PIL import Image
import numpy as np
from typing import List
import onnxruntime as ort
from insightface.app import FaceAnalysis
from catboost import CatBoostClassifier
import shutil

Данные

Для проверки работы алгоритма  я использовал набор данных под название Real World Fasked Face Recognition Dataset (RMFRD). В этих данных содержится 5000 лиц 525 человек в масках и 90 000 лиц без масок. Все изображения находятся в .jpg формате. В данном примере я буду использовать только изображения людей с открытым лицом.

Сразу разделим данные на тренировочные и тестовые.

def get_test_data(dir):
    for subdir in os.listdir(dir):
        path = dir + subdir + '/'
        if len(os.listdir(path))>1:
            filenames = [filename for filename in os.listdir(path) ]
            face_path = path + filenames[0]
            shutil.move(face_path, 'dataset/test_faces/')  

Функция get_test_data()забирает по одному изображению из каждой директории при условии, что оно там не единственное.

Получение эмбеддингов

Библиотека InsightFace уже содержит уже предобученные модели. Одна из таких моделей- buffalo_l. Она подойдет нам как для обнаружении лица на изображении, так и нахождения эмбеддингов. 

app = FaceAnalysis(name="buffalo_l",providers=['CUDAExecutionProvider'])
app.prepare(ctx_id=0, det_size=(256, 256))

Теперь займемся предобработкой данных для подачи их в модель классификации. В этом нам помогут следующие функции.

def extract_face(filename, required_size=(256, 256)):
    image = Image.open(filename)    
    img_arr = np.array(image)  
    im = Image.fromarray((img_arr))
    im = im.resize(required_size)
    rgb_arr = np.array(im.convert('RGB'))   
    emb_res = app.get(rgb_arr)
    try:
      face_array = emb_res[0].embedding
      return face_array     
    except:
      print('no embedding found for this image')

def load_face(dir):
    faces = list()
    for filename in os.listdir(dir):
        path = dir + filename
        face = extract_face(path)
        faces.append(face)
    return faces

def load_dataset(dir):
  X, y = list(), list()
  i = 1
  for subdir in os.listdir(dir):
      path = dir + subdir + '/'
      faces = load_face(path)
      labels = [subdir for i in range(len(faces))]
      print("loaded %d sample for class: %s" % (len(faces),subdir) ) 
      X.extend(faces)
      y.extend(labels)
      i+=1
  return np.array(X), np.array(y)

Пройдемся по всем данным в нашем наборе и получим эмбеддинг для изображений, где модель может обнаружить лицо. С помощью load_dataset() заходим в каждую директорию, а далее load_face() обращается к каждому изображению,
extract_face() возвращает нам необходимый вектор значений соответствующий лицу на том изображении . В случае если нейросеть не сможет обнажить на изображение лица функция extract_face() выдаст сообщение:

no embedding found for this imag.

После обработки всех фото в директории будет выдано сообщение по типу

loaded 137 sample for class: linjunjie

На выходе мы получаем два массива одного размера с эмэддингами изображений и лейблами (имена людей на фото, которые и являются предсказываемым фактором).

def filter_empty_embs(img_set: np.array, img_labels: List[str]):
    good_idx = [i for i,x in enumerate(img_set) if x is not None]
    clean_labels = img_labels[good_idx]
    clean_embs = img_set[good_idx]      
    return clean_embs, clean_labels

Функция filter_empty_embs() удалит пустые значения лэйблов для изображений на которых не были обнаружены лица а следовательно и получены эмбэддинги.

Запустим наши функции для преобразования изображений

get_test_data('dataset/unmasked_users/')
trainX, trainy = load_dataset('dataset/unmasked_users/')

assert len(trainX) == len(trainy)
train_emb, train_labels = filter_empty_embs(trainX, trainy)

assert len(train_emb) == len(train_labels)
print("Train_X size is {} , train_y size is {} ".format(train_emb.shape, train_labels.shape))

Модель классификации

Классификатор будем использовать из библиотеки CatBoost. Набор моделей градиентного бустиннга, которые она в себе содержит, хорошо себя зарекомендовал и применяется для решения реальных задач.

Зададим модель классификатора

clf_model = CatBoostClassifier(iterations=100,
                           task_type="GPU",
                           devices='0:1')

С параметрами классификатора можно поиграться и посмотреть как они влияют на работу модели. Подбор параметров может послужить темой для отдельной статьи. Подробнее про CatBoost можно почитать вот тут.

Обучим нашу модель на эмбенддингах полученных из тренировочной выборки train_emb и именах людей хранящихся в массиве train_labels.

clf_model.fit(np.array(list(train_emb)),
          train_labels,
          verbose=False)

Проверка работы модели

После обучения модели проверим ее на данных которые она еще не видела в наборе test_faces.

preds = []
true_labels = []
for filename in os.listdir('dataset/test_faces/'):
    image = Image.open('dataset/test_faces/'+filename)        
    img_arr = np.array(image)  
    im = Image.fromarray((img_arr))
    required_size=(256, 256)
    im = im.resize(required_size)
    rgb_arr = np.array(im.convert('RGB'))   
    emb_res = app.get(rgb_arr)
    try:
      face_array = emb_res[0].embedding   
    except:
      print('no embedding found for this image')

    predict = clf_model.predict(face_array)

    max_proba = clf_model.predict_proba(img_emb).max()
    if predict[0] in filename:
        true_labels.append(predict[0])
    else:
        true_labels.append(filename)
    preds.append(predict[0])

В качестве метрики качества работы модели будем использовать самую простую - accuracy.

from sklearn.metrics import accuracy_score

print(accuracy_score(true_labels, preds))

На выходе я получил значение 0.8958333333333334, что довольно неплохо.

Заключение

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