winnow/if_curator/image_processing.py

"""Image processing functions for cropping faces and objects."""

import logging
import os

import numpy as np
from PIL import Image

from .config import Config

logger = logging.getLogger(__name__)

# Lazy singleton
_yolo_model = None


def get_yolo_model():
    """Singleton for YOLO model."""
    global _yolo_model
    if _yolo_model is None:
        from ultralytics import YOLO

        logger.info("Loading YOLOv9c model...")
        _yolo_model = YOLO("yolov9c.pt")
    return _yolo_model


def align_face(img: Image.Image, landmarks: list[list[float]] | np.ndarray) -> Image.Image | None:
    """Align face using 5-point landmarks to standard ArcFace input format (112x112).

    This produces a normalized face crop that matches exactly what ArcFace
    was trained on, improving recognition accuracy.

    Args:
        img: Full PIL image containing the face
        landmarks: 5-point facial landmarks [[x,y], ...] (eyes, nose, mouth corners)

    Returns:
        Aligned 112x112 face image, or None if alignment fails
    """
    try:
        from insightface.utils.face_align import norm_crop

        img_np = np.asarray(img)
        lm = np.array(landmarks, dtype=np.float32)
        if lm.shape != (5, 2):
            logger.debug(f"Invalid landmark shape: {lm.shape}, expected (5, 2)")
            return None
        aligned = norm_crop(img_np, lm)
        return Image.fromarray(aligned)
    except ImportError:
        logger.debug("InsightFace not available for face alignment")
        return None
    except Exception as e:
        logger.debug(f"Face alignment failed: {e}")
        return None


def process_face_mode(
    img: Image.Image,
    asset: dict,
    person: dict,
    output_dir: str,
    count: int,
    min_width: int | None = None,
) -> bool:
    """Crop face based on Immich metadata and save to output directory.

    If face alignment is enabled and landmarks are available, produces
    an aligned 112x112 crop. Otherwise falls back to bounding box crop
    with configurable margin.
    """
    min_width = min_width or Config.MIN_FACE_WIDTH

    # Find face metadata for this person
    face_info = None
    for p in asset.get("people", []):
        if p["id"] == person["id"] and (faces := p.get("faces")):
            face_info = faces[0]
            break

    if not face_info:
        logger.debug(f"No face info for {person.get('name')} in asset {asset.get('id')}")
        return False

    img_w, img_h = img.size
    meta_w = face_info.get("imageWidth") or img_w
    meta_h = face_info.get("imageHeight") or img_h

    # Scale bounding box to actual image dimensions
    scale_x, scale_y = img_w / meta_w, img_h / meta_h
    x1 = face_info["boundingBoxX1"] * scale_x
    y1 = face_info["boundingBoxY1"] * scale_y
    x2 = face_info["boundingBoxX2"] * scale_x
    y2 = face_info["boundingBoxY2"] * scale_y

    face_w, face_h = x2 - x1, y2 - y1
    if face_w < min_width or face_h < min_width:
        logger.debug(f"Face too small ({face_w:.1f}x{face_h:.1f})")
        return False

    # Try face alignment if enabled and landmarks available
    if Config.ENABLE_FACE_ALIGNMENT:
        landmarks = face_info.get("landmarks") or face_info.get("landmark")
        if landmarks:
            # Scale landmarks
            scaled_landmarks = [[lm[0] * scale_x, lm[1] * scale_y] for lm in landmarks]
            aligned = align_face(img, scaled_landmarks)
            if aligned is not None:
                aligned.save(os.path.join(output_dir, f"{count}.jpg"), format="JPEG")
                return True

    # Fall back to bounding box crop with configurable margin
    margin = Config.FACE_MARGIN
    margin_x, margin_y = face_w * margin, face_h * margin
    crop_box = (
        max(0, x1 - margin_x),
        max(0, y1 - margin_y),
        min(img_w, x2 + margin_x),
        min(img_h, y2 + margin_y),
    )

    face_crop = img.crop(crop_box)
    face_crop.save(os.path.join(output_dir, f"{count}.jpg"), format="JPEG")
    return True


def process_object_mode(
    img: Image.Image,
    config: dict,
    output_dir: str,
    count: int,
) -> bool:
    """Detect and crop objects using YOLO."""
    try:
        model = get_yolo_model()
        target_class = config.get("object_class", "dog")
        device = "cpu" if os.getenv("FORCE_CPU", "").lower() in ("true", "1", "yes") else None

        results = model(img, verbose=False, device=device)

        found = False
        for idx, (box, cls_id, conf) in enumerate(
            (box, int(box.cls[0]), float(box.conf[0])) for r in results for box in r.boxes
        ):
            if 0 <= cls_id < len(model.names) and model.names[cls_id] == target_class and conf > 0.5:
                x1, y1, x2, y2 = box.xyxy[0].tolist()
                img.crop((x1, y1, x2, y2)).save(
                    os.path.join(output_dir, f"{count}_{idx}.jpg"),
                    format="JPEG",
                )
                found = True

        return found
    except Exception as e:
        logger.error(f"YOLO processing failed: {e}")
        return False


def process_full_mode(img: Image.Image, output_dir: str, count: int) -> bool:
    """Save full image."""
    img.save(os.path.join(output_dir, f"{count}.jpg"), format="JPEG")
    return True