Update app.py

app.py

@@ -1,189 +1,274 @@
 import gradio as gr
-from PIL import Image
-import requests
-import io
-import os
-
-# --- Configuration ---
-# If using Hugging Face Inference API for face detection
-HF_API_TOKEN = os.environ.get("HF_API_TOKEN") # Set this in your Space secrets
-DETECTION_MODEL_URL = "https://api-inference.huggingface.co/models/facebook/detr-resnet-50" # Example
-
-# --- Placeholder for Face Shape Classification ---
-# This is the most challenging part to do with a generic free model.
-# Option 1: Find a face shape classification model on HF Hub (ideal but rare)
-# Option 2: Implement landmark detection + rule-based shape classification (complex)
-# Option 3: Use a very simple heuristic (e.g., aspect ratio of bounding box - very unreliable)
-# Option 4 (Simplest for demo): Let user *select* their face shape or a mock function
-
-def detect_face_api(image_pil):
-    if not HF_API_TOKEN:
-        raise gr.Error("Hugging Face API Token not set in Space secrets!")
-    
-    # Convert PIL Image to bytes
-    img_byte_arr = io.BytesIO()
-    image_pil.save(img_byte_arr, format='PNG')
-    img_byte_arr = img_byte_arr.getvalue()
-
-    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
-    try:
-        response = requests.post(DETECTION_MODEL_URL, headers=headers, data=img_byte_arr)
-        response.raise_for_status()
-        outputs = response.json()
-    except requests.exceptions.RequestException as e:
-        print(f"API Error: {e}")
-        print(f"Response content: {response.content if 'response' in locals() else 'No response'}")
-        return None, "Error calling detection API."
-    except Exception as e:
-        print(f"Other Error during detection: {e}")
-        return None, f"Error processing image: {str(e)}"
-
-
-    # Process outputs to find 'person' or 'face' bounding box
-    # This depends heavily on the model used.
-    # For detr-resnet-50, it returns a list of objects with 'label', 'score', 'box'
-    best_face_box = None
+from PIL import Image, ImageDraw
+import numpy as np
+import torch # Required by transformers
+from transformers import YolosImageProcessor, YolosForObjectDetection
+import mediapipe as mp
+import math # For distance calculations
+
+# --- Model Initialization (Load ONCE when the Space starts) ---
+# 1. Face Detection Model (YOLOS)
+print("Loading face detection model...")
+DETECTION_MODEL_NAME = "hustvl/yolos-tiny" # Smaller model, better for CPU
+try:
+    face_image_processor = YolosImageProcessor.from_pretrained(DETECTION_MODEL_NAME)
+    face_detection_model = YolosForObjectDetection.from_pretrained(DETECTION_MODEL_NAME)
+    print("Face detection model loaded successfully.")
+except Exception as e:
+    print(f"Error loading face detection model: {e}")
+    face_image_processor = None
+    face_detection_model = None
+
+
+# 2. Facial Landmark Model (MediaPipe Face Mesh)
+print("Initializing MediaPipe Face Mesh...")
+try:
+    mp_face_mesh = mp.solutions.face_mesh
+    # static_image_mode=True for processing individual images
+    # max_num_faces=1 as we expect one primary face
+    # min_detection_confidence for robustness
+    face_mesh_detector = mp_face_mesh.FaceMesh(
+        static_image_mode=True,
+        max_num_faces=1,
+        refine_landmarks=True, # Get more detailed landmarks (e.g., iris)
+        min_detection_confidence=0.5)
+    print("MediaPipe Face Mesh initialized successfully.")
+except Exception as e:
+    print(f"Error initializing MediaPipe Face Mesh: {e}")
+    face_mesh_detector = None
+
+# --- Helper Functions ---
+
+def detect_face_local(image_pil):
+    if not face_image_processor or not face_detection_model:
+        return None, "Face detection model not loaded."
+
+    inputs = face_image_processor(images=image_pil, return_tensors="pt")
+    outputs = face_detection_model(**inputs)
+
+    # Post-process to get bounding boxes
+    # target_sizes expects a tensor of [height, width]
+    target_sizes = torch.tensor([image_pil.size[::-1]]) # PIL size is (width, height)
+    results = face_image_processor.post_process_object_detection(outputs, threshold=0.7, target_sizes=target_sizes)[0]
+
+    best_box = None
     max_score = 0
