initialize if_curator project with core modules for image embedding, processing, diversity analysis, and Immich API integration.
This commit is contained in:
Sebastian G
@@ -0,0 +1,4 @@
|
||||
from .cli import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,250 @@
|
||||
"""Interactive CLI for if-curator."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from io import BytesIO
|
||||
|
||||
import requests
|
||||
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 .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 .logging import console, setup_logging
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Strategy presets: (limit, mode_name)
|
||||
STRATEGY_PRESETS = {
|
||||
"1": ("auto", "Auto Diversity"),
|
||||
"2": (30, "Standard (30)"),
|
||||
"3": (100, "Broad (100)"),
|
||||
}
|
||||
|
||||
|
||||
def _get_strategy_choice(has_embedding: bool, entity_type: str) -> tuple[int | str, str]:
|
||||
"""Prompt user for training strategy and return (limit, selection_mode)."""
|
||||
model_name = "InsightFace" if entity_type == "face" else "SigLIP"
|
||||
|
||||
if has_embedding:
|
||||
rprint(" [bold]1.[/bold] Auto (Objective Diversity) [green][Recommended][/green]")
|
||||
rprint(" [dim]• Dynamically selects images until redundancy starts[/dim]")
|
||||
rprint(" [bold]2.[/bold] Standard (30 images)")
|
||||
rprint(" [bold]3.[/bold] Broad (100 images)")
|
||||
rprint(" [bold]4.[/bold] Custom Count")
|
||||
rprint(" [bold]5.[/bold] Skip")
|
||||
|
||||
choice = Prompt.ask("Choice", choices=["1", "2", "3", "4", "5"], default="1")
|
||||
|
||||
if choice == "5":
|
||||
return 0, "skip"
|
||||
if choice == "4":
|
||||
limit = IntPrompt.ask("Enter number of images", default=30)
|
||||
mode = "smart" if Confirm.ask("Use Smart Diversity?", default=True) else "time"
|
||||
return limit, mode
|
||||
if choice in STRATEGY_PRESETS:
|
||||
return STRATEGY_PRESETS[choice][0], "smart"
|
||||
return 30, "smart"
|
||||
|
||||
# Fallback when embedding model not available
|
||||
rprint(f" [yellow]Note: {model_name} not available. Using Time Spread.[/yellow]")
|
||||
rprint(" [bold]1.[/bold] Standard (30 images) [green][Recommended][/green]")
|
||||
rprint(" [bold]2.[/bold] Broad (100 images)")
|
||||
rprint(" [bold]3.[/bold] Custom Count")
|
||||
rprint(" [bold]4.[/bold] Skip")
|
||||
|
||||
choice = Prompt.ask("Choice", choices=["1", "2", "3", "4"], default="1")
|
||||
limits = {"1": 30, "2": 100, "3": IntPrompt.ask("Enter number of images", default=30)}
|
||||
return limits.get(choice, 0), "time" if choice != "4" else "skip"
|
||||
|
||||
|
||||
def 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]
|
||||
name = person["name"]
|
||||
|
||||
console.print(f"\nSelected: [bold green]{name}[/bold green]")
|
||||
|
||||
# Select training mode
|
||||
rprint("\n[bold cyan]Training Mode:[/bold cyan]")
|
||||
rprint(" [bold]1.[/bold] Face (Frigate Face Recognition)")
|
||||
rprint(" [bold]2.[/bold] Object (Frigate Object Classification)")
|
||||
|
||||
mode_choice = Prompt.ask("Choice", choices=["1", "2"], default="1")
|
||||
entity_type = "face" if mode_choice == "1" else "object"
|
||||
|
||||
config = {"name": name, "mode": entity_type}
|
||||
if entity_type == "object":
|
||||
config["object_class"] = Prompt.ask("Enter Object Class (e.g. dog, cat, car)", default="dog")
|
||||
|
||||
# Fetch and filter assets
|
||||
years = IntPrompt.ask("Filter images older than (years)", default=Config.YEARS_FILTER)
|
||||
|
||||
console.print(f"Scanning for {name} ({entity_type})...")
|
||||
with console.status("[bold green]Fetching assets...[/bold green]"):
|
||||
all_assets = fetch_all_assets(person)
|
||||
recent_assets = filter_recent_assets(all_assets, years=years)
|
||||
|
||||
rprint(f" Found [bold]{len(all_assets)}[/bold] total, [bold]{len(recent_assets)}[/bold] in range ({years} years).")
|
||||
|
||||
if not recent_assets:
|
||||
rprint(" [dim]Skipping (0 recent images).[/dim]")
|
||||
return []
|
||||
|
||||
# Strategy selection
|
||||
has_embedding = is_embedding_available(entity_type)
|
||||
rprint(f"\n[bold cyan]Select Training Strategy for {name}:[/bold cyan]")
|
||||
|
||||
limit, selection_mode = _get_strategy_choice(has_embedding, entity_type)
|
||||
if selection_mode == "skip":
|
||||
return []
|
||||
|
||||
# Perform selection
|
||||
selected_assets = _perform_selection(recent_assets, limit, name, selection_mode, entity_type)
|
||||
|
||||
rprint(f" [green]Queued {len(selected_assets)} images for {name}.[/green]")
|
||||
return [{"person": person, "assets": selected_assets, "limit": len(selected_assets), "config": config}]
|
||||
|
||||
|
||||
def _perform_selection(
|
||||
assets: list, limit: int | str, name: str, selection_mode: str, entity_type: str
|
||||
) -> list:
|
||||
"""Run diversity selection with progress display."""
|
||||
if selection_mode == "smart":
|
||||
model_display = "InsightFace (face embeddings)" if entity_type == "face" else "SigLIP (visual embeddings)"
|
||||
rprint(f"\n[cyan]Using {model_display} for diversity analysis...[/cyan]")
|
||||
|
||||
# Pre-load model to avoid interference with progress bar
|
||||
is_embedding_available(entity_type)
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(), TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(), TaskProgressColumn(), console=console,
|
||||
) as progress:
|
||||
task = progress.add_task(f"[cyan]Computing embeddings for {len(assets)} images...", total=None)
|
||||
selected = select_diverse_assets(
|
||||
assets, limit, name,
|
||||
selection_mode=selection_mode,
|
||||
entity_type=entity_type,
|
||||
progress_callback=lambda c, t: progress.update(task, completed=c, total=t),
|
||||
)
|
||||
|
||||
label = f"Auto-diversity selected {len(selected)}" if limit == "auto" else f"Selected {len(selected)}"
|
||||
rprint(f" [green]{label} diverse images.[/green]")
|
||||
return selected
|
||||
|
||||
rprint(f"\n[cyan]Using time-spread selection for {limit} images...[/cyan]")
|
||||
with console.status(f"[bold]Selecting {limit} images evenly distributed over time...[/bold]"):
|
||||
selected = select_diverse_assets(assets, limit, name, selection_mode="time", entity_type=entity_type)
|
||||
rprint(f" [green]Selected {len(selected)} images using time spread.[/green]")
|
||||
return selected
|
||||
|
||||
|
||||
def execute_jobs(jobs: list[dict]) -> None:
|
||||
"""Download and process images for all jobs."""
|
||||
if not jobs:
|
||||
return
|
||||
|
||||
console.rule("[bold blue]Execution Phase")
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(), TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(), TaskProgressColumn(), console=console,
|
||||
) as progress:
|
||||
grand_total = sum(j["limit"] for j in jobs)
|
||||
overall_task = progress.add_task("[green]Overall Progress", total=grand_total)
|
||||
|
||||
for job in jobs:
|
||||
person, assets, config = job["person"], job["assets"], job["config"]
|
||||
name, mode = person["name"], config.get("mode", "face")
|
||||
|
||||
job_task = progress.add_task(f"Processing {name}...", total=len(assets))
|
||||
person_dir = os.path.join(Config.OUTPUT_DIR, name)
|
||||
os.makedirs(person_dir, exist_ok=True)
|
||||
|
||||
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:
|
||||
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)
|
||||
if mode == "face"
|
||||
else process_object_mode(img, config, person_dir, count)
|
||||
if mode == "object"
|
||||
else process_full_mode(img, person_dir, count)
|
||||
)
|
||||
if saved:
|
||||
count += 1
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to process asset {asset['id']}: {e}")
|
||||
|
||||
progress.advance(job_task)
|
||||
progress.advance(overall_task)
|
||||
|
||||
progress.remove_task(job_task)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Entry point for if-curator CLI."""
|
||||
try:
|
||||
setup_logging(verbose=False)
|
||||
|
||||
console.print(r"""
|
||||
[bold blue]if-curator[/bold blue]
|
||||
[dim]Immich -> Frigate Training Data Curator[/dim]
|
||||
""")
|
||||
|
||||
ConfigManager.get().interactive_setup()
|
||||
|
||||
try:
|
||||
Config.validate()
|
||||
except ValueError as e:
|
||||
rprint(f"[bold red]Configuration Error:[/bold red] {e}")
|
||||
return
|
||||
|
||||
rprint(f"Server: [dim]{Config.IMMICH_URL}[/dim]")
|
||||
rprint(f"Output: [dim]{Config.OUTPUT_DIR}[/dim]")
|
||||
|
||||
people = get_people()
|
||||
if not people:
|
||||
rprint("[bold red]Could not fetch people from Immich. Check URL/Key.[/bold red]")
|
||||
return
|
||||
|
||||
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:
|
||||
rprint("[yellow]No jobs configured.[/yellow]")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
rprint("\n[bold red]Aborted by user.[/bold red]")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,101 @@
|
||||
"""Configuration management for if-curator."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import ClassVar
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from rich.prompt import Prompt
|
||||
|
||||
load_dotenv()
|
||||
|
||||
CONFIG_FILE = Path(".immich_config.json")
|
||||
|
||||
|
||||
class Config:
|
||||
"""Singleton configuration with uppercase attribute access for backward compatibility."""
|
||||
|
||||
_instance: ClassVar["Config | None"] = None
|
||||
|
||||
# Configuration values
|
||||
IMMICH_URL: str | None = None
|
||||
API_KEY: str | None = None
|
||||
OUTPUT_DIR: str = "./frigate_train"
|
||||
YEARS_FILTER: int = 10
|
||||
MIN_FACE_WIDTH: int = 50
|
||||
|
||||
def __new__(cls) -> "Config":
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._load()
|
||||
return cls._instance
|
||||
|
||||
def _load(self) -> None:
|
||||
"""Load configuration from environment and config file."""
|
||||
# Load from environment (highest priority)
|
||||
self.IMMICH_URL = os.getenv("IMMICH_URL")
|
||||
self.API_KEY = os.getenv("API_KEY")
|
||||
self.OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./frigate_train")
|
||||
self.YEARS_FILTER = int(os.getenv("YEARS_FILTER", "10"))
|
||||
self.MIN_FACE_WIDTH = int(os.getenv("MIN_FACE_WIDTH", "50"))
|
||||
|
||||
# Fall back to config file for missing values
|
||||
if CONFIG_FILE.exists():
|
||||
try:
|
||||
data = json.loads(CONFIG_FILE.read_text())
|
||||
self.IMMICH_URL = self.IMMICH_URL or data.get("IMMICH_URL")
|
||||
self.API_KEY = self.API_KEY or data.get("API_KEY")
|
||||
if not os.getenv("OUTPUT_DIR"):
|
||||
self.OUTPUT_DIR = data.get("OUTPUT_DIR", self.OUTPUT_DIR)
|
||||
except (json.JSONDecodeError, OSError) as e:
|
||||
logging.warning(f"Failed to load config file: {e}")
|
||||
|
||||
def save(self) -> None:
|
||||
"""Persist configuration to file."""
|
||||
try:
|
||||
CONFIG_FILE.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"IMMICH_URL": self.IMMICH_URL,
|
||||
"API_KEY": self.API_KEY,
|
||||
"OUTPUT_DIR": self.OUTPUT_DIR,
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
logging.info(f"Configuration saved to {CONFIG_FILE}")
|
||||
except OSError as e:
|
||||
logging.error(f"Failed to save config: {e}")
|
||||
|
||||
def interactive_setup(self) -> None:
|
||||
"""Prompt user for missing configuration."""
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
|
||||
if not self.IMMICH_URL:
|
||||
console.print("[yellow]Immich URL not found.[/yellow]")
|
||||
self.IMMICH_URL = Prompt.ask("Enter Immich URL (e.g. http://192.168.1.5:2283)")
|
||||
self.save()
|
||||
|
||||
if not self.API_KEY:
|
||||
console.print("[yellow]Immich API Key not found.[/yellow]")
|
||||
self.API_KEY = Prompt.ask("Enter Immich API Key", password=True)
|
||||
self.save()
|
||||
|
||||
def validate(self) -> None:
|
||||
"""Raise ValueError if required config is missing."""
|
||||
if not self.IMMICH_URL or not self.API_KEY:
|
||||
raise ValueError("Missing Immich URL or API Key.")
|
||||
|
||||
|
||||
# Singleton instance and backward-compatible aliases
|
||||
Config = Config() # type: ignore[misc]
|
||||
ConfigManager = type("ConfigManager", (), {"get": staticmethod(lambda: Config)})
|
||||
|
||||
|
||||
def get_headers() -> dict[str, str]:
|
||||
"""Return HTTP headers for Immich API requests."""
|
||||
return {"x-api-key": Config.API_KEY or "", "Accept": "application/json"}
|
||||
@@ -0,0 +1,181 @@
|
||||
"""
|
||||
Diversity selection for training data curation.
|
||||
|
||||
Uses Farthest Point Sampling (FPS) algorithm with embeddings:
|
||||
- Faces: InsightFace embeddings
|
||||
- Objects: SigLIP embeddings
|
||||
"""
|
||||
|
||||
import logging
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
from .config import Config, get_headers
|
||||
from .embeddings import get_embedding, is_embedding_available
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def select_diverse_assets(
|
||||
assets: list,
|
||||
limit: int | str,
|
||||
entity_name: str,
|
||||
selection_mode: str = "smart",
|
||||
entity_type: str = "face",
|
||||
progress_callback=None,
|
||||
) -> list:
|
||||
"""
|
||||
Select diverse assets using Farthest Point Sampling or time spread.
|
||||
|
||||
Args:
|
||||
assets: List of asset dicts from Immich API
|
||||
limit: Number to select, or "auto" for dynamic selection
|
||||
entity_name: Name of the person/object for logging
|
||||
selection_mode: 'smart' (embedding-based) or 'time' (time spread)
|
||||
entity_type: 'face' or 'object' - determines embedding model
|
||||
progress_callback: Optional callback(current, total) for progress
|
||||
|
||||
Returns:
|
||||
List of selected assets
|
||||
"""
|
||||
# Fast path: fewer assets than limit
|
||||
if limit != "auto" and len(assets) <= limit:
|
||||
return assets
|
||||
|
||||
# Sort by creation time
|
||||
assets = sorted(assets, key=lambda x: x.get("fileCreatedAt", ""))
|
||||
|
||||
if selection_mode != "smart" or not is_embedding_available(entity_type):
|
||||
if selection_mode == "smart":
|
||||
model_name = "InsightFace" if entity_type == "face" else "SigLIP"
|
||||
logger.warning(f"{model_name} unavailable. Falling back to time spread.")
|
||||
return _select_time_spread(assets, limit)
|
||||
|
||||
try:
|
||||
return _select_by_embedding(assets, limit, entity_type, progress_callback)
|
||||
except Exception as e:
|
||||
logger.error(f"Smart Diversity failed: {e}. Falling back to time spread.")
|
||||
return _select_time_spread(assets, limit)
|
||||
|
||||
|
||||
def _fetch_thumbnail(asset_id: str, timeout: int = 10) -> Image.Image | None:
|
||||
"""Fetch thumbnail from Immich API."""
|
||||
try:
|
||||
url = f"{Config.IMMICH_URL}/api/assets/{asset_id}/thumbnail?size=preview&format=JPEG"
|
||||
resp = requests.get(url, headers=get_headers(), timeout=timeout)
|
||||
return Image.open(BytesIO(resp.content)) if resp.ok else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _select_by_embedding(
|
||||
assets: list,
|
||||
limit: int | str,
|
||||
entity_type: str,
|
||||
progress_callback=None,
|
||||
) -> list:
|
||||
"""Select assets using embedding-based Farthest Point Sampling."""
|
||||
# Determine candidate pool (cap at 3000 for performance)
|
||||
effective_limit = 30 if limit == "auto" else limit
|
||||
pool_size = min(3000, max(effective_limit * 20, len(assets)))
|
||||
|
||||
# Subsample if needed (evenly distributed in time)
|
||||
if len(assets) > pool_size:
|
||||
indices = np.linspace(0, len(assets) - 1, pool_size, dtype=int)
|
||||
candidates = [assets[i] for i in indices]
|
||||
else:
|
||||
candidates = assets
|
||||
|
||||
# Compute embeddings
|
||||
embeddings, valid_candidates = [], []
|
||||
for i, asset in enumerate(candidates):
|
||||
if progress_callback:
|
||||
progress_callback(i, len(candidates))
|
||||
|
||||
img = _fetch_thumbnail(asset["id"])
|
||||
if img is None:
|
||||
continue
|
||||
|
||||
emb = get_embedding(img, entity_type)
|
||||
if emb is not None:
|
||||
embeddings.append(emb)
|
||||
valid_candidates.append(asset)
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(len(candidates), len(candidates))
|
||||
|
||||
if not embeddings:
|
||||
logger.warning("No valid embeddings found. Falling back to time spread.")
|
||||
return _select_time_spread(assets, limit)
|
||||
|
||||
if limit != "auto" and len(valid_candidates) < limit:
|
||||
logger.warning(f"Only {len(valid_candidates)} valid embeddings. Returning all.")
|
||||
return valid_candidates
|
||||
|
||||
# Farthest Point Sampling with vectorized distance computation
|
||||
return _farthest_point_sampling(
|
||||
embeddings, valid_candidates, limit, auto_threshold=0.15
|
||||
)
|
||||
|
||||
|
||||
def _farthest_point_sampling(
|
||||
embeddings: list,
|
||||
candidates: list,
|
||||
limit: int | str,
|
||||
auto_threshold: float = 0.15,
|
||||
) -> list:
|
||||
"""Vectorized Farthest Point Sampling."""
|
||||
emb_matrix = np.vstack(embeddings) # (N, D)
|
||||
n = len(emb_matrix)
|
||||
|
||||
# Normalize for cosine distance (cosine_dist = 1 - cosine_sim)
|
||||
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]
|
||||
min_dists = np.full(n, np.inf)
|
||||
|
||||
target = 500 if limit == "auto" else limit
|
||||
|
||||
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
|
||||
|
||||
# Find farthest point
|
||||
best_idx = np.argmax(min_dists)
|
||||
best_dist = min_dists[best_idx]
|
||||
|
||||
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})."
|
||||
)
|
||||
break
|
||||
|
||||
selected.append(best_idx)
|
||||
|
||||
logger.info(f"Smart selection complete. Picked {len(selected)} diverse images.")
|
||||
return [candidates[i] for i in selected]
|
||||
|
||||
|
||||
def _select_time_spread(assets: list, limit: int | str) -> list:
|
||||
"""Select N assets evenly distributed in time."""
|
||||
if limit == "auto":
|
||||
limit = 30
|
||||
|
||||
logger.info(f"Selecting {limit} images using time spread.")
|
||||
|
||||
if len(assets) <= limit:
|
||||
return assets
|
||||
|
||||
indices = np.linspace(0, len(assets) - 1, limit, dtype=int)
|
||||
return [assets[i] for i in np.unique(indices)]
|
||||
@@ -0,0 +1,192 @@
|
||||
"""
|
||||
Unified embedding interface for faces and objects.
|
||||
|
||||
- Faces: InsightFace (ArcFace/Buffalo_L)
|
||||
- Objects: SigLIP (Vision Transformer via transformers)
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Lazy-loaded singletons
|
||||
_insightface_app = None
|
||||
_siglip_model = None
|
||||
_siglip_processor = None
|
||||
|
||||
|
||||
def _is_force_cpu() -> bool:
|
||||
"""Check if CPU mode is forced via environment variable."""
|
||||
return os.getenv("FORCE_CPU", "").lower() in ("true", "1", "yes")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# InsightFace (Faces)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def get_insightface_app():
|
||||
"""Singleton for InsightFace app with automatic GPU/CPU fallback."""
|
||||
global _insightface_app
|
||||
if _insightface_app is not None:
|
||||
return _insightface_app
|
||||
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
from insightface.app import FaceAnalysis
|
||||
|
||||
# Get providers, excluding TensorRT to avoid noisy errors
|
||||
providers = [p for p in ort.get_available_providers() if p != "TensorrtExecutionProvider"]
|
||||
logger.info(f"Available ONNX providers: {providers}")
|
||||
|
||||
# Determine device: 0 for GPU, -1 for CPU
|
||||
gpu_providers = {
|
||||
"CUDAExecutionProvider", "ROCmExecutionProvider",
|
||||
"MPSExecutionProvider", "CoreMLExecutionProvider",
|
||||
}
|
||||
ctx_id = -1 if _is_force_cpu() else (0 if gpu_providers & set(providers) else -1)
|
||||
|
||||
device_str = "GPU" if ctx_id >= 0 else "CPU"
|
||||
logger.info(f"Loading InsightFace Buffalo_L on {device_str} (ctx_id={ctx_id})...")
|
||||
|
||||
# Suppress C-level output during model loading
|
||||
with open(os.devnull, "w") as devnull, contextlib.redirect_stdout(devnull), contextlib.redirect_stderr(devnull):
|
||||
_insightface_app = FaceAnalysis(name="buffalo_l", root="~/.insightface", providers=providers)
|
||||
_insightface_app.prepare(ctx_id=ctx_id, det_size=(640, 640))
|
||||
|
||||
return _insightface_app
|
||||
|
||||
except ImportError:
|
||||
logger.error("InsightFace not installed!")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load InsightFace: {e}")
|
||||
# Retry on CPU if GPU failed
|
||||
if ctx_id == 0:
|
||||
logger.warning("Retrying InsightFace on CPU...")
|
||||
try:
|
||||
from insightface.app import FaceAnalysis
|
||||
_insightface_app = FaceAnalysis(name="buffalo_l", root="~/.insightface")
|
||||
_insightface_app.prepare(ctx_id=-1, det_size=(640, 640))
|
||||
return _insightface_app
|
||||
except Exception as ex:
|
||||
logger.error(f"CPU fallback failed: {ex}")
|
||||
return None
|
||||
|
||||
|
||||
def get_face_embedding(img_pil: Image.Image) -> np.ndarray | None:
|
||||
"""Get embedding of the largest face in a PIL image."""
|
||||
app = get_insightface_app()
|
||||
if not app:
|
||||
return None
|
||||
|
||||
try:
|
||||
# InsightFace expects BGR cv2 image
|
||||
img_bgr = cv2.cvtColor(np.asarray(img_pil), cv2.COLOR_RGB2BGR)
|
||||
faces = app.get(img_bgr)
|
||||
|
||||
if not faces:
|
||||
return None
|
||||
|
||||
# Return embedding of largest face
|
||||
largest = max(faces, key=lambda f: (f.bbox[2] - f.bbox[0]) * (f.bbox[3] - f.bbox[1]))
|
||||
return largest.embedding
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting face embedding: {e}")
|
||||
return None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# SigLIP (Objects)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def get_siglip_model():
|
||||
"""Singleton for SigLIP model and processor with GPU auto-detection."""
|
||||
global _siglip_model, _siglip_processor
|
||||
if _siglip_model is not None:
|
||||
return _siglip_model, _siglip_processor
|
||||
|
||||
try:
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
from transformers import AutoImageProcessor, SiglipVisionModel
|
||||
|
||||
model_name = "google/siglip-base-patch16-224"
|
||||
logger.info(f"Loading SigLIP model ({model_name})...")
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.filterwarnings("ignore", category=FutureWarning)
|
||||
warnings.filterwarnings("ignore", message=".*use_fast.*")
|
||||
_siglip_processor = AutoImageProcessor.from_pretrained(model_name, use_fast=True)
|
||||
_siglip_model = SiglipVisionModel.from_pretrained(model_name)
|
||||
|
||||
_siglip_model.eval()
|
||||
|
||||
# Move to GPU if available
|
||||
if not _is_force_cpu():
|
||||
if torch.cuda.is_available():
|
||||
_siglip_model = _siglip_model.cuda()
|
||||
logger.info("SigLIP running on CUDA GPU")
|
||||
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
_siglip_model = _siglip_model.to("mps")
|
||||
logger.info("SigLIP running on Apple MPS")
|
||||
else:
|
||||
logger.info("SigLIP running on CPU")
|
||||
else:
|
||||
logger.info("FORCE_CPU set. SigLIP running on CPU")
|
||||
|
||||
return _siglip_model, _siglip_processor
|
||||
|
||||
except ImportError as e:
|
||||
logger.error(f"transformers/torch not installed: {e}")
|
||||
return None, None
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load SigLIP: {e}")
|
||||
return None, None
|
||||
|
||||
|
||||
def get_object_embedding(img_pil: Image.Image) -> np.ndarray | None:
|
||||
"""Get 768-dim SigLIP embedding for an image."""
|
||||
model, processor = get_siglip_model()
|
||||
if model is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
inputs = processor(images=img_pil, return_tensors="pt")
|
||||
device = next(model.parameters()).device
|
||||
inputs = {k: v.to(device) for k, v in inputs.items()}
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = model(**inputs)
|
||||
return outputs.pooler_output.squeeze().cpu().numpy()
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting object embedding: {e}")
|
||||
return None
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Unified Interface
|
||||
# =============================================================================
|
||||
|
||||
|
||||
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 is_embedding_available(entity_type: str = "face") -> bool:
|
||||
"""Check if embedding model is available for the given entity type."""
|
||||
if entity_type == "face":
|
||||
return get_insightface_app() is not None
|
||||
model, _ = get_siglip_model()
|
||||
return model is not None
|
||||
@@ -0,0 +1,111 @@
|
||||
"""Image processing functions for cropping faces and objects."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from PIL import Image
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Lazy singleton
|
||||
_yolo_model = None
|
||||
|
||||
|
||||
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 process_face_mode(
|
||||
img: Image.Image,
|
||||
asset: dict,
|
||||
person: dict,
|
||||
output_dir: str,
|
||||
count: int,
|
||||
min_width: int = 50,
|
||||
) -> bool:
|
||||
"""Crop face based on Immich metadata and save to output directory."""
|
||||
# Find face metadata for this person
|
||||
face_info = None
|
||||
for p in asset.get("people", []):
|
||||
if p["id"] == person["id"] and (faces := p.get("faces")):
|
||||
face_info = faces[0]
|
||||
break
|
||||
|
||||
if not face_info:
|
||||
logger.debug(f"No face info for {person.get('name')} in asset {asset.get('id')}")
|
||||
return False
|
||||
|
||||
img_w, img_h = img.size
|
||||
meta_w = face_info.get("imageWidth") or img_w
|
||||
meta_h = face_info.get("imageHeight") or img_h
|
||||
|
||||
# Scale bounding box to actual image dimensions
|
||||
scale_x, scale_y = img_w / meta_w, img_h / meta_h
|
||||
x1 = face_info["boundingBoxX1"] * scale_x
|
||||
y1 = face_info["boundingBoxY1"] * scale_y
|
||||
x2 = face_info["boundingBoxX2"] * scale_x
|
||||
y2 = face_info["boundingBoxY2"] * scale_y
|
||||
|
||||
face_w, face_h = x2 - x1, y2 - y1
|
||||
if face_w < min_width or face_h < min_width:
|
||||
logger.debug(f"Face too small ({face_w:.1f}x{face_h:.1f})")
|
||||
return False
|
||||
|
||||
# Add 10% margin
|
||||
margin_x, margin_y = face_w * 0.10, face_h * 0.10
|
||||
crop_box = (
|
||||
max(0, x1 - margin_x),
|
||||
max(0, y1 - margin_y),
|
||||
min(img_w, x2 + margin_x),
|
||||
min(img_h, y2 + margin_y),
|
||||
)
|
||||
|
||||
face_crop = img.crop(crop_box)
|
||||
face_crop.save(os.path.join(output_dir, f"{count}.jpg"), format="JPEG")
|
||||
return True
|
||||
|
||||
|
||||
def process_object_mode(
|
||||
img: Image.Image,
|
||||
config: dict,
|
||||
output_dir: str,
|
||||
count: int,
|
||||
) -> bool:
|
||||
"""Detect and crop objects using YOLO."""
|
||||
try:
|
||||
model = get_yolo_model()
|
||||
target_class = config.get("object_class", "dog")
|
||||
device = "cpu" if os.getenv("FORCE_CPU", "").lower() in ("true", "1", "yes") else None
|
||||
|
||||
results = model(img, verbose=False, device=device)
|
||||
|
||||
found = False
|
||||
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
|
||||
):
|
||||
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()
|
||||
img.crop((x1, y1, x2, y2)).save(
|
||||
os.path.join(output_dir, f"{count}_{idx}.jpg"),
|
||||
format="JPEG",
|
||||
)
|
||||
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."""
|
||||
img.save(os.path.join(output_dir, f"{count}.jpg"), format="JPEG")
|
||||
return True
|
||||
@@ -0,0 +1,95 @@
|
||||
"""Immich API client for fetching people and assets."""
|
||||
|
||||
import logging
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import requests
|
||||
|
||||
from .config import Config, get_headers
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_people() -> list[dict]:
|
||||
"""Fetch all people from Immich."""
|
||||
try:
|
||||
resp = requests.get(
|
||||
f"{Config.IMMICH_URL}/api/people",
|
||||
headers=get_headers(),
|
||||
timeout=10,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return resp.json().get("people", [])
|
||||
except requests.RequestException as e:
|
||||
logger.error(f"Failed to fetch people: {e}")
|
||||
return []
|
||||
|
||||
|
||||
def fetch_all_assets(person: dict) -> list[dict]:
|
||||
"""Fetch all assets for a person with pagination."""
|
||||
name = person.get("name", "Unknown")
|
||||
person_id = person["id"]
|
||||
url = f"{Config.IMMICH_URL}/api/search/metadata"
|
||||
page_size = 1000
|
||||
|
||||
logger.info(f"Fetching assets for {name}...")
|
||||
|
||||
assets = []
|
||||
for page in range(1, 1000): # Safety limit
|
||||
try:
|
||||
resp = requests.post(
|
||||
url,
|
||||
json={"personIds": [person_id], "size": page_size, "page": page},
|
||||
headers=get_headers(),
|
||||
timeout=30,
|
||||
)
|
||||
|
||||
if not resp.ok:
|
||||
logger.error(f"Error fetching assets for {name} (page {page}): {resp.status_code}")
|
||||
break
|
||||
|
||||
page_assets = resp.json().get("assets", [])
|
||||
if isinstance(page_assets, dict):
|
||||
page_assets = page_assets.get("items", [])
|
||||
|
||||
if not page_assets:
|
||||
break
|
||||
|
||||
assets.extend(page_assets)
|
||||
logger.debug(f"Fetched page {page}, total: {len(assets)}")
|
||||
|
||||
if len(page_assets) < page_size:
|
||||
break
|
||||
|
||||
except requests.RequestException as e:
|
||||
logger.error(f"Exception fetching assets for {name}: {e}")
|
||||
break
|
||||
|
||||
return assets
|
||||
|
||||
|
||||
def filter_recent_assets(assets: list[dict], years: int | None = None) -> list[dict]:
|
||||
"""Filter assets to keep only those from the last N years."""
|
||||
years = years or Config.YEARS_FILTER
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=365 * years)
|
||||
|
||||
logger.debug(f"Filtering assets older than {years} years ({cutoff})")
|
||||
|
||||
recent, skipped = [], 0
|
||||
for asset in assets:
|
||||
created_at_str = asset.get("fileCreatedAt")
|
||||
if not created_at_str:
|
||||
continue
|
||||
|
||||
try:
|
||||
# Handle ISO8601 with 'Z' suffix
|
||||
created_at = datetime.fromisoformat(created_at_str.replace("Z", "+00:00"))
|
||||
if created_at > cutoff:
|
||||
recent.append(asset)
|
||||
else:
|
||||
skipped += 1
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
logger.info(f"Retained {len(recent)} assets (filtered {skipped} old assets).")
|
||||
return recent
|
||||
@@ -0,0 +1,52 @@
|
||||
"""Logging configuration for if-curator."""
|
||||
|
||||
import logging
|
||||
import warnings
|
||||
|
||||
from rich.console import Console
|
||||
from rich.logging import RichHandler
|
||||
|
||||
# Shared console instance - must be the same as used by Progress bars
|
||||
console = Console()
|
||||
|
||||
NOISY_LOGGERS = (
|
||||
"urllib3",
|
||||
"PIL",
|
||||
"ultralytics",
|
||||
"insightface",
|
||||
"onnxruntime",
|
||||
"matplotlib",
|
||||
"transformers",
|
||||
"torch",
|
||||
)
|
||||
|
||||
|
||||
def setup_logging(verbose: bool = False) -> logging.Logger:
|
||||
"""Configure logging with Rich console and file output."""
|
||||
level = logging.DEBUG if verbose else logging.INFO
|
||||
|
||||
# Configure root logger
|
||||
root = logging.getLogger()
|
||||
root.setLevel(level)
|
||||
root.handlers.clear()
|
||||
|
||||
# Rich console handler - uses shared console to avoid breaking progress bars
|
||||
root.addHandler(RichHandler(rich_tracebacks=True, markup=True, console=console))
|
||||
|
||||
# File handler (always debug level)
|
||||
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")
|
||||
)
|
||||
root.addHandler(file_handler)
|
||||
|
||||
# Silence noisy libraries
|
||||
for lib in NOISY_LOGGERS:
|
||||
logging.getLogger(lib).setLevel(logging.WARNING)
|
||||
|
||||
# Suppress Python warnings from ML libraries
|
||||
warnings.filterwarnings("ignore", category=UserWarning, module="onnxruntime")
|
||||
warnings.filterwarnings("ignore", category=FutureWarning, module="transformers")
|
||||
|
||||
return root
|
||||
Reference in New Issue
Block a user