Upload 2 files

app.py

@@ -7,6 +7,12 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 
+try:
+    import mediapipe as mp  # type: ignore
+    HAS_MEDIAPIPE = True
+except Exception:  # pragma: no cover - optional dependency
+    HAS_MEDIAPIPE = False
+
 
 def _ensure_rgb_uint8(image: np.ndarray) -> np.ndarray:
     """Convert an input image array to RGB uint8 format.
@@ -41,6 +47,47 @@ def _central_crop_bbox(width: int, height: int, frac: float = 0.6) -> Tuple[int,
     return x1, y1, x2, y2
 
 
+def _detect_face_bbox_mediapipe(image_rgb: np.ndarray) -> Optional[Tuple[int, int, int, int]]:
+    """Detect a face bounding box using MediaPipe Face Detection and return (x1, y1, x2, y2).
+
+    Returns None if detection fails or mediapipe is unavailable.
+    """
+    if not HAS_MEDIAPIPE:
+        return None
+    height, width = image_rgb.shape[:2]
+    try:
+        with mp.solutions.face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5) as detector:
+            results = detector.process(image_rgb)
+            detections = results.detections or []
+            if not detections:
+                return None
+            # Pick the largest bbox
+            def bbox_area(det):
+                bbox = det.location_data.relative_bounding_box
+                return max(0.0, bbox.width) * max(0.0, bbox.height)
+
+            best = max(detections, key=bbox_area)
+            rb = best.location_data.relative_bounding_box
+            x1 = int(np.clip(rb.xmin * width, 0, width - 1))
+            y1 = int(np.clip(rb.ymin * height, 0, height - 1))
+            x2 = int(np.clip((rb.xmin + rb.width) * width, 0, width))
+            y2 = int(np.clip((rb.ymin + rb.height) * height, 0, height))
+
+            # Expand a bit to include cheeks/forehead
+            pad_x = int(0.08 * width)
+            pad_y = int(0.12 * height)
+            x1 = int(np.clip(x1 - pad_x, 0, width - 1))
+            y1 = int(np.clip(y1 - pad_y, 0, height - 1))
+            x2 = int(np.clip(x2 + pad_x, 0, width))
+            y2 = int(np.clip(y2 + pad_y, 0, height))
+
+            if x2 - x1 < 10 or y2 - y1 < 10:
+                return None
+            return x1, y1, x2, y2
+    except Exception:
+        return None
+
+
 def _binary_open_close(mask: np.ndarray, kernel_size: int = 5, iterations: int = 1) -> np.ndarray:
     """Apply morphological open then close to clean the binary mask."""
     kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
@@ -129,7 +176,7 @@ def _solid_color_image(color_rgb: np.ndarray, size: Tuple[int, int] = (160, 160)
     return swatch
 
 
-def detect_skin_tone(image: np.ndarray, center_focus: bool = True) -> Tuple[str, np.ndarray, np.ndarray]:
+def detect_skin_tone(image: np.ndarray, center_focus: bool = True, use_face_detector: bool = False) -> Tuple[str, np.ndarray, np.ndarray]:
     """Main pipeline: returns (hex_code, color_swatch_image, debug_mask_overlay).
 
     - image: input image as numpy array (H, W, 3) RGB uint8
@@ -138,8 +185,15 @@ def detect_skin_tone(image: np.ndarray, center_focus: bool = True) -> Tuple[str,
     rgb = _ensure_rgb_uint8(image)
     height, width = rgb.shape[:2]
 
-    # Optionally restrict to central crop to avoid background
-    if center_focus:
+    # Optionally restrict to detected face, else center crop, else full image
+    face_bbox: Optional[Tuple[int, int, int, int]] = None
+    if use_face_detector:
+        face_bbox = _detect_face_bbox_mediapipe(rgb)
+
+    if face_bbox is not None:
+        x1, y1, x2, y2 = face_bbox
+        central_rgb = rgb[y1:y2, x1:x2]
+    elif center_focus:
         x1, y1, x2, y2 = _central_crop_bbox(width, height, frac=0.7)
         central_rgb = rgb[y1:y2, x1:x2]
     else:
@@ -218,21 +272,28 @@ with gr.Blocks(title="Skin Tone Detector") as demo:
                 height=360,
             )
             center_focus = gr.Checkbox(value=True, label="Center focus (ignore edges)")
+            use_face_det = gr.Checkbox(value=True if HAS_MEDIAPIPE else False, label="Use face detection (MediaPipe)")
             run_btn = gr.Button("Detect Skin Tone", variant="primary")
 
         with gr.Column():
             hex_output = gr.HTML(label="HEX Color")
             swatch_output = gr.Image(label="Color Swatch", type="numpy")
             debug_output = gr.Image(label="Mask Overlay", type="numpy")
+            if not HAS_MEDIAPIPE:
+                gr.Markdown("MediaPipe not installed or unavailable. Face detection toggle will be ignored.")
 
-    def _run(image: Optional[np.ndarray], center_focus: bool):
+    def _run(image: Optional[np.ndarray], center_focus: bool, use_face_det_flag: bool):
         if image is None:
             return _hex_html("#000000"), np.zeros((160, 160, 3), dtype=np.uint8), None
-        hex_code, swatch, debug = detect_skin_tone(image, center_focus=center_focus)
+        hex_code, swatch, debug = detect_skin_tone(
+            image,
+            center_focus=center_focus,
+            use_face_detector=(use_face_det_flag and HAS_MEDIAPIPE),
+        )
         return _hex_html(hex_code), swatch, debug
 
-    run_btn.click(_run, inputs=[input_image, center_focus], outputs=[hex_output, swatch_output, debug_output])
-    input_image.change(_run, inputs=[input_image, center_focus], outputs=[hex_output, swatch_output, debug_output])
+    run_btn.click(_run, inputs=[input_image, center_focus, use_face_det], outputs=[hex_output, swatch_output, debug_output])
+    input_image.change(_run, inputs=[input_image, center_focus, use_face_det], outputs=[hex_output, swatch_output, debug_output])
 
 
 if __name__ == "__main__":

requirements.txt

@@ -2,4 +2,5 @@ gradio>=4.44.0
 opencv-python-headless>=4.10.0.84
 numpy>=1.26.0
 Pillow>=10.3.0
+mediapipe>=0.10.14