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库,主要叙述了相关的实现过程,实现了对假的百万级人脸检索,哈哈✔😎