+    person_label_id = None # YOLOS typically detects 'person' (label 0 in COCO usually)
+
+    # Find the 'person' class ID if model config is available (or assume it if known)
+    # For general YOLOS, 'person' is often label 0 if trained on COCO.
+    # If your model has specific face labels, adjust this.
+    # For hustvl/yolos-tiny, it's trained on COCO, where "person" is label 0.
+    # Check model.config.id2label if needed
+    # person_label_id = face_detection_model.config.label2id.get("person", 0) # More robust way
+
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        # Assuming 'person' class is the one we want, or if it detects 'face' directly
+        # For YOLOS, it's more likely to detect 'person'. We take the highest score 'person'.
+        # You might need to adjust this if the model has a specific 'face' label
+        if label == 0: # Assuming label 0 is 'person' for COCO-trained YOLOS
+            if score > max_score:
+                max_score = score
+                best_box = box.tolist() # [xmin, ymin, xmax, ymax]
+
+    if best_box:
+        cropped_image = image_pil.crop(best_box)
+        return cropped_image, None # No error message
+    else:
+        return None, "No face/person detected with sufficient confidence."
+
+def get_landmarks_mediapipe(image_pil):
+    if not face_mesh_detector:
+        return None, "MediaPipe Face Mesh not initialized."
+
+    # MediaPipe expects BGR numpy array
+    image_np = np.array(image_pil.convert('RGB'))
+    image_rgb = image_np[:, :, ::-1].copy() # PIL RGB to CV2 BGR (not strictly needed here as MP handles RGB)
+                                          # but good practice if using OpenCV functions later
+    image_rgb_mp = np.array(image_pil.convert('RGB')) # MediaPipe prefers RGB
+
+    results = face_mesh_detector.process(image_rgb_mp)
+
+    if results.multi_face_landmarks:
+        return results.multi_face_landmarks[0], None # Return landmarks for the first face
+    else:
+        return None, "Could not detect facial landmarks."
+
+def _distance(p1, p2):
+    return math.sqrt((p1.x - p2.x)**2 + (p1.y - p2.y)**2 + (p1.z - p2.z)**2)
+
+def _distance_2d(p1, p2, img_width, img_height):
+    # Convert normalized coordinates to pixel coordinates for more intuitive ratios
+    x1, y1 = p1.x * img_width, p1.y * img_height
+    x2, y2 = p2.x * img_width, p2.y * img_height
+    return math.sqrt((x1 - x2)**2 + (y1 - y2)**2)
+
+
+def estimate_face_shape_from_landmarks(landmarks, img_width, img_height):
+    if not landmarks:
+        return "Unknown"
+
+    # Key landmark indices for MediaPipe Face Mesh (468 landmarks total)
+    # These are approximate and might need fine-tuning or using specific contour points
+    # Forehead: e.g., landmark 10 (top of forehead)
+    # Jaw: e.g., landmarks 172, 397 (jaw points), 152 (chin)
+    # Cheekbones: e.g., landmarks 234, 454 (outer cheekbones) or 116, 345 (zygomatic arch)
+    # Face Width: Widest points, often around cheekbones (e.g. 234 to 454)
+    # Face Height: Top of forehead (e.g. 10) to chin (e.g. 152)
+
+    # Example: Use specific points from standard MediaPipe landmark map
+    # (https://github.com/google/mediapipe/blob/master/mediapipe/modules/face_geometry/data/canonical_face_model_uv_visualization.png)
+
+    # Points for measurements (these are just examples, adjust as needed)
+    # These indices are 0-based from the 468 landmarks
+    # Check https://viz.mediapipe.dev/face_ όχιmesh_webgl_demo for interactive map
     
