Update app.py

app.py

@@ -1,265 +1,367 @@
 import gradio as gr
 from PIL import Image, ImageDraw
 import numpy as np
-import torch # Required by transformers
-from transformers import YolosImageProcessor, YolosForObjectDetection
+import torch
+from transformers import YolosImageProcessor, YolosForObjectDetection # Keep for YOLOS architecture
 import mediapipe as mp
-import math # For distance calculations
+import math
 
-# --- Model Initialization (Load ONCE when the Space starts) ---
-# 1. Face Detection Model (YOLOS)
+# --- Model Initialization ---
+# 1. Face Detection Model (Switching to yolos-face if possible)
 print("Loading face detection model...")
-DETECTION_MODEL_NAME = "hustvl/yolos-tiny" # Smaller model, better for CPU
+# Try hustvl/yolos-face (specific for faces)
+# If it fails, we can revert to hustvl/yolos-tiny (person detection)
+DETECTION_MODEL_NAME = "hustvl/yolos-face"
+# DETECTION_MODEL_NAME_FALLBACK = "hustvl/yolos-tiny" # If yolos-face fails
+
 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.")
+    # For yolos-face, the label for "face" is often 0. Check model.config.id2label if unsure.
+    # It seems like hustvl/yolos-face is also trained on a single class "face" (ID 0)
+    FACE_LABEL_ID = 0 # Assuming face is label 0 for this model
+    print(f"Face detection model ({DETECTION_MODEL_NAME}) loaded successfully.")
 except Exception as e:
-    print(f"Error loading face detection model: {e}")
+    print(f"Error loading {DETECTION_MODEL_NAME}: {e}. ")
+    # Fallback or error handling needed if you want to deploy with a working detector
+    # For now, we'll let it fail if the primary model doesn't load to highlight the issue.
     face_image_processor = None
     face_detection_model = None
+    FACE_LABEL_ID = -1 # Indicate model not loaded
 
 
 # 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)
+        refine_landmarks=True,
         min_detection_confidence=0.5)
+    mp_drawing = mp.solutions.drawing_utils # For drawing landmarks
+    drawing_spec = mp_drawing.DrawingSpec(thickness=1, circle_radius=1, color=(0,255,0)) # Green dots
     print("MediaPipe Face Mesh initialized successfully.")
 except Exception as e:
     print(f"Error initializing MediaPipe Face Mesh: {e}")
     face_mesh_detector = None
+    mp_drawing = None
 
-# --- Helper Functions ---
 
+# --- 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."
+    if not face_image_processor or not face_detection_model or FACE_LABEL_ID == -1:
+        return None, "Face detection model not loaded or configured properly."
 
     inputs = face_image_processor(images=image_pil, return_tensors="pt")
-    outputs = face_detection_model(**inputs)
+    with torch.no_grad(): # Important for inference
+        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]
+    target_sizes = torch.tensor([image_pil.size[::-1]])
+    results = face_image_processor.post_process_object_detection(outputs, threshold=0.8, target_sizes=target_sizes)[0] # Increased threshold
 
     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 label == FACE_LABEL_ID: # Check for the specific 'face' label
             if score > max_score:
                 max_score = score
-                best_box = box.tolist() # [xmin, ymin, xmax, ymax]
+                best_box = box.tolist()
 
     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."
+        # Add a small padding to the bounding box to ensure the whole face is included
+        padding_w = (best_box[2] - best_box[0]) * 0.1 # 10% padding width
+        padding_h = (best_box[3] - best_box[1]) * 0.1 # 10% padding height
 
-def get_landmarks_mediapipe(image_pil):
-    if not face_mesh_detector:
-        return None, "MediaPipe Face Mesh not initialized."
+        xmin = max(0, best_box[0] - padding_w)
+        ymin = max(0, best_box[1] - padding_h)
+        xmax = min(image_pil.width, best_box[2] + padding_w)
+        ymax = min(image_pil.height, best_box[3] + padding_h)
 
