initialize if_curator project with core modules for image embedding, processing, diversity analysis, and Immich API integration.
This commit is contained in:
Sebastian G
+235
@@ -0,0 +1,235 @@
|
||||
# Custom
|
||||
frigate_train/
|
||||
runs/
|
||||
yolov9c.pt
|
||||
.insightface/
|
||||
.huggingface/
|
||||
.cache/huggingface
|
||||
|
||||
# Python-generated files
|
||||
__pycache__/
|
||||
*.py[oc]
|
||||
build/
|
||||
dist/
|
||||
wheels/
|
||||
*.egg-info
|
||||
|
||||
# Virtual environments
|
||||
.venv
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[codz]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py.cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
# Pipfile.lock
|
||||
|
||||
# UV
|
||||
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# uv.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
# poetry.lock
|
||||
# poetry.toml
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
|
||||
# https://pdm-project.org/en/latest/usage/project/#working-with-version-control
|
||||
# pdm.lock
|
||||
# pdm.toml
|
||||
.pdm-python
|
||||
.pdm-build/
|
||||
|
||||
# pixi
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
|
||||
# pixi.lock
|
||||
# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
|
||||
# in the .venv directory. It is recommended not to include this directory in version control.
|
||||
.pixi
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# Redis
|
||||
*.rdb
|
||||
*.aof
|
||||
*.pid
|
||||
|
||||
# RabbitMQ
|
||||
mnesia/
|
||||
rabbitmq/
|
||||
rabbitmq-data/
|
||||
|
||||
# ActiveMQ
|
||||
activemq-data/
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.envrc
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
# .idea/
|
||||
|
||||
# Abstra
|
||||
# Abstra is an AI-powered process automation framework.
|
||||
# Ignore directories containing user credentials, local state, and settings.
|
||||
# Learn more at https://abstra.io/docs
|
||||
.abstra/
|
||||
|
||||
# Visual Studio Code
|
||||
# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
|
||||
# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. However, if you prefer,
|
||||
# you could uncomment the following to ignore the entire vscode folder
|
||||
# .vscode/
|
||||
|
||||
# Ruff stuff:
|
||||
.ruff_cache/
|
||||
|
||||
# PyPI configuration file
|
||||
.pypirc
|
||||
|
||||
# Marimo
|
||||
marimo/_static/
|
||||
marimo/_lsp/
|
||||
__marimo__/
|
||||
|
||||
# Streamlit
|
||||
.streamlit/secrets.toml
|
||||
@@ -0,0 +1 @@
|
||||
3.12
|
||||
@@ -0,0 +1,126 @@
|
||||
<div align="center">
|
||||
|
||||
# 🖼️ if-curator
|
||||
### Immich to Frigate Curator
|
||||
|
||||
[](https://python.org)
|
||||
[](https://immich.app)
|
||||
[](https://frigate.video)
|
||||
|
||||
*A specialized tool to extract **high-quality, diverse** training images from your Immich library for Frigate's Face Recognition (ArcFace) and Object/State Classification models.*
|
||||
|
||||
</div>
|
||||
|
||||
> [!WARNING]
|
||||
> **Regarding Object Classification**
|
||||
>
|
||||
> Frigate **does not support** uploading custom images for object classification training via the UI or API.
|
||||
> This tool currently prepares the dataset (crops and categorizes images) for training external models (like YOLO) manually.
|
||||
|
||||
---
|
||||
|
||||
## ⚡ 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.
|
||||
|
||||
---
|
||||
|
||||
## ✨ Features
|
||||
|
||||
### 🎯 Unified Selection Strategies
|
||||
Both Face and Object modes offer the same powerful options:
|
||||
- **Auto (Objective Diversity) [Recommended]**: Dynamically selects images until redundancy starts
|
||||
- **Standard (30 images)**: Balanced set using Smart Diversity
|
||||
- **Broad (100 images)**: Extensive set using Smart Diversity
|
||||
- **Custom Count**: You choose the limit
|
||||
|
||||
### 👤 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
|
||||
|
||||
### 📦 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.*
|
||||
|
||||
---
|
||||
|
||||
## 🚀 Installation
|
||||
|
||||
### Prerequisites
|
||||
- **Python 3.12+**
|
||||
- **[uv](https://astral.sh/uv/)** (highly recommended)
|
||||
- **Immich Server** (v1.106+)
|
||||
|
||||
### Setup
|
||||
|
||||
```bash
|
||||
git clone <repository_url>
|
||||
cd if-curator
|
||||
uv sync
|
||||
```
|
||||
|
||||
### 🏎️ GPU Support (Recommended)
|
||||
For faster embedding computation, install with GPU extras:
|
||||
|
||||
```bash
|
||||
uv sync --extra gpu
|
||||
```
|
||||
*Automatically detects CUDA (NVIDIA), ROCm (AMD), or MPS (macOS).*
|
||||
|
||||
---
|
||||
|
||||
## 💻 Usage
|
||||
|
||||
Run the command-line interface:
|
||||
|
||||
```bash
|
||||
uv run -m 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.
|
||||
|
||||
```text
|
||||
Using SigLIP (visual embeddings) for diversity analysis...
|
||||
Computing embeddings for 69 images... ████████████████ 100%
|
||||
Auto-diversity selected 38 optimally diverse images.
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 🛠️ Configuration
|
||||
|
||||
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 |
|
||||
|
||||
---
|
||||
|
||||
## 🧠 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
|
||||
|
||||
@@ -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
|
||||
@@ -0,0 +1,40 @@
|
||||
[project]
|
||||
name = "if-curator"
|
||||
version = "0.1.0"
|
||||
description = "Immich to Frigate training sets"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
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.optional-dependencies]
|
||||
gpu = [
|
||||
"onnxruntime-gpu>=1.23.2",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"ruff>=0.9.2",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 120
|
||||
target-version = "py312"
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "F", "I"]
|
||||
|
||||
Reference in New Issue
Block a user