def __init__( self, index="datasets/vector_generous_jazz_2021-02-02_11-53-36.index", model=None, ) -> None: super().__init__() self.mtcnn = MTCNN(image_size=160, margin=0, selection_method="probability") self.threshold = 0.001 self.to_tensor = torchvision.transforms.ToTensor() self.indexIDMap = faiss.read_index(index) self.dictionary = LabelNames("data/data.p") if os.uname().machine != "ppc64le": self.preprocessor = FaceAlignmentMTCNN() if not model: self.checkpoint = torch.load( "checkpoints/generous-jazz-275_epoch_19", map_location=torch.device("cpu"), ) self.model = InceptionResnetV1() self.model.load_state_dict(self.checkpoint["model_state_dict"]) else: self.model = model self.model.eval()
def live_detection(database, threshold): device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') print('Running on device: {}'.format(device)) 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='vggface2').eval().to(device) with open(database, "rb") as pkl_in: database = pickle.load(pkl_in) embeddings_set, id_to_name = database cap = cv2.VideoCapture(0) # Check if the webcam is opened correctly if not cap.isOpened(): raise IOError("Cannot open webcam") while True: ret, frame = cap.read() rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) try: face = mtcnn(rgb) c, _ = mtcnn.detect(rgb) c = c.flatten().tolist() embeddings = resnet(face.unsqueeze(0).to(device)).detach() if embeddings is not None: index = (embeddings_set - embeddings).norm(dim=-1).argmin().item() dist = (embeddings_set - embeddings).norm(dim=-1).min().item() if dist < threshold: name = id_to_name[index] else: name = "Unknown" cv2.rectangle(frame, (int(c[0]), int(c[1])), (int(c[2]), int(c[3])), (0, 0, 255), 2) cv2.putText(frame, name, (int(c[0]) - 120, int(c[1]) - 10), cv2.FONT_HERSHEY_TRIPLEX, 1.0, [0, 0, 255], 2, 1) except Exception: pass cv2.imshow('Input', frame) c = cv2.waitKey(1) if c == 27: break cap.release() cv2.destroyAllWindows()
def get_data(): print("getting data...") timing = perf_counter() CLASSES_PER_BATCH = 30 SAMPLES_PER_CLASS = 40 BATCH_SIZE = CLASSES_PER_BATCH * SAMPLES_PER_CLASS # dataset = WebfaceDataset("../../data/CASIA-WebFace") dataset = WebfaceDataset("datasets/CASIA-WebFace") train_loader, _, _ = get_data_loaders( dataset, CLASSES_PER_BATCH, SAMPLES_PER_CLASS, train_proportion=0.01, val_proportion=0.89, test_proportion=0.1, batch_size=2000, ) print("loading model...") if torch.cuda.is_available(): checkpoint = torch.load( "checkpoints/generous-jazz-275_epoch_19", map_location=torch.device("cuda"), ) else: checkpoint = torch.load( "checkpoints/generous-jazz-275_epoch_19", map_location=torch.device("cpu"), ) model = InceptionResnetV1() model.load_state_dict(checkpoint["model_state_dict"]) if torch.cuda.is_available(): model = model.cuda() print("calculating embeddings...") embeddings, targets = extract_embeddings(train_loader, model) print(f"took {perf_counter() - timing} seconds") dump( (embeddings, targets), f"datasets/generous-jazz-275_epoch_19_{datetime.fromtimestamp(time()).strftime('%Y-%m-%d_%H-%M-%S')}.joblib", ) return embeddings, targets
def __init__(self) -> None: super().__init__() self.face_database = {} self.to_tensor = torchvision.transforms.ToTensor() self.threshold = 0.2 self.checkpoint = torch.load("checkpoints/charmed-cosmos-135_epoch_19", map_location=torch.device("cpu")) self.model = InceptionResnetV1() self.model.load_state_dict(self.checkpoint["model_state_dict"]) self.model.eval() print("Starting initaliziation") self.initalize() print("Ended initaliziation")
def __init__(self, pretrained=False): super(FECNet, self).__init__() growth_rate = 64 depth = 100 block_config = [5] efficient = True self.Inc = InceptionResnetV1(pretrained='vggface2', device='cuda').eval() for param in self.Inc.parameters(): param.requires_grad = False self.dense = DenseNet(growth_rate=growth_rate, block_config=block_config, num_classes=16, small_inputs=True, efficient=efficient, num_init_features=512).cuda() if (pretrained): load_weights(self)
def embedding(file_name): import PIL.Image as Image from models.mtcnn import MTCNN from models.inception_resnet_v1 import InceptionResnetV1 workers = 10 #device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') device = torch.device('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) crop_imgs = [] labels = [] with open (file_name, 'rt') as file: for line in file: l = line.strip('\n').split(' ') path = l[0] labels.append(l[1]) x = Image.open(path) x_aligned, prob = mtcnn(x, return_prob=True) crop_imgs.append(x_aligned) print('total', len(crop_imgs),'images read!') embeddings = [] if len(crop_imgs)%100==0 and len(crop_imgs)>=100: for i in range(len(crop_imgs)//100): em = resnet(torch.stack(crop_imgs[i*100:100*(i+1)]).to(device)).detach().cpu() embeddings+=em else: for i in range(len(crop_imgs)): em = resnet(torch.stack(crop_imgs).to(device)).detach().cpu() embeddings+=em return embeddings, labels
# Note that, since MTCNN is a collection of neural nets and other code, the # device must be passed in the following way to enable copying of objects when # needed internally. # See `help(MTCNN)` for details. mtcnn = MTCNN(image_size=160, margin=0, min_face_size=20, thresholds=[0.6, 0.7, 0.7], factor=0.709, prewhiten=True, device=device) # Define Inception Resnet V1 module # Set classify=True for pretrained classifier # See `help(InceptionResnetV1)` for details resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device) # Define a dataset and data loader dataset = datasets.ImageFolder('data/test_images') dataset.idx_to_class = {i: c for c, i in dataset.class_to_idx.items()} loader = DataLoader(dataset, collate_fn=lambda x: x[0]) # Perfom MTCNN facial detection aligned = [] names = [] for x, y in loader: x_aligned, prob = mtcnn(x, return_prob=True) if x_aligned is not None: print('Face detected with probability: {:8f}'.format(prob)) aligned.append(x_aligned) names.append(dataset.idx_to_class[y])
def init_GAN(): """Function used during presentation to initialize GAN and classifier quickly.""" class Generator(nn.Module): def __init__(self, ngpu): super(Generator, self).__init__() self.ngpu = ngpu self.main = nn.Sequential( # input is Z, going into a convolution nn.ConvTranspose2d( nz, ngf * 8, 4, 1, 0, bias=False), nn.BatchNorm2d(ngf * 8), nn.ReLU(True), # state size. (ngf*8) x 4 x 4 nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False), nn.BatchNorm2d(ngf * 4), nn.ReLU(True), # state size. (ngf*4) x 8 x 8 nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 1, bias=False), nn.BatchNorm2d(ngf * 2), nn.ReLU(True), # state size. (ngf*2) x 16 x 16 nn.ConvTranspose2d( ngf * 2, ngf, 4, 3, 3, bias=False), nn.BatchNorm2d(ngf), nn.ReLU(True), # state size. (ngf) x 32 x 32 nn.ConvTranspose2d( ngf, nc, 4, 4, 6, bias=False), nn.Tanh() # state size. (nc) x 160 x 160 ) def forward(self, input): return self.main(input) class Discriminator(nn.Module): def __init__(self, ngpu): super(Discriminator, self).__init__() self.ngpu = ngpu self.main = nn.Sequential( # input is (nc) x 160 x 160 nn.Conv2d(nc, ndf, 4, 4, 6, bias=False), nn.LeakyReLU(0.2, inplace=True), # state size. (ndf) x 32 x 32 nn.Conv2d(ndf, ndf * 2, 4, 3, 3, bias=False), nn.BatchNorm2d(ndf * 2), nn.LeakyReLU(0.2, inplace=True), # state size. (ndf*2) x 16 x 16 nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False), nn.BatchNorm2d(ndf * 4), nn.LeakyReLU(0.2, inplace=True), # state size. (ndf*4) x 8 x 8 nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False), nn.BatchNorm2d(ndf * 8), nn.LeakyReLU(0.2, inplace=True), # state size. (ndf*8) x 4 x 4 nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False), #nn.Sigmoid() ) def forward(self, input): return self.main(input) netD = None; netG = None device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu") bs,nz = 64,100 nc = 3; ndf = 160; ngf = 160 ngpu = 1 netD = Discriminator(ngpu).to(device) netG = Generator(ngpu).to(device) def weights_init(m): classname = m.__class__.__name__ if classname.find('Conv') != -1: nn.init.normal_(m.weight.data, 0.0, 0.02) elif classname.find('BatchNorm') != -1: nn.init.normal_(m.weight.data, 1.0, 0.02) nn.init.constant_(m.bias.data, 0) netD.apply(weights_init) netG.apply(weights_init) # FT class_names = ['Adrien_Brody','Alejandro_Toledo','Angelina_Jolie','Arnold_Schwarzenegger','Carlos_Moya','Charles_Moose','James_Blake','Jennifer_Lopez','Michael_Chaykowsky','Roh_Moo-hyun','Venus_Williams'] model_ft = InceptionResnetV1(pretrained='vggface2', classify=False, num_classes=len(class_names)) layer_list = list(model_ft.children())[-5:] model_ft = nn.Sequential(*list(model_ft.children())[:-5]) class Flatten(nn.Module): def __init__(self): super(Flatten, self).__init__() def forward(self, x): x = x.view(x.size(0), -1) return x class normalize(nn.Module): def __init__(self): super(normalize, self).__init__() def forward(self, x): x = F.normalize(x, p=2, dim=1) return x ## IF TRAINING JUST LAST LAYERS for param in model_ft.parameters(): param.requires_grad = False model_ft.avgpool_1a = nn.AdaptiveAvgPool2d(output_size=1) model_ft.last_linear = nn.Sequential( Flatten(), nn.Linear(in_features=1792, out_features=512, bias=False), normalize() ) model_ft.logits = nn.Linear(layer_list[3].in_features, len(class_names)) model_ft.softmax = nn.Softmax(dim=1) model_ft = model_ft.to(device) return netD, netG, model_ft
if __name__ == "__main__": torch.manual_seed(42) EPOCHS = 200 LEARNING_RATE = 0.001 DROPOUT_PROB = 0.6 SCALE_INCEPTION_A = 0.17 SCALE_INCEPTION_B = 0.10 SCALE_INCEPTION_C = 0.20 MARGIN = 0.2 CLASSES_PER_BATCH = 35 SAMPLES_PER_CLASS = 20 BATCH_SIZE = CLASSES_PER_BATCH * SAMPLES_PER_CLASS model = InceptionResnetV1(DROPOUT_PROB, SCALE_INCEPTION_A, SCALE_INCEPTION_B, SCALE_INCEPTION_C) if torch.cuda.is_available(): model = model.cuda() optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE) scheduler = MultiStepLR(optimizer, milestones=[50, 100], gamma=0.1) triplet_gen = triplet_generator.get_semihard # enable resuming of a training run # os.environ["WANDB_RESUME"] = "must" # os.environ["WANDB_RUN_ID"] = "qpilunyg" wandb.init( project="face-recognition",
def run_AGN(last_layers: str = 'n', ngpu: int = 0, testing: str = 'n', write_vid: str = 'y'): """ Function to run full AGN training from scratch. 1. Train GAN to produce realistic eyeglass frames 2. Finetune facial recognition classifier on images of my face 3. Adversarially train generator to produce eyeglasses that trick the classifier """ device = torch.device("cuda:0" if ( torch.cuda.is_available() and ngpu > 0) else "cpu") # Initialize dataset of eyeglasses tensor_dataset = EyeglassesDataset(csv_file='data/files_sample.csv', root_dir='data/eyeglasses/', transform=transforms.Compose([ transforms.ToPILImage(), transforms.Resize(160), transforms.ToTensor(), Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ])) # Initialize D & G bs = 64 nc, ndf, ngf, nz = 3, 160, 160, 100 netD = Discriminator(ngpu, nc, ndf, ngf, nz).to(device) netG = Generator(ngpu, nc, ndf, ngf, nz).to(device) # Apply weights netD.apply(weights_init) netG.apply(weights_init) # Initialize BCELoss function - BCEWithLogitsLoss is used due to removal of Sigmoid in D criterion = nn.BCEWithLogitsLoss() # Create batch of latent vectors that we will use to visualize the progression of the generator fixed_noise = torch.randn(64, nz, 1, 1, device=device) # Build dataloader of eyeglass frames eye_dl = DataLoader(tensor_dataset, batch_size=64, shuffle=True) # Pre-train the GAN on the eyeglasses img_list, G_losses, D_losses, netG, netD = pretrain_gan(netD, netG, eye_dl, criterion, fixed_noise, nz, testing=testing) # Data augmentation and normalization for training - Just normalization for validation data_transforms = { 'train': transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), 'val': transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]), } # Specify path for location of images of my face data_dir = 'data/test_me' # Build datasets and loaders for both train and validation sets image_datasets = { x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x]) for x in ['train', 'val'] } dataloaders = { x: torch.utils.data.DataLoader(image_datasets[x], batch_size=8, shuffle=True) for x in ['train', 'val'] } dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']} class_names = image_datasets['train'].classes # Get pretrained resnet on vggface2 dataset model_ft = InceptionResnetV1(pretrained='vggface2', classify=False, num_classes=len(class_names)) # Remove the last layers after conv block layer_list = list(model_ft.children())[-5:] # all final layers # All beginning layers model_ft = nn.Sequential(*list(model_ft.children())[:-5]) if last_layers == 'y': for param in model_ft.parameters(): param.requires_grad = False # Re-initialize layers to set requires_grad to True model_ft.avgpool_1a = nn.AdaptiveAvgPool2d(output_size=1) model_ft.last_linear = nn.Sequential( Flatten(), nn.Linear(in_features=1792, out_features=512, bias=False), normalize()) model_ft.logits = nn.Linear(layer_list[3].in_features, len(class_names)) model_ft.softmax = nn.Softmax(dim=1) # Place model on device model_ft = model_ft.to(device) # Define loss function for finetuning criterion = nn.CrossEntropyLoss() # Observe that all parameters are being optimized optimizer_ft = optim.SGD(model_ft.parameters(), lr=1e-2, momentum=0.9) # Decay LR by a factor of *gamma* every *step_size* epochs exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1) # Finetune the facial recognition classifier model_ft, FT_losses = train_ft_model(model_ft, dataloaders, dataset_sizes, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=1) # Save the model torch.save({ 'state_dict': model_ft.state_dict() }, f"tmp/model_ft_loss{FT_losses[-1]}_{dt.datetime.today().strftime('%Y%m%d')}.pth.tar" ) ### AGN training ### # If you have not already saved the frames of the video of you can do that here ### Perform data pre-processing step on frames from video described in readme using command line after this step ### if write_vid == 'y': vidcap = cv2.VideoCapture('data/IMG_2411.MOV') success, image = vidcap.read() count = 0 success = True while success: cv2.imwrite( f"data/agn_me_extras160/Michael_Chaykowsky/Michael_Chaykowsky_{format(count, '04d')}.png", image) success, image = vidcap.read() print('Read a new frame: ', success) count += 1 # Transforms for the images and the coordinates (where eyeglasses go) seperately t_img = transforms.Compose([ transforms.ToPILImage(), transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1), # transforms.RandomHorizontalFlip(), # couldn't get to work # transforms.RandomRotation(5), # couldn't get to work transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) t_landmarks = transforms.Compose([transforms.ToTensor()]) # Custom dataset for images of my face to affix glasses image_datasets_me = MeDataset('data/bboxes_fnames.csv', 'data/agn_me_extras160/Michael_Chaykowsky', bs=1, transform_img=t_img, transform_land=t_landmarks) # Build dataloader for face images dataloader_me = torch.utils.data.DataLoader(image_datasets_me, batch_size=64, shuffle=True) # Pre-process video data of my face to get images to adversarially train the generator coords = vid_preprocess() # Full training of AGN netG, img_list, G_losses, D_losses, d1s, d3s, num_fooled = train_AGN( netG, netD, model_ft, eye_dl, dataloader_me, class_names, nz, num_epochs=1) # Save adversarial generator for testing later torch.save({'state_dict': netG.state_dict()}, f"tmp/adv_G_{dt.datetime.today().strftime('%Y%m%d')}.pth.tar")
import sys from PIL import Image sys.path.append('./') from models.mtcnn import MTCNN # noqa from models.inception_resnet_v1 import InceptionResnetV1 # noqa from models.utils import face_verify # noqa if __name__ == "__main__": mtcnn = MTCNN() resnet = InceptionResnetV1(pretrained='vggface2', classify=False).eval() know_imgs = [ './data/test_data/mingxing/train/chenglong/4d017614bbd293dc.jpg', './data/test_data/mingxing/train/gutianle/3d6fd78e5d9b0703.jpg', './data/test_data/mingxing/train/huge/01e31e7e0d162357.jpg', './data/test_data/mingxing/train/jindong/0ad09ddfefe51c62.jpg', ] test_img = './data/test_data/mingxing/test/huge.jpg' known_encodings = list() for img_id, know_img in enumerate(know_imgs): save_path = './data/test_data/img_{}.jpg'.format(img_id) img = Image.open(know_img) # Get cropped and prewhitened image tensor img_cropped = mtcnn(img, save_path=save_path) # Calculate embedding (unsqueeze to add batch dimension) img_embedding = resnet( img_cropped.unsqueeze(0)).view(-1).detach().numpy() print('embed id:{}'.format(img_id)) known_encodings.append(img_embedding) print('已有图像编码完成,known_encodings:{}'.format(len(known_encodings)))
#### TEST EMBEDDINGS #### expected = [[[0.000000, 1.395957, 0.785551, 1.456866, 1.466266], [1.395957, 0.000000, 1.264742, 0.902874, 0.911210], [0.785551, 1.264742, 0.000000, 1.360339, 1.405513], [1.456866, 0.902874, 1.360339, 0.000000, 1.066445], [1.466266, 0.911210, 1.405513, 1.066445, 0.000000]], [[0.000000, 1.330782, 0.846278, 1.359174, 1.222049], [1.330782, 0.000000, 1.157455, 0.989477, 0.974240], [0.846278, 1.157455, 0.000000, 1.309103, 1.234498], [1.359174, 0.989477, 1.309103, 0.000000, 1.066433], [1.222049, 0.974240, 1.234498, 1.066433, 0.000000]]] for i, ds in enumerate(['vggface2', 'casia-webface']): resnet_pt = InceptionResnetV1(pretrained=ds).eval() start = time() embs = resnet_pt(aligned) print('\nResnet time: {:6f} seconds\n'.format(time() - start)) embs_fromfile = resnet_pt(aligned_fromfile) dists = [[(emb - e).norm().item() for e in embs] for emb in embs] dists_fromfile = [[(emb - e).norm().item() for e in embs_fromfile] for emb in embs_fromfile] print('\nOutput:') print(pd.DataFrame(dists, columns=names, index=names)) print('\nOutput (from file):') print(pd.DataFrame(dists_fromfile, columns=names, index=names))
dataset.samples = [(p, p.replace(data_dir, data_dir + '_cropped')) for p, _ in dataset.samples] loader = DataLoader(dataset, num_workers=workers, batch_size=batch_size, collate_fn=training.collate_pil) for i, (x, y) in enumerate(loader): mtcnn(x, save_path=y) print('\rBatch {} of {}'.format(i + 1, len(loader)), end='') # Remove mtcnn to reduce GPU memory usage del mtcnn resnet = InceptionResnetV1(classify=True, pretrained='vggface2', num_classes=len( dataset.class_to_idx)).to(device) # print('resnet:{}'.format(resnet)) # optimizer = optim.Adam(resnet.parameters(), lr=0.001) print([name for name, param in resnet.named_parameters()]) # 微调时,只重新训练输出层参数 optim_params = [ param for name, param in resnet.named_parameters() if name in {'logits.weight', 'logits.bias'} ] print('optim_params:{}'.format(optim_params)) optimizer = optim.Adam(optim_params) scheduler = MultiStepLR(optimizer, [5, 10]) trans = transforms.Compose( [np.float32,
import tornado.web import json from tornado.escape import json_decode from models.mtcnn import MTCNN from models.inception_resnet_v1 import InceptionResnetV1 from FaceRecognition import cosSimilarity, getVec from utils import getConnection define("port", default=8080, help="run on the given port", type=int) tornado.options.parse_command_line() BASE_DIR = os.path.dirname(__file__) # FaceNet mtcnn = MTCNN(image_size=160, margin=10) resnet = InceptionResnetV1().eval() class MainHandler(tornado.web.RequestHandler): def get(self): self.render("index.html") def post(self): self.render("index.html") class FaceRecognitionHandler(tornado.web.RequestHandler): def get(self): pass def post(self):
import pygame.locals pygame.init() from imports import * from archs import * from models.inception_resnet_v1 import InceptionResnetV1 device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') print('Running on device: {}'.format(device)) # Load Model FT pretrained class_names = [ 'Adrien_Brody','Alejandro_Toledo','Angelina_Jolie','Arnold_Schwarzenegger','Carlos_Moya', 'Charles_Moose','James_Blake','Jennifer_Lopez','Michael_Chaykowsky','Roh_Moo-hyun','Venus_Williams'] model_ft = InceptionResnetV1(pretrained='vggface2', classify=False, num_classes=len(class_names)) layer_list = list(model_ft.children())[-5:] model_ft = nn.Sequential(*list(model_ft.children())[:-5]) model_ft.avgpool_1a = nn.AdaptiveAvgPool2d(output_size=1) model_ft.last_linear = nn.Sequential( Flatten(), nn.Linear(in_features=1792, out_features=512, bias=False), normalize() ) model_ft.logits = nn.Linear(layer_list[3].in_features, len(class_names)) model_ft.softmax = nn.Softmax(dim=1) model_ft = model_ft.to(device) # Build masks imgGlass = cv2.imread("data/glasses_mask.png", -1) r = 160.0 / imgGlass.shape[1]
def tensorflow2pytorch(): lookup_inception_resnet_v1 = { 'conv2d_1a': ['InceptionResnetV1/Conv2d_1a_3x3', load_tf_basicConv2d], 'conv2d_2a': ['InceptionResnetV1/Conv2d_2a_3x3', load_tf_basicConv2d], 'conv2d_2b': ['InceptionResnetV1/Conv2d_2b_3x3', load_tf_basicConv2d], 'conv2d_3b': ['InceptionResnetV1/Conv2d_3b_1x1', load_tf_basicConv2d], 'conv2d_4a': ['InceptionResnetV1/Conv2d_4a_3x3', load_tf_basicConv2d], 'conv2d_4b': ['InceptionResnetV1/Conv2d_4b_3x3', load_tf_basicConv2d], 'repeat_1': ['InceptionResnetV1/Repeat/block35', load_tf_repeat_1], 'mixed_6a': ['InceptionResnetV1/Mixed_6a', load_tf_mixed6a], 'repeat_2': ['InceptionResnetV1/Repeat_1/block17', load_tf_repeat_2], 'mixed_7a': ['InceptionResnetV1/Mixed_7a', load_tf_mixed7a], 'repeat_3': ['InceptionResnetV1/Repeat_2/block8', load_tf_repeat_3], 'block8': ['InceptionResnetV1/Block8', load_tf_block17_8], 'last_linear': ['InceptionResnetV1/Bottleneck/weights', load_tf_linear], 'last_bn': ['InceptionResnetV1/Bottleneck/BatchNorm', load_tf_batchNorm], 'logits': ['Logits', load_tf_linear], } print('\nLoad VGGFace2-trained weights and save\n') mdl = InceptionResnetV1(num_classes=8631).eval() tf_mdl_dir = 'data/20180402-114759' data_name = 'vggface2' load_tf_model_weights(mdl, lookup_inception_resnet_v1, tf_mdl_dir) state_dict = mdl.state_dict() torch.save(state_dict, f'{tf_mdl_dir}-{data_name}.pt') torch.save( { 'logits.weight': state_dict['logits.weight'], 'logits.bias': state_dict['logits.bias'], }, f'{tf_mdl_dir}-{data_name}-logits.pt') state_dict.pop('logits.weight') state_dict.pop('logits.bias') torch.save(state_dict, f'{tf_mdl_dir}-{data_name}-features.pt') print('\nLoad CASIA-Webface-trained weights and save\n') mdl = InceptionResnetV1(num_classes=10575).eval() tf_mdl_dir = 'data/20180408-102900' data_name = 'casia-webface' load_tf_model_weights(mdl, lookup_inception_resnet_v1, tf_mdl_dir) state_dict = mdl.state_dict() torch.save(state_dict, f'{tf_mdl_dir}-{data_name}.pt') torch.save( { 'logits.weight': state_dict['logits.weight'], 'logits.bias': state_dict['logits.bias'], }, f'{tf_mdl_dir}-{data_name}-logits.pt') state_dict.pop('logits.weight') state_dict.pop('logits.bias') torch.save(state_dict, f'{tf_mdl_dir}-{data_name}-features.pt') lookup_pnet = { 'conv1': ['pnet/conv1', load_tf_conv2d_trans], 'prelu1': ['pnet/PReLU1', load_tf_linear], 'conv2': ['pnet/conv2', load_tf_conv2d_trans], 'prelu2': ['pnet/PReLU2', load_tf_linear], 'conv3': ['pnet/conv3', load_tf_conv2d_trans], 'prelu3': ['pnet/PReLU3', load_tf_linear], 'conv4_1': ['pnet/conv4-1', load_tf_conv2d_trans], 'conv4_2': ['pnet/conv4-2', load_tf_conv2d_trans], } lookup_rnet = { 'conv1': ['rnet/conv1', load_tf_conv2d_trans], 'prelu1': ['rnet/prelu1', load_tf_linear], 'conv2': ['rnet/conv2', load_tf_conv2d_trans], 'prelu2': ['rnet/prelu2', load_tf_linear], 'conv3': ['rnet/conv3', load_tf_conv2d_trans], 'prelu3': ['rnet/prelu3', load_tf_linear], 'dense4': ['rnet/conv4', load_tf_linear], 'prelu4': ['rnet/prelu4', load_tf_linear], 'dense5_1': ['rnet/conv5-1', load_tf_linear], 'dense5_2': ['rnet/conv5-2', load_tf_linear], } lookup_onet = { 'conv1': ['onet/conv1', load_tf_conv2d_trans], 'prelu1': ['onet/prelu1', load_tf_linear], 'conv2': ['onet/conv2', load_tf_conv2d_trans], 'prelu2': ['onet/prelu2', load_tf_linear], 'conv3': ['onet/conv3', load_tf_conv2d_trans], 'prelu3': ['onet/prelu3', load_tf_linear], 'conv4': ['onet/conv4', load_tf_conv2d_trans], 'prelu4': ['onet/prelu4', load_tf_linear], 'dense5': ['onet/conv5', load_tf_linear], 'prelu5': ['onet/prelu5', load_tf_linear], 'dense6_1': ['onet/conv6-1', load_tf_linear], 'dense6_2': ['onet/conv6-2', load_tf_linear], 'dense6_3': ['onet/conv6-3', load_tf_linear], } print('\nLoad PNet weights and save\n') tf_mdl_dir = lambda sess: detect_face.create_mtcnn(sess, None) mdl = PNet() data_name = 'pnet' load_tf_model_weights(mdl, lookup_pnet, tf_mdl_dir, is_resnet=False, arg_num=0) torch.save(mdl.state_dict(), f'data/{data_name}.pt') tf.reset_default_graph() with tf.Session() as sess: compare_mtcnn(mdl, tf_mdl_dir, sess, 0, torch.randn(1, 256, 256, 3).detach()) print('\nLoad RNet weights and save\n') mdl = RNet() data_name = 'rnet' load_tf_model_weights(mdl, lookup_rnet, tf_mdl_dir, is_resnet=False, arg_num=1) torch.save(mdl.state_dict(), f'data/{data_name}.pt') tf.reset_default_graph() with tf.Session() as sess: compare_mtcnn(mdl, tf_mdl_dir, sess, 1, torch.randn(1, 24, 24, 3).detach()) print('\nLoad ONet weights and save\n') mdl = ONet() data_name = 'onet' load_tf_model_weights(mdl, lookup_onet, tf_mdl_dir, is_resnet=False, arg_num=2) torch.save(mdl.state_dict(), f'data/{data_name}.pt') tf.reset_default_graph() with tf.Session() as sess: compare_mtcnn(mdl, tf_mdl_dir, sess, 2, torch.randn(1, 48, 48, 3).detach())
img = trans(img) return img trans = transforms.Compose([ transforms.Resize(512) ]) trans_cropped = transforms.Compose([ np.float32, transforms.ToTensor(), prewhiten ]) mtcnn_pt = MTCNN() resnet_pt = InceptionResnetV1(pretrained='vggface2').eval() dataset = datasets.ImageFolder('data/test_images', transform=trans) dataset.idx_to_class = {k: v for v, k in dataset.class_to_idx.items()} loader = DataLoader(dataset, num_workers=8, collate_fn=lambda x: x[0]) names = [] aligned = [] aligned_fromfile = [] for img, idx in loader: name = dataset.idx_to_class[idx] names.append(name) start = time() img_align = mtcnn_pt(img, save_path='data/test_images_aligned/{}/1.png'.format(name)) print('MTCNN time: {:6f} seconds'.format(time() - start)) aligned.append(img_align)
_, train_labels = joblib.load( # "../../data/generous-jazz-275_epoch_19_2021-02-01_14-41-10.joblib" "datasets/generous-jazz-275_epoch_19_2021-02-01_14-41-10.joblib") if torch.cuda.is_available(): checkpoint = torch.load( "checkpoints/generous-jazz-275_epoch_19", map_location=torch.device("cuda"), ) else: checkpoint = torch.load( "checkpoints/generous-jazz-275_epoch_19", map_location=torch.device("cpu"), ) model = InceptionResnetV1() model.load_state_dict(checkpoint["model_state_dict"]) model.eval() dataset = WebfaceDataset("datasets/CASIA-WebFace") # dataset = WebfaceDataset("../../data/CASIA-WebFace") CLASSES_PER_BATCH = 35 SAMPLES_PER_CLASS = 40 BATCH_SIZE = CLASSES_PER_BATCH * SAMPLES_PER_CLASS print("Loading dataset...") train_loader, val_loader, _ = get_data_loaders( dataset, CLASSES_PER_BATCH, SAMPLES_PER_CLASS, train_proportion=0.8, val_proportion=0.1,