Facial Recognition Before Prediction

app.py

@@ -1,16 +1,28 @@
 import gradio as gr
 from fastai.vision.all import *
 from transformers import AutoImageProcessor, AutoModelForImageClassification
-
 from PIL import Image
 import requests
+import face_recognition
 
 learn_inf = load_learner("export.pkl")
 processor = AutoImageProcessor.from_pretrained("dima806/facial_emotions_image_detection")
 model = AutoModelForImageClassification.from_pretrained("dima806/facial_emotions_image_detection")
 
+def extract_face(image)-> Image.Image:
+    # Detect face locations
+    face_locations = face_recognition.face_locations(image)
+
+    # If a face is detected, extract the first one
+    if face_locations:
+        top, right, bottom, left = face_locations[0]
+        face_image = Image.fromarray(image[top:bottom, left:right])
+        return face_image
+    else:
+        return image
+
 def predict(value) -> str:
-    image = Image.fromarray(value)
+    image = extract_face(Image.fromarray(value)).convert("L")
     inputs = processor(images=image, return_tensors="pt")
     outputs = model(**inputs)
     logits = outputs.logits

requirements.txt

@@ -1,2 +1,3 @@
 fastai
-transformers
+transformers
+face_recognition