feat: Implement advanced diversity selection algorithms, comprehensive quality filtering, and caching for improved image curation.
This commit is contained in:
@@ -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
|
||||||
|
|||||||
@@ -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.
|
||||||
@@ -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
|
||||||
|
```
|
||||||
|
|||||||
@@ -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,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
@@ -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
@@ -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
@@ -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":
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
@@ -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"
|
||||||
|
|||||||
Reference in New Issue
Block a user