-    # Try to find "face", then "person"
-    for obj_type in ["face", "person"]: # Some models might label "face" directly
-        for obj in outputs:
-            if obj['label'].lower() == obj_type and obj['score'] > max_score:
-                max_score = obj['score']
-                best_face_box = obj['box']
+    # Face Height: Top of Forehead (10) to Chin (152)
+    p_forehead_top = landmarks.landmark[10]
+    p_chin = landmarks.landmark[152]
+    face_height = _distance_2d(p_forehead_top, p_chin, img_width, img_height)
+
+    # Face Width (approx at cheekbones): Left (234) to Right (454)
+    p_cheek_left = landmarks.landmark[234]
+    p_cheek_right = landmarks.landmark[454]
+    face_width_cheek = _distance_2d(p_cheek_left, p_cheek_right, img_width, img_height)
     
-    if best_face_box:
-        # Crop the image to the detected face
-        # box is {xmin, ymin, xmax, ymax}
-        cropped_image = image_pil.crop((best_face_box['xmin'], best_face_box['ymin'], best_face_box['xmax'], best_face_box['ymax']))
-        return cropped_image, None
-    else:
-        return None, "No face detected with sufficient confidence."
-
-def estimate_face_shape_mock(face_image_pil):
-    """
-    MOCK FUNCTION: In a real app, this would use a model or complex heuristics.
-    For now, let's pretend it analyzes and returns a shape.
-    You could also add a dropdown for the user to select their perceived face shape
-    if AI classification is too hard to implement initially.
-    """
-    # Super simple heuristic (aspect ratio) - VERY UNRELIABLE, just for demo
-    width, height = face_image_pil.size
-    aspect_ratio = height / width
-    if aspect_ratio > 1.25: return "Long"
-    elif 0.9 < aspect_ratio < 1.1: return "Round/Square" # Needs more to differentiate
-    else: return "Oval"
-    # In reality, this needs a dedicated model or landmark-based analysis.
-    # return "Oval" # Or just return a default for now
-
-def get_hairstyle_suggestions(face_shape, gender="neutral"): # Gender could be another input
+    # Forehead Width (approx temples): Left (70) to Right (300) - might be too wide, adjust
+    # Or use points like 54 and 284 for a narrower forehead measure
+    p_forehead_left = landmarks.landmark[54] # More like outer brow
+    p_forehead_right = landmarks.landmark[284] # More like outer brow
+    forehead_width = _distance_2d(p_forehead_left, p_forehead_right, img_width, img_height)
+
+    # Jawline Width (approx): Point near jaw angle left (172) to right (397)
+    # Or closer to chin base: 143 and 372
+    p_jaw_left = landmarks.landmark[132] # Lower jaw points
+    p_jaw_right = landmarks.landmark[361]
+    jaw_width = _distance_2d(p_jaw_left, p_jaw_right, img_width, img_height)
+
+    # Simple Heuristics (these are very basic and need refinement/testing)
+    # Ratios are more reliable than absolute values due to image scale
+    
+    if face_height == 0 or face_width_cheek == 0: return "Unknown (measurement error)"
+
+    ratio_h_w = face_height / face_width_cheek
+
+    # Print measurements for debugging
+    print(f"H: {face_height:.2f}, W_Cheek: {face_width_cheek:.2f}, W_Forehead: {forehead_width:.2f}, W_Jaw: {jaw_width:.2f}")
+    print(f"Ratio H/W: {ratio_h_w:.2f}")
+    print(f"Forehead/Cheek: {forehead_width/face_width_cheek if face_width_cheek else 0:.2f}")
+    print(f"Jaw/Cheek: {jaw_width/face_width_cheek if face_width_cheek else 0:.2f}")
+
+    # These rules are very basic and a starting point
+    if ratio_h_w > 1.25: # Significantly longer than wide
+        if forehead_width > jaw_width and jaw_width < face_width_cheek * 0.85 :
+            return "Heart/Inverted Triangle" # Narrow chin
+        return "Long/Oblong"
+    elif ratio_h_w < 0.95: # Wider than tall or close to it
+        return "Round/Square (Wide)" # Need more to differentiate round vs square (jaw angle)
+    else: # Height and width are somewhat proportional (0.95 to 1.25)
+        # Check relative widths of forehead, cheeks, jaw
+        f_w = forehead_width
+        c_w = face_width_cheek
+        j_w = jaw_width
+
+        if abs(f_w - c_w) < c_w * 0.1 and abs(c_w - j_w) < c_w * 0.1 and abs(f_w - j_w) < f_w * 0.1:
+            # All widths are roughly similar
+            return "Square" # More angular jaw typically
+        elif c_w > f_w and c_w > j_w:
+            return "Diamond" # Widest at cheeks
+        elif f_w > c_w * 0.9 and f_w > j_w and j_w < c_w * 0.9: # Forehead prominent, jaw narrower
+            return "Heart"
+        elif f_w < c_w and j_w < c_w and abs(f_w - j_w) < f_w * 0.15 : # Forehead and jaw narrower than cheeks, but similar to each other
+             return "Oval" # Often considered ideal, balanced
+        else: # Default or fallback
+            return "Oval/Round" # Difficult to distinguish without more rules
+
+    return "Oval (Default)" # Fallback
+
+
+def get_hairstyle_suggestions(face_shape, gender="neutral"):
+    # (Same suggestion dictionary as before - keep it for brevity)
     suggestions = {
-        "Oval": {
-            "hair": [
-                "Most hairstyles work well. Lucky you!",
-                "Consider layers for volume or a sleek bob.",
-                "Side-swept bangs can be flattering."
-            ],
-            "beard": [
-                "Most beard styles are suitable.",
-                "A classic full beard, well-groomed, is excellent.",
-                "Consider a short boxed beard or a Van Dyke."
-            ]
-        },
-        "Round/Square": { # Simplified for demo, ideally separate
-            "hair": [
-                "Add height on top: pompadour, quiff, faux hawk.",
-                "Avoid blunt bobs ending at the chin or very short, round cuts.",
-                "Layers, textured cuts, and off-center parts can soften features.",
-                "For square faces: Softer styles, waves, or curls can balance a strong jaw."
-            ],
-            "beard": [
-                "For round faces: Beards that add length to the chin: goatee, soul patch, or a beard that's shorter on the sides and longer at the chin.",
-                "For square faces: Styles that soften the jawline. A circle beard or a well-trimmed full beard that rounds the chin."
-            ]
-        },
-        "Long": {
-            "hair": [
-                "Add width: Curls, waves, or layered styles with volume at the sides.",
-                "Avoid excessive height on top.",
-                "Bangs (blunt or side-swept) can shorten the face."
-            ],
-            "beard": [
-                "Styles that add width to the face: fuller on the cheeks, like a full beard or mutton chops.",
-                "Avoid long, pointy beards that further elongate the face."
-            ]
-        },
-        "Heart": {
-             "hair": [
-                "Add volume at the jawline: chin-length bobs, layered shoulder-length cuts.",
-                "Side-swept bangs or a textured fringe can balance a wider forehead.",
-                "Avoid too much height on top."
-            ],
-            "beard": [
-                "Fuller beards that add width to the jawline, like a full beard or a Garibaldi.",
-                "Avoid styles that are too narrow at the chin."
-            ]
-        },
-        # Add Diamond, etc.
+        "Oval": {"hair": ["Most hairstyles work well.", "Layers or sleek bob."], "beard": ["Most beard styles.", "Classic full beard."]},
+        "Oval (Default)": {"hair": ["Try versatile styles like layers or a textured crop.", "Side parts can be flattering."], "beard": ["A well-groomed stubble or a short boxed beard often works."]},
+        "Long/Oblong": {"hair": ["Add width: Curls, waves, layers with side volume.", "Bangs can shorten face."], "beard": ["Fuller on cheeks: full beard, mutton chops."]},
+        "Heart": {"hair": ["Add jawline volume: chin-length bobs, layered shoulder cuts.", "Side-swept bangs."], "beard": ["Fuller beards to add jaw width: Garibaldi."]},
+        "Heart/Inverted Triangle": {"hair": ["Add jawline volume: chin-length bobs, layered shoulder cuts.", "Side-swept bangs for forehead."], "beard": ["Fuller beards to add jaw width: Garibaldi, full beard carefully shaped."]},
+        "Square": {"hair": ["Softer styles, waves, curls. Texture to soften angles.", "Avoid sharp, geometric cuts."], "beard": ["Circle beard, rounded full beard."]},
+        "Round/Square (Wide)": {"hair": ["Add height: pompadour, quiff. Layers, off-center parts.", "Avoid blunt bobs at chin."], "beard": ["For round: goatee, beard longer at chin. For square: soften jaw with rounded styles."]},
+        "Diamond": {"hair": ["Soften forehead & jaw: chin bobs, shoulder length with layers.", "Side-swept fringe."], "beard": ["Fuller at chin, possibly some width at jaw but not cheeks: Balbo, shorter full beard."]},
+        "Oval/Round": {"hair": ["Versatile. Add slight height or soft layers.", "Avoid overly round styles if aiming to balance roundness."], "beard": ["Many styles work. A neatly trimmed beard or a Van Dyke can be good."]},
+        "Unknown": {"hair": ["Upload a clearer image for analysis."], "beard": ["Upload a clearer image for analysis."]},
+        "Unknown (measurement error)": {"hair": ["Could not reliably measure face. Try a different pose or lighting."], "beard": ["Could not reliably measure face. Try a different pose or lighting."]},
     }
     if face_shape in suggestions:
         hair_sug = "\n".join([f"- {s}" for s in suggestions[face_shape]["hair"]])
         beard_sug = "\n".join([f"- {s}" for s in suggestions[face_shape]["beard"]])
         return f"**Haircut Suggestions for {face_shape} Face:**\n{hair_sug}\n\n**Beard Style Suggestions for {face_shape} Face:**\n{beard_sug}"
