Привет, Хабр!

С вами Вадим Дарморезов, участник профессионального сообщества NTA.

Сегодня рассмотрю кейс поиска «близнецов» в паспортных данных, которые были размещены в pdf‑файлах, насчитывающих десятки, а порой и сотни страниц.

Введение

Для поиска схожих изображений последовательность шагов обычно следующая:

1. Считывание и обработка изображений: приведение изображений к одному размеру, перевод в градации серого и т. д.

2. Преобразование изображений в вектора.

3. Поиск разницы между векторами изображений и нахождение «близнецов».

В проектах, связанных с распознаванием лиц своеобразными «флагманами» являются библиотеки dlib/face‑recognition и свёрточные нейронные сети. При этом на просторах русскоязычного интернета довольно мало статей о библиотеке insightface. Именно о её использовании хотелось бы поговорить более подробно.

Insightface — open‑source набор инструментов для анализа 2D и 3D изображений, реализованный с помощью фреймворков машинного обучения PyTorch и MXNet. Данная библиотека эффективно реализует широкий спектр современных алгоритмов распознавания/детектирования/выравнивания лиц, которые оптимизированы как для обучения, так и для развертывания.

Приступлю к установке библиотеки. Выполню команду:

pip install -U insightface

Начиная с версии библиотеки 0.2.0, в качестве бэкенда для вычислений используется не MXNet, а onnxruntime. Данная библиотека (нейронная сеть) позволяет в качестве инференса использовать CPU или GPU.

В случае использования CPU, инференс выполняется на логических ядрах процессора, число которых равно числу физических ядер или, при использовании технологии Hyperthreading, увеличено вдвое. Использование CPU на глубоких нейросетях неэффективно из‑за ограниченного обмена данными с ОЗУ, что существенно влияет на скорость работы. Также ограничения на производительность накладываются самой архитектурой — в процессе инференса решаются простые задачи сравнения, которые легко переносятся на параллельные вычисления, но количество параллельных потоков обработки всегда будет ограничено количеством логических ядер CPU.

Инференс с использованием GPU за счет иной архитектуры процессора, наличия высокоскоростной памяти и гибкой системы управления кэш‑памятью гораздо эффективнее, чем инференс на CPU. Плюсом является кардинальное (до 100 раз) ускорение работы и крайне высокая эффективность обучения по сравнению с CPU.

Для установки необходимо выполнить следующие команды:

pip install onnx

После установки необходимых пакетов необходимо выбрать модель, которая будет использоваться для работы. Список моделей, которые могут быть использованы при работе, представлен в таблице ниже:

Далее, возможны два варианта запуска модели.

Запуск модели

1 вариант – запуск модели, с работающим подключением к сети Интернет

Для этого просто запускаю следующую строку:

app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider'])

Все необходимые для работы модели onnx‑файлы будут скачаны и размещены в директории ~/.insightface/models/. В дальнейшем при инициализации модели дополнительные загрузки производиться не будут.

2 вариант – запуск модели в оффлайн-режиме

При отсутствии возможности подключения компьютера к сети интернет для скачивания файлов, необходимо вручную создать следующую структуру директорий ~/.insightface/models/ и разместить туда предварительно скачанные onnx‑файлы модели.

В моем случае была необходима инициализация модели в оффлайн‑режиме, для чего был разработан класс для настройки рабочего окружения (создание необходимых для работы директорий, перемещение onnx‑файлов модели).

import os
import shutil 
class InitialSetup:
    def create_directories(self):
        directories_list = ['pdf', 'model','faces', 'model_result']
        for directory in directories_list:
            if not os.path.isdir(directory):
                os.mkdir(directory)
                print(f'Директория {directory} успешно создана.')
            else:
                print(f'Директория {directory} уже существует.')
        model_directory = r'/home/datalab/.insightface/models/buffalo_l'
        if not os.path.isdir(model_directory):
            os.makedirs(model_directory)
            print(f'Директория {model_directory} успешно создана.')
        else:
            print(f'Директория {model_directory} уже существует.')
