Update app.py

app.py

@@ -1,135 +1,127 @@
-import cv2
 import gradio as gr
+import cv2
 import numpy as np
-from mediapipe import solutions
+import mediapipe as mp
 from sklearn.cluster import KMeans
 
-# Load the mask images based on face shape and skin tone
-mask_dict = {
-    "oval": {
-        "fair": "masks/oval_fair_mask.png",
-        "medium": "masks/oval_medium_mask.png",
-        "dark": "masks/oval_dark_mask.png"
-    },
-    "round": {
-        "fair": "masks/round_fair_mask.png",
-        "medium": "masks/round_medium_mask.png",
-        "dark": "masks/round_dark_mask.png"
-    }
-    # Add more shapes and tones as necessary
-}
-
-# Function to overlay the mask on the face using OpenCV
-def overlay_mask(face_image, mask_path, face_x, face_y, face_w, face_h):
-    # Load the mask
-    mask = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
-    
-    # Resize the mask to fit the face dimensions
-    mask_resized = cv2.resize(mask, (face_w, face_h))
-    
-    # Create a mask and inverse mask (for transparency handling)
-    mask_alpha = mask_resized[:, :, 3] / 255.0  # Alpha channel
-    mask_rgb = mask_resized[:, :, :3]
-    
-    # Define the region of interest (ROI) on the face image
-    roi = face_image[face_y:face_y + face_h, face_x:face_x + face_w]
-    
-    # Blend the mask onto the face using the alpha mask
-    for c in range(0, 3):  # For each color channel
-        roi[:, :, c] = (1 - mask_alpha) * roi[:, :, c] + mask_alpha * mask_rgb[:, :, c]
-    
-    face_image[face_y:face_y + face_h, face_x:face_x + face_w] = roi
-    return face_image
-
-# Function to detect face shape using MediaPipe FaceMesh
-def detect_face_shape(landmarks):
-    # Extract relevant facial landmarks (e.g., eyes, chin, etc.)
-    left_eye = landmarks[33]
-    right_eye = landmarks[133]
-    chin = landmarks[1]
-    jawline = landmarks[7]
-    
-    # Calculate distances and ratios to infer face shape
-    eye_distance = np.linalg.norm(np.array(left_eye) - np.array(right_eye))
-    chin_distance = np.linalg.norm(np.array(chin) - np.array(jawline))
-
-    # Classify face shape based on proportions (this can be adjusted based on your observations)
-    if eye_distance > chin_distance:
-        return "oval"
-    else:
+# Initialize MediaPipe modules
+mp_face_detection = mp.solutions.face_detection
+mp_face_mesh = mp.solutions.face_mesh
+
+face_detector = mp_face_detection.FaceDetection(model_selection=0, min_detection_confidence=0.6)
+face_mesh = mp_face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1, refine_landmarks=True)
+
+def detect_face_shape(landmarks, image_width, image_height):
+    # Extract specific landmarks
+    jaw_left = landmarks[234]
+    jaw_right = landmarks[454]
+    chin = landmarks[152]
+    forehead = landmarks[10]
+
+    x1 = int(jaw_left.x * image_width)
+    x2 = int(jaw_right.x * image_width)
+    y1 = int(chin.y * image_height)
+    y2 = int(forehead.y * image_height)
+
+    face_width = abs(x2 - x1)
+    face_height = abs(y1 - y2)
+
+    ratio = face_width / face_height if face_height != 0 else 0
+
+    if ratio > 1.05:
         return "round"
+    elif 0.95 < ratio <= 1.05:
+        return "square"
+    else:
+        return "oval"
 
-# Function to detect skin tone based on KMeans clustering
-def detect_skin_tone(face_image, face_x, face_y, face_w, face_h):
-    # Extract the face region
-    face_region = face_image[face_y:face_y + face_h, face_x:face_x + face_w]
-    
-    # Convert the face region to RGB (if it's BGR)
-    face_region_rgb = cv2.cvtColor(face_region, cv2.COLOR_BGR2RGB)
-    
-    # Reshape for clustering
-    face_region_flat = face_region_rgb.reshape((-1, 3))
-    
-    # Apply KMeans to classify the skin tone based on the dominant color
-    kmeans = KMeans(n_clusters=1, random_state=0).fit(face_region_flat)
-    
-    # Get the cluster center (dominant color)
-    dominant_color = kmeans.cluster_centers_[0]
-    
-    # Classify skin tone based on dominant color
-    if dominant_color[0] < 100 and dominant_color[1] < 120 and dominant_color[2] < 130:  # darker shades
-        return "dark"
-    elif dominant_color[0] > 120 and dominant_color[1] > 130 and dominant_color[2] > 140:  # fair shades
+def detect_skin_tone(image, x, y, w, h):
+    roi = image[y:y+h, x:x+w]
+    roi_rgb = cv2.cvtColor(roi, cv2.COLOR_BGR2RGB)
+    roi_flat = roi_rgb.reshape((-1, 3))
+
+    kmeans = KMeans(n_clusters=3, n_init=10)
+    kmeans.fit(roi_flat)
+    avg_color = kmeans.cluster_centers_[0]
+    brightness = np.mean(avg_color)
+
+    if brightness > 200:
         return "fair"
