Update app.py

app.py

@@ -2,36 +2,48 @@ import gradio as gr
 from PIL import Image, ImageDraw
 import numpy as np
 import torch
-from transformers import YolosImageProcessor, YolosForObjectDetection # Keep for YOLOS architecture
+from transformers import YolosImageProcessor, YolosForObjectDetection
 import mediapipe as mp
 import math
+import os # For potential future environment variable use
 
 # --- Model Initialization ---
-# 1. Face Detection Model (Switching to yolos-face if possible)
-print("Loading face detection model...")
-# 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
+# 1. Face Detection Model
+print("Attempting to load face detection model...")
+PRIMARY_DETECTION_MODEL_NAME = "hustvl/yolos-face"
+FALLBACK_DETECTION_MODEL_NAME = "hustvl/yolos-tiny" # Detects 'person'
+FACE_LABEL_ID = -1  # Will be set based on which model loads
+
+face_image_processor = None
+face_detection_model = None
 
 try:
-    face_image_processor = YolosImageProcessor.from_pretrained(DETECTION_MODEL_NAME)
-    face_detection_model = YolosForObjectDetection.from_pretrained(DETECTION_MODEL_NAME)
-    # 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.")
+    print(f"Trying primary model: {PRIMARY_DETECTION_MODEL_NAME}")
+    face_image_processor = YolosImageProcessor.from_pretrained(PRIMARY_DETECTION_MODEL_NAME)
+    face_detection_model = YolosForObjectDetection.from_pretrained(PRIMARY_DETECTION_MODEL_NAME)
+    # For hustvl/yolos-face, the label for "face" is 0.
+    FACE_LABEL_ID = 0 # Corresponds to "face"
+    print(f"Successfully loaded primary face detection model: {PRIMARY_DETECTION_MODEL_NAME} (label 'face': {FACE_LABEL_ID})")
 except Exception as 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
-
+    print(f"Error loading primary model {PRIMARY_DETECTION_MODEL_NAME}: {e}")
+    print(f"Attempting to load fallback model: {FALLBACK_DETECTION_MODEL_NAME}")
+    try:
+        face_image_processor = YolosImageProcessor.from_pretrained(FALLBACK_DETECTION_MODEL_NAME)
+        face_detection_model = YolosForObjectDetection.from_pretrained(FALLBACK_DETECTION_MODEL_NAME)
+        # For hustvl/yolos-tiny (trained on COCO), 'person' is label 0.
+        FACE_LABEL_ID = 0 # We will use 'person' (label 0) as a proxy for face
+        print(f"Successfully loaded fallback detection model: {FALLBACK_DETECTION_MODEL_NAME} (using label 'person': {FACE_LABEL_ID})")
+    except Exception as e2:
+        print(f"Error loading fallback model {FALLBACK_DETECTION_MODEL_NAME}: {e2}")
+        print("!!! CRITICAL: Face detection model could not be loaded. The app might not function correctly. !!!")
+        # face_image_processor and face_detection_model will remain None
 
 # 2. Facial Landmark Model (MediaPipe Face Mesh)
 print("Initializing MediaPipe Face Mesh...")
+mp_face_mesh = None
+face_mesh_detector = None
+mp_drawing = None
+drawing_spec = None
 try:
     mp_face_mesh = mp.solutions.face_mesh
     face_mesh_detector = mp_face_mesh.FaceMesh(
@@ -44,8 +56,7 @@ try:
     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
+    # Variables will remain None
 
 
 # --- Helper Functions ---
@@ -53,64 +64,59 @@ def detect_face_local(image_pil):
     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")
-    with torch.no_grad(): # Important for inference
-        outputs = face_detection_model(**inputs)
+    try:
+        inputs = face_image_processor(images=image_pil, return_tensors="pt")
+        with torch.no_grad(): # Important for inference
+            outputs = face_detection_model(**inputs)
 
-    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
+        target_sizes = torch.tensor([image_pil.size[::-1]])
+        results = face_image_processor.post_process_object_detection(outputs, threshold=0.7, target_sizes=target_sizes)[0] # Threshold
 
-    best_box = None
-    max_score = 0
+        best_box = None
+        max_score = 0
 
-    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-        if label == FACE_LABEL_ID: # Check for the specific 'face' label
-            if score > max_score:
-                max_score = score
-                best_box = box.tolist()
+        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+            if label.item() == FACE_LABEL_ID: # Use .item() to get Python int from tensor
+                if score.item() > max_score:
+                    max_score = score.item()
+                    best_box = box.tolist()
 
-    if best_box:
-        # 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
+        if best_box:
+            padding_w = (best_box[2] - best_box[0]) * 0.15 # 15% padding width
+            padding_h = (best_box[3] - best_box[1]) * 0.15 # 15% padding height
 
-        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)
+            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)
+
+            cropped_image = image_pil.crop((xmin, ymin, xmax, ymax))
+            return cropped_image, None
+        else:
+            return None, "No face/person detected with sufficient confidence."
+    except Exception as e:
+        print(f"Error during local face detection: {e}")
+        return None, f"Error during face detection: {str(e)}"
 
