Upload 2 files

app.py

@@ -0,0 +1,23 @@
+import os
+import gradio as gr
+from utils import recognition
+
+workers = 0 if os.name == 'nt' else 4
+device = "cpu"
+
+
+title = "Compare Faces"
+description = "A FaceNet (MTCNN for the bounding boxes and Inception Resnet V1 for the feature extracting, trainning with VGGFace2 dataset) feature extractor computer vision model to classify images of food as pizza, steak or sushi."
+article = "Created by HienK64BKHN."
+
+
+demo = gr.Interface(fn=recognition,
+                    inputs=gr.Image(label="Image"),
+                    outputs=[gr.Textbox(label="Names"),
+                             gr.Image(label="Result")],
+                    title=title,
+                    description=description,
+                    article=article
+                    )
+
+demo.launch(debug=False)

utils.py

@@ -0,0 +1,96 @@
+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 = "facenet-pytorch-master/data/test_images"
+MULTI_FACES_PATH = "facenet-pytorch-master/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)