flan
341c6b0e85
Immich's /api/faces endpoint only returns bounding boxes, not facial landmarks. This meant align_face() never fired and all crops fell back to a plain bbox rectangle — producing partial crops (forehead-only, off-angle faces) when Immich's detection was slightly off. Now, when ENABLE_FACE_ALIGNMENT is true and InsightFace is loaded, execute_jobs() passes the app to process_face_mode(). For each face, it expands the Immich bbox by 50%, crops that search region, runs InsightFace detection within it, and aligns the nearest face to the standard ArcFace 112×112 format using norm_crop(). Falls back to bbox crop if InsightFace finds no face in the search region. The InsightFace model is already in GPU memory from the diversity/ embedding phase, so the singleton lookup adds no load cost.
219 lines
7.5 KiB
Python
219 lines
7.5 KiB
Python
"""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 _save_jpeg(img: Image.Image, path: str) -> None:
|
|
if img.mode != "RGB":
|
|
img = img.convert("RGB")
|
|
img.save(path, format="JPEG")
|
|
|
|
|
|
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,
|
|
insightface_app=None,
|
|
) -> tuple[int, int] | None:
|
|
"""Crop face based on Immich metadata and save to output directory.
|
|
|
|
Returns (width, height) of the saved crop, or None if no crop was saved.
|
|
When insightface_app is provided and ENABLE_FACE_ALIGNMENT is True,
|
|
re-detects the face in the Immich bbox region using InsightFace to get
|
|
precise landmarks for a proper 112x112 aligned crop. Falls back to
|
|
bounding box crop with configurable margin if alignment is unavailable.
|
|
"""
|
|
min_width = min_width or Config.MIN_FACE_WIDTH
|
|
|
|
# Find face metadata for this person
|
|
face_info = None
|
|
people_list = asset.get("people") or []
|
|
for p in people_list:
|
|
if p is None:
|
|
continue
|
|
if p.get("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 None
|
|
|
|
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 None
|
|
|
|
# Re-detect face with InsightFace for landmark-based alignment.
|
|
# Immich's /api/faces endpoint does not include landmarks, so the
|
|
# align_face fallback below never fires without this step.
|
|
if insightface_app is not None and Config.ENABLE_FACE_ALIGNMENT:
|
|
try:
|
|
# Expand the Immich bbox by 50% to give InsightFace enough context
|
|
# for detection and alignment, then search for the face nearest the
|
|
# centre of that region (handles group photos at the boundary).
|
|
pad_x, pad_y = face_w * 0.5, face_h * 0.5
|
|
search_box = (
|
|
max(0, x1 - pad_x),
|
|
max(0, y1 - pad_y),
|
|
min(img_w, x2 + pad_x),
|
|
min(img_h, y2 + pad_y),
|
|
)
|
|
search_crop = img.crop(search_box)
|
|
detected = insightface_app.get(np.asarray(search_crop))
|
|
if detected:
|
|
cx, cy = search_crop.width / 2, search_crop.height / 2
|
|
best = min(
|
|
detected,
|
|
key=lambda f: abs((f.bbox[0] + f.bbox[2]) / 2 - cx)
|
|
+ abs((f.bbox[1] + f.bbox[3]) / 2 - cy),
|
|
)
|
|
kps = getattr(best, "kps", None)
|
|
if kps is not None and np.asarray(kps).shape == (5, 2):
|
|
aligned = align_face(search_crop, kps)
|
|
if aligned is not None:
|
|
_save_jpeg(aligned, os.path.join(output_dir, f"{count}.jpg"))
|
|
return aligned.size
|
|
except Exception as e:
|
|
logger.debug(f"InsightFace re-detection failed for {asset.get('id')}: {e}")
|
|
|
|
# Landmark alignment from Immich metadata (Immich does not currently
|
|
# expose landmarks, so this path is a future-proofing fallback)
|
|
if Config.ENABLE_FACE_ALIGNMENT:
|
|
landmarks = face_info.get("landmarks") or face_info.get("landmark")
|
|
if 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:
|
|
_save_jpeg(aligned, os.path.join(output_dir, f"{count}.jpg"))
|
|
return aligned.size
|
|
|
|
# Final fallback: 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)
|
|
_save_jpeg(face_crop, os.path.join(output_dir, f"{count}.jpg"))
|
|
return face_crop.size
|
|
|
|
|
|
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")
|
|
import torch
|
|
|
|
if os.getenv("FORCE_CPU", "").lower() in ("true", "1", "yes"):
|
|
device = "cpu"
|
|
elif hasattr(torch, "xpu") and torch.xpu.is_available():
|
|
device = "xpu"
|
|
else:
|
|
device = None # YOLO auto-selects (CUDA/ROCm/CPU)
|
|
|
|
results = model(img, verbose=False, device=device)
|
|
|
|
found = False
|
|
class_idx = 0 # Sequential counter per target class (Issue #10)
|
|
for box in (box for r in results for box in r.boxes):
|
|
cls_id = int(box.cls[0])
|
|
conf = float(box.conf[0])
|
|
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()
|
|
_save_jpeg(
|
|
img.crop((x1, y1, x2, y2)),
|
|
os.path.join(output_dir, f"{count}_{class_idx}.jpg"),
|
|
)
|
|
class_idx += 1
|
|
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."""
|
|
_save_jpeg(img, os.path.join(output_dir, f"{count}.jpg"))
|
|
return True
|
|
|