-    # 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
+        cropped_image = image_pil.crop((xmin, ymin, xmax, ymax))
+        return cropped_image, None
+    else:
+        return None, "No face detected with sufficient confidence."
+
+def get_landmarks_and_draw(image_pil):
+    if not face_mesh_detector or not mp_drawing:
+        return None, "MediaPipe Face Mesh not initialized for landmarks.", image_pil # Return original if no detector
 
+    image_rgb_mp = np.array(image_pil.convert('RGB'))
     results = face_mesh_detector.process(image_rgb_mp)
 
+    annotated_image_pil = image_pil.copy() # Start with a copy of the input PIL image
+
     if results.multi_face_landmarks:
-        return results.multi_face_landmarks[0], None # Return landmarks for the first face
+        landmarks = results.multi_face_landmarks[0] # First face
+        
+        # Draw landmarks on the PIL image
+        # Convert PIL to NumPy array for drawing, then back to PIL
+        image_np_to_draw = np.array(annotated_image_pil)
+        for landmark in landmarks.landmark:
+            x = int(landmark.x * image_np_to_draw.shape[1])
+            y = int(landmark.y * image_np_to_draw.shape[0])
+            # Small green circle for each landmark (using PIL Draw directly for simplicity here)
+            # For more complex drawing, use mp_drawing.draw_landmarks on the numpy array
+        
+        # Using mp_drawing.draw_landmarks for better visualization
+        mp_drawing.draw_landmarks(
+            image=image_np_to_draw,
+            landmark_list=landmarks,
+            connections=mp_face_mesh.FACEMESH_TESSELATION, # Shows mesh
+            # connections=mp_face_mesh.FACEMESH_CONTOURS, # Shows contours
+            landmark_drawing_spec=drawing_spec,
+            connection_drawing_spec=drawing_spec)
+        
+        annotated_image_pil = Image.fromarray(image_np_to_draw)
+        return landmarks, None, annotated_image_pil
     else:
-        return None, "Could not detect facial landmarks."
+        return None, "Could not detect facial landmarks.", annotated_image_pil
 
-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 _distance_2d_normalized(p1, p2): # Operates on normalized 0-1 coordinates
+    return math.sqrt((p1.x - p2.x)**2 + (p1.y - p2.y)**2)
 
