Update app.py

app.py

@@ -1,65 +1,120 @@
 import gradio as gr
-import mediapipe as mp
 import numpy as np
 import math
 import json
-from PIL import Image
-from mediapipe.tasks import python as mp_python
-from mediapipe.tasks.python import vision as mp_vision
 import urllib.request
 import os
+from PIL import Image
 
-# ── Download model if not cached ──
-MODEL_PATH = "/tmp/face_landmarker.task"
-MODEL_URL  = "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task"
+# ── Download face landmark TFLite model ──
+MODEL_PATH = "/tmp/face_landmark.tflite"
+DETECTOR_PATH = "/tmp/face_detection.tflite"
 
-if not os.path.exists(MODEL_PATH):
-    print("Downloading face landmarker model...")
-    urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
-    print("Model downloaded.")
+MODEL_URL    = "https://storage.googleapis.com/mediapipe-assets/face_landmark.tflite"
+DETECTOR_URL = "https://storage.googleapis.com/mediapipe-assets/face_detection_short_range.tflite"
 
-# ── Build FaceLandmarker ──
-base_options = mp_python.BaseOptions(
-    model_asset_path=MODEL_PATH,
-    delegate=mp_python.BaseOptions.Delegate.CPU
-)
-options = mp_vision.FaceLandmarkerOptions(
-    base_options=base_options,
-    output_face_blendshapes=False,
-    output_facial_transformation_matrixes=False,
-    num_faces=1,
-    min_face_detection_confidence=0.4,
-    min_face_presence_confidence=0.4
-)
-landmarker = mp_vision.FaceLandmarker.create_from_options(options)
-print("FaceLandmarker ready.")
+for path, url in [(MODEL_PATH, MODEL_URL), (DETECTOR_PATH, DETECTOR_URL)]:
+    if not os.path.exists(path):
+        print(f"Downloading {url}...")
+        urllib.request.urlretrieve(url, path)
+        print(f"Downloaded to {path}")
+
+# ── Load TFLite interpreter (no OpenGL needed) ──
+try:
+    import tflite_runtime.interpreter as tflite
+    print("Using tflite_runtime")
+except ImportError:
+    import tensorflow as tf
+    tflite = tf.lite
+    print("Using tensorflow.lite")
+
+detector_interp = tflite.Interpreter(model_path=DETECTOR_PATH)
+detector_interp.allocate_tensors()
+landmark_interp = tflite.Interpreter(model_path=MODEL_PATH)
+landmark_interp.allocate_tensors()
+
+det_in  = detector_interp.get_input_details()
+det_out = detector_interp.get_output_details()
+lm_in   = landmark_interp.get_input_details()
+lm_out  = landmark_interp.get_output_details()
+
+print(f"Detector input shape:  {det_in[0]['shape']}")
+print(f"Landmark input shape:  {lm_in[0]['shape']}")
+print("Models ready.")
+
+
+def preprocess(image_pil, target_size):
+    img = image_pil.convert("RGB").resize(target_size, Image.LANCZOS)
+    arr = np.array(img, dtype=np.float32) / 255.0
+    return arr[np.newaxis, ...]  # (1, H, W, 3)
 
 
 def analyse_face(image):
     try:
-        # Convert PIL to MediaPipe Image
-        img_rgb  = np.array(image.convert("RGB"))
-        h, w     = img_rgb.shape[:2]
-        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=img_rgb)
+        orig_w, orig_h = image.size
+
+        # ── Step 1: Detect face bounding box ──
+        det_size = (det_in[0]['shape'][2], det_in[0]['shape'][1])  # (W, H)
+        det_input = preprocess(image, det_size)
+        detector_interp.set_tensor(det_in[0]['index'], det_input)
+        detector_interp.invoke()
 
-        result = landmarker.detect(mp_image)
+        # Get detection output — boxes are [y_min, x_min, y_max, x_max] normalised
+        boxes   = detector_interp.get_tensor(det_out[0]['index'])[0]  # (N, 4)
+        scores  = detector_interp.get_tensor(det_out[1]['index'])[0]  # (N,)
 
-        if not result.face_landmarks or len(result.face_landmarks) == 0:
+        best_idx = int(np.argmax(scores))
+        if scores[best_idx] < 0.4:
             return json.dumps({"error": "No face detected. Please upload a clear, front-facing photo."})
 