-    return "Could not determine suggestions for the estimated face shape."
+    return f"Could not determine suggestions for the estimated face shape: {face_shape}."
+
 
-def analyze_face_and_suggest(front_image, side_image_optional):
-    if front_image is None:
+def analyze_face_and_suggest(front_image_input, side_image_optional):
+    if front_image_input is None:
         return None, "Please upload a front-facing photo.", ""
 
-    # Convert Gradio Image (numpy array) to PIL Image
-    img_pil = Image.fromarray(front_image)
+    # Gradio Image input is a NumPy array
+    img_pil = Image.fromarray(front_image_input).convert("RGB") # Ensure RGB
 
-    # 1. Detect Face (using API or local model)
-    cropped_face_pil, error_msg = detect_face_api(img_pil) # Using API
+    # 1. Detect Face (Local YOLOS model)
+    cropped_face_pil, error_msg = detect_face_local(img_pil)
     if error_msg:
         return None, error_msg, ""
-    if cropped_face_pil is None: # Should be caught by error_msg but as a fallback
+    if cropped_face_pil is None:
         return None, "Could not detect a face.", ""
 
-    # --- Placeholder for processing side_image_optional ---
-    # If side_image_optional is provided, you could:
-    # - Run face detection on it too.
-    # - If using landmarks, try to get landmarks from both.
-    # - Use it to refine the face shape estimation (e.g., confirm jawline).
-    # For now, we'll just acknowledge it.
-    side_info = "Side profile not uploaded."
+    # 2. Get Facial Landmarks (MediaPipe)
+    landmarks, error_msg_lm = get_landmarks_mediapipe(cropped_face_pil)
+    if error_msg_lm:
+        # If landmarks fail, still show cropped face but indicate no shape analysis
+        return cropped_face_pil, f"Face detected. Error getting landmarks: {error_msg_lm}", "Cannot suggest hairstyles without landmark analysis."
+
+    # For drawing landmarks (optional visualization)
+    # cropped_face_with_landmarks_pil = cropped_face_pil.copy()
+    # draw = ImageDraw.Draw(cropped_face_with_landmarks_pil)
+    # for landmark in landmarks.landmark:
+    #     x = int(landmark.x * cropped_face_pil.width)
+    #     y = int(landmark.y * cropped_face_pil.height)
+    #     draw.ellipse((x-1, y-1, x+1, y+1), fill='red')
+
+
+    # 3. Estimate Face Shape from Landmarks
+    img_w, img_h = cropped_face_pil.size
+    estimated_shape = estimate_face_shape_from_landmarks(landmarks, img_w, img_h)
+
+    # --- Side profile (acknowledgement, no processing yet) ---
+    side_info = "Side profile not uploaded or not yet processed."
     if side_image_optional is not None:
