【Milvus的人脸检索】

0. 介绍

在上一篇文章中，介绍了milvus提供的以图搜图的样例，这篇文章就在以图搜图样例的基础上进行修改，实现人脸检索。

常见的人脸任务，分为人脸检测、人脸识别、人脸对比和人脸检索，其中人脸检索的含义是：对给定一张人脸照片，和已有人脸库中的N个人脸进行比对，找出最相似的一张脸或多张脸，并给出相似度排序，实现1 : N或M：N搜索。

现如今，大部分云平台都实现了人脸相关的算法应用，并提供相关的服务，如下图为腾讯云人脸页面相关介绍，可以通过点击链接，体验相关的功能。其他商家的链接如下：

• 百度云：https://ai.baidu.com/tech/face/search
• face++：人脸搜索 - 旷视Face⁺⁺人工智能开放平台 

 此外，人脸检索的应用场景也十分广泛，如旷视下图所描述的一样。

本文使用insightface库来实现对图像中人脸的定位和特征提取，基于insightface构建一个人脸特征提取类，然后修改相应的特征提取函数接口，最后将人脸特征插入到milvus和MySQL数据库中，以便后续进行人脸检索。好了，接下来就让我们看看如何基于milvus向量数据库和insightface实现百万级人脸检索。

1. insightface使用

insightface是一个非常高效的人脸分析库，利用insightface使用极少的代码就能实现对图像中人脸定位、人脸gender和age分析、人脸的landmark以及人脸识别等等功能。

1.1 安装insightface

insightface的安装方式非常简单，使用pip即可，安装命令如下：

pip install insightface

安装完成之后，新建文件，粘贴如下代码，即可测试安装是否正确

import cv2
import numpy as np
import insightface
from insightface.app import FaceAnalysis
from insightface.data import get_image as ins_get_image

