feat: Implement advanced diversity selection algorithms, comprehensive quality filtering, and caching for improved image curation.

if_curator/image_processing.py

@@ -3,8 +3,11 @@
 import logging
 import os
 
+import numpy as np
 from PIL import Image
 
+from .config import Config
+
 logger = logging.getLogger(__name__)
 
 # Lazy singleton
@@ -16,20 +19,59 @@ def get_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 = 50,
+    min_width: int | None = None,
 ) -> bool:
-    """Crop face based on Immich metadata and save to output directory."""
+    """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", []):
@@ -57,8 +99,20 @@ def process_face_mode(
         logger.debug(f"Face too small ({face_w:.1f}x{face_h:.1f})")
         return False
 
-    # Add 10% margin
-    margin_x, margin_y = face_w * 0.10, face_h * 0.10
+    # 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),
@@ -87,9 +141,7 @@ def process_object_mode(
 
         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
+            (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()