def init_model(self):
model = fp.InceptionResnetV1(pretrained='vggface2').to(device)
# load pre trained weights
if os.path.exists(FINETUNED_FILENAME):
model.load_state_dict(torch.load(FINETUNED_FILENAME, map_location=device))
model.train()
model = self.freeze_layers(model)
return model
def __init__(self, location=temporary_location, threshold=0.7):
self.device = device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.threshold = threshold
self.csv_location = location
os.makedirs(location.split("/")[0], exist_ok=True)
torch.no_grad()
self.model = fp.InceptionResnetV1(pretrained='vggface2').eval()
if torch.cuda.is_available():
self.model.cuda()
self.ids, self.target_embeddings = read_target_embeddings(self.csv_location)
if len(self.ids) > 0:
self.kNN = self.fit_targets() # we have kNN object after this
def __init__(self, path, recache=False):
self.resnet = fn.InceptionResnetV1(pretrained='vggface2').eval()
self.mtcnn = fn.MTCNN(image_size=160,
margin=0,
min_face_size=20,
thresholds=[0.6, 0.7, 0.7],
factor=0.709,
post_process=True)
cache_path = os.path.join(path, 'faces.db')
if os.path.exists(cache_path) and not recache:
print('Loading from: ', cache_path)
with open(cache_path, 'rb') as f:
self.matcher_db = pickle.load(f)
else:
print('Building Cache...')
self.matcher_db = self.init_matcher_db(path)
def get_embeddings(source_dlr, target_dlr, args):
with torch.no_grad():
resnet = fp.InceptionResnetV1(pretrained='vggface2').eval()
if torch.cuda.is_available():
resnet.cuda()
if args.pretrained != "" and os.path.exists(args.pretrained):
resnet.load_state_dict(torch.load(args.pretrained))
source_embeddings, source_names = map_images_to_embedding(
resnet, source_dlr)
#target_embeddings, target_names = map_images_to_embedding(resnet, target_dlr, args, target=True)
torch.save(source_embeddings, f"{args.source_filename}_embeddings.pt")
naming_list = source_names #+ target_names
naming_dict = {k: v for v, k in enumerate(naming_list)}
pickle.dump((naming_list, naming_dict),
open(f"{args.source_filename}_embedding_names.pt", "wb"))
print("New embeddings created")
def build_facenet_model(latent_space=512):
mtcnn = fp.MTCNN(device=torch.device("cuda"))
facenet = fp.InceptionResnetV1(pretrained='vggface2').eval()
for p in facenet.parameters():
p.requires_grad = False
facenet = facenet.cuda()
fc = nn.Sequential(*[
# nn.Linear(512, 512),
# nn.ReLU(),
# nn.Linear(512, 512),
# nn.ReLU(),
# nn.Linear(512, 512),
# nn.ReLU(),
# nn.Linear(512, 512),
# nn.ReLU(),
nn.Linear(512, latent_space)
])
fc.train()
fc = fc.cuda()
return mtcnn, facenet, fc
super(FaceNet, self).__init__(**kwargs)
def find_faces(self, img):
return self.detect(img)[0]
def draw_rect(self, frame, faces, text):
if faces is None: return frame
img = frame.copy()
for i, box in enumerate(faces):
x, y, h, w = box.astype(int)
cv2.rectangle(img, (x, y), (h, w), (80, 18, 236), 2)
cv2.rectangle(img, (x, y), (h, y - 15), (80, 18, 236), cv2.FILLED)
try:
name = text[i]
except:
name = 'searching'
cv2.putText(img, name, (x + 6, y - 2), cv2.FONT_HERSHEY_DUPLEX,
0.5, (255, 255, 255), 1)
return img
def crop_face(self, img, box, size=160):
faces = []
for face in box:
x, y, h, w = face.astype(int)
crop = cv2.resize(img[y:y + w, x:x + h], (size, size))
faces.append(crop)
return np.array(faces)
net = facenet.InceptionResnetV1(pretrained='vggface2').eval()
def runFacenet():
mtcnn = fp.MTCNN()
resnet = fp.InceptionResnetV1(pretrained='casia-webface').eval()
return mtcnn, resnet