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

This commit is contained in:
Sebastian G
2026-03-02 20:09:45 -05:00
parent 2c4f93bfb9
commit cc7293b66a
15 changed files with 1836 additions and 781 deletions
+1
View File
@@ -5,6 +5,7 @@ yolov9c.pt
.insightface/ .insightface/
.huggingface/ .huggingface/
.cache/huggingface .cache/huggingface
.if_cache/
.immich_config.json .immich_config.json
# Python-generated files # Python-generated files
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2026 Sebastian
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
+107 -50
View File
@@ -11,31 +11,6 @@
</div> </div>
---
> [!WARNING]
> This is 100% vibe-coded.
## ⚡ Why This Tool?
> **"Diversity matters far more than volume."** — *Frigate Developer Tips*
Training AI models on "bulk" data is often harmful. If you feed the model 50 images from the same 10-second video clip, it learns to recognize the *lighting and background*, not the actual *face* or *object*.
`if-curator` solves this using **AI-powered diversity selection**:
| Mode | Embedding Model | Algorithm |
| :--- | :--- | :--- |
| **👤 Face** | InsightFace (ArcFace) | Farthest Point Sampling |
| **🐶 Object** | SigLIP (Vision Transformer) | Farthest Point Sampling |
Both use the same **Farthest Point Sampling (FPS)** algorithm that mathematically selects images until redundancy starts, ensuring optimal diversity whether that's 20 or 150 images.
> [!WARNING] > [!WARNING]
> **Regarding Object Classification** > **Regarding Object Classification**
> >
@@ -44,25 +19,67 @@ Both use the same **Farthest Point Sampling (FPS)** algorithm that mathematicall
--- ---
## ⚡ Why This Tool?
> **"Diversity matters far more than volume."** — *Frigate Developer Tips*
Training AI models on "bulk" data is often harmful. If you feed the model 50 images from the same 10-second video clip, it learns to recognize the *lighting and background*, not the actual *face* or *object*.
`if-curator` solves this using **AI-powered diversity selection** and **quality filtering**:
| Mode | Embedding Model | Algorithm |
| :--- | :--- | :--- |
| **👤 Face** | InsightFace (ArcFace) | K-Medoids Clustering + FPS + Hard Example Weighting |
| **🐶 Object** | SigLIP (Vision Transformer) | K-Medoids Clustering + FPS |
The pipeline uses **K-Medoids clustering** to guarantee coverage of every distinct "look", then fills the remaining budget with **Farthest Point Sampling (FPS)** biased toward **hard examples** (unusual angles, partial occlusions). An **adaptive threshold** stops selection automatically when adding more images becomes redundant.
---
## ✨ Features ## ✨ Features
### 🎯 Unified Selection Strategies ### 🎯 Smart Selection
Both Face and Object modes offer the same powerful options: - **Auto Diversity [Recommended]**: Clusters images by visual similarity, selects representatives from each cluster, then fills with maximally-diverse picks until redundancy starts (capped at 80)
- **Auto (Objective Diversity) [Recommended]**: Dynamically selects images until redundancy starts
- **Standard (30 images)**: Balanced set using Smart Diversity - **Standard (30 images)**: Balanced set using Smart Diversity
- **Broad (100 images)**: Extensive set using Smart Diversity - **Broad (100 images)**: Extensive set using Smart Diversity
- **Custom Count**: You choose the limit - **Custom Count**: You choose the limit
### � Quality Filtering
Bad training data hurts ArcFace models. Images are automatically rejected if they are:
- **Blurry** — Laplacian variance below threshold
- **Grayscale / IR** — ArcFace is trained on color images only
- **Over/Underexposed** — Washed-out or too dark to use
- **Low confidence** — Partial or occluded face detections
- **Too small** — Faces under 100px (configurable) lack features
### 👤 Face Recognition Prep ### 👤 Face Recognition Prep
- Uses **InsightFace** (ArcFace/Buffalo_L) embeddings - Uses **InsightFace** (ArcFace/Buffalo_L) embeddings on **face crops** (not full images — avoids wrong-face in group photos)
- Extracts faces using Immich's metadata - **Hard example prioritization** — unusual angles, sunglasses, and low-confidence detections are biased for selection
- **Auto-Diversity** picks the optimal set size based on visual distinctness - **Face alignment** via InsightFace landmarks (standard 112×112 ArcFace input)
- Downloads **full-resolution** originals for final crops (falls back to JPEG preview for HEIC/RAW)
- Configurable crop margin (default 15%)
### 📦 Object/State Classification Prep ### 📦 Object/State Classification Prep
- Uses **SigLIP** (Vision Transformer) embeddings for semantic diversity - Uses **SigLIP** (Vision Transformer) embeddings for semantic diversity
- **YOLOv9c** to detect and crop specific objects (dogs, cars, etc.) - **YOLOv9c** to detect and crop specific objects (dogs, cars, etc.)
- Captures variation in poses, lighting, and backgrounds - Captures variation in poses, lighting, and backgrounds
- *Note: As mentioned, Frigate upload is pending support.*
### ⚡ Performance
- **Concurrent thumbnail downloads** (8 parallel workers)
- **Batch-capable** SigLIP embeddings for GPU efficiency
- Optional **disk-based embedding cache** for faster re-runs
- **Multi-person batch mode** — process multiple people in one session
### 📋 Preview Before Download
After selection, a summary table shows what will be processed:
```
📋 Training Job Preview
┏━━━━━━━━━━━┳━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ Person ┃ Mode ┃ Images ┃ Date Range ┃
┡━━━━━━━━━━━╇━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ Sebastian │ face │ 80 │ 2021-04-03 → 2026-02-18 │
└───────────┴──────┴────────┴─────────────────────────┘
```
--- ---
@@ -93,22 +110,25 @@ uv sync --extra gpu
## 💻 Usage ## 💻 Usage
Run the command-line interface:
```bash ```bash
uv run -m if_curator uv run if-curator
``` ```
### Interactive Flow ### Interactive Flow
The tool will guide you through: The tool will guide you through:
1. **Select Person/Subject**: Choose from your Immich people. 1. **Select Person** — Choose from your Immich people (supports multi-person batch)
2. **Training Mode**: Face (Recognition) or Object (Classification). 2. **Training Mode** — Face (Recognition) or Object (Classification)
3. **Strategy**: Auto, Standard, Broad, etc. 3. **Strategy** — Auto, Standard, Broad, or Custom
4. **Preview** — Review the selection summary before downloading
5. **Execute** — Downloads and processes images with progress tracking
```text ```text
Using SigLIP (visual embeddings) for diversity analysis... Using InsightFace (face embeddings) for diversity analysis...
Computing embeddings for 69 images... ████████████████ 100% Quality filtering removed 76 images.
Auto-diversity selected 38 optimally diverse images. Adaptive threshold: 0.1721 (median_dist=0.8605, fraction=0.2)
Clustering 223 embeddings into 20 groups (K-Medoids)...
Selected 20 cluster medoids as initial picks.
Selection complete: 80 images (0 hard examples with confidence < 0.85).
``` ```
--- ---
@@ -117,18 +137,55 @@ Auto-diversity selected 38 optimally diverse images.
The tool prompts for your Immich URL and API Key on the first run and saves them to `.immich_config.json`. The tool prompts for your Immich URL and API Key on the first run and saves them to `.immich_config.json`.
| Variable | Description | ### Environment Variables
| :--- | :--- |
| `IMMICH_URL` | Full URL to Immich (e.g. `http://192.168.1.10:2283`) | | Variable | Default | Description |
| `API_KEY` | Your Immich API Key | | :--- | :--- | :--- |
| `FORCE_CPU` | Set to `true` to disable GPU acceleration | | `IMMICH_URL` | — | Full URL to Immich (e.g. `http://192.168.1.10:2283`) |
| `API_KEY` | — | Your Immich API Key |
| `FORCE_CPU` | `false` | Disable GPU acceleration |
| `MIN_FACE_WIDTH` | `100` | Minimum face crop size (pixels) |
| `BLUR_THRESHOLD` | `100.0` | Laplacian variance threshold for blur detection |
| `MIN_CONFIDENCE` | `0.7` | Minimum Immich detection confidence |
| `MAX_AUTO_IMAGES` | `80` | Safety cap for auto-diversity mode |
| `FACE_MARGIN` | `0.15` | Crop margin around face (fraction) |
| `USE_FULL_RESOLUTION` | `true` | Download originals for final crops |
| `ENABLE_FACE_ALIGNMENT` | `true` | Align faces to ArcFace 112×112 format |
| `ENABLE_CACHE` | `false` | Cache embeddings to disk for faster re-runs |
| `CACHE_DIR` | `.if_cache` | Directory for embedding cache |
--- ---
## 🧠 Technical Details ## 🧠 Technical Details
- **InsightFace**: Face detection and embedding (ArcFace) ### Models
- **SigLIP**: Visual embeddings via `transformers` (OpenAI CLIP alternative) - **InsightFace (Buffalo_L)** — Face detection and embedding (ArcFace, 512-d)
- **YOLOv9c**: State-of-the-art object detection for cropping - **SigLIP** — Visual embeddings via `transformers` (google/siglip-base-patch16-224, 768-d)
- **Rich**: Beautiful terminal UI - **YOLOv9c** — Object detection for cropping
### Algorithms
- **K-Medoids Clustering** — Groups embeddings into k clusters using cosine distance, selecting actual data points (medoids) as cluster centers. Guarantees one representative from every distinct "look"
- **Farthest Point Sampling** — After medoid selection, fills remaining budget by iteratively selecting the most distant point from the current set
- **Hard Example Weighting** — Candidates with detection confidence < 0.85 get a 1.2–1.5× distance boost, biasing selection toward challenging images (unusual angles, occlusions)
- **Adaptive Auto-Threshold** — Computed as 20% of the median pairwise cosine distance; stops when the next-best image is too similar
- **Quality Filtering** — Blur (Laplacian), grayscale/IR (channel comparison), exposure (histogram), confidence (Immich metadata)
- **Face Crop Embedding** — Extracts the target person's face (using Immich bbox) before embedding, preventing wrong-face selection in group photos
### Architecture
```
Immich API ─► Fetch Assets by Person ─► Time Filter
│
Concurrent Thumbnail Download (8 workers)
│
Quality Filtering (blur, IR, exposure...)
│
Face Crop Extraction (bbox from Immich metadata)
│
Compute Embeddings (InsightFace / SigLIP)
│
K-Medoids Clustering → FPS + Hard Example Weighting
│
Preview Summary Table
│
Download Full-Res ─► Face Alignment ─► Save
```
+8
View File
@@ -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"
+87
View File
@@ -0,0 +1,87 @@
"""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."""
global _cache
if _cache is None:
_cache = EmbeddingCache(cache_dir)
return _cache
+96 -33
View File
@@ -9,13 +9,15 @@ from PIL import Image
from rich import print as rprint from rich import print as rprint
from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
from rich.prompt import Confirm, IntPrompt, Prompt from rich.prompt import Confirm, IntPrompt, Prompt
from rich.table import Table
from .config import Config, ConfigManager from .config import Config, ConfigManager
from .diversity import select_diverse_assets from .diversity import select_diverse_assets
from .embeddings import is_embedding_available from .embeddings import is_embedding_available
from .image_processing import process_face_mode, process_full_mode, process_object_mode from .image_processing import process_face_mode, process_full_mode, process_object_mode
from .immich_api import fetch_all_assets, filter_recent_assets, get_people from .immich_api import fetch_all_assets, fetch_full_image, filter_recent_assets, get_people
from .logging import console, setup_logging from .logging import console, setup_logging
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Strategy presets: (limit, mode_name) # Strategy presets: (limit, mode_name)
@@ -62,23 +64,9 @@ def _get_strategy_choice(has_embedding: bool, entity_type: str) -> tuple[int | s
return limits.get(choice, 0), "time" if choice != "4" else "skip" return limits.get(choice, 0), "time" if choice != "4" else "skip"
def interactive_configure(people: list[dict]) -> list[dict]: def _configure_person(person: dict, people: list[dict]) -> dict | None:
"""Interactive phase: select person, mode, and configure training strategy.""" """Configure training for a single person. Returns job dict or None."""
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 []
# Select person
console.print("\n[bold cyan]Select Person to Train:[/bold cyan]")
for idx, p in enumerate(valid_people, 1):
console.print(f" [bold]{idx}.[/bold] {p['name']}")
p_choice = IntPrompt.ask("Enter Number", choices=[str(i) for i in range(1, len(valid_people) + 1)])
person = valid_people[p_choice - 1]
name = person["name"] name = person["name"]
console.print(f"\nSelected: [bold green]{name}[/bold green]") console.print(f"\nSelected: [bold green]{name}[/bold green]")
# Select training mode # Select training mode
@@ -105,7 +93,7 @@ def interactive_configure(people: list[dict]) -> list[dict]:
if not recent_assets: if not recent_assets:
rprint(" [dim]Skipping (0 recent images).[/dim]") rprint(" [dim]Skipping (0 recent images).[/dim]")
return [] return None
# Strategy selection # Strategy selection
has_embedding = is_embedding_available(entity_type) has_embedding = is_embedding_available(entity_type)
@@ -113,18 +101,52 @@ def interactive_configure(people: list[dict]) -> list[dict]:
limit, selection_mode = _get_strategy_choice(has_embedding, entity_type) limit, selection_mode = _get_strategy_choice(has_embedding, entity_type)
if selection_mode == "skip": if selection_mode == "skip":
return [] return None
# Perform selection # Perform selection
selected_assets = _perform_selection(recent_assets, limit, name, selection_mode, entity_type) selected_assets = _perform_selection(recent_assets, limit, name, selection_mode, entity_type)
rprint(f" [green]Queued {len(selected_assets)} images for {name}.[/green]") rprint(f" [green]Queued {len(selected_assets)} images for {name}.[/green]")
return [{"person": person, "assets": selected_assets, "limit": len(selected_assets), "config": config}] return {"person": person, "assets": selected_assets, "limit": len(selected_assets), "config": config}
def _perform_selection( def interactive_configure(people: list[dict]) -> list[dict]:
assets: list, limit: int | str, name: str, selection_mode: str, entity_type: str """Interactive phase: select person(s), mode, and configure training strategy.
) -> list:
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 _perform_selection(assets: list, limit: int | str, name: str, selection_mode: str, entity_type: str) -> list:
"""Run diversity selection with progress display.""" """Run diversity selection with progress display."""
if selection_mode == "smart": if selection_mode == "smart":
model_display = "InsightFace (face embeddings)" if entity_type == "face" else "SigLIP (visual embeddings)" model_display = "InsightFace (face embeddings)" if entity_type == "face" else "SigLIP (visual embeddings)"
@@ -134,12 +156,17 @@ def _perform_selection(
is_embedding_available(entity_type) is_embedding_available(entity_type)
with Progress( with Progress(
SpinnerColumn(), TextColumn("[progress.description]{task.description}"), SpinnerColumn(),
BarColumn(), TaskProgressColumn(), console=console, TextColumn("[progress.description]{task.description}"),
BarColumn(),
TaskProgressColumn(),
console=console,
) as progress: ) as progress:
task = progress.add_task(f"[cyan]Computing embeddings for {len(assets)} images...", total=None) task = progress.add_task(f"[cyan]Computing embeddings for {len(assets)} images...", total=None)
selected = select_diverse_assets( selected = select_diverse_assets(
assets, limit, name, assets,
limit,
name,
selection_mode=selection_mode, selection_mode=selection_mode,
entity_type=entity_type, entity_type=entity_type,
progress_callback=lambda c, t: progress.update(task, completed=c, total=t), progress_callback=lambda c, t: progress.update(task, completed=c, total=t),
@@ -156,6 +183,30 @@ def _perform_selection(
return selected return selected
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()
def execute_jobs(jobs: list[dict]) -> None: def execute_jobs(jobs: list[dict]) -> None:
"""Download and process images for all jobs.""" """Download and process images for all jobs."""
if not jobs: if not jobs:
@@ -163,9 +214,14 @@ def execute_jobs(jobs: list[dict]) -> None:
console.rule("[bold blue]Execution Phase") console.rule("[bold blue]Execution Phase")
use_full_res = Config.USE_FULL_RESOLUTION
with Progress( with Progress(
SpinnerColumn(), TextColumn("[progress.description]{task.description}"), SpinnerColumn(),
BarColumn(), TaskProgressColumn(), console=console, TextColumn("[progress.description]{task.description}"),
BarColumn(),
TaskProgressColumn(),
console=console,
) as progress: ) as progress:
grand_total = sum(j["limit"] for j in jobs) grand_total = sum(j["limit"] for j in jobs)
overall_task = progress.add_task("[green]Overall Progress", total=grand_total) overall_task = progress.add_task("[green]Overall Progress", total=grand_total)
@@ -181,17 +237,22 @@ def execute_jobs(jobs: list[dict]) -> None:
count = 0 count = 0
for asset in assets: for asset in assets:
try: try:
# Use full-resolution for final output when configured
if use_full_res:
img = fetch_full_image(asset["id"])
else:
resp = requests.get( resp = requests.get(
f"{Config.IMMICH_URL}/api/assets/{asset['id']}/thumbnail?size=preview&format=JPEG", f"{Config.IMMICH_URL}/api/assets/{asset['id']}/thumbnail?size=preview&format=JPEG",
headers={"x-api-key": Config.API_KEY, "Accept": "application/json"}, headers={"x-api-key": Config.API_KEY, "Accept": "application/json"},
timeout=30, timeout=30,
) )
if not resp.ok: 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]") progress.console.print(f"[red]Failed download {asset['id']}[/red]")
else: else:
img = Image.open(BytesIO(resp.content))
saved = ( saved = (
process_face_mode(img, asset, person, person_dir, count, min_width=Config.MIN_FACE_WIDTH) process_face_mode(img, asset, person, person_dir, count)
if mode == "face" if mode == "face"
else process_object_mode(img, config, person_dir, count) else process_object_mode(img, config, person_dir, count)
if mode == "object" if mode == "object"
@@ -236,10 +297,12 @@ def main() -> None:
jobs = interactive_configure(people) jobs = interactive_configure(people)
if jobs and Confirm.ask(f"\nReady to process {sum(j['limit'] for j in jobs)} images?"): if jobs:
_show_preview(jobs)
if Confirm.ask(f"Ready to process {sum(j['limit'] for j in jobs)} images?"):
execute_jobs(jobs) execute_jobs(jobs)
rprint("\n[bold green]Done! Happy Training.[/bold green]") rprint("\n[bold green]Done! Happy Training.[/bold green]")
elif not jobs: else:
rprint("[yellow]No jobs configured.[/yellow]") rprint("[yellow]No jobs configured.[/yellow]")
except KeyboardInterrupt: except KeyboardInterrupt:
+15 -1
View File
@@ -24,7 +24,21 @@ class Config:
API_KEY: str | None = None API_KEY: str | None = None
OUTPUT_DIR: str = "./frigate_train" OUTPUT_DIR: str = "./frigate_train"
YEARS_FILTER: int = 10 YEARS_FILTER: int = 10
MIN_FACE_WIDTH: int = 50
# Quality filtering
MIN_FACE_WIDTH: int = 100
BLUR_THRESHOLD: float = 100.0
MIN_CONFIDENCE: float = 0.7
MAX_AUTO_IMAGES: int = 80
# 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": def __new__(cls) -> "Config":
if cls._instance is None: if cls._instance is None:
+330 -30
View File
@@ -1,9 +1,12 @@
""" """
Diversity selection for training data curation. Diversity selection for training data curation.
Uses Farthest Point Sampling (FPS) algorithm with embeddings: Selection pipeline:
- Faces: InsightFace embeddings 1. Concurrent thumbnail download
- Objects: SigLIP embeddings 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 import logging
@@ -15,6 +18,7 @@ from PIL import Image
from .config import Config, get_headers from .config import Config, get_headers
from .embeddings import get_embedding, is_embedding_available from .embeddings import get_embedding, is_embedding_available
from .quality import assess_quality
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -28,7 +32,7 @@ def select_diverse_assets(
progress_callback=None, progress_callback=None,
) -> list: ) -> list:
""" """
Select diverse assets using Farthest Point Sampling or time spread. Select diverse assets using cluster-aware FPS or time spread.
Args: Args:
assets: List of asset dicts from Immich API assets: List of asset dicts from Immich API
@@ -61,6 +65,11 @@ def select_diverse_assets(
return _select_time_spread(assets, limit) return _select_time_spread(assets, limit)
# =============================================================================
# Thumbnail & Metadata Helpers
# =============================================================================
def _fetch_thumbnail(asset_id: str, timeout: int = 10) -> Image.Image | None: def _fetch_thumbnail(asset_id: str, timeout: int = 10) -> Image.Image | None:
"""Fetch thumbnail from Immich API.""" """Fetch thumbnail from Immich API."""
try: try:
@@ -71,13 +80,106 @@ def _fetch_thumbnail(asset_id: str, timeout: int = 10) -> Image.Image | None:
return None 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( def _select_by_embedding(
assets: list, assets: list,
limit: int | str, limit: int | str,
entity_type: str, entity_type: str,
progress_callback=None, progress_callback=None,
) -> list: ) -> list:
"""Select assets using embedding-based Farthest Point Sampling.""" """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) # Determine candidate pool (cap at 3000 for performance)
effective_limit = 30 if limit == "auto" else limit effective_limit = 30 if limit == "auto" else limit
pool_size = min(3000, max(effective_limit * 20, len(assets))) pool_size = min(3000, max(effective_limit * 20, len(assets)))
@@ -89,24 +191,68 @@ def _select_by_embedding(
else: else:
candidates = assets candidates = assets
# Compute embeddings # --- Phase 1: Concurrent thumbnail download ---
embeddings, valid_candidates = [], [] from concurrent.futures import ThreadPoolExecutor, as_completed
for i, asset in enumerate(candidates):
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: if progress_callback:
progress_callback(i, len(candidates)) 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
img = _fetch_thumbnail(asset["id"]) # --- 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: if img is None:
continue continue
emb = get_embedding(img, entity_type) 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: if emb is not None:
embeddings.append(emb) embeddings.append(emb)
valid_candidates.append(asset) valid_candidates.append(asset)
confidence_scores.append(confidence)
if progress_callback: if progress_callback:
progress_callback(len(candidates), len(candidates)) progress_callback(len(candidates), len(candidates))
if quality_filtered > 0:
logger.info(f"Quality filtering removed {quality_filtered} images.")
if not embeddings: if not embeddings:
logger.warning("No valid embeddings found. Falling back to time spread.") logger.warning("No valid embeddings found. Falling back to time spread.")
return _select_time_spread(assets, limit) return _select_time_spread(assets, limit)
@@ -115,58 +261,212 @@ def _select_by_embedding(
logger.warning(f"Only {len(valid_candidates)} valid embeddings. Returning all.") logger.warning(f"Only {len(valid_candidates)} valid embeddings. Returning all.")
return valid_candidates return valid_candidates
# Farthest Point Sampling with vectorized distance computation # --- Phase 5: Cluster-aware selection ---
return _farthest_point_sampling( return _cluster_aware_selection(
embeddings, valid_candidates, limit, auto_threshold=0.15 embeddings,
valid_candidates,
limit,
entity_type=entity_type,
confidence_scores=confidence_scores,
) )
def _farthest_point_sampling( # =============================================================================
# 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, embeddings: list,
candidates: list, candidates: list,
limit: int | str, limit: int | str,
auto_threshold: float = 0.15, entity_type: str = "face",
confidence_scores: list | None = None,
) -> list: ) -> list:
"""Vectorized Farthest Point Sampling.""" """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) emb_matrix = np.vstack(embeddings) # (N, D)
n = len(emb_matrix) n = len(emb_matrix)
# Normalize for cosine distance (cosine_dist = 1 - cosine_sim) # Normalize for cosine distance
norms = np.linalg.norm(emb_matrix, axis=1, keepdims=True) norms = np.linalg.norm(emb_matrix, axis=1, keepdims=True)
emb_normed = emb_matrix / np.maximum(norms, 1e-8) emb_normed = emb_matrix / np.maximum(norms, 1e-8)
# Start from median-time sample # Build confidence weight array for hard example boosting
selected = [n // 2] 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) min_dists = np.full(n, np.inf)
target = 500 if limit == "auto" else limit # 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: while len(selected) < target:
# Update min distances with last selected point # Hard example weighting: boost distance for low-confidence candidates
last_emb = emb_normed[selected[-1]] # Confidence < 0.85 gets up to 1.5× distance boost
dists_to_last = 1 - emb_normed @ last_emb # Cosine distance hard_weight = np.where(conf_array < 0.85, 1.0 + (0.85 - conf_array) * 2.0, 1.0)
min_dists = np.minimum(min_dists, dists_to_last) weighted_dists = min_dists * hard_weight
min_dists[selected[-1]] = -np.inf # Exclude already selected
# Find farthest point best_idx = int(np.argmax(weighted_dists))
best_idx = np.argmax(min_dists) best_dist = min_dists[best_idx] # Use unweighted for threshold comparison
best_dist = min_dists[best_idx]
if best_dist == -np.inf: if best_dist == -np.inf:
break # All points selected break # All points selected
if limit == "auto" and best_dist < auto_threshold: if limit == "auto" and best_dist < auto_threshold:
logger.info( logger.info(
f"Auto-stop: Next best image {best_dist:.3f} away (threshold {auto_threshold})." f"Auto-stop: Next best image {best_dist:.3f} away " f"(adaptive threshold {auto_threshold:.4f})."
) )
break break
selected.append(best_idx) 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.")
logger.info(f"Smart selection complete. Picked {len(selected)} diverse images.")
return [candidates[i] for i in selected] return [candidates[i] for i in selected]
# =============================================================================
# Time Spread Fallback
# =============================================================================
def _select_time_spread(assets: list, limit: int | str) -> list: def _select_time_spread(assets: list, limit: int | str) -> list:
"""Select N assets evenly distributed in time.""" """Select N assets evenly distributed in time."""
if limit == "auto": if limit == "auto":
+82 -7
View File
@@ -1,8 +1,9 @@
""" """
Unified embedding interface for faces and objects. Unified embedding interface for faces and objects.
- Faces: InsightFace (ArcFace/Buffalo_L) - Faces: InsightFace (ArcFace/Buffalo_L) — or reuse from Immich
- Objects: SigLIP (Vision Transformer via transformers) - Objects: SigLIP (Vision Transformer via transformers)
- Caching: Disk-based cache avoids recomputation on reruns
""" """
import contextlib import contextlib
@@ -14,6 +15,8 @@ import cv2
import numpy as np import numpy as np
from PIL import Image from PIL import Image
from .cache import get_cache
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Lazy-loaded singletons # Lazy-loaded singletons
@@ -48,8 +51,10 @@ def get_insightface_app():
# Determine device: 0 for GPU, -1 for CPU # Determine device: 0 for GPU, -1 for CPU
gpu_providers = { gpu_providers = {
"CUDAExecutionProvider", "ROCmExecutionProvider", "CUDAExecutionProvider",
"MPSExecutionProvider", "CoreMLExecutionProvider", "ROCmExecutionProvider",
"MPSExecutionProvider",
"CoreMLExecutionProvider",
} }
ctx_id = -1 if _is_force_cpu() else (0 if gpu_providers & set(providers) else -1) ctx_id = -1 if _is_force_cpu() else (0 if gpu_providers & set(providers) else -1)
@@ -73,6 +78,7 @@ def get_insightface_app():
logger.warning("Retrying InsightFace on CPU...") logger.warning("Retrying InsightFace on CPU...")
try: try:
from insightface.app import FaceAnalysis from insightface.app import FaceAnalysis
_insightface_app = FaceAnalysis(name="buffalo_l", root="~/.insightface") _insightface_app = FaceAnalysis(name="buffalo_l", root="~/.insightface")
_insightface_app.prepare(ctx_id=-1, det_size=(640, 640)) _insightface_app.prepare(ctx_id=-1, det_size=(640, 640))
return _insightface_app return _insightface_app
@@ -179,14 +185,83 @@ def get_object_embedding(img_pil: Image.Image) -> np.ndarray | None:
return None 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 # Unified Interface with Caching
# ============================================================================= # =============================================================================
def get_embedding(img_pil: Image.Image, entity_type: str = "face") -> np.ndarray | None: def get_embedding(
"""Get embedding for an image based on entity type ('face' or 'object').""" img_pil: Image.Image,
return get_face_embedding(img_pil) if entity_type == "face" else get_object_embedding(img_pil) 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_embedding_available(entity_type: str = "face") -> bool: def is_embedding_available(entity_type: str = "face") -> bool:
+59 -7
View File
@@ -3,8 +3,11 @@
import logging import logging
import os import os
import numpy as np
from PIL import Image from PIL import Image
from .config import Config
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Lazy singleton # Lazy singleton
@@ -16,20 +19,59 @@ def get_yolo_model():
global _yolo_model global _yolo_model
if _yolo_model is None: if _yolo_model is None:
from ultralytics import YOLO from ultralytics import YOLO
logger.info("Loading YOLOv9c model...") logger.info("Loading YOLOv9c model...")
_yolo_model = YOLO("yolov9c.pt") _yolo_model = YOLO("yolov9c.pt")
return _yolo_model 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( def process_face_mode(
img: Image.Image, img: Image.Image,
asset: dict, asset: dict,
person: dict, person: dict,
output_dir: str, output_dir: str,
count: int, count: int,
min_width: int = 50, min_width: int | None = None,
) -> bool: ) -> bool:
"""Crop face based on Immich metadata and save to output directory.""" """Crop face based on Immich metadata and save to output directory.
If face alignment is enabled and landmarks are available, produces
an aligned 112x112 crop. Otherwise falls back to bounding box crop
with configurable margin.
"""
min_width = min_width or Config.MIN_FACE_WIDTH
# Find face metadata for this person # Find face metadata for this person
face_info = None face_info = None
for p in asset.get("people", []): for p in asset.get("people", []):
@@ -57,8 +99,20 @@ def process_face_mode(
logger.debug(f"Face too small ({face_w:.1f}x{face_h:.1f})") logger.debug(f"Face too small ({face_w:.1f}x{face_h:.1f})")
return False return False
# Add 10% margin # Try face alignment if enabled and landmarks available
margin_x, margin_y = face_w * 0.10, face_h * 0.10 if Config.ENABLE_FACE_ALIGNMENT:
landmarks = face_info.get("landmarks") or face_info.get("landmark")
if landmarks:
# Scale landmarks
scaled_landmarks = [[lm[0] * scale_x, lm[1] * scale_y] for lm in landmarks]
aligned = align_face(img, scaled_landmarks)
if aligned is not None:
aligned.save(os.path.join(output_dir, f"{count}.jpg"), format="JPEG")
return True
# Fall back to bounding box crop with configurable margin
margin = Config.FACE_MARGIN
margin_x, margin_y = face_w * margin, face_h * margin
crop_box = ( crop_box = (
max(0, x1 - margin_x), max(0, x1 - margin_x),
max(0, y1 - margin_y), max(0, y1 - margin_y),
@@ -87,9 +141,7 @@ def process_object_mode(
found = False found = False
for idx, (box, cls_id, conf) in enumerate( for idx, (box, cls_id, conf) in enumerate(
(box, int(box.cls[0]), float(box.conf[0])) (box, int(box.cls[0]), float(box.conf[0])) for r in results for box in r.boxes
for r in results
for box in r.boxes
): ):
if 0 <= cls_id < len(model.names) and model.names[cls_id] == target_class and conf > 0.5: 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() x1, y1, x2, y2 = box.xyxy[0].tolist()
+117 -1
View File
@@ -1,15 +1,30 @@
"""Immich API client for fetching people and assets.""" """Immich API client for fetching people, assets, and face data."""
import logging import logging
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone from datetime import datetime, timedelta, timezone
from io import BytesIO
import numpy as np
import requests import requests
from PIL import Image
from .config import Config, get_headers from .config import Config, get_headers
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@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]: def get_people() -> list[dict]:
"""Fetch all people from Immich.""" """Fetch all people from Immich."""
try: try:
@@ -68,6 +83,107 @@ def fetch_all_assets(person: dict) -> list[dict]:
return assets 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", {}).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 (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 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 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]: def filter_recent_assets(assets: list[dict], years: int | None = None) -> list[dict]:
"""Filter assets to keep only those from the last N years.""" """Filter assets to keep only those from the last N years."""
years = years or Config.YEARS_FILTER years = years or Config.YEARS_FILTER
+1 -3
View File
@@ -36,9 +36,7 @@ def setup_logging(verbose: bool = False) -> logging.Logger:
# File handler (always debug level) # File handler (always debug level)
file_handler = logging.FileHandler("immich_export.log") file_handler = logging.FileHandler("immich_export.log")
file_handler.setLevel(logging.DEBUG) file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter( file_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
)
root.addHandler(file_handler) root.addHandler(file_handler)
# Silence noisy libraries # Silence noisy libraries
+132
View File
@@ -0,0 +1,132 @@
"""Image quality assessment for training data curation.
Filters out images that would hurt Frigate's ArcFace model training:
blur, grayscale/IR, over/underexposure, tiny faces, low-confidence detections.
"""
import logging
from dataclasses import dataclass, field
import cv2
import numpy as np
from PIL import Image
logger = logging.getLogger(__name__)
@dataclass
class QualityResult:
"""Result of quality assessment on a face/image crop."""
passed: bool
reasons: list[str] = field(default_factory=list)
@property
def reason(self) -> str:
return "; ".join(self.reasons) if self.reasons else "OK"
def check_blur(img_np: np.ndarray, threshold: float = 100.0) -> tuple[bool, str]:
"""Detect blur using Laplacian variance.
Lower variance = blurrier image. ArcFace needs clear facial features.
"""
gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY) if img_np.ndim == 3 else img_np
variance = cv2.Laplacian(gray, cv2.CV_64F).var()
if variance < threshold:
return False, f"Blurry (laplacian={variance:.1f}, threshold={threshold})"
return True, ""
def check_grayscale(img_np: np.ndarray) -> tuple[bool, str]:
"""Detect grayscale/IR images by checking channel similarity.
ArcFace is trained on color images; IR/grayscale degrades recognition.
"""
if img_np.ndim != 3 or img_np.shape[2] < 3:
return False, "Grayscale (single channel)"
# Compare channel means — IR/grayscale has nearly identical R, G, B
means = img_np[:, :, :3].mean(axis=(0, 1))
max_diff = max(abs(means[0] - means[1]), abs(means[1] - means[2]), abs(means[0] - means[2]))
if max_diff < 5.0:
return False, f"Grayscale/IR (channel diff={max_diff:.1f})"
return True, ""
def check_exposure(img_np: np.ndarray, lo: float = 30.0, hi: float = 225.0) -> tuple[bool, str]:
"""Check for severe under/overexposure using mean brightness.
Extremely dark or blown-out faces lack usable features.
"""
gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY) if img_np.ndim == 3 else img_np
mean_brightness = gray.mean()
if mean_brightness < lo:
return False, f"Underexposed (brightness={mean_brightness:.1f}, min={lo})"
if mean_brightness > hi:
return False, f"Overexposed (brightness={mean_brightness:.1f}, max={hi})"
return True, ""
def check_face_size(face_width: float, face_height: float, min_px: int = 100) -> tuple[bool, str]:
"""Validate face crop is large enough for meaningful features."""
if face_width < min_px or face_height < min_px:
return False, f"Face too small ({face_width:.0f}x{face_height:.0f}, min={min_px}px)"
return True, ""
def check_confidence(score: float | None, min_conf: float = 0.7) -> tuple[bool, str]:
"""Check detection confidence from Immich.
Low confidence often means partial, occluded, or false-positive faces.
"""
if score is None:
return True, "" # No score available, pass by default
if score < min_conf:
return False, f"Low confidence ({score:.2f}, min={min_conf})"
return True, ""
def assess_quality(
img: Image.Image,
face_bbox: tuple[float, float, float, float] | None = None,
confidence: float | None = None,
blur_threshold: float = 100.0,
min_face_px: int = 100,
min_confidence: float = 0.7,
) -> QualityResult:
"""Run all quality checks on an image.
Args:
img: PIL Image (RGB)
face_bbox: (x1, y1, x2, y2) bounding box, or None to skip face-size check
confidence: Detection confidence score from Immich, or None
blur_threshold: Laplacian variance threshold for blur detection
min_face_px: Minimum face dimension in pixels
min_confidence: Minimum acceptable detection confidence
Returns:
QualityResult with passed=True if all checks pass
"""
img_np = np.asarray(img)
reasons = []
# Run all checks, collect failures
checks = [
check_blur(img_np, blur_threshold),
check_grayscale(img_np),
check_exposure(img_np),
check_confidence(confidence, min_confidence),
]
if face_bbox is not None:
x1, y1, x2, y2 = face_bbox
checks.append(check_face_size(x2 - x1, y2 - y1, min_face_px))
for passed, reason in checks:
if not passed:
reasons.append(reason)
return QualityResult(passed=len(reasons) == 0, reasons=reasons)
+40 -3
View File
@@ -3,24 +3,37 @@ name = "if-curator"
version = "0.1.0" version = "0.1.0"
description = "Immich to Frigate training sets" description = "Immich to Frigate training sets"
readme = "README.md" readme = "README.md"
license = "MIT"
requires-python = ">=3.12" requires-python = ">=3.12"
authors = [{ name = "ds-sebastian" }]
keywords = ["immich", "frigate", "face-recognition", "training-data", "arcface", "insightface"]
classifiers = [
"Development Status :: 3 - Alpha",
"Intended Audience :: Developers",
"License :: OSI Approved :: MIT License",
"Programming Language :: Python :: 3.12",
"Topic :: Scientific/Engineering :: Image Recognition",
]
dependencies = [ dependencies = [
"insightface>=0.7.3", "insightface>=0.7.3",
"numpy>=2.2.6", "numpy>=2.2.6",
"onnxruntime>=1.23.2", "onnxruntime>=1.23.2",
"opencv-python>=4.12.0.88",
"opencv-python-headless>=4.12.0.88", "opencv-python-headless>=4.12.0.88",
"pillow>=12.1.0", "pillow>=12.1.0",
"python-dotenv>=1.2.1", "python-dotenv>=1.2.1",
"requests>=2.32.5", "requests>=2.32.5",
"rich>=14.2.0", "rich>=14.2.0",
"scipy>=1.17.0",
"sentencepiece>=0.2.1",
"torch>=2.9.1", "torch>=2.9.1",
"transformers>=4.57.6", "transformers>=4.57.6",
"ultralytics>=8.3.252", "ultralytics>=8.3.252",
] ]
[project.scripts]
if-curator = "if_curator.cli:main"
[project.urls]
Repository = "https://github.com/ds-sebastian/if_curator"
[project.optional-dependencies] [project.optional-dependencies]
gpu = [ gpu = [
"onnxruntime-gpu>=1.23.2", "onnxruntime-gpu>=1.23.2",
@@ -28,6 +41,7 @@ gpu = [
[dependency-groups] [dependency-groups]
dev = [ dev = [
"pytest>=8.0",
"ruff>=0.9.2", "ruff>=0.9.2",
] ]
@@ -38,3 +52,26 @@ target-version = "py312"
[tool.ruff.lint] [tool.ruff.lint]
select = ["E", "F", "I"] select = ["E", "F", "I"]
[tool.deptry]
pep621_dev_dependency_groups = ["dev"]
[tool.deptry.package_module_name_map]
pillow = "PIL"
opencv-python-headless = "cv2"
python-dotenv = "dotenv"
insightface = "insightface"
numpy = "numpy"
onnxruntime = "onnxruntime"
requests = "requests"
rich = "rich"
torch = "torch"
transformers = "transformers"
ultralytics = "ultralytics"
[tool.deptry.per_rule_ignores]
# onnxruntime-gpu is an optional dep that replaces onnxruntime at runtime
DEP002 = ["onnxruntime-gpu"]
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
Generated
+731 -637
View File
File diff suppressed because it is too large Load Diff