-
-def estimate_face_shape_from_landmarks(landmarks, img_width, img_height):
+def estimate_face_shape_from_landmarks_v2(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
-    
-    # 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)
+        return "Unknown", {}
+
+    # --- Key Anthropometric Ratios for Face Shape ---
+    # Based on MediaPipe landmark indices (https://github.com/google/mediapipe/blob/master/mediapipe/python/solutions/face_mesh_connections.py)
+    # These are illustrative; precise points can vary based on definitions.
+
+    # Face Height: Approx. Trichion (hairline top - 10) to Gnathion (chin bottom - 152)
+    # Since landmark 10 can be on hair, let's use a point slightly lower or average of forehead top points
+    # Top of forehead (10), Midpoint between brows (e.g., 168 or 9)
+    # Chin point (152)
+    p_forehead_top_center = landmarks.landmark[10]
+    p_chin_bottom = landmarks.landmark[152]
+    face_height = abs(p_forehead_top_center.y - p_chin_bottom.y) # Using normalized Y diff
+
+    # Bizygomatic Width (Cheekbone to Cheekbone - widest part of face):
+    # Left Zygion (approx. 234 or 130/133) to Right Zygion (approx. 454 or 359/362)
+    # Let's use standard contour points for face oval: e.g., Left: 234, 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)
+    face_width_cheeks = abs(p_cheek_left.x - p_cheek_right.x) # Normalized X diff
+
+    # Forehead Width (between temporal crests):
+    # E.g., Left: 70, Right: 300 or more specific points like 54 (L) and 284 (R)
+    p_forehead_L = landmarks.landmark[70] # Or 103, 67
+    p_forehead_R = landmarks.landmark[300] # Or 332, 297
+    forehead_width = abs(p_forehead_L.x - p_forehead_R.x)
+
+    # Bigonial Width (Jaw Angle to Jaw Angle):
+    # E.g., Left Gonion (approx. 172 or 137) to Right Gonion (approx. 397 or 366)
+    p_jaw_angle_L = landmarks.landmark[172] # Or 147
+    p_jaw_angle_R = landmarks.landmark[397] # Or 376
+    jaw_width_gonial = abs(p_jaw_angle_L.x - p_jaw_angle_R.x)
+
+    # Chin Width (just above Gnathion, width of the chin prominence)
+    # E.g. landmarks 176 and 400, or slightly higher like 135 and 364
+    p_chin_width_L = landmarks.landmark[143] # Points on the chin body
+    p_chin_width_R = landmarks.landmark[372]
+    chin_width = abs(p_chin_width_L.x - p_chin_width_R.x)
+
+
+    measurements = {
+        "face_height_norm": face_height,
+        "face_width_cheeks_norm": face_width_cheeks,
+        "forehead_width_norm": forehead_width,
+        "jaw_width_gonial_norm": jaw_width_gonial,
+        "chin_width_norm": chin_width
+    }
+    print("Normalized Measurements:", measurements)
+
+    # --- Classification Logic (Needs significant refinement) ---
+    # Ratios are key to normalize for face size in image
+    if face_width_cheeks == 0: return "Unknown (div zero)", measurements
+
+    # Facial Index: Height / Width
+    facial_index = face_height / face_width_cheeks if face_width_cheeks > 0 else 0
+
+    # Relative Widths (compared to cheekbone width as it's often the widest)
+    forehead_to_cheek_ratio = forehead_width / face_width_cheeks
+    jaw_to_cheek_ratio = jaw_width_gonial / face_width_cheeks
+    jaw_to_forehead_ratio = jaw_width_gonial / forehead_width if forehead_width > 0 else 0
+
+    shape = "Unknown"
+
+    # These rules are more like guidelines and need extensive testing and tuning.
+    # Consider a decision tree or a more structured approach.
+    if facial_index > 1.05: # Longer than wide
+        if forehead_to_cheek_ratio > 0.9 and jaw_to_cheek_ratio > 0.9: # All widths similar
+            shape = "Long/Oblong"
+        elif forehead_to_cheek_ratio > jaw_to_cheek_ratio and chin_width < jaw_width_gonial * 0.7:
+            shape = "Heart/Inverted Triangle" # Wider forehead, narrow chin
+        else:
+            shape = "Long"
+    elif facial_index < 0.95: # Wider than long, or close to equal
+        if forehead_to_cheek_ratio > 0.85 and jaw_to_cheek_ratio > 0.85 and abs(forehead_width - jaw_width_gonial) < forehead_width * 0.15:
+             # Similar widths, angular jaw suggests Square, rounded suggests Round
+             # Differentiating Square and Round needs jawline CURVE analysis, which is harder.
+             # Let's use jaw_width vs cheek_width: if jaw is nearly as wide as cheeks -> Square tendency
+            if jaw_width_gonial > face_width_cheeks * 0.85:
+                shape = "Square"
+            else:
+                shape = "Round"
+        else: # If widths are not all similar
+            shape = "Round" # Default for wider faces
+    else: # facial_index between 0.95 and 1.05 (balanced height/width)
+        # Oval: Forehead slightly narrower than cheeks, jawline tapers smoothly, narrower than forehead.
+        # Diamond: Widest at cheeks, forehead and jaw narrower.
+        if face_width_cheeks > forehead_width and face_width_cheeks > jaw_width_gonial and chin_width < jaw_width_gonial * 0.8:
+            shape = "Diamond"
+        elif forehead_width > jaw_width_gonial and face_width_cheeks > jaw_width_gonial and chin_width < jaw_width_gonial * 0.75:
+            # forehead_to_cheek_ratio slightly less than 1 (e.g. 0.8-0.98)
+            # jaw_to_cheek_ratio less than forehead_to_cheek
+            if 0.8 < forehead_to_cheek_ratio < 1.0 and jaw_to_cheek_ratio < forehead_to_cheek_ratio * 0.95:
+                 shape = "Oval"
+            else: # If forehead is widest and tapers
+                 shape = "Heart"
+
+        elif abs(forehead_width - jaw_width_gonial) < forehead_width * 0.15 and abs(face_width_cheeks - forehead_width) < forehead_width * 0.1:
+            shape = "Square" # All widths relatively similar
+        else:
+            shape = "Oval" # Fallback for balanced, if not Diamond or strong Square
+
+    if shape == "Unknown": # If no specific rules matched strongly
+        if 0.95 <= facial_index <= 1.05 and forehead_to_cheek_ratio < 1 and jaw_to_cheek_ratio < forehead_to_cheek_ratio:
+            shape = "Oval (Default)"
+        else:
+            shape = "Round (Default)"
+            
+    return shape, measurements
+
+
+def get_side_profile_assessment(side_image_pil):
+    if not side_image_pil:
+        return "Not provided", None
+
+    side_image_pil = side_image_pil.convert("RGB")
+    # For simplicity, assume face is prominent in side profile.
+    # In a real app, run detection here too.
+    landmarks, error_msg_lm, _ = get_landmarks_and_draw(side_image_pil) # We don't need drawn image here
+
+    if error_msg_lm or not landmarks:
+        return f"Could not analyze ({error_msg_lm or 'no landmarks'})", None
+
+    # --- Assess Jawline from Side Profile ---
+    # Points for jaw angle (simplified):
+    # Point near ear lobe (e.g., landmark 127, 234 can be temple for side)
+    # Let's try specific side profile landmarks if they differ, or use consistent ones.
+    # For jaw angle: 172 (Gonion area), 152 (Chin tip/Pogonion), a point up along the ramus (e.g. 177 or 34)
     
-    # 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
+    # This requires good landmark stability on side profiles, which can be tricky.
+    # A very simple proxy: horizontal prominence of chin vs. a point higher on jaw.
+    p_chin_tip = landmarks.landmark[152]
+    p_jaw_angle_approx = landmarks.landmark[172] # Approximate gonion from front view set
+    p_upper_jaw_point = landmarks.landmark[135] # A point higher on the mandible body
+
+    # We need to consider the orientation. Let's assume face is looking left or right.
+    # A very rough heuristic:
+    # If chin (p_chin_tip.x) is significantly "forward" (more extreme x value, depending on orientation)
+    # than the jaw_angle_approx.x, it might suggest a stronger jaw.
+    # This is highly dependent on head rotation and landmark stability.
+
+    # A more robust method would involve angles, but that requires careful landmark selection.
+    # For now, let's just acknowledge if landmarks were found.
+    # In future, one could calculate the angle formed by landmarks such as:
+    # - A point on the ear (e.g. 127)
+    # - The gonion (jaw angle, e.g. 172 from frontal set, or a side-specific one)
+    # - The pogonion (chin tip, 152)
+    # A smaller angle (more acute) might indicate a sharper jawline.
     
-    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.", "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."]},
+    # Example using 3 points to form an angle: A=ear_pt, B=jaw_angle_pt, C=chin_pt
+    # vec_BA = (A.x-B.x, A.y-B.y)
+    # vec_BC = (C.x-B.x, C.y-B.y)
+    # dot_product = vec_BA[0]*vec_BC[0] + vec_BA[1]*vec_BC[1]
+    # mag_BA = math.sqrt(vec_BA[0]**2 + vec_BA[1]**2)
+    # mag_BC = math.sqrt(vec_BC[0]**2 + vec_BC[1]**2)
+    # angle_rad = math.acos(dot_product / (mag_BA * mag_BC))
+    # angle_deg = math.degrees(angle_rad)
+    # This is sensitive to landmark choice and head pose.
+
+    return "Analyzed (details TBD)", landmarks # Placeholder
+
+def get_hairstyle_suggestions_v2(face_shape, side_profile_info=""):
+    # (Expanded suggestion dictionary - keep it outside for brevity if very long)
+    # This needs to be more granular based on the new shapes from estimate_face_shape_v2
+    base_suggestions = {
+        "Oval": {"hair": ["Most styles work. Consider layers, textured crops, or side parts."], "beard": ["Versatile. Classic full beard, short boxed, or stubble."]},
+        "Oval (Default)": {"hair": ["Versatile. Try layers or a textured crop. Side parts can be flattering."], "beard": ["Well-groomed stubble or a short boxed beard."]},
+        "Long/Oblong": {"hair": ["Add width: Curls, waves, shoulder-length with layers. Bangs (blunt/side-swept). Avoid height."], "beard": ["Fuller on cheeks: full beard, mutton chops. Avoid long, pointy beards."]},
+        "Long": {"hair": ["Add width: Curls, waves, shoulder-length with layers. Bangs (blunt/side-swept). Avoid height."], "beard": ["Fuller on cheeks: full beard, mutton chops. Avoid long, pointy beards."]},
+        "Heart": {"hair": ["Add jawline volume: chin-length bobs, layered shoulder cuts. Side-swept bangs/textured fringe for forehead."], "beard": ["Fuller beards to add jaw width: Garibaldi, full beard carefully shaped."]},
+        "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 shaped."]},
+        "Square": {"hair": ["Softer styles: waves, curls, layers. Textured cuts, off-center parts. Avoid sharp, geometric cuts if aiming to soften."], "beard": ["Circle beard, rounded full beard. Stubble can highlight jaw if desired."]},
+        "Round": {"hair": ["Add height and length: pompadour, quiff, faux hawk, side part. Layers. Avoid blunt bobs at chin or very short, round cuts."], "beard": ["Add length to chin: goatee, soul patch, beard shorter on sides & longer at chin (ducktail)."]},
+        "Diamond": {"hair": ["Soften forehead & jaw: chin bobs, shoulder length with layers, textured fringe. Side-swept bangs."], "beard": ["Fuller at chin, possibly some width at jaw but not cheeks: Balbo, shorter full beard."]},
         "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."]},
