【向量数据库】相似向量检索Faiss数据库的安装及余弦相似度计算（python）

简介

在使用 Faiss 库进行人脸余弦相似度计算的应用中，通常会遵循以下步骤：

特征提取：首先，需要使用人脸识别模型对人脸图像进行特征提取。这些特征通常是高维度的向量，代表了每张人脸图像的抽象特征。

构建 Faiss 索引：将提取出的人脸特征向量构建成 Faiss 可以处理的索引结构。通常会选择使用内积（dot product）作为相似度度量方式，因为 Faiss 默认使用内积进行近似的余弦相似度计算。

相似度查询：当有新的人脸图像需要进行相似度匹配时，将这张人脸图像提取的特征向量作为查询向量，使用 Faiss 进行相似度查询，找出在数据库中与查询向量最相似的人脸特征向量。得到的相似度得分可以用来衡量两张人脸图像之间的相似程度。

应用场景：基于人脸相似度匹配的应用场景包括人脸识别、人脸搜索、人脸聚类等。例如，在人脸识别中，可以通过计算查询人脸与数据库中存储的人脸特征向量之间的余弦相似度来识别或验证用户身份。

代码

注意,faiss只提供了内积搜索，使用下面的程序计算余弦相似度时，需要先将各向量转化为单位向量，单位向量的内积，才是余弦值。

import faiss  # 导入 faiss 库
import numpy as np
class FaceDatabase:
    def __init__(self, dimension=512):
        """
        初始化 FaceDatabase 类
        
        Parameters:
        dimension (int): 人脸向量的维度，默认为 512
        """
        self.dimension = dimension  # 保存人脸向量的维度
        self.index = faiss.IndexFlatIP(dimension)  # 创建 Faiss 的余弦相似度索引
        self.ids = []  # 保存人脸向量对应的 id

    def addFace(self, face_id, face_vector):
        """
        将人脸向量添加进数据库
        
        Parameters:
        face_id: 人脸 id
        face_vector: 人脸向量
        """
        if face_vector.shape[0] != self.dimension:
            raise ValueError("Face vector dimension does not match the database dimension")  # 抛出维度不匹配的异常
        self.index.add(np.expand_dims(face_vector, axis=0))  # 向 Faiss 索引中添加人脸向量
        self.ids.append(face_id)  # 记录人脸向量的 id

    def searchSimilarFaces(self, query_vector, threshold):
        """
        查询相似的人脸向量
            
        Parameters:
        query_vector: 查询向量
        threshold: 相似度阈值
        Returns:
        tuple or None: 返回超过阈值的最相似人脸向量和对应的距离，如果没有超过阈值的，则返回 None
        """
        distances, indices = self.index.search(np.expand_dims(query_vector, axis=0), len(self.ids))  # 使用 Faiss 进行搜索4
        print(distances)
        similar_faces = [(self.ids[i], distances[0][i]) for i in range(len(self.ids)) if distances[0][i] > threshold]  # 保存超过阈值的最相似人脸向量
        if similar_faces:
            most_similar_face = max(similar_faces, key=lambda x: x[1])  # 找到最相似的人脸
            return most_similar_face  # 返回最相似的人脸和对应的距离
        else:
            return None  # 如果没有超过阈值的则返回 None

    def removeFaceById(self, face_id):
        """
        根据人脸 id 删除数据库内的人脸向量
        
        Parameters:
        face_id: 待删除的人脸 id
        """
        #先删除高索引的元素，再删除低索引的元素，避免索引错位的问题。
        index_to_remove = [i for i, stored_id in enumerate(self.ids) if stored_id == face_id]  # 找到需要删除的人脸向量的索引
        for i in sorted(index_to_remove, reverse=True):
            self.index.remove_ids(np.array([i]))  # 从 Faiss 索引中删除对应的人脸向量
            del self.ids[i]  # 从 id 列表中删除对应的 id

    def getNumFaces(self):
        """
        获取数据库内存储的人脸向量数量
        
        Returns:
        int: 人脸向量数量
        """
        return len(self.ids)  # 返回保存的人脸向量数量

    def clearDatabase(self):
        """
        清空数据库，删除所有的人脸向量
        """
        self.index.reset()  # 重置 Faiss 索引
        self.ids.clear()  # 清空 id 列表
        
if __name__ == '__main__':
   # 创建一个 FaceDatabase 实例
    face_db = FaceDatabase()

    # 添加人脸向量和对应的 id
    id1 = '001'
    vector1 = np.random.rand(512).astype('float32')  # 随机生成一个512维向量
    face_db.addFace(id1, vector1)

    id2 = '002'
    vector2 = np.random.rand(512).astype('float32')
    face_db.addFace(id2, vector2)

    id3 = '003'
    vector3 = np.random.rand(512).astype('float32')
    face_db.addFace(id3, vector3)

    # 搜索相似人脸
    query_vector = np.random.rand(512).astype('float32')  # 随机生成一个512维向量作为查询向量
    threshold = -1000000000000000000
    result = face_db.searchSimilarFaces(query_vector, threshold)
    if result:
        similar_id, distance = result
        print(f"The most similar face is {similar_id} with distance {distance}")
    else:
        print("No similar face found")

    # 删除人脸
    face_db.removeFaceById(id2)
    print(f"Number of faces in the database: {face_db.getNumFaces()}")

    # 清空数据库
    face_db.clearDatabase()
    print(f"Number of faces in the database: {face_db.getNumFaces()}")