Spaces:
Running
Running
| import gradio as gr | |
| from PIL import Image | |
| import os | |
| import cv2 | |
| import numpy as np | |
| import mediapipe as mp | |
| # Load hairstyles from 'hairstyles' folder | |
| def load_hairstyles(): | |
| folder = "hairstyles" | |
| if not os.path.exists(folder): | |
| return [] | |
| return [ | |
| Image.open(os.path.join(folder, f)).convert("RGBA") | |
| for f in sorted(os.listdir(folder)) if f.endswith(".png") | |
| ] | |
| hairstyles = load_hairstyles() | |
| mp_face_detection = mp.solutions.face_detection | |
| # 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 | |
| # 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] | |
| 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)) | |
| if not results.detections: | |
| return user_img # No face detected, return original | |
| # 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) | |
| # Estimate head top | |
| top_y = max(y - int(h * 0.6), 0) | |
| # 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") | |
| style_slider = gr.Slider(0, max(len(hairstyles)-1, 0), step=1, label="🎨 Select Hairstyle") | |
| apply_btn = gr.Button("✨ Apply Hairstyle") | |
| with gr.Column(): | |
| result_output = gr.Image(label="🔍 Result Preview") | |
| apply_btn.click( | |
| fn=apply_hairstyle, | |
| inputs=[image_input, style_slider], | |
| outputs=result_output | |
| ) | |
| demo.launch() | |