Update app.py

app.py

@@ -8,6 +8,7 @@ import gradio as gr
 import numpy as np
 from datasets import load_dataset
 from main import DinoRegressionHeteroImages
+import mediapipe as mp
 
 # ============================
 # Model definition
@@ -84,15 +85,46 @@ imgtransform = transforms.Compose([
                          std=[0.229, 0.224, 0.225]),
 ])
 
+
+mp_face_detection = mp.solutions.face_detection
+mp_drawing = mp.solutions.drawing_utils
+
+def crop_face(img: Image.Image):
+    """Detects the largest face and crops it, fallback to full image."""
+    img_rgb = np.array(img.convert("RGB"))
+    h, w, _ = img_rgb.shape
+
+    with mp_face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5) as face_det:
+        results = face_det.process(img_rgb)
+
+        if results.detections:
+            # take the first (largest) face
+            bbox = results.detections[0].location_data.relative_bounding_box
+            x1 = max(int(bbox.xmin * w), 0)
+            y1 = max(int(bbox.ymin * h), 0)
+            x2 = min(int((bbox.xmin + bbox.width) * w), w)
+            y2 = min(int((bbox.ymin + bbox.height) * h), h)
+
+            face = img.crop((x1, y1, x2, y2))
+            return face
+
+    # fallback: return original if no face found
+    return img
+
+
 # ============================
 # Prediction function
 # ============================
 def predict(img):
+    # Crop face
+    face = crop_face(img)
+    
     x = imgtransform(img).unsqueeze(0).to(device)
     with torch.no_grad():
         mu, logvar = model(x)
     pred_years = mu.item() * lifespan_std + lifespan_mean
-    return f"Estimated remaining lifespan: {pred_years:.1f} years"
+    
+    return f"<div style='font-size:32px; font-weight:bold; color:#2E86AB;'>Estimated remaining lifespan: {pred_years:.1f} years</div>"
 
 # ============================
 # Gradio UI (Webcam + Upload)
@@ -105,7 +137,7 @@ img.save("example.jpg")
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil", sources=["upload", "webcam"], label="Face Photo"),
-    outputs=gr.Textbox(label="Predicted Lifespan"),
+    outputs=gr.HTML(label="Prediction"),
     live=False,
     title="Lifespan Predictor (Demo)",
     description="Upload or take a face photo. Model predicts remaining lifespan (years).",