app = FaceAnalysis(providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
app.prepare(ctx_id=0, det_size=(640, 640))
img = ins_get_image('t1')
faces = app.get(img)
rimg = app.draw_on(img, faces)
cv2.imwrite("./t1_output.jpg", rimg)

 上述代码，构建FaceAnalysis对象，完成人脸定位、人脸属性分析和人脸特征提取等工作，并将结果保存到t1_output.jpg文件中。

1.2 安装onnxruntime-gpu

如下图所示，当insightface的版本大于等于0.2时，采用onnxruntime作为推理框架，默认使用CPU进行推理，因此，为了加快模型运行速度，可以安装onnxruntime-gpu来利用GPU资源加速推理。

安装onnxruntime-gpu库的命令非常简单，使用pip安装即可，命令如下：

pip install onnxruntime-gpu

 但是，由于cuda环境和onnxruntime-gpu版本的不匹配，会导致无法利用GPU。因此，在安装onnxruntime-gpu之前，需要查看当前设备cuda和cudnn的版本，从而安装正确版本的onnxruntime-gpu。如我的设备的cuda版本为11.4，按照下表安装合适的版本。大家可以通过该链接查询，需要安装的版本。

 通过下面代码可以简单测试onnxruntime-gpu是否安装正确。

import onnxruntime as ort

print(ort.get_device())
# print(ort.get_all_providers())
print(ort.get_available_providers())

输出如下所示，当打印的available_providers中包含CUDAExecutionProvider即可。

GPU
['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']

2. 构建特征提取网络

基于insightface，可以使用极少的代码构建人脸的特征提取网络，具体代码查考如下。

from insightface.app import FaceAnalysis
import insightface
# import os
import cv2
# from tqdm import tqdm
# import pickle
assert insightface.__version__>='0.3'


class FaceRecognition():
    def __init__(self) -> None:
        self.app = FaceAnalysis(allowed_modules=['detection', 'recognition',"genderage"], providers=['CUDAExecutionProvider'])
        # detection network input size
        self.app.prepare(ctx_id=0, det_size=(640, 640))

    def extract_face_features(self, img_path):
        try:
            img_data = cv2.imread(img_path)
            feats = self.app.get(img_data)
        except Exception as e:
                return 
        return feats        # list

 通过上述代码，就能完成对一张图像中每个人脸特征的提取，具体来说，包括人脸的位置信息、关键点信息、人脸对应的性别和年龄，以及人脸的特征信息。

3. milvus的使用

以milvus中以图搜图代码为基础，对其中的文件内容进行修改来实现人脸检索功能。

3.1 encode

使用第2节中的特征提取代码替换encode.py文件中的内容。

3.2 load操作

由于特征提取接口返回内容的结果发生变化，因此需要对load.py文件中extract_features函数进行相应的修改，代码内容如下：

# Get the vector of images
def extract_features(img_dir, model):
    try:
        cache = Cache('./tmp')
        face_embeddings = list()
        face_properites = list()
        img_list = get_imgs(img_dir)
        total = len(img_list)
        cache['total'] = total
        for i, img_path in enumerate(img_list):
            try:
                # path_encoded_list,norm_feat = model.batch_extract_feat(img_dir)
                faces = model.extract_face_features(img_path)
                for face in faces:

                    face_embeddings.append(face.normed_embedding)
                    face_dict = {
                        "image_path": img_path,
                        "gender": face["gender"],
                        "age": face["age"]
                    }
                    face_properites.append(face_dict)
                    cache['current'] = i+1
                    print(f"Extracting feature from image No. {i + 1} , {total} images in total")
            except Exception as e:
                LOGGER.error(f"Error with extracting feature from image {e}")
                continue
        return face_embeddings, face_properites
    except Exception as e:
        LOGGER.error(f"Error with extracting feature from image {e}")
        sys.exit(1)

此外插入到MySQL中的内容也发生了变化，因此，插入数据的形式的format_data函数内容，修改成如下：

# Combine the id of the vector and the name of the image into a list
def format_data(ids, properites):
    data = []
    for i in range(len(ids)):
        value = (str(ids[i]), properites[i]["image_path"].encode(), properites[i]["gender"], properites[i]["age"])
        data.append(value)
    return data

对应的MySQL创建table和插入数据的语句也需要修改成如下：

    def create_mysql_table(self, table_name):
        # Create mysql table if not exists
        self.test_connection()
        sql = "create table if not exists " + table_name + "(milvus_id TEXT, image_path TEXT, gender BOOLEAN, age INT );"
        try:
            self.cursor.execute(sql)
            LOGGER.debug(f"MYSQL create table: {table_name} with sql: {sql}")
        except Exception as e:
            LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
            sys.exit(1)
    def load_data_to_mysql(self, table_name, data):
        # Batch insert (Milvus_ids, img_path) to mysql
        self.test_connection()
        sql = "insert into " + table_name + " (milvus_id,image_path, gender, age) values (%s,%s, %s, %s);"
        try:
            self.cursor.executemany(sql, data)
            self.conn.commit()
            LOGGER.debug(f"MYSQL loads data to table: {table_name} successfully")
        except Exception as e:
            LOGGER.error(f"MYSQL ERROR: {e} with sql: {sql}")
            sys.exit(1)

3.3 配置文件

由于特征提取网络发生了变换，默认情况下，人脸采用arcfac-resnet50提取特征维度为512，因此，需要对VECTOR_DIMENSION进行修改。同样，为了与以图搜图分开，将DEFAULT_TABLE修改为milvus_face_search，具体修改如下图所示。

 至此，我们完成了代码的相关修改，接下来只需要启动服务，插入数据样本库，执行搜索即可。

4. 启动服务

如上一篇文章一样，使用如下命令，启动服务。

uvicorn main:app --reload

 在浏览器中输入127.0.0.1：8000/docs进入FastAPI - Swagger UI，得到如下页面内容。

 进入/img/load条目，输入table名称和数据样本库的路径，构建图像的向量样本库用于后续的检索

 进入/embedding/load条目，输入table名称，将数据加载到内存中。

 进入/img/search条目，输入table名称、图像文件和topK值，执行检索

返回topK个与目标图像相似的文件路径。

 5. 启动客户端

如上一篇文章一样，启动milvus提供的 milvusbootcamp/img-search-client:1.0   容器，在浏览器中输入127.0.0.1：80001即可进入客户端。

 从上图中可以看到当前的样本库大小为8000余个数据，由于这里只是做一个demo，所以没有插入上百万级数据，俺也没有这么多数据库。

祭出神仙姐姐作为目标样本，得到如下的搜索结果。

 6. 总结

在本文中，基于milvus和insightface库，主要叙述了相关的实现过程，实现了对假的百万级人脸检索，哈哈✔