-        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
+    if not face_mesh_detector or not mp_drawing or not drawing_spec:
+        return None, "MediaPipe Face Mesh not initialized for landmarks.", image_pil
 
-    image_rgb_mp = np.array(image_pil.convert('RGB'))
+    image_rgb_mp = np.array(image_pil.convert('RGB')) # MediaPipe prefers RGB
     results = face_mesh_detector.process(image_rgb_mp)
 
-    annotated_image_pil = image_pil.copy() # Start with a copy of the input PIL image
+    annotated_image_pil = image_pil.copy()
 
     if results.multi_face_landmarks:
-        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
+        landmarks = results.multi_face_landmarks[0]
         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
+
+        # Draw landmarks using MediaPipe's utility
         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)
         
@@ -120,51 +126,33 @@ def get_landmarks_and_draw(image_pil):
         return None, "Could not detect facial landmarks.", annotated_image_pil
 
 
-def _distance_2d_normalized(p1, p2): # Operates on normalized 0-1 coordinates
+def _distance_2d_normalized(p1, p2):
     return math.sqrt((p1.x - p2.x)**2 + (p1.y - p2.y)**2)
 
 def estimate_face_shape_from_landmarks_v2(landmarks, img_width, img_height):
     if not landmarks:
         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
+    face_height = abs(p_forehead_top_center.y - p_chin_bottom.y)
 
-    # 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_cheeks = abs(p_cheek_left.x - p_cheek_right.x) # Normalized X diff
+    face_width_cheeks = abs(p_cheek_left.x - p_cheek_right.x)
 
-    # 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
+    p_forehead_L = landmarks.landmark[70]
+    p_forehead_R = landmarks.landmark[300]
     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
+    p_jaw_angle_L = landmarks.landmark[172]
+    p_jaw_angle_R = landmarks.landmark[397]
     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_L = landmarks.landmark[143]
     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,
@@ -172,65 +160,51 @@ def estimate_face_shape_from_landmarks_v2(landmarks, img_width, img_height):
         "jaw_width_gonial_norm": jaw_width_gonial,
         "chin_width_norm": chin_width
     }
-    print("Normalized Measurements:", measurements)
+    # print("Normalized Measurements:", {k: round(v,3) for k,v in measurements.items()})
 
-    # --- 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
+        if forehead_to_cheek_ratio > 0.85 and jaw_to_cheek_ratio > 0.85 and abs(forehead_width - jaw_width_gonial) < forehead_width * 0.20 :
+             shape = "Long/Oblong" # All widths relatively similar but face is long
+        elif forehead_width > jaw_width_gonial and chin_width < jaw_width_gonial * 0.85:
+            shape = "Heart/Inverted Triangle"
         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:
+    elif facial_index < 0.95: # Wider than long, or close to equal width/height and not distinctly Diamond/Heart
+        if forehead_to_cheek_ratio > 0.85 and jaw_to_cheek_ratio > 0.85 and abs(forehead_width - jaw_width_gonial) < forehead_width * 0.20:
+            if jaw_width_gonial > face_width_cheeks * 0.88: # Strong jaw compared to cheeks
                 shape = "Square"
             else:
                 shape = "Round"
-        else: # If widths are not all similar
-            shape = "Round" # Default for wider faces
+        else: # If widths are not all similar, default to Round for wider faces
+            shape = "Round"
     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:
+        if face_width_cheeks > forehead_width and face_width_cheeks > jaw_width_gonial and chin_width < jaw_width_gonial * 0.85:
             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:
+        elif forehead_width > jaw_width_gonial and face_width_cheeks > jaw_width_gonial and chin_width < jaw_width_gonial * 0.8:
+            if 0.80 < 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
+            else:
                  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
+        elif abs(forehead_width - jaw_width_gonial) < forehead_width * 0.15 and abs(face_width_cheeks - forehead_width) < forehead_width * 0.15 :
+            shape = "Square"
         else:
-            shape = "Oval" # Fallback for balanced, if not Diamond or strong Square
+            shape = "Oval" # General fallback for balanced faces not matching other criteria
 
     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:
+        if 0.95 <= facial_index <= 1.05 and forehead_to_cheek_ratio < 1.0 and jaw_to_cheek_ratio < forehead_to_cheek_ratio:
             shape = "Oval (Default)"
+        elif facial_index < 0.95:
+             shape = "Round (Default)"
         else:
-            shape = "Round (Default)"
+             shape = "Long (Default)"
             
     return shape, measurements
 
@@ -239,68 +213,37 @@ def get_side_profile_assessment(side_image_pil):
     if not side_image_pil:
         return "Not provided", None
 