-    else:
+    elif brightness > 100:
         return "medium"
+    else:
+        return "dark"
+
+def overlay_mask(image, face_shape, skin_tone, x, y, w, h):
+    overlay = image.copy()
+    mask = np.zeros_like(image, dtype=np.uint8)
+
+    color_dict = {
+        "fair": (255, 182, 193),    # light pink
+        "medium": (0, 191, 255),    # deep sky blue
+        "dark": (138, 43, 226)      # blue violet
+    }
+    color = color_dict.get(skin_tone, (255, 255, 255))
+
+    if face_shape == "oval":
+        center = (x + w // 2, y + h // 2)
+        axes = (w // 2, h // 2)
+        cv2.ellipse(mask, center, axes, 0, 0, 360, color, -1)
+    elif face_shape == "round":
+        radius = min(w, h) // 2
+        center = (x + w // 2, y + h // 2)
+        cv2.circle(mask, center, radius, color, -1)
+    elif face_shape == "square":
+        cv2.rectangle(mask, (x, y), (x + w, y + h), color, -1)
+    else:
+        cv2.rectangle(mask, (x, y), (x + w, y + h), color, -1)
+
+    alpha = 0.4
+    blended = cv2.addWeighted(mask, alpha, image, 1 - alpha, 0)
+
+    cv2.putText(blended, f"{face_shape}, {skin_tone}", (x, y - 10),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
+
+    return blended
 
-# Function to process the image: detect face, determine shape and skin tone, and apply mask
 def process_image(image):
-    # Convert to grayscale for face detection
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    
-    # Initialize face detection
-    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-    faces = face_cascade.detectMultiScale(gray, 1.1, 4)
-    
-    # Initialize MediaPipe FaceMesh for face shape detection
-    mp_face_mesh = solutions.face_mesh.FaceMesh(static_image_mode=True)
-    mp_results = mp_face_mesh.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
-
-    # Get face shape (assuming landmarks are available from MediaPipe FaceMesh)
-    face_shape = "oval"  # Default value; it should be determined by landmark analysis
-    skin_tone = "fair"   # Default value; it should be determined by skin tone detection
-    
-    if mp_results.multi_face_landmarks:
-        for landmarks in mp_results.multi_face_landmarks:
-            face_shape = detect_face_shape(landmarks)
-    
-    # Assume we are getting the first face (you can extend this if there are multiple faces)
-    for (x, y, w, h) in faces:
-        # Detect skin tone for the face region
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    ih, iw, _ = image.shape
+
+    results = face_detector.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+    if not results.detections:
+        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    for detection in results.detections:
+        bboxC = detection.location_data.relative_bounding_box
+        x = int(bboxC.xmin * iw)
+        y = int(bboxC.ymin * ih)
+        w = int(bboxC.width * iw)
+        h = int(bboxC.height * ih)
+        x, y = max(x, 0), max(y, 0)
+
+        # Detect mesh
+        results_mesh = face_mesh.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+        if results_mesh.multi_face_landmarks:
+            landmarks = results_mesh.multi_face_landmarks[0].landmark
+            face_shape = detect_face_shape(landmarks, iw, ih)
+        else:
+            face_shape = "oval"
+
         skin_tone = detect_skin_tone(image, x, y, w, h)
-        
-        # Select mask based on face shape and skin tone
-        mask_path = mask_dict[face_shape][skin_tone]
-        
-        # Apply the mask
-        image = overlay_mask(image, mask_path, x, y, w, h)
-    
-    return image
-
-# Gradio interface
-def gradio_interface(image):
-    processed_image = process_image(image)
-    return processed_image
-
-# Create Gradio interface
-iface = gr.Interface(
-    fn=gradio_interface,
-    inputs=gr.inputs.Image(type="numpy", label="Upload Image or Use Webcam"),
-    outputs=gr.outputs.Image(type="numpy", label="Image with Mask"),
-    live=True
+        image = overlay_mask(image, face_shape, skin_tone, x, y, w, h)
+
+    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+# Gradio UI
+demo = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="numpy", label="Upload or Snap Image"),
+    outputs=gr.Image(label="Face Shape + Skin Tone + Mask Overlay"),
+    live=True,
+    title="Face Shape & Skin Tone Analyzer",
+    description="This app detects face shape & skin tone and overlays a dynamic mask using OpenCV."
 )
 
-# Launch the Gradio interface
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()