-        side_img_pil = Image.fromarray(side_image_optional)
-        # You could try to detect face/landmarks on side_img_pil here
-        side_info = "Side profile uploaded (analysis can be enhanced in future versions)."
-        # For a more advanced system, you'd combine info from front and side.
-
-    # 2. Estimate Face Shape (the hard part)
-    # For demo, using a mock function. Replace with actual logic.
-    estimated_shape = estimate_face_shape_mock(cropped_face_pil)
-    
-    # 3. Get Suggestions
+        side_info = "Side profile uploaded (analysis can be enhanced in future versions to refine shape)."
+        # Potentially run detection + landmarks on side_image_optional here too
+        # And combine information for a more robust `estimated_shape`
+
+    # 4. Get Suggestions
     suggestions_text = get_hairstyle_suggestions(estimated_shape)
 
     return cropped_face_pil, f"Estimated Face Shape: **{estimated_shape}**\n{side_info}", suggestions_text
 
 # --- Gradio Interface ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# ✂️ AI Hairstyle & Beard Suggester 🧔")
+    gr.Markdown("# ✂️ AI Hairstyle & Beard Suggester (Local Models) 🧔")
     gr.Markdown(
-        "Upload a clear, front-facing photo of a person. "
-        "Optionally, upload a side profile for potentially better results (future enhancement)."
-        "\n*Disclaimer: This is a demo using simplified logic for face shape estimation. Suggestions are general.*"
+        "Upload a clear, front-facing photo. "
+        "Optionally, upload a side profile (currently not used for analysis but can be added)."
+        "\n*Disclaimer: This app uses local AI models for face detection and landmark-based shape estimation. Suggestions are based on general heuristics and may not be perfect.*"
     )
 
     with gr.Row():
         with gr.Column(scale=1):
