Update app.py

app.py

@@ -1,8 +1,11 @@
 import gradio as gr
 from PIL import Image
 import os
+import cv2
+import numpy as np
+import mediapipe as mp
 
-# Load all hairstyle images from folder
+# Load hairstyles from 'hairstyles' folder
 def load_hairstyles():
     folder = "hairstyles"
     if not os.path.exists(folder):
@@ -13,36 +16,54 @@ def load_hairstyles():
     ]
 
 hairstyles = load_hairstyles()
+mp_face_detection = mp.solutions.face_detection
 
-# Automatically apply hairstyle to uploaded image
+# Apply hairstyle using face detection to auto-align on head
 def apply_hairstyle(user_img, style_index):
     if user_img is None or not hairstyles:
         return None
 
-    user_img = user_img.convert("RGBA")
-    base_w, base_h = user_img.size
-    hairstyle = hairstyles[style_index]
+    # Convert to CV2 image for detection
+    img_cv2 = np.array(user_img.convert("RGB"))
+    img_cv2 = cv2.cvtColor(img_cv2, cv2.COLOR_RGB2BGR)
+    h_img, w_img = img_cv2.shape[:2]
 
-    # Resize hairstyle to match width of head image
-    scale_factor = 0.9  # Adjust if needed
-    new_width = int(base_w * scale_factor)
-    new_height = int(hairstyle.height * (new_width / hairstyle.width))
-    resized_hairstyle = hairstyle.resize((new_width, new_height))
+    with mp_face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5) as face_detection:
+        results = face_detection.process(cv2.cvtColor(img_cv2, cv2.COLOR_BGR2RGB))
 
-    # Automatically center it near the top
-    paste_x = (base_w - new_width) // 2
-    paste_y = int(base_h * 0.05)  # Adjust this offset if needed
+        if not results.detections:
+            return user_img  # No face detected, return original
 
-    composite = Image.new("RGBA", user_img.size)
-    composite.paste(resized_hairstyle, (paste_x, paste_y), resized_hairstyle)
+        # Use first face
+        detection = results.detections[0]
+        bbox = detection.location_data.relative_bounding_box
+        x = int(bbox.xmin * w_img)
+        y = int(bbox.ymin * h_img)
+        w = int(bbox.width * w_img)
+        h = int(bbox.height * h_img)
 
-    final = Image.alpha_composite(user_img, composite)
-    return final.convert("RGB")
+        # Estimate head top
+        top_y = max(y - int(h * 0.6), 0)
 
-# Gradio Interface (no sliders shown to user)
-with gr.Blocks() as demo:
-    gr.Markdown("## 💇 Salon Virtual Hairstyle Try-On (Auto Fit – No Adjustment Needed)")
+        # Load and resize hairstyle
+        hairstyle = hairstyles[style_index]
+        new_width = int(w * 1.1)  # Slightly wider than face
+        new_height = int(hairstyle.height * (new_width / hairstyle.width))
+        resized_hair = hairstyle.resize((new_width, new_height))
+
+        # Create output image
+        user_img = user_img.convert("RGBA")
+        composite = Image.new("RGBA", user_img.size)
+        paste_x = x - int((new_width - w) / 2)
+        paste_y = top_y
+        composite.paste(resized_hair, (paste_x, paste_y), resized_hair)
+
+        final = Image.alpha_composite(user_img, composite)
+        return final.convert("RGB")
 
+# Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("## 💇 Salon Virtual Hairstyle Try-On (Face-Aligned, No Adjustment Needed)")
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type="pil", label="📷 Upload Your Image")
@@ -58,3 +79,4 @@ with gr.Blocks() as demo:
     )
 
 demo.launch()
+