-        lm = result.face_landmarks[0]   # list of NormalizedLandmark
+        y1, x1, y2, x2 = boxes[best_idx]
+        # Add 20% padding
+        pad = 0.20
+        bw  = x2 - x1
+        bh  = y2 - y1
+        x1  = max(0.0, x1 - bw * pad)
+        y1  = max(0.0, y1 - bh * pad)
+        x2  = min(1.0, x2 + bw * pad)
+        y2  = min(1.0, y2 + bh * pad)
+
+        crop = image.crop((int(x1 * orig_w), int(y1 * orig_h),
+                           int(x2 * orig_w), int(y2 * orig_h)))
+        crop_w, crop_h = crop.size
+
+        # ── Step 2: Run landmark model on cropped face ──
+        lm_size   = (lm_in[0]['shape'][2], lm_in[0]['shape'][1])
+        lm_input  = preprocess(crop, lm_size)
+        landmark_interp.set_tensor(lm_in[0]['index'], lm_input)
+        landmark_interp.invoke()
+
+        raw_lm = landmark_interp.get_tensor(lm_out[0]['index'])  # shape varies
+        raw_lm = raw_lm.reshape(-1, 3)  # (468, 3) — x, y, z all in [0, lm_size]
+
+        lm_w, lm_h = lm_size
+        # Normalise to [0,1] relative to original image
+        landmarks = []
+        for pt in raw_lm:
+            # pt in landmark input coords → map to crop → map to original
+            nx = (pt[0] / lm_w) * (x2 - x1) + x1
+            ny = (pt[1] / lm_h) * (y2 - y1) + y1
+            nz = float(pt[2]) / lm_w
+            landmarks.append({"x": float(nx), "y": float(ny), "z": nz})
+
+        # ── Score calculation (same geometry as before) ──
+        h, w = orig_h, orig_w
 
         def dist(a, b):
             return math.sqrt(
-                ((lm[a].x - lm[b].x) * w) ** 2 +
-                ((lm[a].y - lm[b].y) * h) ** 2
+                ((landmarks[a]['x'] - landmarks[b]['x']) * w) ** 2 +
+                ((landmarks[a]['y'] - landmarks[b]['y']) * h) ** 2
             )
 
+        lm = [type('L', (), {'x': p['x'], 'y': p['y'], 'z': p['z']})() for p in landmarks]
+
         face_ref = dist(33, 263)
-        if face_ref == 0:
-            return json.dumps({"error": "Could not measure face. Please try a different photo."})
+        if face_ref < 1:
+            return json.dumps({"error": "Face too small. Please use a closer photo."})
 
-        # ── Age estimation ──
         face_height  = dist(10, 152)
         lower_face_h = dist(168, 152)
         lower_ratio  = lower_face_h / face_height if face_height > 0 else 0.52
@@ -85,14 +140,10 @@ def analyse_face(image):
         age_mid   = max(18, min(72, round(age_raw)))
         age_range = f"{max(18, age_mid - 4)}\u2013{age_mid + 4}"
 
-        # ── Wrinkle score ──
-        wrinkle = round(max(1.0, min(9.9, 1 + texture * 18 + (age_mid - 18) * 0.10)), 1)
-
-        # ── Elasticity score ──
+        wrinkle    = round(max(1.0, min(9.9, 1 + texture * 18 + (age_mid - 18) * 0.10)), 1)
         cheek_sag  = ((dist(116, 61) + dist(345, 291)) / 2) / face_ref
         elasticity = round(max(1.0, min(9.9, 10 - (cheek_sag - 0.55) * 18 - (age_mid - 18) * 0.10)), 1)
 
-        # ── Jawline score ──
         jaw_pts = [234,93,132,58,172,136,150,149,176,148,152,377,400,378,379,365,397,288,361,323,454]
         jaw_dev = 0.0
         for i in range(1, len(jaw_pts) - 1):
@@ -108,8 +159,6 @@ def analyse_face(image):
         age_factor    = round(max(0.0, min(1.0, (age_mid - 18) / 54)), 3)
         years_younger = max(3, round(age_factor * 14 + 2))
 
-        landmarks_out = [{"x": float(l.x), "y": float(l.y), "z": float(l.z)} for l in lm]
-
         return json.dumps({
             "age_range":     age_range,
             "age_mid":       age_mid,
@@ -118,7 +167,7 @@ def analyse_face(image):
             "jawline":       jawline,
             "years_younger": years_younger,
             "age_factor":    age_factor,
-            "landmarks":     landmarks_out,
+            "landmarks":     landmarks,
             "image_width":   w,
             "image_height":  h
         })