def extract_feature(model_path, backbone_net, face_scrub_path, megaface_path, batch_size=32, gpus='0', do_norm=False): if backbone_net == 'MobileFace': net = mobilefacenet.MobileFaceNet() elif backbone_net == 'Res50_IR': net = cbam.CBAMResNet_IR(50, feature_dim=args.feature_dim, mode='ir') elif backbone_net == 'SERes50_IR': net = cbam.CBAMResNet_IR(50, feature_dim=args.feature_dim, mode='se_ir') elif backbone_net == 'CBAMRes50_IR': net = cbam.CBAMResNet_IR(50, feature_dim=args.feature_dim, mode='cbam_ir') elif backbone_net == 'Res100_IR': net = cbam.CBAMResNet_IR(100, feature_dim=args.feature_dim, mode='ir') elif backbone_net == 'SERes100_IR': net = cbam.CBAMResNet_IR(100, feature_dim=args.feature_dim, mode='se_ir') elif backbone_net == 'CBAMRes100_IR': net = cbam.CBAMResNet_IR(100, feature_dim=args.feature_dim, mode='cbam_ir') else: print(args.backbone, ' is not available!') multi_gpus = False if len(gpus.split(',')) > 1: multi_gpus = True os.environ['CUDA_VISIBLE_DEVICES'] = gpus device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') net.load_state_dict(torch.load(model_path)['net_state_dict']) if multi_gpus: net = DataParallel(net).to(device) else: net = net.to(device) net.eval() transform = transforms.Compose([ transforms.ToTensor(), # range [0, 255] -> [0.0,1.0] transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) # range [0.0, 1.0] -> [-1.0,1.0] ]) megaface_dataset = MegaFace(face_scrub_path, megaface_path, transform=transform) megaface_loader = torch.utils.data.DataLoader(megaface_dataset, batch_size=batch_size, shuffle=False, num_workers=12, drop_last=False) for data in megaface_loader: img, img_path= data[0].to(device), data[1] with torch.no_grad(): output = net(img).data.cpu().numpy() if do_norm is False: for i in range(len(img_path)): abs_path = img_path[i] + '.feat' write_mat(abs_path, output[i]) print('extract 1 batch...without feature normalization') else: for i in range(len(img_path)): abs_path = img_path[i] + '.feat' feat = output[i] feat = feat / np.sqrt((np.dot(feat, feat))) write_mat(abs_path, feat) print('extract 1 batch...with feature normalization') print('all images have been processed!')
def loadModel(data_root, file_list, backbone_net, gpus='0', resume=None): if backbone_net == 'MobileFace': net = mobilefacenet.MobileFaceNet() elif backbone_net == 'Res50_IR': net = cbam.CBAMResNet_IR(50, feature_dim=args.feature_dim, mode='ir') elif backbone_net == 'SERes50_IR': net = cbam.CBAMResNet_IR(50, feature_dim=args.feature_dim, mode='se_ir') elif backbone_net == 'CBAMRes50_IR': net = cbam.CBAMResNet_IR(50, feature_dim=args.feature_dim, mode='cbam_ir') elif backbone_net == 'Res100_IR': net = cbam.CBAMResNet_IR(100, feature_dim=args.feature_dim, mode='ir') elif backbone_net == 'SERes100_IR': net = cbam.CBAMResNet_IR(100, feature_dim=args.feature_dim, mode='se_ir') elif backbone_net == 'CBAMRes100_IR': net = cbam.CBAMResNet_IR(100, feature_dim=args.feature_dim, mode='cbam_ir') else: print(args.backbone, ' is not available!') # gpu init multi_gpus = False if len(gpus.split(',')) > 1: multi_gpus = True os.environ['CUDA_VISIBLE_DEVICES'] = gpus device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') net.load_state_dict(torch.load(resume)['net_state_dict']) if multi_gpus: net = DataParallel(net).to(device) else: net = net.to(device) transform = transforms.Compose([ transforms.ToTensor(), # range [0, 255] -> [0.0,1.0] transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) # range [0.0, 1.0] -> [-1.0,1.0] ]) cfp_dataset = CFP_FP(data_root, file_list, transform=transform) cfp_loader = torch.utils.data.DataLoader(cfp_dataset, batch_size=128, shuffle=False, num_workers=4, drop_last=False) return net.eval(), device, cfp_dataset, cfp_loader