133 lines
4.3 KiB
Python
133 lines
4.3 KiB
Python
"""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)
|