特征库制作完整代码

import os
from a_Net_module import *
import torch
from PIL import Image
from detect_img import Detector
import cv2
import numpy as np


class FaceDatabase():
    def __init__(self, personName_path, params_path,database_path):

        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        net = FaceNet(500).to(device)
        net.load_state_dict(torch.load(params_path))
        net.eval()

        self.database = {}
        for personName in os.listdir(personName_path):
            person_feats = []
            for personImg_name in os.listdir(os.path.join(personName_path, personName)):
                img = Image.open(os.path.join(personName_path, personName, personImg_name))
                person_feat = net.encode((torch.unsqueeze(transf(img), dim=0)).to(device))
                person_feats.append(person_feat)
            person_feats = torch.cat(person_feats, dim=0)
            self.database[personName] = person_feats

        with open(database_path,'w',encoding='utf8') as f:
            torch.save(self.database,database_path)
        f.close()



if __name__ == '__main__':
    personName_path = r'./face_database'
    params_path = r'./params/params_1.pth'
    database_path=r'./params/face_database.txt'
    database = FaceDatabase(personName_path, params_path,database_path)
    data=torch.load(database_path)
    print(data['何钢'].shape)