"""Image quality assessment for training data curation. Filters out images that would hurt Frigate's ArcFace model training: blur, grayscale/IR, over/underexposure, tiny faces, low-confidence detections. """ import logging from dataclasses import dataclass, field import cv2 import numpy as np from PIL import Image logger = logging.getLogger(__name__) @dataclass class QualityResult: """Result of quality assessment on a face/image crop.""" passed: bool reasons: list[str] = field(default_factory=list) @property def reason(self) -> str: return "; ".join(self.reasons) if self.reasons else "OK" def check_blur(img_np: np.ndarray, threshold: float = 100.0) -> tuple[bool, str]: """Detect blur using Laplacian variance. Lower variance = blurrier image. ArcFace needs clear facial features. """ gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY) if img_np.ndim == 3 else img_np variance = cv2.Laplacian(gray, cv2.CV_64F).var() if variance < threshold: return False, f"Blurry (laplacian={variance:.1f}, threshold={threshold})" return True, "" def check_grayscale(img_np: np.ndarray) -> tuple[bool, str]: """Detect grayscale/IR images by checking channel similarity. ArcFace is trained on color images; IR/grayscale degrades recognition. """ if img_np.ndim != 3 or img_np.shape[2] < 3: return False, "Grayscale (single channel)" # Compare channel means — IR/grayscale has nearly identical R, G, B means = img_np[:, :, :3].mean(axis=(0, 1)) max_diff = max(abs(means[0] - means[1]), abs(means[1] - means[2]), abs(means[0] - means[2])) if max_diff < 5.0: return False, f"Grayscale/IR (channel diff={max_diff:.1f})" return True, "" def check_exposure(img_np: np.ndarray, lo: float = 30.0, hi: float = 225.0) -> tuple[bool, str]: """Check for severe under/overexposure using mean brightness. Extremely dark or blown-out faces lack usable features. """ gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY) if img_np.ndim == 3 else img_np mean_brightness = gray.mean() if mean_brightness < lo: return False, f"Underexposed (brightness={mean_brightness:.1f}, min={lo})" if mean_brightness > hi: return False, f"Overexposed (brightness={mean_brightness:.1f}, max={hi})" return True, "" def check_face_size(face_width: float, face_height: float, min_px: int = 100) -> tuple[bool, str]: """Validate face crop is large enough for meaningful features.""" if face_width < min_px or face_height < min_px: return False, f"Face too small ({face_width:.0f}x{face_height:.0f}, min={min_px}px)" return True, "" def check_confidence(score: float | None, min_conf: float = 0.7) -> tuple[bool, str]: """Check detection confidence from Immich. Low confidence often means partial, occluded, or false-positive faces. """ if score is None: return True, "" # No score available, pass by default if score < min_conf: return False, f"Low confidence ({score:.2f}, min={min_conf})" return True, "" def assess_quality( img: Image.Image, face_bbox: tuple[float, float, float, float] | None = None, confidence: float | None = None, blur_threshold: float = 100.0, min_face_px: int = 100, min_confidence: float = 0.7, ) -> QualityResult: """Run all quality checks on an image. Args: img: PIL Image (RGB) face_bbox: (x1, y1, x2, y2) bounding box, or None to skip face-size check confidence: Detection confidence score from Immich, or None blur_threshold: Laplacian variance threshold for blur detection min_face_px: Minimum face dimension in pixels min_confidence: Minimum acceptable detection confidence Returns: QualityResult with passed=True if all checks pass """ img_np = np.asarray(img) reasons = [] # Run all checks, collect failures checks = [ check_blur(img_np, blur_threshold), check_grayscale(img_np), check_exposure(img_np), check_confidence(confidence, min_confidence), ] if face_bbox is not None: x1, y1, x2, y2 = face_bbox checks.append(check_face_size(x2 - x1, y2 - y1, min_face_px)) for passed, reason in checks: if not passed: reasons.append(reason) return QualityResult(passed=len(reasons) == 0, reasons=reasons)