detect.py

import logging
import math

import cv2
import numpy as np

logger = logging.getLogger(__name__)


def _expand_box(x, y, w, h, img_w, img_h, scale=1.2):
    pad_w = int(w * (scale - 1) / 2)
    pad_h = int(h * (scale - 1) / 2)
    x1 = max(0, x - pad_w)
    y1 = max(0, y - pad_h)
    x2 = min(img_w, x + w + pad_w)
    y2 = min(img_h, y + h + pad_h)
    return x1, y1, x2, y2


def _crop_to_face(img):
    cv_img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
    gray = cv2.cvtColor(cv_img, cv2.COLOR_BGR2GRAY)

    face_cascade_path = cv2.data.haarcascades + "haarcascade_frontalface_alt2.xml"
    face_cascade = cv2.CascadeClassifier(face_cascade_path)
    if face_cascade.empty():
        face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + "haarcascade_frontalface_default.xml")

    faces = face_cascade.detectMultiScale(gray, 1.3, 5)
    if len(faces) == 0:
        logger.info("No face detected; using full image for classification.")
        return img

    x, y, w, h = max(faces, key=lambda f: f[2] * f[3])
    img_h, img_w = gray.shape[:2]
    x1, y1, x2, y2 = _expand_box(x, y, w, h, img_w, img_h, scale=1.25)
    return img.crop((x1, y1, x2, y2))


def _gaussian_score(value, target, sigma):
    return math.exp(-((value - target) ** 2) / (2 * sigma ** 2))


def _shape_probabilities_from_geometry(landmarks):
    from geometry import extract_features

    feats = extract_features(landmarks)
    lw_ratio = feats.get("lw_ratio", 0)
    jaw_ratio = feats.get("jaw_ratio", 0)

    if lw_ratio <= 0 or jaw_ratio <= 0:
        return None

    targets = {
        "Round": {"lw": 1.05, "jaw": 0.88},
        "Square": {"lw": 1.12, "jaw": 0.95},
        "Oval": {"lw": 1.25, "jaw": 0.86},
        "Oblong": {"lw": 1.35, "jaw": 0.86},
        "Heart": {"lw": 1.22, "jaw": 0.75},
        "Diamond": {"lw": 1.20, "jaw": 0.70},
    }

    scores = {}
    for shape, target in targets.items():
        lw_score = _gaussian_score(lw_ratio, target["lw"], 0.1)
        jaw_score = _gaussian_score(jaw_ratio, target["jaw"], 0.07)
        scores[shape] = lw_score * jaw_score

    total = sum(scores.values())
    if total <= 0:
        return None

    probabilities = {k: round(v / total, 4) for k, v in scores.items()}
    return dict(sorted(probabilities.items(), key=lambda item: item[1], reverse=True))


def _blend_probabilities(primary, secondary, alpha=0.55):
    labels = set(primary.keys()) | set(secondary.keys())
    blended = {}
    for label in labels:
        blended[label] = alpha * primary.get(label, 0) + (1 - alpha) * secondary.get(label, 0)
    total = sum(blended.values())
    if total <= 0:
        return primary
    blended = {k: round(v / total, 4) for k, v in blended.items()}
    return dict(sorted(blended.items(), key=lambda item: item[1], reverse=True))



def detect_face_shape(image_path):
    """
    Detects face shape using PIL and trained SVM model.
    """
    from PIL import Image
    from landmarks import get_landmarks
    from classifier import classify_face_shape

    try:
        # Load image with PIL
        if isinstance(image_path, str):
            img = Image.open(image_path)
        else:
            # Assume it's already a PIL image or numpy array
            img = image_path

        if img is None:
            raise ValueError("Could not load image")

        # Crop to detected face region to improve accuracy
        if isinstance(image_path, str):
            img = img.convert("RGB")
            img = _crop_to_face(img)

        geom_probs = None
        if isinstance(image_path, str):
            try:
                landmarks = get_landmarks(image_path)
                geom_probs = _shape_probabilities_from_geometry(landmarks)
            except Exception as e:
                logger.info(f"Landmark-based detection failed, falling back: {e}")

        # Classify directly (classifier handles resizing/grayscale)
        shape_probs = classify_face_shape(img)

        if geom_probs:
            geom_values = list(geom_probs.values())
            top = geom_values[0] if geom_values else 0
            runner_up = geom_values[1] if len(geom_values) > 1 else 0
            if top >= 0.45 and (top - runner_up) >= 0.08:
                return geom_probs
            return _blend_probabilities(geom_probs, shape_probs, alpha=0.55)
        
        # Get the top shape
        best_shape = list(shape_probs.keys())[0]
        logger.info(f"Detected face shape: {best_shape} for {image_path}")
        
        return shape_probs
        
    except Exception as e:
        logger.error(f"Detection failed: {e}")
        return {"Error": 1.0}