+        "Unknown (div zero)": {"hair": ["Measurement error. Try different image."], "beard": ["Measurement error. Try different image."]},
+        "Round (Default)": {"hair": ["Add height and length: pompadour, quiff, faux hawk, side part. Layers. Avoid blunt bobs at chin or very short, round cuts."], "beard": ["Add length to chin: goatee, soul patch, beard shorter on sides & longer at chin (ducktail)."]},
     }
-    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 f"Could not determine suggestions for the estimated face shape: {face_shape}."
+    
+    sugg = base_suggestions.get(face_shape, {"hair": ["General advice: consult a professional stylist."], "beard": ["Experiment with styles that you feel confident in."]})
+    
+    hair_sug = "\n".join([f"- {s}" for s in sugg["hair"]])
+    beard_sug = "\n".join([f"- {s}" for s in sugg["beard"]])
+    
+    # Add side profile note if relevant
+    if "Analyzed" in side_profile_info:
+        side_note = "\n\n*Side profile was analyzed. Future versions might use this for more tailored advice (e.g., jawline definition).*"
+    else:
+        side_note = ""
+
+    return f"**Haircut Suggestions for {face_shape} Face:**\n{hair_sug}\n\n**Beard Style Suggestions for {face_shape} Face:**\n{beard_sug}{side_note}"
 
 
-def analyze_face_and_suggest(front_image_input, side_image_optional):
+def analyze_face_and_suggest_v2(front_image_input, side_image_input_optional):
     if front_image_input is None:
