facenet + fiass 实现人脸识别

1. from facenet_pytorch import MTCNN, InceptionResnetV1
   from PIL import Image, ImageDraw
   import cv2
   import time
   import numpy as np
   import torch
   import faiss
   
   
   class FaissDB(object):
       def __init__(self, d=512):
           """创建相似性搜索库"""
           self.index = faiss.IndexFlatL2(d)  # build the index
   
       def add(self, data):
           """添加向量"""
           # index.add(1000004, xb)  # add vectors to the index
           self.index.add(data)
   
   
       def search(self, target, k=1):
           """
           查询faiss向量数据库获取欧氏距离最小的特征向量
           Args:
               target: 目标向量
               k: 返回最相近的特征向量数量
   
           Returns:
   
           """
           # 从索引中得到前K近个向量， search方法返回值D（欧氏距离值）， I（索引值）
           D, I = self.index.search(target, k)  # sanity check
           # val = filter(lambda x: x < 0.6, D)
           return I
   
   
   class Video(object):
       def __init__(self, mtcnn, resnet):
           self.mtcnn = mtcnn
           self.resnet = resnet
   
       def detection_face(self):
           my_faiss = FaissDB()
           # 特征向量的尺寸
           d = 512  # dimension
           # 生成10万条随机特征向量
           xb = torch.randn((100000, d), dtype=torch.float)
           my_faiss.add(xb)
           fcj = torch.tensor(np.load("/Users/wb-fcj414969/project/facenet-pytorch-master/test/img/fcj.npy"))
           my_faiss.add(fcj)
   
           before_detection_time = 0
           # cv2.resizeWindow("cap", (480, 480))
           capture = cv2.VideoCapture(0)  # 0为电脑内置摄像头
           # 构架视频窗口，cv2.WINDOW_NORMAL标识窗口的大小
           cv2.namedWindow("frame", cv2.WINDOW_NORMAL)
           # cv2.WND_PROP_FULLSCREEN标识全屏
           # cv2.setWindowProperty("frame", cv2.WND_PROP_FULLSCREEN, cv2.WND_PROP_FULLSCREEN)
           # 设置窗口FPS，不生效
           # capture.set(cv2.CAP_PROP_FPS, 10)
           capture.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
           capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
   
           start_time = time.time()
           while True:
               print("FPS:%s" %cv2.CAP_PROP_FPS)
               time1 = time.time()
               # 摄像头读取,ret为是否成功打开摄像头,true,false。 frame为视频的每一帧图像
               ret, frame = capture.read()
               # 格式转变，BGRtoRGB
               frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
               frame = cv2.flip(frame, 1)  # 摄像头是和人对立的，将图像左右调换回来正常显示。
   
               # 人脸检测，显示人脸框，显示特征点
               boxes, probs, points = self.mtcnn.detect(frame, landmarks=True)
               if boxes is not None:
                   max_boxe, amx_probs, max_point = self.mtcnn.select_boxes(boxes, probs, points, frame)
                   img = Image.fromarray(frame)  # ndarry类型转Image类型
                   draw = ImageDraw.Draw(img)
                   draw.rectangle((tuple(max_boxe[0][:2].tolist()), tuple(max_boxe[0][2:].tolist())), width=1)
                   for item in max_point[0]:
                       draw.point(tuple(item), fill=(255, 0, 0))
   
   
                   # 获取人脸张量信息, 并存储
                   img_cropped = self.mtcnn(frame)
                   # Calculate embedding (unsqueeze to add batch dimension)
                   img_embedding = torch.tensor(self.resnet(img_cropped).detach().numpy())
   
                   # 查询向量数据库
                   data = my_faiss.search(img_embedding)
   
                   if data[0]:
                       # 此处通过检测到的人脸索引值获取人脸特征查询数据库获取人员信息，并显示在屏幕中
                       draw.text(tuple(max_boxe[0][:2].tolist()), "fjc")
   
                   # Image类型转ndarry类型后，显示处理后的图片
                   frame = np.array(img)
   
               # RGBtoBGR满足opencv显示格式
               frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
   
               cv2.imshow("frame", frame)
               time2 = time.time()
               print("人脸框检测时间：%s" % str(time2 - time1))
   
               # 判断是否点击了退出键
               c = cv2.waitKey(100)
               if c == 27:
                   break
   
           cv2.destroyAllWindows()
   
   
   if __name__ == "__main__":
       # 116.398663,39.919528
       # 116.408149,39.929599
   
       # a = [[(i, j) for j in range(39919528, 39929599)] for i in range(116398663, 116408149)]
       #
       # b = torch.tensor(a)
       # b = b.view([-1, 2])
       #
       # # 测试faiss速度
       # my_faiss = FaissDB(d=2)
       # my_faiss.add(b)
       # predict = torch.tensor([[116401754, 39925781]])
       # predict = predict.repeat(1000, 1)
       # time1 = time.time()
       # result = my_faiss.search(predict)
       # time2 = time.time()
       # print(time2-time1)
   
       image_size = 160
       mtcnn = MTCNN(image_size=160, margin=0, keep_all=True)
       mtcnn.eval()
       # Create an inception resnet (in eval mode):
       resnet = InceptionResnetV1(pretrained='vggface2')
       resnet.eval()
   
       video = Video(mtcnn, resnet)
       video.detection_face()