【Milvus的人脸检索】

最新推荐文章于 2024-09-13 07:37:52 发布

hello_dear_you

最新推荐文章于 2024-09-13 07:37:52 发布

阅读量2.6k

点赞数 5

分类专栏： # Milvus 文章标签： milvus insightface 人脸检索

本文链接：https://blog.csdn.net/hello_dear_you/article/details/127976638

版权

Milvus 专栏收录该内容

5 篇文章 9 订阅

订阅专栏

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)