def move_model_files(self):
        #список необходимых для работы модели onnx-файлов
        insightface_work_files = (
            'genderage.onnx',
            'w600k_r50.onnx',
            'det_10g.onnx',
            '2d106det.onnx',
            '1k3d68.onnx'
        )
        # определяю список onnx-файлов в необходимой для работы модели директории
        insightface_model_directory = r'/home/datalab/.insightface/models/buffalo_l'
        insightface_files = set(os.listdir(insightface_model_directory))
        #Проверяю, есть ли необходимые для работы модели файлы в необходимой директории
        if insightface_files==insightface_work_files:
            print('Все необходимые для работы модели onnx-файлы размещены.')
        else:
            #выгружаю список onnx-файлов модели
            model_files = [os.path.join('model', file) for file in os.listdir('model')]
            clear_model_files = set([file.split('/')[-1] for file in model_files])
            print(clear_model_files)
            if clear_model_files == insightface_work_files:
                for model_file in model_files:
                    shutil.copy(model_file, insightface_model_directory)
                print('Перемещение необходимых файлов прошло успешно.')
            else:
                print('Проверьте список onnx-файлов для перемещения.')

После формирования рабочего окружения можно приступать к обработке pdf‑файлов и обработке изображений.

Обработка pdf-файлов и обработка изображений

Начну с импорта необходимых библиотек:

import glob
import numpy as np
import matplotlib.pyplot as plt
import cv2
from pathlib import Path
from tqdm import tqdm
import pandas as pd
import fitz
import traceback
import onnxruntime as ort
from insightface.app import FaceAnalysis
from sklearn.neighbors import NearestNeighbors
from numpy.linalg import norm
from PIL import Image

После импорта библиотек посмотрю, как можно извлечь изображение из pdf‑файлов. Для этой задачи была выбрана библиотека fitz. Для корректной работы данной библиотеки необходимо установить пакет pymupdf.

doc = fitz.open(путь к pdf-файлу)
for i in range(len(doc)):
    for img in doc.get_page_images(i):
        xref = img[0]
        pix = fitz.Pixmap(doc, xref)
        if pix.n > 4:
            pix = fitz.Pixmap(fitz.csRGB, pix)
        img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)
        try:
            img = np.ascontiguousarray(img[...,[2,1,0]])
        except IndexError:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
            img = np.ascontiguousarray(img[...,[2,1,0]])

Далее, рассмотрю, как работает поиск лиц на изображениях. Инициализирую модель:

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

В качестве примера я хотел бы использовать знаменитое селфи Эллен Дедженерес с церемонии Оскар-2014:

Данное фото выбрано не только потому, что на нем представлен каст выдающихся актеров Голливуда, но и по следующим причинам:

• на фото представлены как мужчины, так и женщины;

• представлены люди различных возрастов;

• лица некоторых звезд видны не полностью — закрыты волосами соседей, руками и т. д.

Код для распознавания достаточно прост:

image = cv2.imread('test.jpg') # считываю изображение 
faces = app.get(image) # произвожу распознавание лиц
rimg = app.draw_on(image, faces) # отрисовка области с лицами
cv2.imwrite('res.jpg', rimg) # сохранение результата

Посмотрю, что получилось:

С детекцией лиц на изображении модель справилась отлично, были распознаны даже частично прикрытые лица, а вот с определением пола и возраста ситуация не так однозначна. Модели не удается точно определять пол, в случае затрудненной видимости анализируемого объекта, как получилось в случае Анджелины Джоли (порядка 50 процентов лица скрыто), так же при определении возраста возникают значительные погрешности (например, возраст Брэдли Купера на момент снимка — 39 лет, Дженнифер Лоуренс — 24 года).

Посмотрю, какие значения хранятся в переменной faces на примере одного лица:

{'bbox': array([1048.6523,477.87848, 1427.735,1018.7425 ], dtype=float32),
 'kps': array([[1109.926,676.95526],
        [1291.9822,678.36816],
        [1178.8099,779.84735],
        [1122.0046 ,848.871  ],
        [1304.9967,849.50073]], dtype=float32),
 'det_score': 0.9315067,
 'landmark_3d_68': array([[ 1.0514039e+03,  6.8547186e+02,  3.2998676e+02],
        ***
        [ 1.1808237e+03,  8.8022034e+02,  4.9980091e+01]], dtype=float32),
 'pose': array([ -2.3625553, -11.447153 ,  -1.7689382], dtype=float32),
 'landmark_2d_106': array([[1219.4696 , 1027.5748 ],
        ***
        [1340.001  ,  621.7045 ]], dtype=float32),
 'gender': 1,
 'age': 57,
 'embedding': array([ 6.67082489e-01, -7.11157694e-02,  9.92161810e-01, -1.89440691e+00,
        ***
         1.37064215e-02, -7.82325566e-02,  5.46212256e-01, -6.86526656e-01],
       dtype=float32)}

Необходимые для дальнейшего анализа переменные:

• bbox — хранит в себе координаты точек, ограничивающих область лица;

• gender — пол человека, которому принадлежит обнаруженное лицо;

