feat: Implement advanced diversity selection algorithms, comprehensive quality filtering, and caching for improved image curation.
This commit is contained in:
Sebastian G
@@ -5,6 +5,7 @@ yolov9c.pt
|
||||
.insightface/
|
||||
.huggingface/
|
||||
.cache/huggingface
|
||||
.if_cache/
|
||||
|
||||
.immich_config.json
|
||||
# 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>
|
||||
|
||||
|
||||
---
|
||||
|
||||
> [!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]
|
||||
> **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
|
||||
|
||||
### 🎯 Unified Selection Strategies
|
||||
Both Face and Object modes offer the same powerful options:
|
||||
- **Auto (Objective Diversity) [Recommended]**: Dynamically selects images until redundancy starts
|
||||
### 🎯 Smart Selection
|
||||
- **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)
|
||||
- **Standard (30 images)**: Balanced set using Smart Diversity
|
||||
- **Broad (100 images)**: Extensive set using Smart Diversity
|
||||
- **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
|
||||
- Uses **InsightFace** (ArcFace/Buffalo_L) embeddings
|
||||
- Extracts faces using Immich's metadata
|
||||
- **Auto-Diversity** picks the optimal set size based on visual distinctness
|
||||
- Uses **InsightFace** (ArcFace/Buffalo_L) embeddings on **face crops** (not full images — avoids wrong-face in group photos)
|
||||
- **Hard example prioritization** — unusual angles, sunglasses, and low-confidence detections are biased for selection
|
||||
- **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
|
||||
- Uses **SigLIP** (Vision Transformer) embeddings for semantic diversity
|
||||
- **YOLOv9c** to detect and crop specific objects (dogs, cars, etc.)
|
||||
- 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
|
||||
|
||||
Run the command-line interface:
|
||||
|
||||
```bash
|
||||
uv run -m if_curator
|
||||
uv run if-curator
|
||||
```
|
||||
|
||||
### Interactive Flow
|
||||
The tool will guide you through:
|
||||
1. **Select Person/Subject**: Choose from your Immich people.
|
||||
2. **Training Mode**: Face (Recognition) or Object (Classification).
|
||||
3. **Strategy**: Auto, Standard, Broad, etc.
|
||||
1. **Select Person** — Choose from your Immich people (supports multi-person batch)
|
||||
2. **Training Mode** — Face (Recognition) or Object (Classification)
|
||||
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
|
||||
Using SigLIP (visual embeddings) for diversity analysis...
|
||||
Computing embeddings for 69 images... ████████████████ 100%
|
||||
Auto-diversity selected 38 optimally diverse images.
|
||||
Using InsightFace (face embeddings) for diversity analysis...
|
||||
Quality filtering removed 76 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`.
|
||||
|
||||
| Variable | Description |
|
||||
| :--- | :--- |
|
||||
| `IMMICH_URL` | Full URL to Immich (e.g. `http://192.168.1.10:2283`) |
|
||||
| `API_KEY` | Your Immich API Key |
|
||||
| `FORCE_CPU` | Set to `true` to disable GPU acceleration |
|
||||
### Environment Variables
|
||||
|
||||
| Variable | Default | Description |
|
||||
| :--- | :--- | :--- |
|
||||
| `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
|
||||
|
||||
- **InsightFace**: Face detection and embedding (ArcFace)
|
||||
- **SigLIP**: Visual embeddings via `transformers` (OpenAI CLIP alternative)
|
||||
- **YOLOv9c**: State-of-the-art object detection for cropping
|
||||
- **Rich**: Beautiful terminal UI
|
||||
### Models
|
||||
- **InsightFace (Buffalo_L)** — Face detection and embedding (ArcFace, 512-d)
|
||||
- **SigLIP** — Visual embeddings via `transformers` (google/siglip-base-patch16-224, 768-d)
|
||||
- **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
|
||||
+103
-40
@@ -9,13 +9,15 @@ from PIL import Image
|
||||
from rich import print as rprint
|
||||
from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
|
||||
from rich.prompt import Confirm, IntPrompt, Prompt
|
||||
from rich.table import Table
|
||||
|
||||
from .config import Config, ConfigManager
|
||||
from .diversity import select_diverse_assets
|
||||
from .embeddings import is_embedding_available
|
||||
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
|
||||
|
||||
logger = logging.getLogger(__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"
|
||||
|
||||
|
||||
def interactive_configure(people: list[dict]) -> list[dict]:
|
||||
"""Interactive phase: select person, mode, and configure training strategy."""
|
||||
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]
|
||||
def _configure_person(person: dict, people: list[dict]) -> dict | None:
|
||||
"""Configure training for a single person. Returns job dict or None."""
|
||||
name = person["name"]
|
||||
|
||||
console.print(f"\nSelected: [bold green]{name}[/bold green]")
|
||||
|
||||
# Select training mode
|
||||
@@ -105,7 +93,7 @@ def interactive_configure(people: list[dict]) -> list[dict]:
|
||||
|
||||
if not recent_assets:
|
||||
rprint(" [dim]Skipping (0 recent images).[/dim]")
|
||||
return []
|
||||
return None
|
||||
|
||||
# Strategy selection
|
||||
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)
|
||||
if selection_mode == "skip":
|
||||
return []
|
||||
return None
|
||||
|
||||
# Perform selection
|
||||
selected_assets = _perform_selection(recent_assets, limit, name, selection_mode, entity_type)
|
||||
|
||||
rprint(f" [green]Queued {len(selected_assets)} images for {name}.[/green]")
|
||||
return [{"person": person, "assets": selected_assets, "limit": len(selected_assets), "config": config}]
|
||||
return {"person": person, "assets": selected_assets, "limit": len(selected_assets), "config": config}
|
||||
|
||||
|
||||
def _perform_selection(
|
||||
assets: list, limit: int | str, name: str, selection_mode: str, entity_type: str
|
||||
) -> list:
|
||||
def interactive_configure(people: list[dict]) -> list[dict]:
|
||||
"""Interactive phase: select person(s), mode, and configure training strategy.
|
||||
|
||||
Supports multi-person batch mode — after configuring one person,
|
||||
prompts to add another.
|
||||
"""
|
||||
valid_people = sorted([p for p in people if p.get("name")], key=lambda x: x["name"])
|
||||
|
||||
if not valid_people:
|
||||
rprint("[red]No people found with names in Immich.[/red]")
|
||||
return []
|
||||
|
||||
jobs = []
|
||||
|
||||
while True:
|
||||
# Select person
|
||||
console.print("\n[bold cyan]Select Person to Train:[/bold cyan]")
|
||||
for idx, p in enumerate(valid_people, 1):
|
||||
# Mark already-queued people
|
||||
marker = " [dim](queued)[/dim]" if any(j["person"]["id"] == p["id"] for j in jobs) else ""
|
||||
console.print(f" [bold]{idx}.[/bold] {p['name']}{marker}")
|
||||
|
||||
p_choice = IntPrompt.ask("Enter Number", choices=[str(i) for i in range(1, len(valid_people) + 1)])
|
||||
person = valid_people[p_choice - 1]
|
||||
|
||||
job = _configure_person(person, valid_people)
|
||||
if job:
|
||||
jobs.append(job)
|
||||
|
||||
# Multi-person: ask to add another
|
||||
if not Confirm.ask("\nAdd another person?", default=False):
|
||||
break
|
||||
|
||||
return jobs
|
||||
|
||||
|
||||
def _perform_selection(assets: list, limit: int | str, name: str, selection_mode: str, entity_type: str) -> list:
|
||||
"""Run diversity selection with progress display."""
|
||||
if selection_mode == "smart":
|
||||
model_display = "InsightFace (face embeddings)" if entity_type == "face" else "SigLIP (visual embeddings)"
|
||||
@@ -134,12 +156,17 @@ def _perform_selection(
|
||||
is_embedding_available(entity_type)
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(), TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(), TaskProgressColumn(), console=console,
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
console=console,
|
||||
) as progress:
|
||||
task = progress.add_task(f"[cyan]Computing embeddings for {len(assets)} images...", total=None)
|
||||
selected = select_diverse_assets(
|
||||
assets, limit, name,
|
||||
assets,
|
||||
limit,
|
||||
name,
|
||||
selection_mode=selection_mode,
|
||||
entity_type=entity_type,
|
||||
progress_callback=lambda c, t: progress.update(task, completed=c, total=t),
|
||||
@@ -156,6 +183,30 @@ def _perform_selection(
|
||||
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:
|
||||
"""Download and process images for all jobs."""
|
||||
if not jobs:
|
||||
@@ -163,9 +214,14 @@ def execute_jobs(jobs: list[dict]) -> None:
|
||||
|
||||
console.rule("[bold blue]Execution Phase")
|
||||
|
||||
use_full_res = Config.USE_FULL_RESOLUTION
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(), TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(), TaskProgressColumn(), console=console,
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
console=console,
|
||||
) as progress:
|
||||
grand_total = sum(j["limit"] for j in jobs)
|
||||
overall_task = progress.add_task("[green]Overall Progress", total=grand_total)
|
||||
@@ -181,17 +237,22 @@ def execute_jobs(jobs: list[dict]) -> None:
|
||||
count = 0
|
||||
for asset in assets:
|
||||
try:
|
||||
resp = requests.get(
|
||||
f"{Config.IMMICH_URL}/api/assets/{asset['id']}/thumbnail?size=preview&format=JPEG",
|
||||
headers={"x-api-key": Config.API_KEY, "Accept": "application/json"},
|
||||
timeout=30,
|
||||
)
|
||||
if not resp.ok:
|
||||
# Use full-resolution for final output when configured
|
||||
if use_full_res:
|
||||
img = fetch_full_image(asset["id"])
|
||||
else:
|
||||
resp = requests.get(
|
||||
f"{Config.IMMICH_URL}/api/assets/{asset['id']}/thumbnail?size=preview&format=JPEG",
|
||||
headers={"x-api-key": Config.API_KEY, "Accept": "application/json"},
|
||||
timeout=30,
|
||||
)
|
||||
img = Image.open(BytesIO(resp.content)) if resp.ok else None
|
||||
|
||||
if img is None:
|
||||
progress.console.print(f"[red]Failed download {asset['id']}[/red]")
|
||||
else:
|
||||
img = Image.open(BytesIO(resp.content))
|
||||
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"
|
||||
else process_object_mode(img, config, person_dir, count)
|
||||
if mode == "object"
|
||||
@@ -236,10 +297,12 @@ def main() -> None:
|
||||
|
||||
jobs = interactive_configure(people)
|
||||
|
||||
if jobs and Confirm.ask(f"\nReady to process {sum(j['limit'] for j in jobs)} images?"):
|
||||
execute_jobs(jobs)
|
||||
rprint("\n[bold green]Done! Happy Training.[/bold green]")
|
||||
elif not jobs:
|
||||
if jobs:
|
||||
_show_preview(jobs)
|
||||
if Confirm.ask(f"Ready to process {sum(j['limit'] for j in jobs)} images?"):
|
||||
execute_jobs(jobs)
|
||||
rprint("\n[bold green]Done! Happy Training.[/bold green]")
|
||||
else:
|
||||
rprint("[yellow]No jobs configured.[/yellow]")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
|
||||
+15
-1
@@ -24,7 +24,21 @@ class Config:
|
||||
API_KEY: str | None = None
|
||||
OUTPUT_DIR: str = "./frigate_train"
|
||||
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":
|
||||
if cls._instance is None:
|
||||
|
||||
+332
-32
@@ -1,9 +1,12 @@
|
||||
"""
|
||||
Diversity selection for training data curation.
|
||||
|
||||
Uses Farthest Point Sampling (FPS) algorithm with embeddings:
|
||||
- Faces: InsightFace embeddings
|
||||
- Objects: SigLIP embeddings
|
||||
Selection pipeline:
|
||||
1. Concurrent thumbnail download
|
||||
2. Quality filtering (blur, IR, exposure, confidence, face size)
|
||||
3. Face crop extraction (embed person's face, not full image)
|
||||
4. Embedding computation (InsightFace or SigLIP)
|
||||
5. Cluster-aware selection (K-Medoids + FPS with hard example weighting)
|
||||
"""
|
||||
|
||||
import logging
|
||||
@@ -15,6 +18,7 @@ from PIL import Image
|
||||
|
||||
from .config import Config, get_headers
|
||||
from .embeddings import get_embedding, is_embedding_available
|
||||
from .quality import assess_quality
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -28,7 +32,7 @@ def select_diverse_assets(
|
||||
progress_callback=None,
|
||||
) -> list:
|
||||
"""
|
||||
Select diverse assets using Farthest Point Sampling or time spread.
|
||||
Select diverse assets using cluster-aware FPS or time spread.
|
||||
|
||||
Args:
|
||||
assets: List of asset dicts from Immich API
|
||||
@@ -61,6 +65,11 @@ def select_diverse_assets(
|
||||
return _select_time_spread(assets, limit)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Thumbnail & Metadata Helpers
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _fetch_thumbnail(asset_id: str, timeout: int = 10) -> Image.Image | None:
|
||||
"""Fetch thumbnail from Immich API."""
|
||||
try:
|
||||
@@ -71,13 +80,106 @@ def _fetch_thumbnail(asset_id: str, timeout: int = 10) -> Image.Image | 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(
|
||||
assets: list,
|
||||
limit: int | str,
|
||||
entity_type: str,
|
||||
progress_callback=None,
|
||||
) -> 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)
|
||||
effective_limit = 30 if limit == "auto" else limit
|
||||
pool_size = min(3000, max(effective_limit * 20, len(assets)))
|
||||
@@ -89,24 +191,68 @@ def _select_by_embedding(
|
||||
else:
|
||||
candidates = assets
|
||||
|
||||
# Compute embeddings
|
||||
embeddings, valid_candidates = [], []
|
||||
for i, asset in enumerate(candidates):
|
||||
if progress_callback:
|
||||
progress_callback(i, len(candidates))
|
||||
# --- Phase 1: Concurrent thumbnail download ---
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
img = _fetch_thumbnail(asset["id"])
|
||||
thumbnail_map: dict[str, Image.Image] = {}
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
futures = {pool.submit(_fetch_thumbnail, a["id"]): a for a in candidates}
|
||||
for i, future in enumerate(as_completed(futures)):
|
||||
if progress_callback:
|
||||
progress_callback(i, len(candidates))
|
||||
asset = futures[future]
|
||||
try:
|
||||
img = future.result()
|
||||
if img is not None:
|
||||
thumbnail_map[asset["id"]] = img
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
# --- Phase 2-4: Quality filter → Crop → Embed ---
|
||||
embeddings, valid_candidates, confidence_scores = [], [], []
|
||||
quality_filtered = 0
|
||||
|
||||
for asset in candidates:
|
||||
img = thumbnail_map.get(asset["id"])
|
||||
if img is None:
|
||||
continue
|
||||
|
||||
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:
|
||||
embeddings.append(emb)
|
||||
valid_candidates.append(asset)
|
||||
confidence_scores.append(confidence)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(len(candidates), len(candidates))
|
||||
|
||||
if quality_filtered > 0:
|
||||
logger.info(f"Quality filtering removed {quality_filtered} images.")
|
||||
|
||||
if not embeddings:
|
||||
logger.warning("No valid embeddings found. Falling back to time spread.")
|
||||
return _select_time_spread(assets, limit)
|
||||
@@ -115,58 +261,212 @@ def _select_by_embedding(
|
||||
logger.warning(f"Only {len(valid_candidates)} valid embeddings. Returning all.")
|
||||
return valid_candidates
|
||||
|
||||
# Farthest Point Sampling with vectorized distance computation
|
||||
return _farthest_point_sampling(
|
||||
embeddings, valid_candidates, limit, auto_threshold=0.15
|
||||
# --- Phase 5: Cluster-aware selection ---
|
||||
return _cluster_aware_selection(
|
||||
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,
|
||||
candidates: list,
|
||||
limit: int | str,
|
||||
auto_threshold: float = 0.15,
|
||||
entity_type: str = "face",
|
||||
confidence_scores: list | None = None,
|
||||
) -> 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)
|
||||
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)
|
||||
emb_normed = emb_matrix / np.maximum(norms, 1e-8)
|
||||
|
||||
# Start from median-time sample
|
||||
selected = [n // 2]
|
||||
# Build confidence weight array for hard example boosting
|
||||
conf_array = np.ones(n)
|
||||
if confidence_scores and entity_type == "face":
|
||||
for i, c in enumerate(confidence_scores):
|
||||
if c is not None:
|
||||
conf_array[i] = c
|
||||
|
||||
# Compute adaptive threshold for auto mode
|
||||
auto_threshold = _compute_adaptive_threshold(emb_normed, entity_type) if limit == "auto" else 0.0
|
||||
target = Config.MAX_AUTO_IMAGES if limit == "auto" else limit
|
||||
|
||||
# --- Stage 1: K-Medoids clustering ---
|
||||
k = min(max(5, target // 4), n // 3, n) # e.g., 5-20 clusters
|
||||
logger.info(f"Clustering {n} embeddings into {k} groups (K-Medoids)...")
|
||||
|
||||
# Compute full cosine distance matrix
|
||||
dist_matrix = 1 - emb_normed @ emb_normed.T
|
||||
|
||||
medoid_indices, cluster_labels = _kmedoids(dist_matrix, k)
|
||||
selected = list(medoid_indices)
|
||||
selected_set = set(selected)
|
||||
|
||||
logger.info(f"Selected {len(selected)} cluster medoids as initial picks.")
|
||||
|
||||
# --- Stage 2: FPS with hard example weighting ---
|
||||
min_dists = np.full(n, np.inf)
|
||||
|
||||
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:
|
||||
# Update min distances with last selected point
|
||||
last_emb = emb_normed[selected[-1]]
|
||||
dists_to_last = 1 - emb_normed @ last_emb # Cosine distance
|
||||
min_dists = np.minimum(min_dists, dists_to_last)
|
||||
min_dists[selected[-1]] = -np.inf # Exclude already selected
|
||||
# Hard example weighting: boost distance for low-confidence candidates
|
||||
# Confidence < 0.85 gets up to 1.5× distance boost
|
||||
hard_weight = np.where(conf_array < 0.85, 1.0 + (0.85 - conf_array) * 2.0, 1.0)
|
||||
weighted_dists = min_dists * hard_weight
|
||||
|
||||
# Find farthest point
|
||||
best_idx = np.argmax(min_dists)
|
||||
best_dist = min_dists[best_idx]
|
||||
best_idx = int(np.argmax(weighted_dists))
|
||||
best_dist = min_dists[best_idx] # Use unweighted for threshold comparison
|
||||
|
||||
if best_dist == -np.inf:
|
||||
break # All points selected
|
||||
|
||||
if limit == "auto" and best_dist < auto_threshold:
|
||||
logger.info(
|
||||
f"Auto-stop: Next best image {best_dist:.3f} away (threshold {auto_threshold})."
|
||||
f"Auto-stop: Next best image {best_dist:.3f} away " f"(adaptive threshold {auto_threshold:.4f})."
|
||||
)
|
||||
break
|
||||
|
||||
selected.append(best_idx)
|
||||
selected_set.add(best_idx)
|
||||
|
||||
# Update min distances
|
||||
dists_to_new = dist_matrix[best_idx]
|
||||
min_dists = np.minimum(min_dists, dists_to_new)
|
||||
min_dists[best_idx] = -np.inf
|
||||
|
||||
# Log hard example stats
|
||||
if entity_type == "face":
|
||||
selected_conf = [conf_array[i] for i in selected if conf_array[i] < 1.0]
|
||||
hard_count = sum(1 for c in selected_conf if c < 0.85)
|
||||
logger.info(
|
||||
f"Selection complete: {len(selected)} images " f"({hard_count} hard examples with confidence < 0.85)."
|
||||
)
|
||||
else:
|
||||
logger.info(f"Selection complete: {len(selected)} diverse images.")
|
||||
|
||||
logger.info(f"Smart selection complete. Picked {len(selected)} diverse images.")
|
||||
return [candidates[i] for i in selected]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Time Spread Fallback
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _select_time_spread(assets: list, limit: int | str) -> list:
|
||||
"""Select N assets evenly distributed in time."""
|
||||
if limit == "auto":
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
"""
|
||||
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)
|
||||
- Caching: Disk-based cache avoids recomputation on reruns
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
@@ -14,6 +15,8 @@ import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from .cache import get_cache
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Lazy-loaded singletons
|
||||
@@ -48,8 +51,10 @@ def get_insightface_app():
|
||||
|
||||
# Determine device: 0 for GPU, -1 for CPU
|
||||
gpu_providers = {
|
||||
"CUDAExecutionProvider", "ROCmExecutionProvider",
|
||||
"MPSExecutionProvider", "CoreMLExecutionProvider",
|
||||
"CUDAExecutionProvider",
|
||||
"ROCmExecutionProvider",
|
||||
"MPSExecutionProvider",
|
||||
"CoreMLExecutionProvider",
|
||||
}
|
||||
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...")
|
||||
try:
|
||||
from insightface.app import FaceAnalysis
|
||||
|
||||
_insightface_app = FaceAnalysis(name="buffalo_l", root="~/.insightface")
|
||||
_insightface_app.prepare(ctx_id=-1, det_size=(640, 640))
|
||||
return _insightface_app
|
||||
@@ -179,14 +185,83 @@ def get_object_embedding(img_pil: Image.Image) -> np.ndarray | 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:
|
||||
"""Get embedding for an image based on entity type ('face' or 'object')."""
|
||||
return get_face_embedding(img_pil) if entity_type == "face" else get_object_embedding(img_pil)
|
||||
def get_embedding(
|
||||
img_pil: Image.Image,
|
||||
entity_type: str = "face",
|
||||
asset_id: str | None = None,
|
||||
immich_embedding: np.ndarray | None = None,
|
||||
) -> np.ndarray | None:
|
||||
"""Get embedding for an image based on entity type.
|
||||
|
||||
Priority:
|
||||
1. Pre-fetched Immich embedding (if provided)
|
||||
2. Disk cache (if enabled and asset_id provided)
|
||||
3. Local model computation (InsightFace or SigLIP)
|
||||
|
||||
Args:
|
||||
img_pil: The image to embed
|
||||
entity_type: 'face' or 'object'
|
||||
asset_id: Optional asset ID for cache lookup
|
||||
immich_embedding: Optional pre-fetched embedding from Immich API
|
||||
"""
|
||||
from .config import Config
|
||||
|
||||
use_cache = Config.ENABLE_CACHE and asset_id is not None
|
||||
cache = get_cache(Config.CACHE_DIR) if use_cache else None
|
||||
model_key = "immich" if entity_type == "face" else "siglip"
|
||||
|
||||
# 1. Use Immich embedding if provided
|
||||
if immich_embedding is not None:
|
||||
if cache:
|
||||
cache.put(asset_id, immich_embedding, model_key)
|
||||
return immich_embedding
|
||||
|
||||
# 2. Check disk cache
|
||||
if cache:
|
||||
cached = cache.get(asset_id, model_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
# 3. Compute locally
|
||||
if entity_type == "face":
|
||||
emb = get_face_embedding(img_pil)
|
||||
model_key = "insightface"
|
||||
else:
|
||||
emb = get_object_embedding(img_pil)
|
||||
|
||||
# Cache the result
|
||||
if emb is not None and cache:
|
||||
cache.put(asset_id, emb, model_key)
|
||||
|
||||
return emb
|
||||
|
||||
|
||||
def is_embedding_available(entity_type: str = "face") -> bool:
|
||||
|
||||
@@ -3,8 +3,11 @@
|
||||
import logging
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from .config import Config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Lazy singleton
|
||||
@@ -16,20 +19,59 @@ def get_yolo_model():
|
||||
global _yolo_model
|
||||
if _yolo_model is None:
|
||||
from ultralytics import YOLO
|
||||
|
||||
logger.info("Loading YOLOv9c model...")
|
||||
_yolo_model = YOLO("yolov9c.pt")
|
||||
return _yolo_model
|
||||
|
||||
|
||||
def align_face(img: Image.Image, landmarks: list[list[float]] | np.ndarray) -> Image.Image | None:
|
||||
"""Align face using 5-point landmarks to standard ArcFace input format (112x112).
|
||||
|
||||
This produces a normalized face crop that matches exactly what ArcFace
|
||||
was trained on, improving recognition accuracy.
|
||||
|
||||
Args:
|
||||
img: Full PIL image containing the face
|
||||
landmarks: 5-point facial landmarks [[x,y], ...] (eyes, nose, mouth corners)
|
||||
|
||||
Returns:
|
||||
Aligned 112x112 face image, or None if alignment fails
|
||||
"""
|
||||
try:
|
||||
from insightface.utils.face_align import norm_crop
|
||||
|
||||
img_np = np.asarray(img)
|
||||
lm = np.array(landmarks, dtype=np.float32)
|
||||
if lm.shape != (5, 2):
|
||||
logger.debug(f"Invalid landmark shape: {lm.shape}, expected (5, 2)")
|
||||
return None
|
||||
aligned = norm_crop(img_np, lm)
|
||||
return Image.fromarray(aligned)
|
||||
except ImportError:
|
||||
logger.debug("InsightFace not available for face alignment")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.debug(f"Face alignment failed: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def process_face_mode(
|
||||
img: Image.Image,
|
||||
asset: dict,
|
||||
person: dict,
|
||||
output_dir: str,
|
||||
count: int,
|
||||
min_width: int = 50,
|
||||
min_width: int | None = None,
|
||||
) -> 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
|
||||
face_info = None
|
||||
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})")
|
||||
return False
|
||||
|
||||
# Add 10% margin
|
||||
margin_x, margin_y = face_w * 0.10, face_h * 0.10
|
||||
# Try face alignment if enabled and landmarks available
|
||||
if Config.ENABLE_FACE_ALIGNMENT:
|
||||
landmarks = face_info.get("landmarks") or face_info.get("landmark")
|
||||
if landmarks:
|
||||
# Scale landmarks
|
||||
scaled_landmarks = [[lm[0] * scale_x, lm[1] * scale_y] for lm in landmarks]
|
||||
aligned = align_face(img, scaled_landmarks)
|
||||
if aligned is not None:
|
||||
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 = (
|
||||
max(0, x1 - margin_x),
|
||||
max(0, y1 - margin_y),
|
||||
@@ -87,9 +141,7 @@ def process_object_mode(
|
||||
|
||||
found = False
|
||||
for idx, (box, cls_id, conf) in enumerate(
|
||||
(box, int(box.cls[0]), float(box.conf[0]))
|
||||
for r in results
|
||||
for box in r.boxes
|
||||
(box, int(box.cls[0]), float(box.conf[0])) 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:
|
||||
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
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
from .config import Config, get_headers
|
||||
|
||||
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]:
|
||||
"""Fetch all people from Immich."""
|
||||
try:
|
||||
@@ -68,6 +83,107 @@ def fetch_all_assets(person: dict) -> list[dict]:
|
||||
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]:
|
||||
"""Filter assets to keep only those from the last N years."""
|
||||
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 = logging.FileHandler("immich_export.log")
|
||||
file_handler.setLevel(logging.DEBUG)
|
||||
file_handler.setFormatter(
|
||||
logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
|
||||
)
|
||||
file_handler.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
|
||||
root.addHandler(file_handler)
|
||||
|
||||
# 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"
|
||||
description = "Immich to Frigate training sets"
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
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 = [
|
||||
"insightface>=0.7.3",
|
||||
"numpy>=2.2.6",
|
||||
"onnxruntime>=1.23.2",
|
||||
"opencv-python>=4.12.0.88",
|
||||
"opencv-python-headless>=4.12.0.88",
|
||||
"pillow>=12.1.0",
|
||||
"python-dotenv>=1.2.1",
|
||||
"requests>=2.32.5",
|
||||
"rich>=14.2.0",
|
||||
"scipy>=1.17.0",
|
||||
"sentencepiece>=0.2.1",
|
||||
"torch>=2.9.1",
|
||||
"transformers>=4.57.6",
|
||||
"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]
|
||||
gpu = [
|
||||
"onnxruntime-gpu>=1.23.2",
|
||||
@@ -28,6 +41,7 @@ gpu = [
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"pytest>=8.0",
|
||||
"ruff>=0.9.2",
|
||||
]
|
||||
|
||||
@@ -38,3 +52,26 @@ target-version = "py312"
|
||||
[tool.ruff.lint]
|
||||
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