-        return None, "Please upload a front-facing photo.", ""
+        return None, "Please upload a front-facing photo.", "", {}
 
-    # Gradio Image input is a NumPy array
-    img_pil = Image.fromarray(front_image_input).convert("RGB") # Ensure RGB
+    img_pil = Image.fromarray(front_image_input).convert("RGB")
 
-    # 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:
-        return None, "Could not detect a face.", ""
+    # 1. Detect Face
+    cropped_face_pil, error_msg_detect = detect_face_local(img_pil)
+    if error_msg_detect:
+        return None, error_msg_detect, "", {}
+    if cropped_face_pil is None: # Should be caught by error_msg but as a fallback
+        return None, "Could not detect a face.", "", {}
 
-    # 2. Get Facial Landmarks (MediaPipe)
-    landmarks, error_msg_lm = get_landmarks_mediapipe(cropped_face_pil)
+    # 2. Get Facial Landmarks and Draw them
+    landmarks, error_msg_lm, face_with_landmarks_pil = get_landmarks_and_draw(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')
+        return face_with_landmarks_pil, f"Face detected. Error getting landmarks: {error_msg_lm}", "Cannot suggest hairstyles without landmark analysis.", {}
 
+    # 3. Estimate Face Shape
+    img_w, img_h = cropped_face_pil.size # Use cropped face dimensions
+    estimated_shape, measurements = estimate_face_shape_from_landmarks_v2(landmarks, img_w, img_h)
+    
+    measurements_str = "\n".join([f"- {k.replace('_norm',' (norm.)')}: {v:.2f}" for k,v in measurements.items()])
+    analysis_text = f"Estimated Face Shape: **{estimated_shape}**\n\nNormalized Measurements:\n{measurements_str}"
 
-    # 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_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. Analyze Side Profile (Basic)
+    side_profile_status = "Not provided or not analyzed"
+    side_profile_data = None
+    if side_image_input_optional is not None:
+        side_pil = Image.fromarray(side_image_input_optional).convert("RGB")
+        side_profile_status, side_profile_data = get_side_profile_assessment(side_pil)
+        analysis_text += f"\n\nSide Profile: {side_profile_status}"
+        # Future: Modify 'estimated_shape' or suggestions based on side_profile_data
 
-    # 4. Get Suggestions
-    suggestions_text = get_hairstyle_suggestions(estimated_shape)
+    # 5. Get Suggestions
+    suggestions_text = get_hairstyle_suggestions_v2(estimated_shape, side_profile_status)
 
-    return cropped_face_pil, f"Estimated Face Shape: **{estimated_shape}**\n{side_info}", suggestions_text
+    return face_with_landmarks_pil, analysis_text, suggestions_text
 
 # --- Gradio Interface ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# ✂️ AI Hairstyle & Beard Suggester (Local Models) 🧔")
+    gr.Markdown("# ✂️ AI Hairstyle & Beard Suggester (Enhanced Local Models) 🧔")
     gr.Markdown(
-        "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.*"
+        "Upload a clear, front-facing photo. Optionally, upload a side profile."
+        "\n*Disclaimer: This app uses local AI models for face detection and landmark-based shape estimation. Suggestions are general and based on heuristics.*"
     )
 
     with gr.Row():
@@ -268,20 +370,20 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             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 (or Cropped with Landmarks)")
-            output_shape_info = gr.Markdown(label="Face Analysis")
+            output_image_landmarks = gr.Image(label="Detected Face with Landmarks")
+            output_analysis_info = gr.Markdown(label="Face Analysis & Measurements")
             output_suggestions = gr.Markdown(label="Suggestions")
 
     submit_btn.click(
-        analyze_face_and_suggest,
+        analyze_face_and_suggest_v2,
         inputs=[front_image_input, side_image_input],
-        outputs=[output_image, output_shape_info, output_suggestions]
+        outputs=[output_image_landmarks, output_analysis_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.Markdown("--- \n ### Note on Face Shape Estimation: \n The face shape estimation is based on ratios of distances between facial landmarks detected by MediaPipe. The categories (Oval, Round, Square, etc.) and the rules to classify them are experimental and may require further refinement for high accuracy. Landmark visualization shows the points used.")
 
 
 if __name__ == "__main__":
-    if face_detection_model and face_mesh_detector: # Only launch if models loaded
+    if face_detection_model and face_mesh_detector:
         demo.launch()
     else:
-        print("Gradio app not launched due to model loading errors.")
+        print("Gradio app not launched due to model loading errors. Check face detection model name and availability.")