+    # Convert Gradio Image (numpy array) to PIL Image if it's not already
+    if isinstance(side_image_pil, np.ndarray):
+        side_image_pil = Image.fromarray(side_image_pil)
+    
     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
+    landmarks, error_msg_lm, _ = get_landmarks_and_draw(side_image_pil)
 
     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)
-    
-    # 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.
     
-    # 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
+    # Basic assessment placeholder
+    # E.g. Chin prominence (landmark 152's x vs jaw angle 172's x)
+    # This is highly dependent on consistent side view and requires careful calibration
+    # For now, just acknowledge landmarks were found
+    return "Analyzed (basic landmark detection)", landmarks
 
 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."]},
+        "Long (Default)": {"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)."]},
+        "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)."]},
         "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 (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)."]},
     }
     
     sugg = base_suggestions.get(face_shape, {"hair": ["General advice: consult a professional stylist."], "beard": ["Experiment with styles that you feel confident in."]})
@@ -308,57 +251,58 @@ def get_hairstyle_suggestions_v2(face_shape, side_profile_info=""):
     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
+    side_note = ""
     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 = ""
+        side_note = "\n\n*Side profile analyzed. Future versions could use this for more tailored advice (e.g., jawline definition).*"
+    elif "Not provided" not in side_profile_info and side_profile_info: # If there was an attempt but it failed
+        side_note = f"\n\n*Side profile: {side_profile_info}*"
+
 
     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_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.", ""
+
+    # Ensure models are loaded
+    if not face_detection_model or not face_mesh_detector:
+        error_msg = []
+        if not face_detection_model: error_msg.append("Face detector not loaded.")
+        if not face_mesh_detector: error_msg.append("Landmark detector not loaded.")
+        return None, " ".join(error_msg) + " Please check Space logs.", ""
 
     img_pil = Image.fromarray(front_image_input).convert("RGB")
 
-    # 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.", "", {}
+        return None, error_msg_detect, "" # No measurements if face detection fails
+    if cropped_face_pil is None:
+        return None, "Could not detect a face.", ""
 
-    # 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:
-        return face_with_landmarks_pil, f"Face detected. Error getting landmarks: {error_msg_lm}", "Cannot suggest hairstyles without landmark analysis.", {}
+        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
+    img_w, img_h = cropped_face_pil.size
     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()])
+    measurements_str = "\n".join([f"- {k.replace('_norm',' (norm. ratio)'):<25}: {v:.3f}" for k,v in measurements.items()])
     analysis_text = f"Estimated Face Shape: **{estimated_shape}**\n\nNormalized Measurements:\n{measurements_str}"
 
-    # 4. Analyze Side Profile (Basic)
-    side_profile_status = "Not provided or not analyzed"
-    side_profile_data = None
+    side_profile_status = "Not provided"
     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)
+        # Pass the numpy array directly
+        side_profile_status, _ = get_side_profile_assessment(side_image_input_optional)
         analysis_text += f"\n\nSide Profile: {side_profile_status}"
-        # Future: Modify 'estimated_shape' or suggestions based on side_profile_data
 
-    # 5. Get Suggestions
     suggestions_text = get_hairstyle_suggestions_v2(estimated_shape, side_profile_status)
 
     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 (Enhanced Local Models) 🧔")
+    gr.Markdown("# ✂️ AI Hairstyle & Beard Suggester 🧔")
     gr.Markdown(
         "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.*"
@@ -379,11 +323,18 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         inputs=[front_image_input, side_image_input],
         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 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.")
+    gr.Markdown("--- \n ### Notes: \n - **Face Shape Estimation:** Based on ratios of distances between facial landmarks (MediaPipe). The categories (Oval, Round, etc.) and classification rules are experimental. \n - **Landmark Visualization:** Green mesh shows detected facial landmarks. \n - **Model Loading:** Tries `hustvl/yolos-face` first, then `hustvl/yolos-tiny` (person detection) as fallback. Check Space logs for details.")
 
 
 if __name__ == "__main__":
-    if face_detection_model and face_mesh_detector:
+    # Only launch if at least the fallback detection model and mediapipe loaded
+    if (face_detection_model and face_image_processor and FACE_LABEL_ID != -1) and \
+       (face_mesh_detector and mp_drawing and drawing_spec):
+        print("Launching Gradio App...")
         demo.launch()
     else:
-        print("Gradio app not launched due to model loading errors. Check face detection model name and availability.")
+        print("Gradio app not launched due to critical model loading errors. Please check the logs.")
+        if not (face_detection_model and face_image_processor and FACE_LABEL_ID != -1):
+            print("-> Face detection model failed to load.")
+        if not (face_mesh_detector and mp_drawing and drawing_spec):
+            print("-> MediaPipe landmark model failed to initialize.")