face_detection.py

import os
import cv2
import numpy as np

MODEL_DIR = os.path.join(os.path.dirname(__file__), "models")
MODEL_PATH = os.path.join(MODEL_DIR, "blaze_face_short_range.tflite")
_detector = None


def _get_detector():
    from mediapipe.tasks.python import BaseOptions
    from mediapipe.tasks.python.vision.face_detector import FaceDetector, FaceDetectorOptions
    from mediapipe.tasks.python.vision.core.image import Image, ImageFormat
    global _detector
    if _detector is None:
        base = BaseOptions(model_asset_path=MODEL_PATH)
        opts = FaceDetectorOptions(base_options=base, min_detection_confidence=0.5)
        _detector = FaceDetector.create_from_options(opts)
    return _detector


#  Skin color fallback 
def _detect_skin_region(image):
    """Detect skin-colored region using HSV. Returns (x, y, w, h) or None."""
    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    lower1 = np.array([0, 20, 70], dtype=np.uint8)
    upper1 = np.array([20, 255, 255], dtype=np.uint8)
    lower2 = np.array([160, 20, 70], dtype=np.uint8)
    upper2 = np.array([180, 255, 255], dtype=np.uint8)

    mask1 = cv2.inRange(hsv, lower1, upper1)
    mask2 = cv2.inRange(hsv, lower2, upper2)
    mask = cv2.bitwise_or(mask1, mask2)

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

    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None

    largest = max(contours, key=cv2.contourArea)
    img_area = image.shape[0] * image.shape[1]
    if cv2.contourArea(largest) < img_area * 0.05:
        return None

    x, y, bw, bh = cv2.boundingRect(largest)
    return (x, y, bw, bh)


#  Helper 
def _crop_region(image, x, y, bw, bh, padding=0.2, shift_up=0.0):
    """Crop region with padding, clamped to image bounds.
    
    shift_up: fraction of face height to shift crop upward 
              (keeps same crop size, moves window up)
    """
    h, w = image.shape[:2]
    pad_x = int(bw * padding)
    pad_y = int(bh * padding)
    shift_px = int(bh * shift_up)

    crop_h = bh + 2 * pad_y  # desired crop height
    crop_w = bw + 2 * pad_x

    # Center on face, then shift up
    cx = x + bw // 2
    cy = y + bh // 2 - shift_px

    x1 = max(0, cx - crop_w // 2)
    y1 = max(0, cy - crop_h // 2)
    x2 = min(w, x1 + crop_w)
    y2 = min(h, y1 + crop_h)

    # Re-adjust if clipped at top boundary
    if y1 == 0 and y2 - y1 < crop_h:
        y2 = min(h, crop_h)

    return image[y1:y2, x1:x2]


#  Orientation 
def _classify_orientation(face_result, detection=None):
    """Klasifikasi orientasi wajah: 'frontal' atau 'side_profile'.
    
    Menggunakan keypoints MediaPipe (mata/telinga) jika tersedia,
    fallback ke aspect ratio bounding box.
    """
    x, y, w, h = face_result["bounds"]
    aspect = h / max(w, 1)

    # Jika ada keypoints dari MediaPipe
    if detection is not None and hasattr(detection, 'keypoints') and len(detection.keypoints) >= 6:
        kp = detection.keypoints
        # keypoints: 0=right_eye, 1=left_eye, 2=nose, 3=mouth, 4=right_ear, 5=left_ear
        reye, leye = kp[0], kp[1]
        rear, lear = kp[4], kp[5]
        # Jarak horizontal mata (normalized 0-1)
        eye_dist = abs(leye.x - reye.x)
        # Frontal: kedua mata terpisah lebar, ear di sisi luar
        if eye_dist > 0.12:
            return "frontal"
        # Side profile: mata berdekatan, atau ear mendekati center
        if eye_dist < 0.08 or abs(rear.x - lear.x) < 0.03:
            return "side_profile"

    # Fallback: aspect ratio
    # Frontal ~1.0-1.6, side profile ~1.6-2.5
    return "side_profile" if aspect > 1.7 else "frontal"


#  Public API 
def detect_face(image):
    """Detect face in image and return info dict or None.

    Returns:
        dict with keys:
            - bounds: (x, y, w, h) face bounding box on original image
            - score: confidence score (0-1)
            - method: "mediapipe" | "skin"
            - orientation: "frontal" | "side_profile" (mediapipe only)
        or None if no face found.
    """
    # Step 1: MediaPipe Face Detection
    try:
        from mediapipe.tasks.python.vision.core.image import Image, ImageFormat
        detector = _get_detector()
        rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        mp_img = Image(image_format=ImageFormat.SRGB, data=rgb)
        result = detector.detect(mp_img)
        if result.detections:
            best = max(result.detections, key=lambda d: d.categories[0].score)
            bb = best.bounding_box
            x = int(bb.origin_x)
            y = int(bb.origin_y)
            fw = int(bb.width)
            fh = int(bb.height)
            face_result = {
                "bounds": (x, y, fw, fh),
                "score": float(best.categories[0].score),
                "method": "mediapipe",
            }
            face_result["orientation"] = _classify_orientation(face_result, detection=best)
            return face_result
    except Exception as e:
        print(f"[face_detection] MediaPipe failed: {e}")

    # Step 2: Skin color fallback
    skin = _detect_skin_region(image)
    if skin is not None:
        x, y, fw, fh = int(skin[0]), int(skin[1]), int(skin[2]), int(skin[3])
        return {"bounds": (x, y, fw, fh), "score": 1.0, "method": "skin", "orientation": "frontal"}

    return None


def draw_face_box(image, face_result, color=(0, 255, 0), thickness=3):
    """Draw face bounding box on a copy of the image."""
    result = image.copy()
    if face_result is None:
        return result
    x, y, w, h = face_result["bounds"]
    cv2.rectangle(result, (x, y), (x + w, y + h), color, thickness)

    orientation = face_result.get("orientation", "?")
    label = f"face ({face_result['method']}) {orientation} {face_result['score']:.2f}"
    (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
    cv2.rectangle(result, (x, y - th - 6), (x + tw + 4, y), color, -1)
    cv2.putText(result, label, (x + 2, y - 3),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)

    return result


def crop_face(image, padding=0.2, shift_up=0.15,
              padding_side=0.30, shift_up_side=0.15):
    """Detect face area and crop image with orientation-aware settings.

    Args:
        image: numpy array (BGR image from cv2.imread)
        padding: face padding for frontal faces
        shift_up: face shift for frontal faces
        padding_side: face padding for side profiles
        shift_up_side: face shift for side profiles

    Returns:
        Cropped face image, or original image if nothing detected.
    """
    face_result = detect_face(image)
    if face_result is None:
        return image

    orientation = face_result.get("orientation", "frontal")
    if orientation == "side_profile":
        padding = padding_side
        shift_up = shift_up_side

    x, y, fw, fh = face_result["bounds"]
    return _crop_region(image, x, y, fw, fh, padding, shift_up=shift_up)