-            front_image_input = gr.Image(type="numpy", label="Front Face Photo (Required)")
-            side_image_input = gr.Image(type="numpy", label="Side Profile Photo (Optional)")
+            front_image_input = gr.Image(type="numpy", label="Front Face Photo (Required)", sources=["upload", "webcam"])
+            side_image_input = gr.Image(type="numpy", label="Side Profile Photo (Optional)", sources=["upload", "webcam"])
             submit_btn = gr.Button("Get Suggestions", variant="primary")
         with gr.Column(scale=2):
-            output_image = gr.Image(label="Detected Face")
+            output_image = gr.Image(label="Detected Face (or Cropped with Landmarks)")
             output_shape_info = gr.Markdown(label="Face Analysis")
             output_suggestions = gr.Markdown(label="Suggestions")
 
@@ -192,18 +277,11 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         inputs=[front_image_input, side_image_input],
         outputs=[output_image, output_shape_info, output_suggestions]
     )
+    gr.Markdown("--- \n ### Note on Face Shape Estimation: \n The face shape estimation is based on ratios of distances between facial landmarks. The categories (Oval, Round, Square, etc.) and the rules to classify them are simplified. For more accurate results, a dedicated face shape classification model or more complex geometric analysis would be needed. The landmark points used are: \n - **Height:** Top of Forehead (MP Landmark 10) to Chin (MP 152) \n - **Cheek Width:** Left Cheek (MP 234) to Right Cheek (MP 454) \n - **Forehead Width:** Left Outer Brow (MP 54) to Right Outer Brow (MP 284) \n - **Jaw Width:** Left Lower Jaw (MP 132) to Right Lower Jaw (MP 361)")
 
-    gr.Examples(
-        examples=[
-            # Add paths to example images if you have them in your Space
-            # ["path/to/example_front.jpg", "path/to/example_side.jpg"],
-            # ["another_front.png", None],
-        ],
-        inputs=[front_image_input, side_image_input],
-        outputs=[output_image, output_shape_info, output_suggestions],
-        fn=analyze_face_and_suggest,
-        cache_examples=False # or True if your function is deterministic and inputs are fixed
-    )
 
 if __name__ == "__main__":
-    demo.launch()
+    if face_detection_model and face_mesh_detector: # Only launch if models loaded
+        demo.launch()
+    else:
+        print("Gradio app not launched due to model loading errors.")