flan/winnow
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: remove object mode pipeline (YOLO, SigLIP, TRAINING_MODE, OBJECT_CLASS)

Browse Source 
winnow is a face recognition training tool. Object mode required manual
file placement with no Frigate API, pulled in torch/torchvision/transformers/
ultralytics (~2 GB), and was architecturally misaligned with the project goal.

Removed:
- process_object_mode (YOLO inference), get_yolo_model
- SigLIP model stack (get_siglip_model, get_object_embedding, batch variant)
- entity_type branching throughout diversity, embeddings, jobs, executor
- TRAINING_MODE and OBJECT_CLASS env vars
- torch, torchvision, transformers, ultralytics dependencies
- pytorch index entries from pyproject.toml
- Object mode from README (Modes section, env var table, How It Works)
This commit is contained in:
flan
2026-06-14 17:29:25 +00:00
parent 3105b15beb
commit 835016e0e3
7 changed files with 91 additions and 421 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+14 -28
View File
@@ -39,8 +39,7 @@ Immich library
│
▼
4. Compute embeddings from the same preview thumbnails
• Faces → InsightFace (ArcFace / Buffalo_L) → 512-dim vector
• Objects → SigLIP (Vision Transformer) → 768-dim vector
• InsightFace (ArcFace / Buffalo_L) → 512-dim vector
│
▼
5. Near-duplicate removal — greedy cosine-distance pass drops burst shots
@@ -61,35 +60,25 @@ Immich library
7. Download full-resolution originals from Immich
│
▼
8. Crop and process
• Face mode: EXIF-corrected, landmark-aligned 112×112 crop (ArcFace format)
• Object mode: YOLOv9c detection → one crop per matched instance
8. Crop and process — EXIF-corrected, landmark-aligned 112×112 crop (ArcFace format)
│
▼
9. Deliver
• Face mode: upload crops to Frigate's face registration API
↳ below MAX_AUTO_IMAGES — upload, unless the novelty gate
(FRIGATE_SCORE_CEILING) determines the candidate is already
covered by the current training set
↳ at cap + QUALITY_REPLACEMENT=true — with Frigate scoring active,
swap the most redundant tracked image (highest pre-upload recognize
score) if the candidate is more novel (lower score); falling back to
blur-score comparison when no Frigate scores are available; manually
added files are never touched
↳ at cap + QUALITY_REPLACEMENT=false — skip this person
• Object mode: save crops to disk → place into your Frigate data directory
9. Deliver — upload crops to Frigate's face registration API
↳ below MAX_AUTO_IMAGES — upload, unless the novelty gate
(FRIGATE_SCORE_CEILING) determines the candidate is already
covered by the current training set
↳ at cap + QUALITY_REPLACEMENT=true — with Frigate scoring active,
swap the most redundant tracked image (highest pre-upload recognize
score) if the candidate is more novel (lower score); falling back to
blur-score comparison when no Frigate scores are available; manually
added files are never touched
↳ at cap + QUALITY_REPLACEMENT=false — skip this person
```
Uploaded and rejected asset IDs are persisted across runs. The same image is never processed twice; rejected assets are permanently skipped unless `RETRY_REJECTED=true`.
---
## Modes
**Face mode** (default) — extracts face crops using Immich's bounding box metadata, applies EXIF orientation correction, and aligns them to ArcFace's standard 112×112 format using 5-point facial landmarks. Crops are uploaded directly to Frigate's face registration API.
**Object mode** — runs each full-resolution image through YOLOv9c to detect instances of a target class (dog, cat, car, etc.), crops each detection, and saves it to the output directory. Frigate has no API for uploading object training data; place the crops into your Frigate data directory manually.
---
## Running in Docker
@@ -116,7 +105,7 @@ services:
- FRIGATE_URL=http://192.168.1.10:5000
- CRON_SCHEDULE=0 3 * * 0
volumes:
- /path/to/models:/models
- /path/to/models:/models # INSIGHTFACE_HOME — persists Buffalo_L model (~300 MB)
- /path/to/cache:/app/.if_cache
- /path/to/output:/app/frigate_train
deploy:
@@ -162,7 +151,7 @@ See [compose.yml](compose.yml) for the full annotated example with all options.
| *(empty string)* | Stay alive, run nothing — trigger manually with `docker exec -it winnow winnow` |
| Cron expression | Run on startup, then repeat on schedule |
In scheduled mode the process (and loaded models) stays resident between runs. The first run after a fresh install downloads the embedding models (~1–2 GB); subsequent runs use the cached models from the mounted volume.
In scheduled mode the process (and loaded models) stays resident between runs. The first run after a fresh install downloads InsightFace Buffalo_L (~300 MB); subsequent runs use the cached model from the mounted volume.
---
@@ -180,10 +169,8 @@ In scheduled mode the process (and loaded models) stays resident between runs. T
| Variable | Default | Description |
| :--- | :--- | :--- |
| `TRAINING_MODE` | `face` | `face` — upload crops to Frigate; `object` — save crops to disk |
| `STRATEGY` | `adaptive` | `adaptive` — embedding-based diversity selection, stops when candidates become redundant; `standard` — fixed 30 images; `broad` — fixed 100 images |
| `LIMIT` | *(unset)* | Exact image count — overrides `STRATEGY` |
| `OBJECT_CLASS` | `dog` | Target class for object mode (any YOLO class: `dog`, `cat`, `car`, etc.) |
| `AUTO_MODE` | *(auto)* | Skip interactive prompts and process all people unattended — auto-detected when no TTY is present (Docker, cron); set `true` to force in a terminal |
| `VERBOSE` | `false` | Enable DEBUG-level console output (log file is always DEBUG) |
@@ -227,7 +214,6 @@ These defaults are tuned for Frigate's ArcFace requirements. winnow will warn on
| `OPENVINO_DEVICE` | `CPU` | Intel variant only: set `GPU` to use Arc or iGPU; default runs on CPU |
| `ENABLE_CACHE` | `true` | Cache computed embeddings to disk (speeds up re-runs on the same library) |
| `CACHE_DIR` | `.if_cache` | Path for embedding cache and upload tracker files |
| `HF_HOME` | *(system)* | HuggingFace model cache path (SigLIP) |
| `INSIGHTFACE_HOME` | *(system)* | InsightFace model cache path (Buffalo_L) |
### Output
pyproject.toml
+1 -29
View File
@@ -1,7 +1,7 @@
[project]
name = "winnow"
version = "0.4.10"
description = "Selects diverse, high-quality photos from Immich as training data for Frigate face recognition and object classification."
description = "Selects diverse, high-quality photos from Immich as training data for Frigate face recognition."
license = "AGPL-3.0-or-later"
requires-python = ">=3.13"
authors = [{ name = "Holden Salomon", email = "holden@arch.fyi" }]
@@ -27,10 +27,6 @@ dependencies = [
"python-dotenv>=1.2.1",
"requests>=2.32.5",
"rich>=14.2.0",
"torch>=2.12.0",
"torchvision>=0.27.0",
"transformers>=5.12.0",
"ultralytics>=8.4.66",
]
[project.scripts]
@@ -53,27 +49,6 @@ required-environments = [
"sys_platform == 'linux' and platform_machine == 'aarch64'",
]
[tool.uv.sources]
torch = [
{ index = "pytorch-cu126", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
{ index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
]
torchvision = [
{ index = "pytorch-cu126", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
{ index = "pytorch-cpu", marker = "sys_platform != 'linux'" },
]
[[tool.uv.index]]
name = "pytorch-cu126"
url = "https://download.pytorch.org/whl/cu126"
explicit = true
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
[dependency-groups]
dev = [
@@ -104,9 +79,6 @@ nvidia-cudnn-cu12 = "nvidia.cudnn"
onnxruntime-gpu = "onnxruntime"
requests = "requests"
rich = "rich"
torch = "torch"
transformers = "transformers"
ultralytics = "ultralytics"
[tool.deptry.per_rule_ignores]
DEP002 = ["onnxruntime-gpu", "nvidia-cudnn-cu12"]
winnow/diversity.py
+32 -55
View File
@@ -5,7 +5,7 @@ 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)
4. Embedding computation (InsightFace)
5. Cluster-aware selection (K-Medoids + FPS with hard example weighting)
"""
@@ -28,7 +28,6 @@ def select_diverse_assets(
limit: int | str,
entity_name: str,
selection_mode: str = "smart",
entity_type: str = "face",
person_id: str | None = None,
progress_callback=None,
) -> list:
@@ -38,9 +37,8 @@ def select_diverse_assets(
Args:
assets: List of asset dicts from Immich API
limit: Number to select, or "auto" for dynamic selection
entity_name: Name of the person/object for logging
entity_name: Name of the person for logging
selection_mode: 'smart' (embedding-based) or 'time' (time spread)
entity_type: 'face' or 'object' - determines embedding model
progress_callback: Optional callback(current, total) for progress
Returns:
@@ -53,14 +51,13 @@ def select_diverse_assets(
# Sort by creation time
assets = sorted(assets, key=lambda x: x.get("fileCreatedAt", ""))
if selection_mode != "smart" or not is_embedding_available(entity_type):
if selection_mode != "smart" or not is_embedding_available():
if selection_mode == "smart":
model_name = "InsightFace" if entity_type == "face" else "SigLIP"
logger.warning(f"{model_name} unavailable. Falling back to time spread.")
logger.warning("InsightFace unavailable. Falling back to time spread.")
return _select_time_spread(assets, limit)
try:
return _select_by_embedding(assets, limit, entity_type, person_id, progress_callback)
return _select_by_embedding(assets, limit, person_id, progress_callback)
except Exception as e:
logger.error(f"Smart Diversity failed: {e}. Falling back to time spread.")
return _select_time_spread(assets, limit)
@@ -179,7 +176,6 @@ def _crop_face_from_thumbnail(
def _select_by_embedding(
assets: list,
limit: int | str,
entity_type: str,
person_id: str | None = None,
progress_callback=None,
) -> list:
@@ -188,7 +184,7 @@ def _select_by_embedding(
Pipeline:
1. Concurrent thumbnail download
2. Quality filtering
3. Face crop extraction (face mode only)
3. Face crop extraction
4. Embedding computation
5. Cluster-aware selection with hard example weighting
"""
@@ -243,28 +239,25 @@ def _select_by_embedding(
confidence = _get_face_confidence(asset, person_id=person_id)
if entity_type == "face":
face_bbox = _get_face_bbox(asset, person_id=person_id)
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
face_bbox = _get_face_bbox(asset, person_id=person_id)
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
asset["quality_score"] = quality.blur_score
face_crop = _crop_face_from_thumbnail(img, asset, person_id=person_id)
embed_img = face_crop if face_crop is not None else img
else:
embed_img = img
asset["quality_score"] = quality.blur_score
face_crop = _crop_face_from_thumbnail(img, asset, person_id=person_id)
embed_img = face_crop if face_crop is not None else img
emb = get_embedding(embed_img, entity_type, asset_id=asset["id"])
emb = get_embedding(embed_img, asset_id=asset["id"])
if emb is not None:
embeddings.append(emb)
valid_candidates.append(asset)
@@ -300,7 +293,6 @@ def _select_by_embedding(
embeddings,
valid_candidates,
limit,
entity_type=entity_type,
confidence_scores=confidence_scores,
)
@@ -429,11 +421,11 @@ def _kmedoids(dist_matrix: np.ndarray, k: int, max_iter: int = 50) -> tuple[list
# =============================================================================
def _compute_adaptive_threshold(emb_normed: np.ndarray, entity_type: str) -> float:
def _compute_adaptive_threshold(emb_normed: np.ndarray) -> 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.
the threshold as 20% of the median pairwise distance.
"""
n = len(emb_normed)
sample_size = min(200, n)
@@ -441,22 +433,14 @@ def _compute_adaptive_threshold(emb_normed: np.ndarray, entity_type: str) -> flo
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)]
if len(upper_tri) == 0:
return 0.05
median_dist = float(np.median(upper_tri))
threshold = max(0.05, median_dist * 0.20)
# 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.debug(
f"Adaptive threshold: {threshold:.4f} "
f"(median_dist={median_dist:.4f}, fraction={fraction}, type={entity_type})"
)
logger.debug(f"Adaptive threshold: {threshold:.4f} (median_dist={median_dist:.4f})")
return threshold
@@ -464,7 +448,6 @@ def _cluster_aware_selection(
embeddings: list,
candidates: list,
limit: int | str,
entity_type: str = "face",
confidence_scores: list | None = None,
) -> list:
"""Two-stage selection: K-Medoids clustering → FPS with hard example weighting.
@@ -484,13 +467,13 @@ def _cluster_aware_selection(
# Build confidence weight array for hard example boosting
conf_array = np.ones(n)
if confidence_scores and entity_type == "face":
if confidence_scores:
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
auto_threshold = _compute_adaptive_threshold(emb_normed) if limit == "auto" else 0.0
target = Config.MAX_AUTO_IMAGES if limit == "auto" else limit
# --- Stage 1: K-Medoids clustering ---
@@ -542,15 +525,9 @@ def _cluster_aware_selection(
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.")
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 ({hard_count} hard examples with confidence < 0.85).")
return [candidates[i] for i in selected]
winnow/embeddings.py
+18 -168
View File
@@ -1,8 +1,7 @@
"""
Unified embedding interface for faces and objects.
Embedding interface for face diversity selection.
- Faces: InsightFace (ArcFace/Buffalo_L) — or reuse from Immich
- Objects: SigLIP (Vision Transformer via transformers)
- Caching: Disk-based cache avoids recomputation on reruns
"""
@@ -42,12 +41,9 @@ def _suppress_output():
os.close(saved_err)
# Lazy-loaded singletons
# Lazy-loaded singleton
_insightface_app = None
_insightface_loaded = False
_siglip_model = None
_siglip_processor = None
_siglip_loaded = False
def _is_force_cpu() -> bool:
@@ -202,169 +198,41 @@ def get_face_embedding(img_pil: Image.Image) -> np.ndarray | None:
# =============================================================================
# SigLIP (Objects)
# =============================================================================
def get_siglip_model():
"""Singleton for SigLIP model and processor with GPU auto-detection."""
global _siglip_model, _siglip_processor, _siglip_loaded
if _siglip_loaded:
return _siglip_model, _siglip_processor
_siglip_loaded = True
try:
import warnings
import torch
from transformers import AutoImageProcessor, SiglipVisionModel
model_name = "google/siglip-base-patch16-224"
# Disk cache check — path derived from model_name using HuggingFace's slug convention
hf_home = os.environ.get("HF_HOME", os.path.join(os.path.expanduser("~"), ".cache", "huggingface"))
cache_slug = "models--" + model_name.replace("/", "--")
model_cache = Path(hf_home) / "hub" / cache_slug
if model_cache.exists() and any(model_cache.iterdir()):
logger.info(f"SigLIP {model_name}: found in model cache")
else:
logger.info(f"SigLIP {model_name}: not cached — downloading now (~380 MB)")
logger.info(f"SigLIP {model_name}: loading into memory...")
t0 = time.time()
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", message=".*use_fast.*")
_siglip_processor = AutoImageProcessor.from_pretrained(model_name, use_fast=True)
_siglip_model = SiglipVisionModel.from_pretrained(model_name)
_siglip_model.eval()
# Move to GPU if available (ROCm builds expose torch.cuda.is_available() == True)
if not _is_force_cpu():
if torch.cuda.is_available():
_siglip_model = _siglip_model.cuda()
device_name = "CUDA GPU"
elif hasattr(torch, "xpu") and torch.xpu.is_available():
_siglip_model = _siglip_model.to("xpu")
device_name = "Intel XPU"
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
_siglip_model = _siglip_model.to("mps")
device_name = "Apple MPS"
else:
device_name = "CPU"
else:
device_name = "CPU (FORCE_CPU)"
logger.info(f"SigLIP {model_name}: ready on {device_name} ({time.time() - t0:.1f}s)")
return _siglip_model, _siglip_processor
except ImportError as e:
logger.error(f"transformers/torch not installed: {e}")
return None, None
except Exception as e:
logger.error(f"Failed to load SigLIP: {e}")
return None, None
def get_object_embedding(img_pil: Image.Image) -> np.ndarray | None:
"""Get 768-dim SigLIP embedding for an image."""
model, processor = get_siglip_model()
if model is None:
return None
try:
import torch
inputs = processor(images=img_pil, return_tensors="pt")
device = next(model.parameters()).device
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
return outputs.pooler_output.squeeze().cpu().numpy()
except Exception as e:
logger.error(f"Error getting object embedding: {e}")
return None
def get_object_embeddings_batch(images: list[Image.Image]) -> list[np.ndarray | None]:
"""Get SigLIP embeddings for a batch of images (GPU-efficient)."""
model, processor = get_siglip_model()
if model is None:
return [None] * len(images)
try:
import torch
inputs = processor(images=images, return_tensors="pt", padding=True)
device = next(model.parameters()).device
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
embeddings = outputs.pooler_output.cpu().numpy()
return [embeddings[i] for i in range(len(embeddings))]
except Exception as e:
logger.error(f"Error in batch embedding: {e}")
# Fall back to individual computation
return [get_object_embedding(img) for img in images]
# =============================================================================
# Unified Interface with Caching
# Embedding Interface with Caching
# =============================================================================
def get_embedding(
img_pil: Image.Image,
entity_type: str = "face",
asset_id: str | None = None,
immich_embedding: np.ndarray | None = None,
) -> np.ndarray | None:
"""Get embedding for an image based on entity type.
"""Get embedding for a face image.
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
3. Local InsightFace computation
"""
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
# Use a single consistent cache key per model so lookups and stores always match.
# "immich" was previously used as the face key on the lookup path but "insightface"
# on the store path — meaning the cache was never hit for locally-computed embeddings.
cache_key = "insightface" 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, cache_key)
cache.put(asset_id, immich_embedding, "insightface")
return immich_embedding
# 2. Check disk cache
if cache:
cached = cache.get(asset_id, cache_key)
cached = cache.get(asset_id, "insightface")
if cached is not None:
return cached
# 3. Compute locally
if entity_type == "face":
emb = get_face_embedding(img_pil)
else:
emb = get_object_embedding(img_pil)
emb = get_face_embedding(img_pil)
if emb is not None and cache:
cache.put(asset_id, emb, cache_key)
cache.put(asset_id, emb, "insightface")
return emb
@@ -378,36 +246,18 @@ def _is_module_available(module_name: str) -> bool:
return False
def is_embedding_available(entity_type: str = "face", *, load: bool = False) -> bool:
"""Check if embedding model is available for the given entity type.
def is_embedding_available(*, load: bool = False) -> bool:
"""Check if InsightFace is available.
By default this performs a lightweight import-check only (no model loading).
Pass ``load=True`` to actually load the model (expensive, hundreds of MB).
Args:
entity_type: 'face' or 'object'
load: If True, fully load the model to verify. If False (default),
only check that the required packages are importable.
Pass ``load=True`` to actually load the model (expensive, ~300 MB).
"""
if load:
if entity_type == "face":
return get_insightface_app() is not None
model, _ = get_siglip_model()
return model is not None
# Lightweight check: just verify the packages are importable
if entity_type == "face":
return _is_module_available("insightface") and _is_module_available("onnxruntime")
return _is_module_available("transformers") and _is_module_available("torch")
def load_embedding_model(entity_type: str = "face") -> bool:
"""Explicitly load the embedding model for the given entity type.
Returns True if the model loaded successfully.
"""
if entity_type == "face":
return get_insightface_app() is not None
model, _ = get_siglip_model()
return model is not None
return _is_module_available("insightface") and _is_module_available("onnxruntime")
def load_embedding_model() -> bool:
"""Explicitly load InsightFace. Returns True if the model loaded successfully."""
return get_insightface_app() is not None
winnow/executor.py
+10 -37
View File
@@ -19,7 +19,7 @@ from .frigate_api import (
get_frigate_person_files,
recognize_face,
)
from .image_processing import process_face_mode, process_full_mode, process_object_mode
from .image_processing import process_face_mode
from .immich_api import fetch_face_data, fetch_full_image
from .log_config import console
from .quality import assess_quality
@@ -250,26 +250,21 @@ def execute_jobs(jobs: list[dict]) -> None:
if img is None:
progress.console.print(f"[red]Failed download {asset['id']}[/red]")
else:
saved = (
process_face_mode(img, asset, person, person_dir, count, insightface_app=insightface_app)
if mode == "face"
else process_object_mode(img, config, person_dir, count)
if mode == "object"
else process_full_mode(img, person_dir, count)
saved = process_face_mode(
img, asset, person, person_dir, count, insightface_app=insightface_app
)
if saved:
# Record which asset produced which output file
filename = f"{count}.jpg"
asset_map[filename] = asset["id"]
score_map[filename] = asset.get("quality_score")
if mode == "face" and isinstance(saved, tuple):
if isinstance(saved, tuple):
dims_map[filename] = saved
# Time-spread path: compute blur score from the downloaded
# image. Cap at 1440px so the scale matches the preview
# thumbnails the embedding path uses for scoring — Laplacian
# variance grows with resolution, making full-res and
# thumbnail scores incomparable if left uncapped.
if mode == "face" and score_map[filename] is None:
if score_map[filename] is None:
try:
score_img = img.convert("RGB") if img.mode != "RGB" else img
if score_img.width > 1440 or score_img.height > 1440:
@@ -279,12 +274,6 @@ def execute_jobs(jobs: list[dict]) -> None:
except Exception as exc:
logger.debug(f"Quality score fallback for {asset['id']}: {exc}")
score_map[filename] = 0.0 # unknown quality — treat as lowest
# Also record object-mode variant filenames
if mode == "object":
for f in sorted(os.listdir(person_dir)):
if f.startswith(f"{count}_") and f not in asset_map:
asset_map[f] = asset["id"]
score_map[f] = asset.get("face_confidence")
count += 1
else:
@@ -312,29 +301,13 @@ def execute_jobs(jobs: list[dict]) -> None:
def upload_to_frigate(jobs: list[dict]) -> None:
"""Upload processed face crops to Frigate via API with detailed logging.
Only runs for face-mode jobs. Object-mode crops are saved to the output
directory as the deliverable and must be copied to Frigate manually.
After each successful upload, records the Immich asset ID in the
upload tracker so it is skipped on future runs.
"""
face_jobs = [j for j in jobs if j["config"].get("mode", "face") == "face"]
if not face_jobs:
rprint("[dim]No face-mode jobs to upload.[/dim]")
if not jobs:
rprint("[dim]No jobs to upload.[/dim]")
return
# Notify user about object-mode jobs that were skipped
object_jobs = [j for j in jobs if j["config"].get("mode") == "object"]
for job in object_jobs:
name = job["person"]["name"]
try:
person_dir = _safe_person_dir(Config.OUTPUT_DIR, name)
except ValueError as e:
logger.error(str(e))
continue
rprint(f" [dim]📁 {name} (object): crops saved to {person_dir} — copy to Frigate manually[/dim]")
frigate_url = os.environ.get("FRIGATE_URL", "")
if not frigate_url:
rprint("[yellow]⚠️ FRIGATE_URL not set, skipping upload.[/yellow]")
@@ -347,7 +320,7 @@ def upload_to_frigate(jobs: list[dict]) -> None:
# from the asset_map stored on each job during execute_jobs()
filename_to_asset_id: dict[str, dict[str, str]] = {}
total_files = 0
for job in face_jobs:
for job in jobs:
name = job["person"]["name"]
asset_map = job.get("asset_map", {})
filename_to_asset_id[name] = asset_map
@@ -357,7 +330,7 @@ def upload_to_frigate(jobs: list[dict]) -> None:
rprint(" [yellow]No images found to upload.[/yellow]")
return
rprint(f" People: [bold]{len(face_jobs)}[/bold], Total images: [bold]{total_files}[/bold]")
rprint(f" People: [bold]{len(jobs)}[/bold], Total images: [bold]{total_files}[/bold]")
uploaded, failed = 0, 0
max_retries = 2
@@ -375,7 +348,7 @@ def upload_to_frigate(jobs: list[dict]) -> None:
) as progress:
upload_task = progress.add_task("[green]Uploading to Frigate", total=total_files)
for job in face_jobs:
for job in jobs:
name = job["person"]["name"]
# URL-encode the name for the API (handles spaces, special chars)
encoded_name = quote(name, safe="")
winnow/image_processing.py
+1 -60
View File
@@ -1,4 +1,4 @@
"""Image processing functions for cropping faces and objects."""
"""Image processing functions for cropping faces."""
import logging
import os
@@ -10,9 +10,6 @@ from .config import Config
logger = logging.getLogger(__name__)
# Lazy singleton
_yolo_model = None
def _save_jpeg(img: Image.Image, path: str) -> None:
if img.mode != "RGB":
@@ -20,17 +17,6 @@ def _save_jpeg(img: Image.Image, path: str) -> None:
img.save(path, format="JPEG")
def get_yolo_model():
"""Singleton for YOLO model."""
global _yolo_model
if _yolo_model is None:
from ultralytics import YOLO
logger.info("Loading YOLOv9c model...")
_yolo_model = YOLO("yolov9c.pt")
return _yolo_model
def align_face(img: Image.Image, landmarks: list[list[float]] | np.ndarray) -> Image.Image | None:
"""Align face using 5-point landmarks to standard ArcFace input format (112x112).
@@ -170,49 +156,4 @@ def process_face_mode(
return face_crop.size
def process_object_mode(
img: Image.Image,
config: dict,
output_dir: str,
count: int,
) -> bool:
"""Detect and crop objects using YOLO."""
try:
model = get_yolo_model()
target_class = config.get("object_class", "dog")
import torch
if os.getenv("FORCE_CPU", "").lower() in ("true", "1", "yes"):
device = "cpu"
elif hasattr(torch, "xpu") and torch.xpu.is_available():
device = "xpu"
else:
device = None # YOLO auto-selects (CUDA/ROCm/CPU)
results = model(img, verbose=False, device=device)
found = False
class_idx = 0 # Sequential counter per target class (Issue #10)
for box in (box for r in results for box in r.boxes):
cls_id = int(box.cls[0])
conf = float(box.conf[0])
if 0 <= cls_id < len(model.names) and model.names[cls_id] == target_class and conf > 0.5:
x1, y1, x2, y2 = box.xyxy[0].tolist()
_save_jpeg(
img.crop((x1, y1, x2, y2)),
os.path.join(output_dir, f"{count}_{class_idx}.jpg"),
)
class_idx += 1
found = True
return found
except Exception as e:
logger.error(f"YOLO processing failed: {e}")
return False
def process_full_mode(img: Image.Image, output_dir: str, count: int) -> bool:
"""Save full image."""
_save_jpeg(img, os.path.join(output_dir, f"{count}.jpg"))
return True
winnow/jobs.py
+15 -44
View File
@@ -26,10 +26,8 @@ STRATEGY_PRESETS = {
}
def _get_strategy_choice(has_embedding: bool, entity_type: str) -> tuple[int | str, str]:
def _get_strategy_choice(has_embedding: bool) -> tuple[int | str, str]:
"""Prompt user for training strategy and return (limit, selection_mode)."""
model_name = "InsightFace" if entity_type == "face" else "SigLIP"
if has_embedding:
rprint(" [bold]1.[/bold] Adaptive Diversity [green][Recommended][/green]")
rprint(" [dim]• Dynamically selects images until redundancy starts[/dim]")
@@ -51,7 +49,7 @@ def _get_strategy_choice(has_embedding: bool, entity_type: str) -> tuple[int | s
return 30, "smart"
# Fallback when embedding model not available
rprint(f" [yellow]Note: {model_name} not available. Using Time Spread.[/yellow]")
rprint(" [yellow]Note: InsightFace not available. Using Time Spread.[/yellow]")
rprint(" [bold]1.[/bold] Standard (30 images) [green][Recommended][/green]")
rprint(" [bold]2.[/bold] Broad (100 images)")
rprint(" [bold]3.[/bold] Custom Count")
@@ -86,15 +84,13 @@ def _resolve_strategy(strategy: str, has_embedding: bool) -> tuple[int | str, st
def _perform_selection(
assets: list, limit: int | str, name: str, selection_mode: str, entity_type: str, person_id: str | None = None
assets: list, limit: int | str, name: str, selection_mode: str, person_id: str | None = None
) -> list:
"""Run diversity selection with progress display."""
if selection_mode == "smart":
model_display = "InsightFace (face embeddings)" if entity_type == "face" else "SigLIP (visual embeddings)"
rprint(f"\n[cyan]Using {model_display} for diversity analysis...[/cyan]")
rprint("\n[cyan]Using InsightFace (face embeddings) for diversity analysis...[/cyan]")
# Pre-load model explicitly (separate from availability check)
load_embedding_model(entity_type)
load_embedding_model()
with Progress(
SpinnerColumn(),
@@ -109,7 +105,6 @@ def _perform_selection(
limit,
name,
selection_mode=selection_mode,
entity_type=entity_type,
person_id=person_id,
progress_callback=lambda c, t: progress.update(task, completed=c, total=t),
)
@@ -120,9 +115,7 @@ def _perform_selection(
rprint(f"\n[cyan]Using time-spread selection for {limit} images...[/cyan]")
with console.status(f"[bold]Selecting {limit} images evenly distributed over time...[/bold]"):
selected = select_diverse_assets(
assets, limit, name, selection_mode="time", entity_type=entity_type, person_id=person_id
)
selected = select_diverse_assets(assets, limit, name, selection_mode="time", person_id=person_id)
rprint(f" [green]Selected {len(selected)} images using time spread.[/green]")
return selected
@@ -132,22 +125,12 @@ def _configure_person(person: dict, people: list[dict]) -> dict | None:
name = person["name"]
console.print(f"\nSelected: [bold green]{name}[/bold green]")
# Select training mode
rprint("\n[bold cyan]Training Mode:[/bold cyan]")
rprint(" [bold]1.[/bold] Face (Frigate Face Recognition)")
rprint(" [bold]2.[/bold] Object (Frigate Object Classification)")
mode_choice = Prompt.ask("Choice", choices=["1", "2"], default="1")
entity_type = "face" if mode_choice == "1" else "object"
config = {"name": name, "mode": entity_type, "quality_replacement": Config.QUALITY_REPLACEMENT}
if entity_type == "object":
config["object_class"] = Prompt.ask("Enter Object Class (e.g. dog, cat, car)", default="dog")
config = {"name": name, "mode": "face", "quality_replacement": Config.QUALITY_REPLACEMENT}
# Fetch and filter assets
years = IntPrompt.ask("Filter images older than (years)", default=Config.YEARS_FILTER)
console.print(f"Scanning for {name} ({entity_type})...")
console.print(f"Scanning for {name}...")
with console.status("[bold green]Fetching assets...[/bold green]"):
all_assets = fetch_all_assets(person)
recent_assets = filter_recent_assets(all_assets, years=years)
@@ -171,17 +154,15 @@ def _configure_person(person: dict, people: list[dict]) -> dict | None:
return None
# Strategy selection
has_embedding = is_embedding_available(entity_type)
has_embedding = is_embedding_available()
rprint(f"\n[bold cyan]Select Training Strategy for {name}:[/bold cyan]")
limit, selection_mode = _get_strategy_choice(has_embedding, entity_type)
limit, selection_mode = _get_strategy_choice(has_embedding)
if selection_mode == "skip":
return None
# Perform selection
selected_assets = _perform_selection(
recent_assets, limit, name, selection_mode, entity_type, person_id=person["id"]
)
selected_assets = _perform_selection(recent_assets, limit, name, selection_mode, person_id=person["id"])
rprint(f" [green]Queued {len(selected_assets)} images for {name}.[/green]")
return {"person": person, "assets": selected_assets, "limit": len(selected_assets), "config": config}
@@ -231,7 +212,6 @@ def auto_configure(people: list[dict]) -> list[dict]:
rprint("[red]No people found with names in Immich.[/red]")
return []
mode = os.environ.get("TRAINING_MODE", "face")
strategy = os.environ.get("STRATEGY", "auto")
skip = os.environ.get("SKIP_PEOPLE", "").split(",") if os.environ.get("SKIP_PEOPLE") else []
only = os.environ.get("ONLY_PEOPLE", "").split(",") if os.environ.get("ONLY_PEOPLE") else []
@@ -260,11 +240,7 @@ def auto_configure(people: list[dict]) -> list[dict]:
jobs = []
for person in valid_people:
name = person["name"]
entity_type = mode
config = {"name": name, "mode": entity_type}
if entity_type == "object":
config["object_class"] = os.environ.get("OBJECT_CLASS", "dog")
config = {"name": name, "mode": "face"}
all_assets = fetch_all_assets(person)
recent_assets = filter_recent_assets(all_assets, years=Config.YEARS_FILTER)
@@ -307,7 +283,7 @@ def auto_configure(people: list[dict]) -> list[dict]:
config["quality_replacement"] = quality_replacement_only or Config.QUALITY_REPLACEMENT
has_embedding = is_embedding_available(entity_type)
has_embedding = is_embedding_available()
limit, selection_mode = _resolve_strategy(strategy, has_embedding)
# Cap selection to remaining capacity (no cap when replacement-only — executor
@@ -323,9 +299,7 @@ def auto_configure(people: list[dict]) -> list[dict]:
if selection_mode == "skip":
continue
selected_assets = _perform_selection(
recent_assets, limit, name, selection_mode, entity_type, person_id=person["id"]
)
selected_assets = _perform_selection(recent_assets, limit, name, selection_mode, person_id=person["id"])
if auto_cap is not None:
selected_assets = selected_assets[:auto_cap]
@@ -340,20 +314,17 @@ 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)
table.add_row(name, count, date_range)
console.print()
console.print(table)