Face_Recognition / utils.py
HienK64BKHN's picture
Update utils.py
0fbdd31 verified
raw
history blame
2.89 kB
import torch
import torchvision
import numpy as np
import pandas as pd
import os
from PIL import Image
import cv2
from facenet_pytorch import MTCNN, InceptionResnetV1
device = "cpu"
SINGLE_FACE_PATH = "data/test_images"
MULTI_FACES_PATH = "data/test_images_2"
THRESHOLD = 0.8
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,
keep_all=True
)
resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device)
def extract_face(box, img, margin=20):
img_size = (len(img), len(img[0]))
face_size = 160
margin = [
margin * (box[2] - box[0]) / (face_size - margin),
margin * (box[3] - box[1]) / (face_size - margin),
] #tạo margin bao quanh box cũ
box = [
int(max(box[0] - margin[0] / 2, 0)),
int(max(box[1] - margin[1] / 2, 0)),
int(min(box[2] + margin[0] / 2, img_size[0])),
int(min(box[3] + margin[1] / 2, img_size[1])),
]
img = img[box[1]:box[3], box[0]:box[2]]
if img.size != 0:
face = cv2.resize(img,(face_size, face_size), interpolation=cv2.INTER_AREA)
face = Image.fromarray(face)
return face
return False
def trans(img):
transform = torchvision.transforms.Compose([
torchvision.transforms.ToTensor()
])
return transform(img)
def find_and_compare(model, face):
face = trans(face)
face_embedded = model(face.unsqueeze(dim=0).to(device))
for usr in os.listdir(SINGLE_FACE_PATH):
usr_embedding = torch.load(SINGLE_FACE_PATH + "/" + usr + "/" + usr + ".pt").to(device)
if (face_embedded - usr_embedding).norm().item() < THRESHOLD:
print("Found " + usr + " in the image!")
return usr
return "unknown"
def recognition(img):
if len(img) < 1000 and len(img[0]) < 1000:
box_thickness = 3
text_thickness = 1
text_scale = 0.75
else:
box_thickness = 6
text_thickness = 2
text_scale = 2
img = np.copy(np.asarray(img))
names = []
boxes, _ = mtcnn.detect(img)
for box in boxes:
face = extract_face(box, img)
if face != False:
name = find_and_compare(model=resnet, face=face)
img = cv2.rectangle(img, (int(box[0]),int(box[1])), (int(box[2]),int(box[3])), (255,0,0), box_thickness)
img = cv2.putText(img, name, (int(box[0]),int(box[1])), cv2.FONT_HERSHEY_DUPLEX, text_scale, (0,255,0), text_thickness, cv2.LINE_8)
if name != "unknown" and not (name in names):
names.append(name)
print(names)
if len(names) > 0:
result = "Found "
for name in names:
result += name
if names[-1] != name:
result += ", "
else:
result = "Found no one!"
return result, Image.fromarray(img)