skin_detection.py

"""
Skin detection and extraction logic.
"""

import cv2
import numpy as np
import mediapipe as mp
import os

# Suppress TensorFlow logging caused by MediaPipe
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

mp_face = mp.solutions.face_detection


def detect_face(image):
    """
    Detects the primary face in an image.
    Returns a valid face crop or None.
    """

    h, w, _ = image.shape

    with mp_face.FaceDetection(
        model_selection=1,
        min_detection_confidence=0.5
    ) as detector:
        results = detector.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

        if not results.detections:
            return None

        bbox = results.detections[0].location_data.relative_bounding_box

        x1 = int(bbox.xmin * w)
        y1 = int(bbox.ymin * h)
        x2 = int((bbox.xmin + bbox.width) * w)
        y2 = int((bbox.ymin + bbox.height) * h)

        # -------------------------
        # CLAMP TO IMAGE BOUNDS
        # -------------------------
        x1 = max(0, x1)
        y1 = max(0, y1)
        x2 = min(w, x2)
        y2 = min(h, y2)

        # -------------------------
        # VALIDATE CROP
        # -------------------------
        if x2 <= x1 or y2 <= y1:
            return None

        face = image[y1:y2, x1:x2]

        if face.size == 0:
            return None

        return face


def skin_color_mask(face_img):
    """
    Generates a binary skin mask using HSV + YCrCb thresholds.
    """
    hsv = cv2.cvtColor(face_img, cv2.COLOR_BGR2HSV)
    ycrcb = cv2.cvtColor(face_img, cv2.COLOR_BGR2YCrCb)

    hsv_mask = cv2.inRange(
        hsv,
        (0, 40, 60),
        (25, 255, 255)
    )

    ycrcb_mask = cv2.inRange(
        ycrcb,
        (0, 135, 85),
        (255, 180, 135)
    )

    mask = cv2.bitwise_and(hsv_mask, ycrcb_mask)

    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

    return mask


def extract_skin(image_path):
    """
    Extracts skin-only image and mask.
    Works for:
    - Face images (selfies)
    - Partial skin images
    - Skin-only datasets
    """
    # Configuration ------------------------------
    USE_FACE_DETECTION = True
    FACE_REQUIRED = False
    KEEP_ONLY_LARGEST_SKIN_COMPONENT = True
    MAX_SKIN_CONFIDENCE_FILTERING = False
    MAX_SKIN_CONFIDENCE_RATIO = 0.90
    # --------------------------------------------

    # Read image
    image = cv2.imread(image_path)
    if image is None:
        return None, None

    # Try face detection first
    face = detect_face(image)

    if (face is not None) and USE_FACE_DETECTION:
        region = face
    else:
        if FACE_REQUIRED:
            return None, None
        # Fallback to full image
        region = image
    
    # Reject if no skin detected
    if region is None or region.size == 0:
        return None, None

    # Generate skin mask
    mask = skin_color_mask(region)

    # -------------------------------
    # KEEP ONLY LARGEST SKIN COMPONENT
    # -------------------------------
    if KEEP_ONLY_LARGEST_SKIN_COMPONENT:
        num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(
            mask, connectivity=8
        )

        # If only background detected
        if num_labels <= 1:
            return None, None

        # Ignore background label 0
        largest_label = 1 + stats[1:, cv2.CC_STAT_AREA].argmax()
        mask = np.uint8(labels == largest_label) * 255
    # -------------------------------

    # Reject images with insufficient skin pixels
    skin_pixel_ratio = np.sum(mask > 0) / mask.size
    if skin_pixel_ratio < 0.05:
        return None, None
    # Reject images with too many skin pixels
    if MAX_SKIN_CONFIDENCE_FILTERING:
        if skin_pixel_ratio > MAX_SKIN_CONFIDENCE_RATIO:
            # Probably overexposed / false positive
            return None, None

    skin_only = cv2.bitwise_and(region, region, mask=mask)

    return skin_only, mask