def gen_data_set(root=None): """ return: trainset, valset, {idx:class} """ if root is None: raise ValueError("must set root!") trainset = ImageFolder(root + "train", transform=transform) valset = ImageFolder(root + "val", transform=transform) trainset.idx_to_class = {v: k for k, v in trainset.class_to_idx.items()} valset.idx_to_class = {v: k for k, v in valset.class_to_idx.items()} #a = wulandata.__getitem__(500) if trainset.idx_to_class != valset.idx_to_class: raise ValueError("train and val folder net equal.") return trainset, valset, trainset.idx_to_class
def init_base(root, net, mtcnn=None, K=100): # obtain single face in sample => keep_all=False mtcnn = set_mtcnn(mtcnn, all=False, largest=True) data = ImageFolder(root); workers = 0 if os.name=='nt' else 4 data.idx_to_class = {v:k for k,v in data.class_to_idx.items()} collate_fn = lambda x: x[0] # for batch=1: just 1st in list # collate_fn() used to handle list of (img,label) sample: # loader = DataLoader(data, batch_size=2, collate_fn=collate_fn) # for i in loader: print(i) # Try: lambda x: list(zip(*x)) OR x loader = DataLoader(data, collate_fn=collate_fn, num_workers=workers) faces, names, feat = [], [], [] for img, idx in loader: # keep_all=False: (C,H,W) face, box, prob = mtcnn(img, landmarks=False) if face is None or box is None: continue faces.append(face); names.append(data.idx_to_class[idx]) faces = torch.stack(faces) # stack: add dim->(N,C,H,W) img = cat_face(faces, names, sh='face'); cv2.waitKey() cv2.imwrite(f'{root}/{root}.png', img) # save cat_face net = load_net(net); faces = Norm(faces) for i in range(0, len(faces), K): # feat: (n,512) feat.append(net(faces[i:i+K]).detach().cpu()) base = [torch.cat(feat), names] # cat: merge dim torch.save(base, f'{root}/base.pt') # save base cv2.destroyAllWindows(); return base
# This is for our re-trained model model = InceptionResnetV1(pretrained="vggface2", classify=True) def collate_fn(x): return x[0] path = "../../data/vggface2-80" preprocess = transforms.Compose( [transforms.Resize((224, 224)), transforms.ToTensor()]) dataset = ImageFolder(path) dataset.idx_to_class = {i: c for c, i in dataset.class_to_idx.items()} batch_size = 1 loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn) max_iterations = len(loader) num_samples = len(loader) num_images = max_iterations * batch_size print(f"there are {num_images} images") test_correct = 0 model.eval() mtcnn = MTCNN(image_size=160, margin=0,
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') mtcnn = MTCNN( image_size=160, margin=0, min_face_size=20, thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True, device=device ) resnet = InceptionResnetV1(pretrained='casia-webface').eval().to(device) ds = ImageFolder( root=args['in_path'], ) def collate_fn(x): return list(zip(*x)) ds.idx_to_class = {i: c for c, i in ds.class_to_idx.items()} dl = DataLoader( ds, batch_size=args['batch_size'], shuffle=False, num_workers=args['num_workers'], collate_fn=collate_fn, ) for imgs, classes in tqdm(dl): faces, boxes = align_image(imgs, mtcnn) for i, (img, cls, face, box) in enumerate(zip(imgs, classes, faces, boxes)): if face is None: continue boxed_img = draw_boxes(img, [box])