refactor: rename project to winnow
- Rename Python package directory if_curator/ → winnow/ - Update all imports, entry points, and CLI references - Update pyproject.toml: name, scripts, package list, repository URL - Update Dockerfile, compose.yml, entrypoint.sh, scheduler.py - Update GitHub Actions workflow image names - Update README Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Claude Sonnet 4.6
@@ -0,0 +1,8 @@
|
||||
"""Immich to Frigate training set curator.
|
||||
|
||||
AI-powered tool to extract high-quality, diverse training images from your
|
||||
Immich library for Frigate's Face Recognition (ArcFace) and Object/State
|
||||
Classification models.
|
||||
"""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
@@ -0,0 +1,4 @@
|
||||
from winnow.cli import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,94 @@
|
||||
"""Disk-based embedding cache.
|
||||
|
||||
Caches embeddings keyed by (asset_id, model_version) to avoid
|
||||
recomputing on reruns. Uses numpy binary format for fast I/O.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Model versions — bump these when the upstream model changes
|
||||
MODEL_VERSIONS = {
|
||||
"insightface": "buffalo_l_v1",
|
||||
"siglip": "siglip-base-patch16-224_v1",
|
||||
"immich": "immich_buffalo_l_v1",
|
||||
}
|
||||
|
||||
|
||||
class EmbeddingCache:
|
||||
"""Simple disk-based embedding cache.
|
||||
|
||||
Embeddings are stored as .npy files in a flat directory,
|
||||
keyed by a hash of (asset_id, model_version).
|
||||
"""
|
||||
|
||||
def __init__(self, cache_dir: str = ".if_cache") -> None:
|
||||
self.cache_dir = cache_dir
|
||||
self._ensured = False
|
||||
|
||||
def _ensure_dir(self) -> None:
|
||||
if not self._ensured:
|
||||
os.makedirs(self.cache_dir, exist_ok=True)
|
||||
self._ensured = True
|
||||
|
||||
@staticmethod
|
||||
def _key(asset_id: str, model: str) -> str:
|
||||
version = MODEL_VERSIONS.get(model, model)
|
||||
raw = f"{asset_id}:{version}"
|
||||
return hashlib.sha256(raw.encode()).hexdigest()[:16]
|
||||
|
||||
def _path(self, asset_id: str, model: str) -> str:
|
||||
return os.path.join(self.cache_dir, f"{self._key(asset_id, model)}.npy")
|
||||
|
||||
def get(self, asset_id: str, model: str = "insightface") -> np.ndarray | None:
|
||||
"""Retrieve cached embedding, or None if not cached."""
|
||||
path = self._path(asset_id, model)
|
||||
if os.path.exists(path):
|
||||
try:
|
||||
return np.load(path)
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
def put(self, asset_id: str, embedding: np.ndarray, model: str = "insightface") -> None:
|
||||
"""Store an embedding in the cache."""
|
||||
self._ensure_dir()
|
||||
try:
|
||||
np.save(self._path(asset_id, model), embedding)
|
||||
except Exception as e:
|
||||
logger.debug(f"Cache write failed for {asset_id}: {e}")
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Delete all cached embeddings."""
|
||||
if not os.path.isdir(self.cache_dir):
|
||||
return
|
||||
count = 0
|
||||
for f in os.listdir(self.cache_dir):
|
||||
if f.endswith(".npy"):
|
||||
os.remove(os.path.join(self.cache_dir, f))
|
||||
count += 1
|
||||
logger.info(f"Cleared {count} cached embeddings.")
|
||||
|
||||
|
||||
# Singleton instance
|
||||
_cache: EmbeddingCache | None = None
|
||||
|
||||
|
||||
def get_cache(cache_dir: str = ".if_cache") -> EmbeddingCache:
|
||||
"""Get or create the singleton cache instance.
|
||||
|
||||
Note: The ``cache_dir`` parameter is only used when creating the
|
||||
singleton for the first time. Subsequent calls return the existing
|
||||
instance regardless of ``cache_dir``. If you need a cache with a
|
||||
different directory, instantiate ``EmbeddingCache`` directly.
|
||||
"""
|
||||
global _cache
|
||||
if _cache is None:
|
||||
_cache = EmbeddingCache(cache_dir)
|
||||
return _cache
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
"""Interactive CLI for winnow."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from rich import print as rprint
|
||||
from rich.prompt import Confirm
|
||||
|
||||
from .config import Config, ConfigManager
|
||||
from .executor import execute_jobs, upload_to_frigate
|
||||
from .immich_api import get_people
|
||||
from .jobs import _show_preview, auto_configure, interactive_configure
|
||||
from .logging import console, setup_logging
|
||||
from .upload_tracker import get_person_summary, reset_person
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Entry point for winnow CLI."""
|
||||
try:
|
||||
setup_logging(verbose=False)
|
||||
|
||||
console.print(r"""
|
||||
[bold blue]winnow[/bold blue]
|
||||
[dim]Immich -> Frigate Training Data Curator[/dim]
|
||||
""")
|
||||
|
||||
ConfigManager.get().interactive_setup()
|
||||
|
||||
try:
|
||||
Config.validate()
|
||||
except ValueError as e:
|
||||
rprint(f"[bold red]Configuration Error:[/bold red] {e}")
|
||||
return
|
||||
|
||||
rprint(f"Server: [dim]{Config.IMMICH_URL}[/dim]")
|
||||
rprint(f"Output: [dim]{Config.OUTPUT_DIR}[/dim]")
|
||||
|
||||
# Handle RESET_PERSON before anything else
|
||||
reset_person_name = os.environ.get("RESET_PERSON", "").strip()
|
||||
if reset_person_name:
|
||||
reset_person(reset_person_name)
|
||||
rprint(f"[bold yellow]Reset tracking data for: {reset_person_name}[/bold yellow]")
|
||||
|
||||
# Show per-person tracker summary if data exists
|
||||
summary = get_person_summary()
|
||||
if summary:
|
||||
rprint("\n[dim]Tracker summary:[/dim]")
|
||||
for person_name, counts in summary.items():
|
||||
rprint(f" [dim]{person_name}: {counts['uploaded']} uploaded, {counts['rejected']} rejected[/dim]")
|
||||
|
||||
people = get_people()
|
||||
if not people:
|
||||
rprint("[bold red]Could not fetch people from Immich. Check URL/Key.[/bold red]")
|
||||
return
|
||||
|
||||
# Check for non-interactive mode
|
||||
auto_mode = os.environ.get("AUTO_MODE", "false").lower() == "true"
|
||||
dry_run = os.environ.get("DRY_RUN", "false").lower() in ("true", "1", "yes")
|
||||
|
||||
if dry_run:
|
||||
rprint("[bold yellow]DRY RUN — no images will be downloaded or uploaded[/bold yellow]")
|
||||
|
||||
if auto_mode:
|
||||
rprint("[bold cyan]Running in AUTO mode (non-interactive)[/bold cyan]")
|
||||
jobs = auto_configure(people)
|
||||
else:
|
||||
jobs = interactive_configure(people)
|
||||
|
||||
if jobs:
|
||||
_show_preview(jobs)
|
||||
if dry_run:
|
||||
rprint("\n[bold yellow]Dry run complete — skipping execute and upload.[/bold yellow]")
|
||||
elif auto_mode or Confirm.ask(f"Ready to process {sum(j['limit'] for j in jobs)} images?"):
|
||||
execute_jobs(jobs)
|
||||
upload_to_frigate(jobs)
|
||||
rprint("\n[bold green]Done! Happy Training.[/bold green]")
|
||||
else:
|
||||
rprint("[yellow]No jobs configured.[/yellow]")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
rprint("\n[bold red]Aborted by user.[/bold red]")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,168 @@
|
||||
"""Configuration management for winnow."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import ClassVar
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from rich.prompt import Prompt
|
||||
|
||||
load_dotenv()
|
||||
|
||||
CONFIG_FILE = Path(".immich_config.json")
|
||||
|
||||
|
||||
class _Config:
|
||||
"""Singleton configuration with uppercase attribute access for backward compatibility."""
|
||||
|
||||
_instance: ClassVar["_Config | None"] = None
|
||||
|
||||
# Configuration values
|
||||
IMMICH_URL: str | None = None
|
||||
API_KEY: str | None = None
|
||||
OUTPUT_DIR: str = "./frigate_train"
|
||||
YEARS_FILTER: int = 10
|
||||
|
||||
# Quality filtering
|
||||
MIN_FACE_WIDTH: int = 50
|
||||
BLUR_THRESHOLD: float = 100.0
|
||||
MIN_CONFIDENCE: float = 0.7
|
||||
MAX_AUTO_IMAGES: int = 80
|
||||
|
||||
# People filtering
|
||||
MIN_FACE_COUNT: int = 0
|
||||
|
||||
# Output quality
|
||||
FACE_MARGIN: float = 0.15
|
||||
USE_FULL_RESOLUTION: bool = True
|
||||
ENABLE_FACE_ALIGNMENT: bool = True
|
||||
|
||||
# Caching (opt-in to avoid unexpected files)
|
||||
ENABLE_CACHE: bool = False
|
||||
CACHE_DIR: str = ".if_cache"
|
||||
|
||||
def __new__(cls) -> "_Config":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._load()
|
||||
return cls._instance
|
||||
|
||||
def _load(self) -> None:
|
||||
"""Load configuration from environment and config file."""
|
||||
# Load from environment (highest priority)
|
||||
self.IMMICH_URL = os.getenv("IMMICH_URL")
|
||||
self.API_KEY = os.getenv("API_KEY")
|
||||
self.OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./frigate_train")
|
||||
self.YEARS_FILTER = int(os.getenv("YEARS_FILTER", "10"))
|
||||
self.MIN_FACE_WIDTH = int(os.getenv("MIN_FACE_WIDTH", "50"))
|
||||
self.MIN_FACE_COUNT = int(os.getenv("MIN_FACE_COUNT", "0"))
|
||||
self.BLUR_THRESHOLD = float(os.getenv("BLUR_THRESHOLD", "100.0"))
|
||||
self.MIN_CONFIDENCE = float(os.getenv("MIN_CONFIDENCE", "0.7"))
|
||||
self.MAX_AUTO_IMAGES = int(os.getenv("MAX_AUTO_IMAGES", "80"))
|
||||
self.FACE_MARGIN = float(os.getenv("FACE_MARGIN", "0.15"))
|
||||
self.USE_FULL_RESOLUTION = os.getenv("USE_FULL_RESOLUTION", "true").lower() in ("true", "1", "yes")
|
||||
self.ENABLE_FACE_ALIGNMENT = os.getenv("ENABLE_FACE_ALIGNMENT", "true").lower() in ("true", "1", "yes")
|
||||
self.ENABLE_CACHE = os.getenv("ENABLE_CACHE", "false").lower() in ("true", "1", "yes")
|
||||
self.CACHE_DIR = os.getenv("CACHE_DIR", ".if_cache")
|
||||
|
||||
# Fall back to config file for missing values
|
||||
if CONFIG_FILE.exists():
|
||||
try:
|
||||
data = json.loads(CONFIG_FILE.read_text())
|
||||
self.IMMICH_URL = self.IMMICH_URL or data.get("IMMICH_URL")
|
||||
self.API_KEY = self.API_KEY or data.get("API_KEY")
|
||||
if not os.getenv("OUTPUT_DIR"):
|
||||
self.OUTPUT_DIR = data.get("OUTPUT_DIR", self.OUTPUT_DIR)
|
||||
except (json.JSONDecodeError, OSError) as e:
|
||||
logging.warning(f"Failed to load config file: {e}")
|
||||
|
||||
@classmethod
|
||||
def reset(cls) -> None:
|
||||
"""Reset the singleton — mainly useful for testing or delayed env setup."""
|
||||
cls._instance = None
|
||||
|
||||
def save(self) -> None:
|
||||
"""Persist configuration to file."""
|
||||
try:
|
||||
CONFIG_FILE.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"IMMICH_URL": self.IMMICH_URL,
|
||||
"API_KEY": self.API_KEY,
|
||||
"OUTPUT_DIR": self.OUTPUT_DIR,
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
logging.info(f"Configuration saved to {CONFIG_FILE}")
|
||||
except OSError as e:
|
||||
logging.error(f"Failed to save config: {e}")
|
||||
|
||||
def interactive_setup(self) -> None:
|
||||
"""Prompt user for missing configuration."""
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
if not self.IMMICH_URL:
|
||||
console.print("[yellow]Immich URL not found.[/yellow]")
|
||||
self.IMMICH_URL = Prompt.ask("Enter Immich URL (e.g. http://192.168.1.5:2283)")
|
||||
self.save()
|
||||
|
||||
if not self.API_KEY:
|
||||
console.print("[yellow]Immich API Key not found.[/yellow]")
|
||||
self.API_KEY = Prompt.ask("Enter Immich API Key", password=True)
|
||||
self.save()
|
||||
|
||||
def validate(self) -> None:
|
||||
"""Raise ValueError if required config is missing."""
|
||||
if not self.IMMICH_URL or not self.API_KEY:
|
||||
raise ValueError("Missing Immich URL or API Key.")
|
||||
|
||||
|
||||
# Singleton instance — use a lazy property pattern to avoid import-time side effects
|
||||
# when env vars aren't yet set. Call Config.instance() or just access attributes on
|
||||
# the module-level `Config` (which delegates to the singleton).
|
||||
class _ConfigAccessor:
|
||||
"""Lazy accessor that defers singleton creation until first attribute access.
|
||||
|
||||
This avoids reading .env and config files at import time, so environment
|
||||
variables set after importing the module are properly picked up.
|
||||
"""
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
return getattr(_Config(), name)
|
||||
|
||||
def __setattr__(self, name: str, value):
|
||||
if name.startswith("_"):
|
||||
super().__setattr__(name, value)
|
||||
else:
|
||||
setattr(_Config(), name, value)
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset the underlying singleton."""
|
||||
_Config.reset()
|
||||
|
||||
def interactive_setup(self) -> None:
|
||||
"""Delegate to the singleton."""
|
||||
_Config().interactive_setup()
|
||||
|
||||
def validate(self) -> None:
|
||||
"""Delegate to the singleton."""
|
||||
_Config().validate()
|
||||
|
||||
def save(self) -> None:
|
||||
"""Delegate to the singleton."""
|
||||
_Config().save()
|
||||
|
||||
|
||||
Config = _ConfigAccessor()
|
||||
ConfigManager = type("ConfigManager", (), {"get": staticmethod(lambda: _Config())})
|
||||
|
||||
|
||||
def get_headers() -> dict[str, str]:
|
||||
"""Return HTTP headers for Immich API requests."""
|
||||
return {"x-api-key": Config.API_KEY or "", "Accept": "application/json"}
|
||||
|
||||
@@ -0,0 +1,481 @@
|
||||
"""
|
||||
Diversity selection for training data curation.
|
||||
|
||||
Selection pipeline:
|
||||
1. Concurrent thumbnail download
|
||||
2. Quality filtering (blur, IR, exposure, confidence, face size)
|
||||
3. Face crop extraction (embed person's face, not full image)
|
||||
4. Embedding computation (InsightFace or SigLIP)
|
||||
5. Cluster-aware selection (K-Medoids + FPS with hard example weighting)
|
||||
"""
|
||||
|
||||
import logging
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
from .config import Config, get_headers
|
||||
from .embeddings import get_embedding, is_embedding_available
|
||||
from .quality import assess_quality
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def select_diverse_assets(
|
||||
assets: list,
|
||||
limit: int | str,
|
||||
entity_name: str,
|
||||
selection_mode: str = "smart",
|
||||
entity_type: str = "face",
|
||||
progress_callback=None,
|
||||
) -> list:
|
||||
"""
|
||||
Select diverse assets using cluster-aware FPS or time spread.
|
||||
|
||||
Args:
|
||||
assets: List of asset dicts from Immich API
|
||||
limit: Number to select, or "auto" for dynamic selection
|
||||
entity_name: Name of the person/object for logging
|
||||
selection_mode: 'smart' (embedding-based) or 'time' (time spread)
|
||||
entity_type: 'face' or 'object' - determines embedding model
|
||||
progress_callback: Optional callback(current, total) for progress
|
||||
|
||||
Returns:
|
||||
List of selected assets
|
||||
"""
|
||||
# Fast path: fewer assets than limit
|
||||
if limit != "auto" and len(assets) <= limit:
|
||||
return assets
|
||||
|
||||
# Sort by creation time
|
||||
assets = sorted(assets, key=lambda x: x.get("fileCreatedAt", ""))
|
||||
|
||||
if selection_mode != "smart" or not is_embedding_available(entity_type):
|
||||
if selection_mode == "smart":
|
||||
model_name = "InsightFace" if entity_type == "face" else "SigLIP"
|
||||
logger.warning(f"{model_name} unavailable. Falling back to time spread.")
|
||||
return _select_time_spread(assets, limit)
|
||||
|
||||
try:
|
||||
return _select_by_embedding(assets, limit, entity_type, progress_callback)
|
||||
except Exception as e:
|
||||
logger.error(f"Smart Diversity failed: {e}. Falling back to time spread.")
|
||||
return _select_time_spread(assets, limit)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Thumbnail & Metadata Helpers
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _fetch_thumbnail(asset_id: str, timeout: int = 10) -> Image.Image | None:
|
||||
"""Fetch thumbnail from Immich API."""
|
||||
try:
|
||||
url = f"{Config.IMMICH_URL}/api/assets/{asset_id}/thumbnail?size=preview&format=JPEG"
|
||||
resp = requests.get(url, headers=get_headers(), timeout=timeout)
|
||||
return Image.open(BytesIO(resp.content)) if resp.ok else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _get_face_bbox(asset: dict) -> tuple[float, float, float, float] | None:
|
||||
"""Extract face bounding box from asset metadata if available."""
|
||||
for person in asset.get("people", []):
|
||||
faces = person.get("faces", [])
|
||||
if faces:
|
||||
f = faces[0]
|
||||
return (
|
||||
f.get("boundingBoxX1", 0),
|
||||
f.get("boundingBoxY1", 0),
|
||||
f.get("boundingBoxX2", 0),
|
||||
f.get("boundingBoxY2", 0),
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _get_face_confidence(asset: dict) -> float | None:
|
||||
"""Extract face detection confidence from asset metadata if available."""
|
||||
for person in asset.get("people", []):
|
||||
faces = person.get("faces", [])
|
||||
if faces:
|
||||
return faces[0].get("score") or faces[0].get("confidence")
|
||||
return None
|
||||
|
||||
|
||||
def _crop_face_from_thumbnail(
|
||||
img: Image.Image,
|
||||
asset: dict,
|
||||
margin: float = 0.25,
|
||||
) -> Image.Image | None:
|
||||
"""Crop the face region from a thumbnail using Immich bbox metadata.
|
||||
|
||||
By cropping before embedding, we guarantee InsightFace embeds the
|
||||
correct person's face (not the largest face in a group photo).
|
||||
|
||||
Args:
|
||||
img: Full preview thumbnail
|
||||
asset: Asset dict with people/faces metadata
|
||||
margin: Extra margin around the bbox (fraction, default 25%)
|
||||
|
||||
Returns:
|
||||
Cropped face PIL image, or None if no face metadata available
|
||||
"""
|
||||
bbox = _get_face_bbox(asset)
|
||||
if bbox is None:
|
||||
return None
|
||||
|
||||
x1, y1, x2, y2 = bbox
|
||||
img_w, img_h = img.size
|
||||
|
||||
# Get metadata dimensions to scale bbox
|
||||
for person in asset.get("people", []):
|
||||
faces = person.get("faces", [])
|
||||
if faces:
|
||||
meta_w = faces[0].get("imageWidth") or img_w
|
||||
meta_h = faces[0].get("imageHeight") or img_h
|
||||
scale_x, scale_y = img_w / meta_w, img_h / meta_h
|
||||
x1, y1 = x1 * scale_x, y1 * scale_y
|
||||
x2, y2 = x2 * scale_x, y2 * scale_y
|
||||
break
|
||||
|
||||
face_w, face_h = x2 - x1, y2 - y1
|
||||
|
||||
# Add margin so InsightFace's internal alignment has context
|
||||
mx, my = face_w * margin, face_h * margin
|
||||
crop = img.crop(
|
||||
(
|
||||
max(0, x1 - mx),
|
||||
max(0, y1 - my),
|
||||
min(img_w, x2 + mx),
|
||||
min(img_h, y2 + my),
|
||||
)
|
||||
)
|
||||
|
||||
# Skip if too small for meaningful embedding
|
||||
if crop.width < 30 or crop.height < 30:
|
||||
return None
|
||||
|
||||
return crop
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Embedding Collection
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _select_by_embedding(
|
||||
assets: list,
|
||||
limit: int | str,
|
||||
entity_type: str,
|
||||
progress_callback=None,
|
||||
) -> list:
|
||||
"""Select assets using embedding-based cluster-aware FPS.
|
||||
|
||||
Pipeline:
|
||||
1. Concurrent thumbnail download
|
||||
2. Quality filtering
|
||||
3. Face crop extraction (face mode only)
|
||||
4. Embedding computation
|
||||
5. Cluster-aware selection with hard example weighting
|
||||
"""
|
||||
# Determine candidate pool (cap at 3000 for performance)
|
||||
effective_limit = 30 if limit == "auto" else limit
|
||||
pool_size = min(3000, max(effective_limit * 20, len(assets)))
|
||||
|
||||
# Subsample if needed (evenly distributed in time)
|
||||
if len(assets) > pool_size:
|
||||
indices = np.linspace(0, len(assets) - 1, pool_size, dtype=int)
|
||||
candidates = [assets[i] for i in indices]
|
||||
else:
|
||||
candidates = assets
|
||||
|
||||
# --- Phase 1: Concurrent thumbnail download ---
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
thumbnail_map: dict[str, Image.Image] = {}
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
futures = {pool.submit(_fetch_thumbnail, a["id"]): a for a in candidates}
|
||||
for i, future in enumerate(as_completed(futures)):
|
||||
if progress_callback:
|
||||
progress_callback(i, len(candidates))
|
||||
asset = futures[future]
|
||||
try:
|
||||
img = future.result()
|
||||
if img is not None:
|
||||
thumbnail_map[asset["id"]] = img
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
# --- Phase 2-4: Quality filter → Crop → Embed ---
|
||||
embeddings, valid_candidates, confidence_scores = [], [], []
|
||||
quality_filtered = 0
|
||||
|
||||
for asset in candidates:
|
||||
img = thumbnail_map.get(asset["id"])
|
||||
if img is None:
|
||||
continue
|
||||
|
||||
confidence = _get_face_confidence(asset)
|
||||
|
||||
# Quality gate: filter before expensive embedding computation
|
||||
if entity_type == "face":
|
||||
face_bbox = _get_face_bbox(asset)
|
||||
quality = assess_quality(
|
||||
img,
|
||||
face_bbox=face_bbox,
|
||||
confidence=confidence,
|
||||
blur_threshold=Config.BLUR_THRESHOLD,
|
||||
min_face_px=Config.MIN_FACE_WIDTH,
|
||||
min_confidence=Config.MIN_CONFIDENCE,
|
||||
)
|
||||
if not quality.passed:
|
||||
quality_filtered += 1
|
||||
logger.debug(f"Quality filtered {asset['id']}: {quality.reason}")
|
||||
continue
|
||||
|
||||
# Crop the target person's face before embedding
|
||||
face_crop = _crop_face_from_thumbnail(img, asset)
|
||||
embed_img = face_crop if face_crop is not None else img
|
||||
else:
|
||||
embed_img = img
|
||||
|
||||
emb = get_embedding(embed_img, entity_type, asset_id=asset["id"])
|
||||
if emb is not None:
|
||||
embeddings.append(emb)
|
||||
valid_candidates.append(asset)
|
||||
confidence_scores.append(confidence)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(len(candidates), len(candidates))
|
||||
|
||||
if quality_filtered > 0:
|
||||
logger.info(f"Quality filtering removed {quality_filtered} images.")
|
||||
|
||||
if not embeddings:
|
||||
logger.warning("No valid embeddings found. Falling back to time spread.")
|
||||
return _select_time_spread(assets, limit)
|
||||
|
||||
if limit != "auto" and len(valid_candidates) < limit:
|
||||
logger.warning(f"Only {len(valid_candidates)} valid embeddings. Returning all.")
|
||||
return valid_candidates
|
||||
|
||||
# --- Phase 5: Cluster-aware selection ---
|
||||
return _cluster_aware_selection(
|
||||
embeddings,
|
||||
valid_candidates,
|
||||
limit,
|
||||
entity_type=entity_type,
|
||||
confidence_scores=confidence_scores,
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# K-Medoids (Lightweight Implementation)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _kmedoids(dist_matrix: np.ndarray, k: int, max_iter: int = 50) -> tuple[list[int], np.ndarray]:
|
||||
"""Lightweight K-Medoids clustering using cosine distance matrix.
|
||||
|
||||
Args:
|
||||
dist_matrix: (N, N) pairwise distance matrix
|
||||
k: Number of clusters
|
||||
max_iter: Maximum iterations for swap step
|
||||
|
||||
Returns:
|
||||
(medoid_indices, cluster_labels) tuple
|
||||
"""
|
||||
n = dist_matrix.shape[0]
|
||||
rng = np.random.default_rng(42)
|
||||
|
||||
# Initialize medoids: first = most central point, rest = farthest from chosen
|
||||
total_dist = dist_matrix.sum(axis=1)
|
||||
medoids = [int(np.argmin(total_dist))]
|
||||
|
||||
for _ in range(k - 1):
|
||||
dists_to_chosen = dist_matrix[:, medoids].min(axis=1)
|
||||
dists_to_chosen[medoids] = -np.inf
|
||||
medoids.append(int(np.argmax(dists_to_chosen)))
|
||||
|
||||
# Iterative swap step
|
||||
medoids = list(medoids)
|
||||
labels = np.argmin(dist_matrix[:, medoids], axis=1)
|
||||
cost = sum(dist_matrix[i, medoids[labels[i]]] for i in range(n))
|
||||
|
||||
for _ in range(max_iter):
|
||||
improved = False
|
||||
# Try swapping each medoid with a random non-medoid
|
||||
non_medoids = [i for i in range(n) if i not in medoids]
|
||||
if not non_medoids:
|
||||
break
|
||||
|
||||
for m_idx in range(k):
|
||||
candidates = rng.choice(non_medoids, size=min(10, len(non_medoids)), replace=False)
|
||||
for cand in candidates:
|
||||
new_medoids = medoids.copy()
|
||||
new_medoids[m_idx] = cand
|
||||
new_labels = np.argmin(dist_matrix[:, new_medoids], axis=1)
|
||||
new_cost = sum(dist_matrix[i, new_medoids[new_labels[i]]] for i in range(n))
|
||||
if new_cost < cost:
|
||||
medoids = new_medoids
|
||||
labels = new_labels
|
||||
cost = new_cost
|
||||
improved = True
|
||||
break
|
||||
if improved:
|
||||
break
|
||||
|
||||
if not improved:
|
||||
break
|
||||
|
||||
return medoids, labels
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Cluster-Aware Selection (K-Medoids + FPS Hybrid)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _compute_adaptive_threshold(emb_normed: np.ndarray, entity_type: str) -> float:
|
||||
"""Compute adaptive FPS stop threshold based on actual embedding distribution.
|
||||
|
||||
Instead of a hardcoded threshold, samples pairwise distances and sets
|
||||
the threshold as a fraction of the median pairwise distance.
|
||||
"""
|
||||
n = len(emb_normed)
|
||||
sample_size = min(200, n)
|
||||
rng = np.random.default_rng(42)
|
||||
indices = rng.choice(n, sample_size, replace=False) if n > sample_size else np.arange(n)
|
||||
sample = emb_normed[indices]
|
||||
|
||||
# Compute pairwise cosine distances for the sample
|
||||
pairwise = 1 - sample @ sample.T
|
||||
upper_tri = pairwise[np.triu_indices(len(sample), k=1)]
|
||||
median_dist = float(np.median(upper_tri))
|
||||
|
||||
# Faces: 20% of median (tighter — want fewer, more distinct images)
|
||||
# Objects: 10% of median (wider — want more diversity)
|
||||
fraction = 0.20 if entity_type == "face" else 0.10
|
||||
threshold = max(0.05, median_dist * fraction)
|
||||
|
||||
logger.info(
|
||||
f"Adaptive threshold: {threshold:.4f} "
|
||||
f"(median_dist={median_dist:.4f}, fraction={fraction}, type={entity_type})"
|
||||
)
|
||||
return threshold
|
||||
|
||||
|
||||
def _cluster_aware_selection(
|
||||
embeddings: list,
|
||||
candidates: list,
|
||||
limit: int | str,
|
||||
entity_type: str = "face",
|
||||
confidence_scores: list | None = None,
|
||||
) -> list:
|
||||
"""Two-stage selection: K-Medoids clustering → FPS with hard example weighting.
|
||||
|
||||
Stage 1: Cluster embeddings into k groups, select medoids as initial picks.
|
||||
This guarantees at least one representative from every distinct "look".
|
||||
|
||||
Stage 2: Fill remaining budget with FPS across cluster boundaries,
|
||||
biasing toward hard examples (low-confidence candidates).
|
||||
"""
|
||||
emb_matrix = np.vstack(embeddings) # (N, D)
|
||||
n = len(emb_matrix)
|
||||
|
||||
# Normalize for cosine distance
|
||||
norms = np.linalg.norm(emb_matrix, axis=1, keepdims=True)
|
||||
emb_normed = emb_matrix / np.maximum(norms, 1e-8)
|
||||
|
||||
# Build confidence weight array for hard example boosting
|
||||
conf_array = np.ones(n)
|
||||
if confidence_scores and entity_type == "face":
|
||||
for i, c in enumerate(confidence_scores):
|
||||
if c is not None:
|
||||
conf_array[i] = c
|
||||
|
||||
# Compute adaptive threshold for auto mode
|
||||
auto_threshold = _compute_adaptive_threshold(emb_normed, entity_type) if limit == "auto" else 0.0
|
||||
target = Config.MAX_AUTO_IMAGES if limit == "auto" else limit
|
||||
|
||||
# --- Stage 1: K-Medoids clustering ---
|
||||
k = min(max(5, target // 4), n // 3, n) # e.g., 5-20 clusters
|
||||
logger.info(f"Clustering {n} embeddings into {k} groups (K-Medoids)...")
|
||||
|
||||
# Compute full cosine distance matrix
|
||||
dist_matrix = 1 - emb_normed @ emb_normed.T
|
||||
|
||||
medoid_indices, cluster_labels = _kmedoids(dist_matrix, k)
|
||||
selected = list(medoid_indices)
|
||||
selected_set = set(selected)
|
||||
|
||||
logger.info(f"Selected {len(selected)} cluster medoids as initial picks.")
|
||||
|
||||
# --- Stage 2: FPS with hard example weighting ---
|
||||
min_dists = np.full(n, np.inf)
|
||||
|
||||
# Initialize min distances from all medoids
|
||||
for idx in selected:
|
||||
dists = dist_matrix[idx]
|
||||
min_dists = np.minimum(min_dists, dists)
|
||||
for idx in selected:
|
||||
min_dists[idx] = -np.inf
|
||||
|
||||
while len(selected) < target:
|
||||
# Hard example weighting: boost distance for low-confidence candidates
|
||||
# Confidence < 0.85 gets up to 1.5× distance boost
|
||||
hard_weight = np.where(conf_array < 0.85, 1.0 + (0.85 - conf_array) * 2.0, 1.0)
|
||||
weighted_dists = min_dists * hard_weight
|
||||
|
||||
best_idx = int(np.argmax(weighted_dists))
|
||||
best_dist = min_dists[best_idx] # Use unweighted for threshold comparison
|
||||
|
||||
if best_dist == -np.inf:
|
||||
break # All points selected
|
||||
|
||||
if limit == "auto" and best_dist < auto_threshold:
|
||||
logger.info(
|
||||
f"Auto-stop: Next best image {best_dist:.3f} away " f"(adaptive threshold {auto_threshold:.4f})."
|
||||
)
|
||||
break
|
||||
|
||||
selected.append(best_idx)
|
||||
selected_set.add(best_idx)
|
||||
|
||||
# Update min distances
|
||||
dists_to_new = dist_matrix[best_idx]
|
||||
min_dists = np.minimum(min_dists, dists_to_new)
|
||||
min_dists[best_idx] = -np.inf
|
||||
|
||||
# Log hard example stats
|
||||
if entity_type == "face":
|
||||
selected_conf = [conf_array[i] for i in selected if conf_array[i] < 1.0]
|
||||
hard_count = sum(1 for c in selected_conf if c < 0.85)
|
||||
logger.info(
|
||||
f"Selection complete: {len(selected)} images " f"({hard_count} hard examples with confidence < 0.85)."
|
||||
)
|
||||
else:
|
||||
logger.info(f"Selection complete: {len(selected)} diverse images.")
|
||||
|
||||
return [candidates[i] for i in selected]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Time Spread Fallback
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _select_time_spread(assets: list, limit: int | str) -> list:
|
||||
"""Select N assets evenly distributed in time."""
|
||||
if limit == "auto":
|
||||
limit = 30
|
||||
|
||||
logger.info(f"Selecting {limit} images using time spread.")
|
||||
|
||||
if len(assets) <= limit:
|
||||
return assets
|
||||
|
||||
indices = np.linspace(0, len(assets) - 1, limit, dtype=int)
|
||||
return [assets[i] for i in np.unique(indices)]
|
||||
@@ -0,0 +1,339 @@
|
||||
"""
|
||||
Unified embedding interface for faces and objects.
|
||||
|
||||
- Faces: InsightFace (ArcFace/Buffalo_L) — or reuse from Immich
|
||||
- Objects: SigLIP (Vision Transformer via transformers)
|
||||
- Caching: Disk-based cache avoids recomputation on reruns
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import importlib
|
||||
import logging
|
||||
import os
|
||||
import warnings
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from .cache import get_cache
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Lazy-loaded singletons
|
||||
_insightface_app = None
|
||||
_insightface_loaded = False
|
||||
_siglip_model = None
|
||||
_siglip_processor = None
|
||||
_siglip_loaded = False
|
||||
|
||||
|
||||
def _is_force_cpu() -> bool:
|
||||
"""Check if CPU mode is forced via environment variable."""
|
||||
return os.getenv("FORCE_CPU", "").lower() in ("true", "1", "yes")
|
||||
|
||||
|
||||
def _preload_cuda_libs() -> None:
|
||||
"""Preload CUDA/cuDNN DLLs so onnxruntime-gpu registers CUDAExecutionProvider.
|
||||
|
||||
Starting with onnxruntime-gpu 1.19+, CUDA/cuDNN libraries are no longer
|
||||
bundled inside the ORT package. They must be loaded from the nvidia-*
|
||||
pip packages (nvidia-cuda-runtime-cu12, nvidia-cudnn-cu12) before any
|
||||
InferenceSession is created.
|
||||
|
||||
Calling preload_dlls() with directory="" searches NVIDIA site-packages
|
||||
directories automatically.
|
||||
"""
|
||||
try:
|
||||
import onnxruntime
|
||||
if hasattr(onnxruntime, "preload_dlls"):
|
||||
onnxruntime.preload_dlls(cuda=True, cudnn=True, directory="")
|
||||
logger.info("Preloaded CUDA/cuDNN DLLs for onnxruntime-gpu")
|
||||
else:
|
||||
logger.debug("onnxruntime.preload_dlls() not available (ORT < 1.21)")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to preload CUDA/cuDNN DLLs: {e}")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# InsightFace (Faces)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def get_insightface_app():
|
||||
"""Singleton for InsightFace app with automatic GPU/CPU fallback."""
|
||||
global _insightface_app, _insightface_loaded
|
||||
if _insightface_loaded:
|
||||
return _insightface_app
|
||||
_insightface_loaded = True
|
||||
|
||||
# Preload CUDA/cuDNN DLLs BEFORE any ORT InferenceSession is created
|
||||
_preload_cuda_libs()
|
||||
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
from insightface.app import FaceAnalysis
|
||||
|
||||
# Get providers, excluding TensorRT to avoid noisy errors
|
||||
providers = [p for p in ort.get_available_providers() if p != "TensorrtExecutionProvider"]
|
||||
logger.info(f"Available ONNX providers: {providers}")
|
||||
|
||||
# Determine device: 0 for GPU, -1 for CPU
|
||||
gpu_providers = {
|
||||
"CUDAExecutionProvider",
|
||||
"ROCmExecutionProvider",
|
||||
"MPSExecutionProvider",
|
||||
"CoreMLExecutionProvider",
|
||||
}
|
||||
ctx_id = -1 if _is_force_cpu() else (0 if gpu_providers & set(providers) else -1)
|
||||
|
||||
device_str = "GPU" if ctx_id >= 0 else "CPU"
|
||||
logger.info(f"Loading InsightFace Buffalo_L on {device_str} (ctx_id={ctx_id})...")
|
||||
|
||||
# Suppress C-level output during model loading
|
||||
with open(os.devnull, "w") as devnull, contextlib.redirect_stdout(devnull), contextlib.redirect_stderr(devnull):
|
||||
insightface_home = os.environ.get("INSIGHTFACE_HOME", os.path.expanduser("~/.insightface"))
|
||||
_insightface_app = FaceAnalysis(name="buffalo_l", root=insightface_home, providers=providers)
|
||||
_insightface_app.prepare(ctx_id=ctx_id, det_size=(640, 640))
|
||||
|
||||
return _insightface_app
|
||||
|
||||
except ImportError:
|
||||
logger.error("InsightFace not installed!")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load InsightFace: {e}")
|
||||
# Retry on CPU if GPU failed
|
||||
if ctx_id == 0:
|
||||
logger.warning("Retrying InsightFace on CPU...")
|
||||
try:
|
||||
from insightface.app import FaceAnalysis
|
||||
|
||||
_insightface_app = FaceAnalysis(name="buffalo_l", root=insightface_home)
|
||||
_insightface_app.prepare(ctx_id=-1, det_size=(640, 640))
|
||||
return _insightface_app
|
||||
except Exception as ex:
|
||||
logger.error(f"CPU fallback failed: {ex}")
|
||||
return None
|
||||
|
||||
|
||||
def get_face_embedding(img_pil: Image.Image) -> np.ndarray | None:
|
||||
"""Get embedding of the largest face in a PIL image."""
|
||||
app = get_insightface_app()
|
||||
if not app:
|
||||
return None
|
||||
|
||||
try:
|
||||
# InsightFace expects BGR cv2 image
|
||||
img_bgr = cv2.cvtColor(np.asarray(img_pil), cv2.COLOR_RGB2BGR)
|
||||
|
||||
# Suppress scikit-image FutureWarning from InsightFace's face_align.py
|
||||
with warnings.catch_warnings():
|
||||
warnings.filterwarnings("ignore", message=".*estimate.*is deprecated", category=FutureWarning)
|
||||
faces = app.get(img_bgr)
|
||||
|
||||
if not faces:
|
||||
return None
|
||||
|
||||
# Return embedding of largest face
|
||||
largest = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
|
||||
return largest.embedding
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting face embedding: {e}")
|
||||
return None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# SigLIP (Objects)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def get_siglip_model():
|
||||
"""Singleton for SigLIP model and processor with GPU auto-detection."""
|
||||
global _siglip_model, _siglip_processor, _siglip_loaded
|
||||
if _siglip_loaded:
|
||||
return _siglip_model, _siglip_processor
|
||||
_siglip_loaded = True
|
||||
|
||||
try:
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
from transformers import AutoImageProcessor, SiglipVisionModel
|
||||
|
||||
model_name = "google/siglip-base-patch16-224"
|
||||
logger.info(f"Loading SigLIP model ({model_name})...")
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.filterwarnings("ignore", category=FutureWarning)
|
||||
warnings.filterwarnings("ignore", message=".*use_fast.*")
|
||||
_siglip_processor = AutoImageProcessor.from_pretrained(model_name, use_fast=True)
|
||||
_siglip_model = SiglipVisionModel.from_pretrained(model_name)
|
||||
|
||||
_siglip_model.eval()
|
||||
|
||||
# Move to GPU if available
|
||||
if not _is_force_cpu():
|
||||
if torch.cuda.is_available():
|
||||
_siglip_model = _siglip_model.cuda()
|
||||
logger.info("SigLIP running on CUDA GPU")
|
||||
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
_siglip_model = _siglip_model.to("mps")
|
||||
logger.info("SigLIP running on Apple MPS")
|
||||
else:
|
||||
logger.info("SigLIP running on CPU")
|
||||
else:
|
||||
logger.info("FORCE_CPU set. SigLIP running on CPU")
|
||||
|
||||
return _siglip_model, _siglip_processor
|
||||
|
||||
except ImportError as e:
|
||||
logger.error(f"transformers/torch not installed: {e}")
|
||||
return None, None
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load SigLIP: {e}")
|
||||
return None, None
|
||||
|
||||
|
||||
def get_object_embedding(img_pil: Image.Image) -> np.ndarray | None:
|
||||
"""Get 768-dim SigLIP embedding for an image."""
|
||||
model, processor = get_siglip_model()
|
||||
if model is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
inputs = processor(images=img_pil, return_tensors="pt")
|
||||
device = next(model.parameters()).device
|
||||
inputs = {k: v.to(device) for k, v in inputs.items()}
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = model(**inputs)
|
||||
return outputs.pooler_output.squeeze().cpu().numpy()
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting object embedding: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def get_object_embeddings_batch(images: list[Image.Image]) -> list[np.ndarray | None]:
|
||||
"""Get SigLIP embeddings for a batch of images (GPU-efficient)."""
|
||||
model, processor = get_siglip_model()
|
||||
if model is None:
|
||||
return [None] * len(images)
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
inputs = processor(images=images, return_tensors="pt", padding=True)
|
||||
device = next(model.parameters()).device
|
||||
inputs = {k: v.to(device) for k, v in inputs.items()}
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = model(**inputs)
|
||||
embeddings = outputs.pooler_output.cpu().numpy()
|
||||
return [embeddings[i] for i in range(len(embeddings))]
|
||||
except Exception as e:
|
||||
logger.error(f"Error in batch embedding: {e}")
|
||||
# Fall back to individual computation
|
||||
return [get_object_embedding(img) for img in images]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Unified Interface with Caching
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def get_embedding(
|
||||
img_pil: Image.Image,
|
||||
entity_type: str = "face",
|
||||
asset_id: str | None = None,
|
||||
immich_embedding: np.ndarray | None = None,
|
||||
) -> np.ndarray | None:
|
||||
"""Get embedding for an image based on entity type.
|
||||
|
||||
Priority:
|
||||
1. Pre-fetched Immich embedding (if provided)
|
||||
2. Disk cache (if enabled and asset_id provided)
|
||||
3. Local model computation (InsightFace or SigLIP)
|
||||
|
||||
Args:
|
||||
img_pil: The image to embed
|
||||
entity_type: 'face' or 'object'
|
||||
asset_id: Optional asset ID for cache lookup
|
||||
immich_embedding: Optional pre-fetched embedding from Immich API
|
||||
"""
|
||||
from .config import Config
|
||||
|
||||
use_cache = Config.ENABLE_CACHE and asset_id is not None
|
||||
cache = get_cache(Config.CACHE_DIR) if use_cache else None
|
||||
model_key = "immich" if entity_type == "face" else "siglip"
|
||||
|
||||
# 1. Use Immich embedding if provided
|
||||
if immich_embedding is not None:
|
||||
if cache:
|
||||
cache.put(asset_id, immich_embedding, model_key)
|
||||
return immich_embedding
|
||||
|
||||
# 2. Check disk cache
|
||||
if cache:
|
||||
cached = cache.get(asset_id, model_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
# 3. Compute locally
|
||||
if entity_type == "face":
|
||||
emb = get_face_embedding(img_pil)
|
||||
model_key = "insightface"
|
||||
else:
|
||||
emb = get_object_embedding(img_pil)
|
||||
|
||||
# Cache the result
|
||||
if emb is not None and cache:
|
||||
cache.put(asset_id, emb, model_key)
|
||||
|
||||
return emb
|
||||
|
||||
|
||||
def _is_module_available(module_name: str) -> bool:
|
||||
"""Check if a Python module is importable without importing it fully."""
|
||||
try:
|
||||
importlib.util.find_spec(module_name)
|
||||
return True
|
||||
except (ModuleNotFoundError, ValueError):
|
||||
return False
|
||||
|
||||
|
||||
def is_embedding_available(entity_type: str = "face", *, load: bool = False) -> bool:
|
||||
"""Check if embedding model is available for the given entity type.
|
||||
|
||||
By default this performs a lightweight import-check only (no model loading).
|
||||
Pass ``load=True`` to actually load the model (expensive, hundreds of MB).
|
||||
|
||||
Args:
|
||||
entity_type: 'face' or 'object'
|
||||
load: If True, fully load the model to verify. If False (default),
|
||||
only check that the required packages are importable.
|
||||
"""
|
||||
if load:
|
||||
if entity_type == "face":
|
||||
return get_insightface_app() is not None
|
||||
model, _ = get_siglip_model()
|
||||
return model is not None
|
||||
|
||||
# Lightweight check: just verify the packages are importable
|
||||
if entity_type == "face":
|
||||
return _is_module_available("insightface") and _is_module_available("onnxruntime")
|
||||
return _is_module_available("transformers") and _is_module_available("torch")
|
||||
|
||||
|
||||
def load_embedding_model(entity_type: str = "face") -> bool:
|
||||
"""Explicitly load the embedding model for the given entity type.
|
||||
|
||||
Returns True if the model loaded successfully.
|
||||
"""
|
||||
if entity_type == "face":
|
||||
return get_insightface_app() is not None
|
||||
model, _ = get_siglip_model()
|
||||
return model is not None
|
||||
|
||||
@@ -0,0 +1,339 @@
|
||||
"""Execution phase: image processing and Frigate upload."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
from io import BytesIO
|
||||
from urllib.parse import quote
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
from rich import print as rprint
|
||||
from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
|
||||
|
||||
from .config import Config, get_headers
|
||||
from .image_processing import process_face_mode, process_full_mode, process_object_mode
|
||||
from .immich_api import fetch_face_data, fetch_full_image
|
||||
from .logging import console
|
||||
from .upload_tracker import mark_rejected, mark_uploaded
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _enrich_asset_with_face_data(asset: dict, person: dict) -> dict:
|
||||
"""Enrich an asset dict with face bounding box data from the Immich faces API.
|
||||
|
||||
The search/metadata endpoint does not include face bounding box data,
|
||||
so we fetch it from GET /api/faces?id={asset_id} and inject it into
|
||||
the asset's "people" field so process_face_mode can find it.
|
||||
|
||||
Returns the enriched asset dict (modifies in place and returns it).
|
||||
"""
|
||||
person_id = person["id"]
|
||||
face_data = fetch_face_data(asset["id"], person_id=person_id)
|
||||
|
||||
if face_data is None:
|
||||
logger.debug(f"No face data returned for {person.get('name')} in asset {asset.get('id')}")
|
||||
# Clean any None entries from the people list (can come from Immich API)
|
||||
if "people" in asset:
|
||||
asset["people"] = [p for p in asset["people"] if p is not None]
|
||||
return asset
|
||||
|
||||
# Skip zero-area bounding boxes (face detection failed or no face found)
|
||||
if face_data.bbox == (0, 0, 0, 0):
|
||||
logger.debug(f"Zero-area bounding box for {person.get('name')} in asset {asset.get('id')}")
|
||||
# Clean any None entries from the people list (can come from Immich API)
|
||||
if "people" in asset:
|
||||
asset["people"] = [p for p in asset["people"] if p is not None]
|
||||
return asset
|
||||
|
||||
face_info = {
|
||||
"boundingBoxX1": face_data.bbox[0],
|
||||
"boundingBoxY1": face_data.bbox[1],
|
||||
"boundingBoxX2": face_data.bbox[2],
|
||||
"boundingBoxY2": face_data.bbox[3],
|
||||
"imageWidth": face_data.image_width,
|
||||
"imageHeight": face_data.image_height,
|
||||
}
|
||||
|
||||
# Inject into asset so process_face_mode can find it via asset["people"]
|
||||
asset["people"] = [{"id": person_id, "faces": [face_info]}]
|
||||
return asset
|
||||
|
||||
|
||||
def execute_jobs(jobs: list[dict]) -> None:
|
||||
"""Download and process images for all jobs.
|
||||
|
||||
Builds an asset_map per job (filename → Immich asset ID) so that
|
||||
upload_to_frigate() can mark assets as uploaded after success.
|
||||
"""
|
||||
if not jobs:
|
||||
return
|
||||
|
||||
console.rule("[bold blue]Execution Phase")
|
||||
|
||||
use_full_res = Config.USE_FULL_RESOLUTION
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
console=console,
|
||||
) as progress:
|
||||
grand_total = sum(j["limit"] for j in jobs)
|
||||
overall_task = progress.add_task("[green]Overall Progress", total=grand_total)
|
||||
|
||||
for job in jobs:
|
||||
person, assets, config = job["person"], job["assets"], job["config"]
|
||||
name, mode = person["name"], config.get("mode", "face")
|
||||
|
||||
job_task = progress.add_task(f"Processing {name}...", total=len(assets))
|
||||
person_dir = os.path.join(Config.OUTPUT_DIR, name)
|
||||
# Face crops are transient (uploaded then discarded); wipe before each run.
|
||||
# Object crops are the deliverable; preserve them across runs.
|
||||
if mode == "face" and os.path.isdir(person_dir):
|
||||
shutil.rmtree(person_dir)
|
||||
os.makedirs(person_dir, exist_ok=True)
|
||||
|
||||
# Track filename → asset_id mapping for upload dedup
|
||||
asset_map: dict[str, str] = {}
|
||||
|
||||
count = 0
|
||||
for asset in assets:
|
||||
try:
|
||||
# For face mode, enrich the asset with face bounding box data
|
||||
# from the Immich faces API (not included in search/metadata results)
|
||||
if mode == "face":
|
||||
asset = _enrich_asset_with_face_data(asset, person)
|
||||
|
||||
# Use full-resolution for final output when configured
|
||||
if use_full_res:
|
||||
img = fetch_full_image(asset["id"])
|
||||
else:
|
||||
resp = requests.get(
|
||||
f"{Config.IMMICH_URL}/api/assets/{asset['id']}/thumbnail?size=preview&format=JPEG",
|
||||
headers=get_headers(),
|
||||
timeout=30,
|
||||
)
|
||||
img = Image.open(BytesIO(resp.content)) if resp.ok else None
|
||||
|
||||
if img is None:
|
||||
progress.console.print(f"[red]Failed download {asset['id']}[/red]")
|
||||
else:
|
||||
saved = (
|
||||
process_face_mode(img, asset, person, person_dir, count)
|
||||
if mode == "face"
|
||||
else process_object_mode(img, config, person_dir, count)
|
||||
if mode == "object"
|
||||
else process_full_mode(img, person_dir, count)
|
||||
)
|
||||
if saved:
|
||||
# Record which asset produced which output file
|
||||
filename = f"{count}.jpg"
|
||||
asset_map[filename] = asset["id"]
|
||||
# Also record object-mode variant filenames
|
||||
if mode == "object":
|
||||
for f in sorted(os.listdir(person_dir)):
|
||||
if f.startswith(f"{count}_") and f not in asset_map:
|
||||
asset_map[f] = asset["id"]
|
||||
|
||||
count += 1
|
||||
else:
|
||||
progress.console.print(
|
||||
f"[yellow]Skipped {asset['id']} (no usable face data)[/yellow]"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to process asset {asset['id']}: {e}")
|
||||
|
||||
progress.advance(job_task)
|
||||
progress.advance(overall_task)
|
||||
|
||||
# Store asset_map on the job so upload_to_frigate can use it
|
||||
job["asset_map"] = asset_map
|
||||
|
||||
progress.remove_task(job_task)
|
||||
|
||||
# Log how many images were actually saved vs selected
|
||||
if count < len(assets):
|
||||
logger.info(f"{name}: saved {count}/{len(assets)} selected images")
|
||||
|
||||
|
||||
def upload_to_frigate(jobs: list[dict]) -> None:
|
||||
"""Upload processed face crops to Frigate via API with detailed logging.
|
||||
|
||||
Only runs for face-mode jobs. Object-mode crops are saved to the output
|
||||
directory as the deliverable and must be copied to Frigate manually.
|
||||
|
||||
After each successful upload, records the Immich asset ID in the
|
||||
upload tracker so it is skipped on future runs.
|
||||
"""
|
||||
face_jobs = [j for j in jobs if j["config"].get("mode", "face") == "face"]
|
||||
|
||||
if not face_jobs:
|
||||
rprint("[dim]No face-mode jobs to upload.[/dim]")
|
||||
return
|
||||
|
||||
# Notify user about object-mode jobs that were skipped
|
||||
object_jobs = [j for j in jobs if j["config"].get("mode") == "object"]
|
||||
for job in object_jobs:
|
||||
name = job["person"]["name"]
|
||||
person_dir = os.path.join(Config.OUTPUT_DIR, name)
|
||||
rprint(f" [dim]📁 {name} (object): crops saved to {person_dir} — copy to Frigate manually[/dim]")
|
||||
|
||||
frigate_url = os.environ.get("FRIGATE_URL", "")
|
||||
if not frigate_url:
|
||||
rprint("[yellow]⚠️ FRIGATE_URL not set, skipping upload.[/yellow]")
|
||||
return
|
||||
|
||||
rprint("\n[bold cyan]📤 Uploading to Frigate[/bold cyan]")
|
||||
rprint(f" Target: [dim]{frigate_url}[/dim]")
|
||||
|
||||
# Build a mapping of output filenames → Immich asset IDs
|
||||
# from the asset_map stored on each job during execute_jobs()
|
||||
filename_to_asset_id: dict[str, dict[str, str]] = {}
|
||||
total_files = 0
|
||||
for job in face_jobs:
|
||||
name = job["person"]["name"]
|
||||
asset_map = job.get("asset_map", {})
|
||||
filename_to_asset_id[name] = asset_map
|
||||
total_files += len(asset_map)
|
||||
|
||||
if total_files == 0:
|
||||
rprint(" [yellow]No images found to upload.[/yellow]")
|
||||
return
|
||||
|
||||
rprint(f" People: [bold]{len(face_jobs)}[/bold], Total images: [bold]{total_files}[/bold]")
|
||||
|
||||
uploaded, failed = 0, 0
|
||||
max_retries = 2
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
console=console,
|
||||
) as progress:
|
||||
upload_task = progress.add_task("[green]Uploading to Frigate", total=total_files)
|
||||
|
||||
for job in face_jobs:
|
||||
name = job["person"]["name"]
|
||||
# URL-encode the name for the API (handles spaces, special chars)
|
||||
encoded_name = quote(name, safe="")
|
||||
if " " in name:
|
||||
progress.console.print(f" ℹ️ URL-encoded name for Frigate API: '{name}' → '{encoded_name}'")
|
||||
|
||||
person_dir = os.path.join(Config.OUTPUT_DIR, name)
|
||||
if not os.path.isdir(person_dir):
|
||||
progress.console.print(f" [dim]⏭️ {name}: no output directory, skipping[/dim]")
|
||||
continue
|
||||
|
||||
asset_map = filename_to_asset_id.get(name, {})
|
||||
person_files = sorted(asset_map.keys())
|
||||
|
||||
if not person_files:
|
||||
progress.console.print(f" [dim]⏭️ {name}: no images found[/dim]")
|
||||
continue
|
||||
|
||||
progress.console.print(f" 📁 {name}: uploading {len(person_files)} image(s)...")
|
||||
person_uploaded = 0
|
||||
person_failed = 0
|
||||
|
||||
for fname in person_files:
|
||||
fpath = os.path.join(person_dir, fname)
|
||||
for attempt in range(1, max_retries + 1):
|
||||
try:
|
||||
with open(fpath, "rb") as f:
|
||||
resp = requests.post(
|
||||
f"{frigate_url}/api/faces/{encoded_name}/register",
|
||||
files={"file": (fname, f, "image/jpeg")},
|
||||
timeout=30,
|
||||
)
|
||||
if resp.status_code == 200:
|
||||
uploaded += 1
|
||||
person_uploaded += 1
|
||||
|
||||
# Mark this asset as uploaded so it's skipped on future runs
|
||||
asset_id = asset_map.get(fname)
|
||||
if asset_id:
|
||||
mark_uploaded(asset_id, person_name=name)
|
||||
|
||||
break
|
||||
else:
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"Upload attempt {attempt}/{max_retries} for {fname}:"
|
||||
f" HTTP {resp.status_code}, retrying..."
|
||||
)
|
||||
continue
|
||||
failed += 1
|
||||
person_failed += 1
|
||||
progress.console.print(
|
||||
f" [red]✗ {fname}: HTTP {resp.status_code} (after {max_retries} attempts)[/red]"
|
||||
)
|
||||
try:
|
||||
error_detail = resp.json().get("message", resp.text[:100])
|
||||
progress.console.print(f" [dim]{error_detail}[/dim]")
|
||||
except Exception:
|
||||
error_detail = resp.text[:100]
|
||||
progress.console.print(f" [dim]{error_detail}[/dim]")
|
||||
if resp.status_code == 400 and "face" in error_detail.lower():
|
||||
asset_id = asset_map.get(fname)
|
||||
if asset_id:
|
||||
mark_rejected(asset_id, person_name=name)
|
||||
except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as exc:
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"Upload attempt {attempt}/{max_retries} for {fname}:"
|
||||
f" {type(exc).__name__}, retrying..."
|
||||
)
|
||||
continue
|
||||
failed += 1
|
||||
person_failed += 1
|
||||
label = (
|
||||
"Connection refused"
|
||||
if isinstance(exc, requests.exceptions.ConnectionError)
|
||||
else "Request timed out (30s)"
|
||||
)
|
||||
progress.console.print(
|
||||
f" [red]✗ {fname}: {label} (after {max_retries} attempts)[/red]"
|
||||
)
|
||||
except Exception as e:
|
||||
if attempt < max_retries:
|
||||
logger.warning(
|
||||
f"Upload attempt {attempt}/{max_retries} for {fname}:"
|
||||
f" {type(e).__name__}, retrying..."
|
||||
)
|
||||
continue
|
||||
failed += 1
|
||||
person_failed += 1
|
||||
progress.console.print(
|
||||
f" [red]✗ {fname}: {type(e).__name__} - {e} (after {max_retries} attempts)[/red]"
|
||||
)
|
||||
|
||||
progress.advance(upload_task)
|
||||
|
||||
# Per-person summary
|
||||
if person_failed == 0:
|
||||
progress.console.print(
|
||||
f" ✅ {name}: {person_uploaded}/{person_uploaded} uploaded"
|
||||
)
|
||||
else:
|
||||
progress.console.print(
|
||||
f" ⚠️ {name}: {person_uploaded} succeeded, {person_failed} failed"
|
||||
)
|
||||
|
||||
# Grand summary
|
||||
rprint("\n [bold]Frigate Upload Summary:[/bold]")
|
||||
rprint(f" ✅ Succeeded: [green]{uploaded}[/green]")
|
||||
if failed:
|
||||
rprint(f" ❌ Failed: [red]{failed}[/red]")
|
||||
else:
|
||||
rprint(" ❌ Failed: 0")
|
||||
|
||||
if failed > 0:
|
||||
rprint(" [yellow]Check logs above for per-file error details.[/yellow]")
|
||||
|
||||
if failed == total_files and total_files > 0:
|
||||
rprint(" [bold red]All uploads failed. Verify FRIGATE_URL is reachable and API is enabled.[/bold red]")
|
||||
@@ -0,0 +1,175 @@
|
||||
"""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,
|
||||
) -> 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
|
||||
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 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:
|
||||
_save_jpeg(aligned, os.path.join(output_dir, f"{count}.jpg"))
|
||||
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)
|
||||
_save_jpeg(face_crop, os.path.join(output_dir, f"{count}.jpg"))
|
||||
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
|
||||
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
|
||||
|
||||
@@ -0,0 +1,217 @@
|
||||
"""Immich API client for fetching people, assets, and face data."""
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
from PIL import Image, ImageOps
|
||||
|
||||
from .config import Config, get_headers
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MAX_PAGES = 1000 # Safety limit for pagination
|
||||
|
||||
|
||||
@dataclass
|
||||
class FaceData:
|
||||
"""Pre-computed face data from Immich."""
|
||||
|
||||
embedding: np.ndarray | None
|
||||
bbox: tuple[float, float, float, float] # (x1, y1, x2, y2)
|
||||
confidence: float | None
|
||||
image_width: int
|
||||
image_height: int
|
||||
|
||||
|
||||
def get_people() -> list[dict]:
|
||||
"""Fetch all people from Immich."""
|
||||
try:
|
||||
resp = requests.get(
|
||||
f"{Config.IMMICH_URL}/api/people",
|
||||
headers=get_headers(),
|
||||
timeout=10,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return resp.json().get("people", [])
|
||||
except (requests.RequestException, ValueError) as e:
|
||||
logger.error(f"Failed to fetch people: {e}")
|
||||
return []
|
||||
|
||||
|
||||
def fetch_all_assets(person: dict) -> list[dict]:
|
||||
"""Fetch all assets for a person with pagination."""
|
||||
name = person.get("name", "Unknown")
|
||||
person_id = person["id"]
|
||||
url = f"{Config.IMMICH_URL}/api/search/metadata"
|
||||
page_size = 1000
|
||||
|
||||
logger.info(f"Fetching assets for {name}...")
|
||||
|
||||
assets = []
|
||||
for page in range(1, MAX_PAGES + 1):
|
||||
try:
|
||||
resp = requests.post(
|
||||
url,
|
||||
json={"personIds": [person_id], "size": page_size, "page": page},
|
||||
headers=get_headers(),
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
if not resp.ok:
|
||||
logger.error(f"Error fetching assets for {name} (page {page}): {resp.status_code}")
|
||||
break
|
||||
|
||||
page_assets = resp.json().get("assets", [])
|
||||
if isinstance(page_assets, dict):
|
||||
page_assets = page_assets.get("items", [])
|
||||
|
||||
if not page_assets:
|
||||
break
|
||||
|
||||
assets.extend(page_assets)
|
||||
logger.debug(f"Fetched page {page}, total: {len(assets)}")
|
||||
|
||||
if len(page_assets) < page_size:
|
||||
break
|
||||
|
||||
except (requests.RequestException, ValueError) as e:
|
||||
logger.error(f"Exception fetching assets for {name}: {e}")
|
||||
break
|
||||
|
||||
return assets
|
||||
|
||||
|
||||
def fetch_face_data(asset_id: str, person_id: str | None = None) -> FaceData | None:
|
||||
"""Fetch pre-computed face data (embedding, bbox, confidence) from Immich.
|
||||
|
||||
Queries GET /api/faces?id={asset_id} to retrieve face detection results
|
||||
that Immich already computed using InsightFace Buffalo_L.
|
||||
|
||||
Args:
|
||||
asset_id: The asset to get face data for
|
||||
person_id: Optional person ID to match the specific face
|
||||
|
||||
Returns:
|
||||
FaceData with embedding, bbox, and confidence, or None if unavailable
|
||||
"""
|
||||
try:
|
||||
resp = requests.get(
|
||||
f"{Config.IMMICH_URL}/api/faces",
|
||||
params={"id": asset_id},
|
||||
headers=get_headers(),
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
if not resp.ok:
|
||||
logger.debug(f"Face data endpoint returned {resp.status_code} for {asset_id}")
|
||||
return None
|
||||
|
||||
faces = resp.json()
|
||||
if not faces:
|
||||
return None
|
||||
|
||||
# Match the target person if specified
|
||||
face = None
|
||||
if person_id:
|
||||
face = next(
|
||||
(f for f in faces if (f.get("person") or {}).get("id") == person_id),
|
||||
None,
|
||||
)
|
||||
if face is None:
|
||||
face = faces[0] # Fall back to first/largest face
|
||||
|
||||
# Extract embedding if available
|
||||
embedding = None
|
||||
if "embedding" in face:
|
||||
embedding = np.array(face["embedding"], dtype=np.float32)
|
||||
|
||||
# Extract bounding box
|
||||
bbox = (
|
||||
face.get("boundingBoxX1", 0),
|
||||
face.get("boundingBoxY1", 0),
|
||||
face.get("boundingBoxX2", 0),
|
||||
face.get("boundingBoxY2", 0),
|
||||
)
|
||||
|
||||
return FaceData(
|
||||
embedding=embedding,
|
||||
bbox=bbox,
|
||||
confidence=face.get("score") or face.get("confidence"),
|
||||
image_width=face.get("imageWidth", 0),
|
||||
image_height=face.get("imageHeight", 0),
|
||||
)
|
||||
|
||||
except requests.RequestException as e:
|
||||
logger.debug(f"Failed to fetch face data for {asset_id}: {e}")
|
||||
return None
|
||||
except (AttributeError, KeyError, TypeError, ValueError) as e:
|
||||
logger.debug(f"Failed to parse face data for {asset_id}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def fetch_full_image(asset_id: str, timeout: int = 60) -> Image.Image | None:
|
||||
"""Fetch full-resolution image from Immich, falling back to preview thumbnail.
|
||||
|
||||
The /original endpoint may return HEIC, RAW, or video files that PIL
|
||||
cannot open directly. In that case, we fall back to the JPEG thumbnail.
|
||||
"""
|
||||
# Try original first
|
||||
try:
|
||||
resp = requests.get(
|
||||
f"{Config.IMMICH_URL}/api/assets/{asset_id}/original",
|
||||
headers=get_headers(),
|
||||
timeout=timeout,
|
||||
)
|
||||
if resp.ok:
|
||||
try:
|
||||
return ImageOps.exif_transpose(Image.open(BytesIO(resp.content)))
|
||||
except Exception:
|
||||
logger.debug(f"PIL can't open original for {asset_id}, falling back to preview")
|
||||
except requests.RequestException:
|
||||
logger.debug(f"Original request failed for {asset_id}, falling back to preview")
|
||||
|
||||
# Fall back to preview thumbnail (always JPEG)
|
||||
try:
|
||||
resp = requests.get(
|
||||
f"{Config.IMMICH_URL}/api/assets/{asset_id}/thumbnail?size=preview&format=JPEG",
|
||||
headers=get_headers(),
|
||||
timeout=30,
|
||||
)
|
||||
if resp.ok:
|
||||
return ImageOps.exif_transpose(Image.open(BytesIO(resp.content)))
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to fetch image {asset_id}: {e}")
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def filter_recent_assets(assets: list[dict], years: int | None = None) -> list[dict]:
|
||||
"""Filter assets to keep only those from the last N years."""
|
||||
years = years or Config.YEARS_FILTER
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=365 * years)
|
||||
|
||||
logger.debug(f"Filtering assets older than {years} years ({cutoff})")
|
||||
|
||||
recent, skipped = [], 0
|
||||
for asset in assets:
|
||||
created_at_str = asset.get("fileCreatedAt")
|
||||
if not created_at_str:
|
||||
continue
|
||||
|
||||
try:
|
||||
# Handle ISO8601 with 'Z' suffix
|
||||
created_at = datetime.fromisoformat(created_at_str.replace("Z", "+00:00"))
|
||||
if created_at > cutoff:
|
||||
recent.append(asset)
|
||||
else:
|
||||
skipped += 1
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
logger.info(f"Retained {len(recent)} assets (filtered {skipped} old assets).")
|
||||
return recent
|
||||
|
||||
+302
@@ -0,0 +1,302 @@
|
||||
"""Configuration phase: strategy selection and job building."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from rich import print as rprint
|
||||
from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
|
||||
from rich.prompt import Confirm, IntPrompt, Prompt
|
||||
from rich.table import Table
|
||||
|
||||
from .config import Config
|
||||
from .diversity import select_diverse_assets
|
||||
from .embeddings import is_embedding_available, load_embedding_model
|
||||
from .immich_api import fetch_all_assets, filter_recent_assets
|
||||
from .logging import console
|
||||
from .upload_tracker import filter_already_uploaded
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Strategy presets: (limit, mode_name)
|
||||
STRATEGY_PRESETS = {
|
||||
"1": ("auto", "Auto Diversity"),
|
||||
"2": (30, "Standard (30)"),
|
||||
"3": (100, "Broad (100)"),
|
||||
}
|
||||
|
||||
|
||||
def _get_strategy_choice(has_embedding: bool, entity_type: str) -> tuple[int | str, str]:
|
||||
"""Prompt user for training strategy and return (limit, selection_mode)."""
|
||||
model_name = "InsightFace" if entity_type == "face" else "SigLIP"
|
||||
|
||||
if has_embedding:
|
||||
rprint(" [bold]1.[/bold] Auto (Objective Diversity) [green][Recommended][/green]")
|
||||
rprint(" [dim]• Dynamically selects images until redundancy starts[/dim]")
|
||||
rprint(" [bold]2.[/bold] Standard (30 images)")
|
||||
rprint(" [bold]3.[/bold] Broad (100 images)")
|
||||
rprint(" [bold]4.[/bold] Custom Count")
|
||||
rprint(" [bold]5.[/bold] Skip")
|
||||
|
||||
choice = Prompt.ask("Choice", choices=["1", "2", "3", "4", "5"], default="1")
|
||||
|
||||
if choice == "5":
|
||||
return 0, "skip"
|
||||
if choice == "4":
|
||||
limit = IntPrompt.ask("Enter number of images", default=30)
|
||||
mode = "smart" if Confirm.ask("Use Smart Diversity?", default=True) else "time"
|
||||
return limit, mode
|
||||
if choice in STRATEGY_PRESETS:
|
||||
return STRATEGY_PRESETS[choice][0], "smart"
|
||||
return 30, "smart"
|
||||
|
||||
# Fallback when embedding model not available
|
||||
rprint(f" [yellow]Note: {model_name} not available. Using Time Spread.[/yellow]")
|
||||
rprint(" [bold]1.[/bold] Standard (30 images) [green][Recommended][/green]")
|
||||
rprint(" [bold]2.[/bold] Broad (100 images)")
|
||||
rprint(" [bold]3.[/bold] Custom Count")
|
||||
rprint(" [bold]4.[/bold] Skip")
|
||||
|
||||
choice = Prompt.ask("Choice", choices=["1", "2", "3", "4"], default="1")
|
||||
limits = {"1": 30, "2": 100, "3": IntPrompt.ask("Enter number of images", default=30)}
|
||||
return limits.get(choice, 0), "time" if choice != "4" else "skip"
|
||||
|
||||
|
||||
def _resolve_strategy(strategy: str, has_embedding: bool) -> tuple[int | str, str]:
|
||||
"""Resolve env var strategy to (limit, selection_mode) without prompts."""
|
||||
custom_limit = os.environ.get("LIMIT", "").strip()
|
||||
|
||||
if not has_embedding:
|
||||
limit = int(custom_limit) if custom_limit else 30
|
||||
return limit, "time"
|
||||
|
||||
if custom_limit:
|
||||
return int(custom_limit), "smart"
|
||||
|
||||
strategy_map = {
|
||||
"auto": ("auto", "smart"),
|
||||
"standard": (30, "smart"),
|
||||
"broad": (100, "smart"),
|
||||
}
|
||||
return strategy_map.get(strategy, ("auto", "smart"))
|
||||
|
||||
|
||||
def _perform_selection(assets: list, limit: int | str, name: str, selection_mode: str, entity_type: str) -> list:
|
||||
"""Run diversity selection with progress display."""
|
||||
if selection_mode == "smart":
|
||||
model_display = "InsightFace (face embeddings)" if entity_type == "face" else "SigLIP (visual embeddings)"
|
||||
rprint(f"\n[cyan]Using {model_display} for diversity analysis...[/cyan]")
|
||||
|
||||
# Pre-load model explicitly (separate from availability check)
|
||||
load_embedding_model(entity_type)
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
console=console,
|
||||
) as progress:
|
||||
task = progress.add_task(f"[cyan]Computing embeddings for {len(assets)} images...", total=None)
|
||||
selected = select_diverse_assets(
|
||||
assets,
|
||||
limit,
|
||||
name,
|
||||
selection_mode=selection_mode,
|
||||
entity_type=entity_type,
|
||||
progress_callback=lambda c, t: progress.update(task, completed=c, total=t),
|
||||
)
|
||||
|
||||
label = f"Auto-diversity selected {len(selected)}" if limit == "auto" else f"Selected {len(selected)}"
|
||||
rprint(f" [green]{label} diverse images.[/green]")
|
||||
return selected
|
||||
|
||||
rprint(f"\n[cyan]Using time-spread selection for {limit} images...[/cyan]")
|
||||
with console.status(f"[bold]Selecting {limit} images evenly distributed over time...[/bold]"):
|
||||
selected = select_diverse_assets(assets, limit, name, selection_mode="time", entity_type=entity_type)
|
||||
rprint(f" [green]Selected {len(selected)} images using time spread.[/green]")
|
||||
return selected
|
||||
|
||||
|
||||
def _configure_person(person: dict, people: list[dict]) -> dict | None:
|
||||
"""Configure training for a single person. Returns job dict or None."""
|
||||
name = person["name"]
|
||||
console.print(f"\nSelected: [bold green]{name}[/bold green]")
|
||||
|
||||
# Select training mode
|
||||
rprint("\n[bold cyan]Training Mode:[/bold cyan]")
|
||||
rprint(" [bold]1.[/bold] Face (Frigate Face Recognition)")
|
||||
rprint(" [bold]2.[/bold] Object (Frigate Object Classification)")
|
||||
|
||||
mode_choice = Prompt.ask("Choice", choices=["1", "2"], default="1")
|
||||
entity_type = "face" if mode_choice == "1" else "object"
|
||||
|
||||
config = {"name": name, "mode": entity_type}
|
||||
if entity_type == "object":
|
||||
config["object_class"] = Prompt.ask("Enter Object Class (e.g. dog, cat, car)", default="dog")
|
||||
|
||||
# Fetch and filter assets
|
||||
years = IntPrompt.ask("Filter images older than (years)", default=Config.YEARS_FILTER)
|
||||
|
||||
console.print(f"Scanning for {name} ({entity_type})...")
|
||||
with console.status("[bold green]Fetching assets...[/bold green]"):
|
||||
all_assets = fetch_all_assets(person)
|
||||
recent_assets = filter_recent_assets(all_assets, years=years)
|
||||
|
||||
rprint(f" Found [bold]{len(all_assets)}[/bold] total, [bold]{len(recent_assets)}[/bold] in range ({years} years).")
|
||||
|
||||
# Filter out assets already uploaded to Frigate
|
||||
retry_rejected = os.environ.get("RETRY_REJECTED", "false").lower() in ("true", "1", "yes")
|
||||
before_dedup = len(recent_assets)
|
||||
new_asset_ids = set(filter_already_uploaded([a["id"] for a in recent_assets], retry_rejected=retry_rejected))
|
||||
recent_assets = [a for a in recent_assets if a["id"] in new_asset_ids]
|
||||
skipped = before_dedup - len(recent_assets)
|
||||
if skipped:
|
||||
rprint(f" [dim]Skipped {skipped} assets already uploaded to Frigate.[/dim]")
|
||||
|
||||
if not recent_assets:
|
||||
rprint(" [dim]Skipping (0 new images after dedup).[/dim]")
|
||||
return None
|
||||
|
||||
# Strategy selection
|
||||
has_embedding = is_embedding_available(entity_type)
|
||||
rprint(f"\n[bold cyan]Select Training Strategy for {name}:[/bold cyan]")
|
||||
|
||||
limit, selection_mode = _get_strategy_choice(has_embedding, entity_type)
|
||||
if selection_mode == "skip":
|
||||
return None
|
||||
|
||||
# Perform selection
|
||||
selected_assets = _perform_selection(recent_assets, limit, name, selection_mode, entity_type)
|
||||
|
||||
rprint(f" [green]Queued {len(selected_assets)} images for {name}.[/green]")
|
||||
return {"person": person, "assets": selected_assets, "limit": len(selected_assets), "config": config}
|
||||
|
||||
|
||||
def interactive_configure(people: list[dict]) -> list[dict]:
|
||||
"""Interactive phase: select person(s), mode, and configure training strategy.
|
||||
|
||||
Supports multi-person batch mode — after configuring one person,
|
||||
prompts to add another.
|
||||
"""
|
||||
valid_people = sorted([p for p in people if p.get("name")], key=lambda x: x["name"])
|
||||
|
||||
if not valid_people:
|
||||
rprint("[red]No people found with names in Immich.[/red]")
|
||||
return []
|
||||
|
||||
jobs = []
|
||||
|
||||
while True:
|
||||
# Select person
|
||||
console.print("\n[bold cyan]Select Person to Train:[/bold cyan]")
|
||||
for idx, p in enumerate(valid_people, 1):
|
||||
# Mark already-queued people
|
||||
marker = " [dim](queued)[/dim]" if any(j["person"]["id"] == p["id"] for j in jobs) else ""
|
||||
console.print(f" [bold]{idx}.[/bold] {p['name']}{marker}")
|
||||
|
||||
p_choice = IntPrompt.ask("Enter Number", choices=[str(i) for i in range(1, len(valid_people) + 1)])
|
||||
person = valid_people[p_choice - 1]
|
||||
|
||||
job = _configure_person(person, valid_people)
|
||||
if job:
|
||||
jobs.append(job)
|
||||
|
||||
# Multi-person: ask to add another
|
||||
if not Confirm.ask("\nAdd another person?", default=False):
|
||||
break
|
||||
|
||||
return jobs
|
||||
|
||||
|
||||
def auto_configure(people: list[dict]) -> list[dict]:
|
||||
"""Non-interactive: configure jobs for all named people automatically."""
|
||||
valid_people = sorted([p for p in people if p.get("name")], key=lambda x: x["name"])
|
||||
|
||||
if not valid_people:
|
||||
rprint("[red]No people found with names in Immich.[/red]")
|
||||
return []
|
||||
|
||||
mode = os.environ.get("TRAINING_MODE", "face")
|
||||
strategy = os.environ.get("STRATEGY", "auto")
|
||||
skip = os.environ.get("SKIP_PEOPLE", "").split(",") if os.environ.get("SKIP_PEOPLE") else []
|
||||
only = os.environ.get("ONLY_PEOPLE", "").split(",") if os.environ.get("ONLY_PEOPLE") else []
|
||||
|
||||
if only:
|
||||
valid_people = [p for p in valid_people if p["name"] in only]
|
||||
if skip:
|
||||
valid_people = [p for p in valid_people if p["name"] not in skip]
|
||||
|
||||
# Filter by minimum face count (Issue #6: previously unimplemented)
|
||||
min_face_count = Config.MIN_FACE_COUNT
|
||||
if min_face_count > 0:
|
||||
valid_people = [p for p in valid_people if p.get("assetCount", 0) >= min_face_count]
|
||||
if valid_people:
|
||||
rprint(
|
||||
f" Filtered to {len(valid_people)} people with"
|
||||
f" ≥{min_face_count} assets (MIN_FACE_COUNT={min_face_count})"
|
||||
)
|
||||
|
||||
jobs = []
|
||||
for person in valid_people:
|
||||
name = person["name"]
|
||||
entity_type = mode
|
||||
|
||||
config = {"name": name, "mode": entity_type}
|
||||
if entity_type == "object":
|
||||
config["object_class"] = os.environ.get("OBJECT_CLASS", "dog")
|
||||
|
||||
all_assets = fetch_all_assets(person)
|
||||
recent_assets = filter_recent_assets(all_assets, years=Config.YEARS_FILTER)
|
||||
|
||||
rprint(f" {name}: {len(all_assets)} total, {len(recent_assets)} recent")
|
||||
|
||||
# Filter out assets already uploaded to Frigate
|
||||
retry_rejected = os.environ.get("RETRY_REJECTED", "false").lower() in ("true", "1", "yes")
|
||||
before_dedup = len(recent_assets)
|
||||
new_asset_ids = set(filter_already_uploaded([a["id"] for a in recent_assets], retry_rejected=retry_rejected))
|
||||
recent_assets = [a for a in recent_assets if a["id"] in new_asset_ids]
|
||||
skipped = before_dedup - len(recent_assets)
|
||||
if skipped:
|
||||
rprint(f" [dim]Skipped {skipped} assets already uploaded to Frigate.[/dim]")
|
||||
|
||||
if not recent_assets:
|
||||
rprint(f" [dim]Skipping {name} (0 new images after dedup).[/dim]")
|
||||
continue
|
||||
|
||||
has_embedding = is_embedding_available(entity_type)
|
||||
limit, selection_mode = _resolve_strategy(strategy, has_embedding)
|
||||
|
||||
if selection_mode == "skip":
|
||||
continue
|
||||
|
||||
selected_assets = _perform_selection(recent_assets, limit, name, selection_mode, entity_type)
|
||||
|
||||
if selected_assets:
|
||||
rprint(f" [green]Queued {len(selected_assets)} images for {name}.[/green]")
|
||||
jobs.append({"person": person, "assets": selected_assets, "limit": len(selected_assets), "config": config})
|
||||
|
||||
return jobs
|
||||
|
||||
|
||||
def _show_preview(jobs: list[dict]) -> None:
|
||||
"""Show a summary table of all queued jobs before execution."""
|
||||
table = Table(title="📋 Training Job Preview", show_header=True, header_style="bold cyan")
|
||||
table.add_column("Person", style="bold")
|
||||
table.add_column("Mode", style="dim")
|
||||
table.add_column("Images", justify="right")
|
||||
table.add_column("Date Range", style="dim")
|
||||
|
||||
for job in jobs:
|
||||
name = job["person"]["name"]
|
||||
mode = job["config"].get("mode", "face")
|
||||
count = str(job["limit"])
|
||||
|
||||
# Date range
|
||||
dates = sorted(a.get("fileCreatedAt", "")[:10] for a in job["assets"] if a.get("fileCreatedAt"))
|
||||
date_range = f"{dates[0]} → {dates[-1]}" if len(dates) >= 2 else (dates[0] if dates else "—")
|
||||
|
||||
table.add_row(name, mode, count, date_range)
|
||||
|
||||
console.print()
|
||||
console.print(table)
|
||||
console.print()
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Logging configuration for winnow."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import warnings
|
||||
|
||||
from rich.console import Console
|
||||
from rich.logging import RichHandler
|
||||
|
||||
# Shared console instance - must be the same as used by Progress bars
|
||||
console = Console()
|
||||
|
||||
NOISY_LOGGERS = (
|
||||
"urllib3",
|
||||
"PIL",
|
||||
"ultralytics",
|
||||
"insightface",
|
||||
"onnxruntime",
|
||||
"matplotlib",
|
||||
"transformers",
|
||||
"torch",
|
||||
)
|
||||
|
||||
|
||||
def setup_logging(verbose: bool = False) -> logging.Logger:
|
||||
"""Configure logging with Rich console and file output."""
|
||||
level = logging.DEBUG if verbose else logging.INFO
|
||||
|
||||
# Configure root logger
|
||||
root = logging.getLogger()
|
||||
root.setLevel(level)
|
||||
root.handlers.clear()
|
||||
|
||||
# Rich console handler - uses shared console to avoid breaking progress bars
|
||||
root.addHandler(RichHandler(rich_tracebacks=True, markup=True, console=console))
|
||||
|
||||
# File handler (always debug level) — log file respects OUTPUT_DIR if set
|
||||
log_dir = os.environ.get("OUTPUT_DIR", ".")
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
log_path = os.path.join(log_dir, "immich_export.log")
|
||||
file_handler = logging.FileHandler(log_path)
|
||||
file_handler.setLevel(logging.DEBUG)
|
||||
file_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
|
||||
root.addHandler(file_handler)
|
||||
|
||||
# Silence noisy libraries
|
||||
for lib in NOISY_LOGGERS:
|
||||
logging.getLogger(lib).setLevel(logging.WARNING)
|
||||
|
||||
# Suppress Python warnings from ML libraries
|
||||
warnings.filterwarnings("ignore", category=UserWarning, module="onnxruntime")
|
||||
warnings.filterwarnings("ignore", category=FutureWarning, module="transformers")
|
||||
|
||||
return root
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
"""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 = 50) -> 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 = 50,
|
||||
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)
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
"""Persistent tracker for Immich asset IDs already uploaded/rejected by Frigate.
|
||||
|
||||
Two separate JSON files in CACHE_DIR:
|
||||
frigate_uploaded_ids.json — successfully uploaded assets
|
||||
frigate_rejected_ids.json — assets Frigate rejected (e.g. no face detected)
|
||||
|
||||
Both are excluded from future candidate pools. To reset:
|
||||
- All: delete both files
|
||||
- One person: call reset_person("Name") or set RESET_PERSON=Name
|
||||
- Rejects only: delete frigate_rejected_ids.json, or set RETRY_REJECTED=true
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
UPLOAD_TRACKER_FILE = "frigate_uploaded_ids.json"
|
||||
REJECT_TRACKER_FILE = "frigate_rejected_ids.json"
|
||||
|
||||
|
||||
def _tracker_path(filename: str) -> Path:
|
||||
try:
|
||||
from .config import Config
|
||||
return Path(Config.CACHE_DIR) / filename
|
||||
except (ImportError, AttributeError):
|
||||
return Path(filename)
|
||||
|
||||
|
||||
def _load(filename: str) -> dict:
|
||||
path = _tracker_path(filename)
|
||||
if not path.exists():
|
||||
return {}
|
||||
try:
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
except (json.JSONDecodeError, OSError) as e:
|
||||
logger.warning(f"Could not load tracker {filename}: {e}")
|
||||
return {}
|
||||
|
||||
|
||||
def _save(filename: str, data: dict) -> None:
|
||||
path = _tracker_path(filename)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(path, "w") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
|
||||
|
||||
def _flat_key(filename: str) -> str:
|
||||
return "uploaded_asset_ids" if "uploaded" in filename else "rejected_asset_ids"
|
||||
|
||||
|
||||
def _load_flat(filename: str) -> set[str]:
|
||||
return set(_load(filename).get(_flat_key(filename), []))
|
||||
|
||||
|
||||
def _mark(filename: str, asset_id: str, person_name: str | None) -> None:
|
||||
data = _load(filename)
|
||||
flat_key = _flat_key(filename)
|
||||
flat = set(data.get(flat_key, []))
|
||||
flat.add(asset_id)
|
||||
data[flat_key] = sorted(flat)
|
||||
if person_name:
|
||||
by_person = data.setdefault("by_person", {})
|
||||
person_ids = set(by_person.get(person_name, []))
|
||||
person_ids.add(asset_id)
|
||||
by_person[person_name] = sorted(person_ids)
|
||||
_save(filename, data)
|
||||
|
||||
|
||||
# ── Public API ────────────────────────────────────────────────────────────────
|
||||
|
||||
def load_uploaded_ids() -> set[str]:
|
||||
return _load_flat(UPLOAD_TRACKER_FILE)
|
||||
|
||||
|
||||
def load_rejected_ids() -> set[str]:
|
||||
return _load_flat(REJECT_TRACKER_FILE)
|
||||
|
||||
|
||||
def mark_uploaded(asset_id: str, person_name: str | None = None) -> None:
|
||||
_mark(UPLOAD_TRACKER_FILE, asset_id, person_name)
|
||||
logger.debug(f"Marked {asset_id} as uploaded ({person_name})")
|
||||
|
||||
|
||||
def mark_rejected(asset_id: str, person_name: str | None = None) -> None:
|
||||
_mark(REJECT_TRACKER_FILE, asset_id, person_name)
|
||||
logger.debug(f"Marked {asset_id} as rejected ({person_name})")
|
||||
|
||||
|
||||
def reset_person(person_name: str) -> None:
|
||||
"""Remove all uploaded and rejected records for a given person."""
|
||||
for filename in (UPLOAD_TRACKER_FILE, REJECT_TRACKER_FILE):
|
||||
data = _load(filename)
|
||||
flat_key = _flat_key(filename)
|
||||
by_person = data.get("by_person", {})
|
||||
person_ids = set(by_person.pop(person_name, []))
|
||||
if person_ids:
|
||||
flat = set(data.get(flat_key, [])) - person_ids
|
||||
data[flat_key] = sorted(flat)
|
||||
data["by_person"] = by_person
|
||||
_save(filename, data)
|
||||
logger.info(f"Reset tracking data for {person_name}")
|
||||
|
||||
|
||||
def get_person_summary() -> dict[str, dict[str, int]]:
|
||||
"""Return {person_name: {uploaded: N, rejected: N}} for display."""
|
||||
uploaded_by = _load(UPLOAD_TRACKER_FILE).get("by_person", {})
|
||||
rejected_by = _load(REJECT_TRACKER_FILE).get("by_person", {})
|
||||
names = set(uploaded_by) | set(rejected_by)
|
||||
return {
|
||||
name: {
|
||||
"uploaded": len(uploaded_by.get(name, [])),
|
||||
"rejected": len(rejected_by.get(name, [])),
|
||||
}
|
||||
for name in sorted(names)
|
||||
}
|
||||
|
||||
|
||||
def filter_already_uploaded(
|
||||
asset_ids: list[str],
|
||||
retry_rejected: bool = False,
|
||||
) -> list[str]:
|
||||
"""Return asset IDs not yet uploaded (and not rejected, unless retry_rejected)."""
|
||||
exclude = load_uploaded_ids()
|
||||
if not retry_rejected:
|
||||
exclude |= load_rejected_ids()
|
||||
new_ids = [aid for aid in asset_ids if aid not in exclude]
|
||||
skipped = len(asset_ids) - len(new_ids)
|
||||
if skipped:
|
||||
logger.info(f"Skipping {skipped} assets already uploaded or rejected by Frigate")
|
||||
return new_ids
|
||||
Reference in New Issue
Block a user