Adjusted image processing

app.py

@@ -16,11 +16,25 @@ import subprocess
 import os
 
 
+from facenet_pytorch import MTCNN
 
-def resizeImage(image): 
-    resized = image.resize((112, 112))  
-    return resized
+mtcnn = MTCNN(keep_all=False)
+
+
+
+def crop_face_to_112x112(image: Image.Image):
+    if image.size == (112, 112):
+        return image
 
+    boxes, _ = mtcnn.detect(image)
+
+    if boxes is None:
+        raise ValueError("No face detected.")
+
+    x1, y1, x2, y2 = map(int, boxes[0])
+    cropped = image.crop((x1, y1, x2, y2))
+    resized = cropped.resize((112, 112), Image.BILINEAR)
+    return resized
 
 
 
@@ -154,7 +168,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             image_input_enroll = gr.Image(label="Upload a reference facial image.", type="pil", sources="upload")
-            image_input_enroll.change(fn=resizeImage, inputs=image_input_enroll, outputs=image_input_enroll)
+            image_input_enroll.change(fn=crop_face_to_112x112, inputs=image_input_enroll, outputs=image_input_enroll)
         with gr.Column():
             example_gallery = gr.Gallery(value=example_images, columns=3)
         with gr.Column():
@@ -231,7 +245,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             image_input_auth = gr.Image(label="Upload a facial image.", type="pil", sources="upload")
-            image_input_auth.change(fn=resizeImage, inputs=image_input_auth, outputs=image_input_auth)
+            image_input_auth.change(fn=crop_face_to_112x112, inputs=image_input_auth, outputs=image_input_auth)
         with gr.Column():
             example_gallery = gr.Gallery(value=example_images_auth, columns=3)
         with gr.Column():

requirements.txt

@@ -4,4 +4,5 @@ torch
 timm
 opencv-python
 pillow
-torchvision
+torchvision
+facenet-pytorch