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,