• age — возраст человека, которому принадлежит лицо;

• embedding — векторное представление обнаруженного лица.

Объединю полученные знания и коды:

class FaceWorker:
    
    def __init__(self):
        self.app = FaceAnalysis(name="buffalo_l", providers=['CUDAExecutionProvider'])
        self.app.prepare(ctx_id=0, det_size=(256,256))
        self.knn = NearestNeighbors(metric='cosine', algorithm='brute')
        
    def extract_faces_from_pdf(self,files_paths, result_images_directory='faces'):
        errors_count = 0
        try:
            with open('completed_files.csv','a+') as file:
                for file_path in tqdm(files_paths):
                    file_name = Path(file_path).stem
                    doc = fitz.open(file_path)
                    for i in range(len(doc)):
                        for img in doc.get_page_images(i):
                            xref = img[0]
                            pix = fitz.Pixmap(doc, xref)
                            if pix.n > 4:
                                pix = fitz.Pixmap(fitz.csRGB, pix)
                            img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)
                            try:
                                img = np.ascontiguousarray(img[...,[2,1,0]])
                            except IndexError:
                                img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
                                img = np.ascontiguousarray(img[...,[2,1,0]])
                            faces = self.app.get(img)
                            if len(faces)>0:
                                for j,face in enumerate(faces):
                                    try:
                                        bbox = face.bbox
                                        x1,y1,x2,y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
                                        crop_img = img[ y1:y2,x1:x2]
                                        face_directory  = os.path.join(result_images_directory, f'{file_name}_face_{j}.png')
                                        cv2.imwrite(face_directory, crop_img)
                                    except cv2.error as error:
                                        errors_count +=1
                                        continue
                    end_time = dt.now().strftime('%d-%m-%Y %H:%M')
                    file.write(f'{file_path}|{end_time}\n')
        except:
            error = traceback.format_exc()
            print(f'При попытке поиска лиц в pdf-файлах произошла ошибка:\n{error}\nПоследний обработанный файл записан в completed_files.csv\n')
        finally:
            print(f'Ошибок записи cv2.error - {errors_count}')

    def face_vectorizer(self, face_path):
        try:
            image = cv2.imread(face_path)
            faces = self.app.get(image)
            if len(faces)>0:
                return faces[0].embedding
        except:
            error = traceback.format_exc()
            print(error)

Применю полученный код для поиска лиц в pdf‑файлах и их преобразования в векторное представление:

fw = FaceWorker()
pdfs = glob.glob('pdf/*.pdf')
print(f'Количество pdf-файлов для обработки - {len(pdfs)}')
fw.extract_faces_from_pdf(pdfs)
search_faces = glob.glob('faces/*.png')
vectors_dict = {
    'images_paths':[],
    'images_vectors':[]
}
for search_face in tqdm(search_faces):
    vector = fw.face_vectorizer(search_face)
    if vector is not None:
        vectors_dict['images_paths'].append(search_face)
        vectors_dict['images_vectors'].append(vector)
print('Лица преобразованы в вектора.')

Для поиска похожих изображений, представленных в векторном виде, буду использовать метод ближайших соседей из библиотеки sklearn, где в качестве метрики близости векторов будет выступать косинусное расстояние (данный подход не является единственно верным, существует множество методов расчета близости векторов).

Применю реализованный метод:

similar_faces = fw.search_similar_faces(vectors_dict, 30, 0.7)
print('Сформирован список схожих лиц.')
all_similar_images = []
for cluster in similar_faces:
    similar_images = [element[0] for element in cluster]
    all_similar_images.append(similar_images)
filtered_similar_images = []
for i,element in enumerate(all_similar_images):
    if set(element) not in filtered_similar_images:
        filtered_similar_images.append(set(element))
print('Отфильтрованы все возможные комбинации одних и тех же изображений.')

В первом тестовом примере модель не справилась с полным распознаванием Анджелины Джоли, будет логичным протестировать готовый код на датасете известных актеров с целью найти её «близнецов». Результат сравнения представлен ниже:

Заключение

Мне удалось реализовать систему для детектирования лиц в pdf-документах и поиска похожих людей с помощью библиотеки Insightface.Также хотелось бы отметить, что возможна гибкая настройка множества участков данной системы (от извлечения изображений из pdf-документов до методов расчета сходства изображений), что может позволить ускорить не только скорость обработки данных, но и качество распознавания и поиска дублирующихся и похожих лиц. Библиотека insightface богата на различные методы обработки лиц и может быть использована не только для их выявления и сравнения.