diff --git a/compose.yml b/compose.yml index 038a2c9..a048ffb 100644 --- a/compose.yml +++ b/compose.yml @@ -18,6 +18,11 @@ services: # - SKIP_PEOPLE=Unknown # - MIN_FACE_COUNT=5 + # ── Tracker overrides (one-shot, remove after use) ── + # - DRY_RUN=true # Preview selection without downloading/uploading + # - RETRY_REJECTED=true # Re-attempt previously rejected images + # - RESET_PERSON=John # Clear uploaded+rejected history for one person + # ── Scheduling ── # Cron expression (unset = run once and exit) - CRON_SCHEDULE=0 3 * * 0 # Every Sunday at 3 AM diff --git a/if_curator/cli.py b/if_curator/cli.py index 4821ac2..022eb8f 100644 --- a/if_curator/cli.py +++ b/if_curator/cli.py @@ -18,7 +18,7 @@ from .embeddings import is_embedding_available, load_embedding_model from .image_processing import process_face_mode, process_full_mode, process_object_mode from .immich_api import fetch_all_assets, fetch_face_data, fetch_full_image, filter_recent_assets, get_people from .logging import console, setup_logging -from .upload_tracker import filter_already_uploaded, mark_uploaded +from .upload_tracker import filter_already_uploaded, get_person_summary, mark_rejected, mark_uploaded, reset_person logger = logging.getLogger(__name__) @@ -94,8 +94,9 @@ def _configure_person(person: dict, people: list[dict]) -> dict | None: rprint(f" Found [bold]{len(all_assets)}[/bold] total, [bold]{len(recent_assets)}[/bold] in range ({years} years).") # Filter out assets already uploaded to Frigate + retry_rejected = os.environ.get("RETRY_REJECTED", "false").lower() in ("true", "1", "yes") before_dedup = len(recent_assets) - new_asset_ids = set(filter_already_uploaded([a["id"] for a in recent_assets])) + new_asset_ids = set(filter_already_uploaded([a["id"] for a in recent_assets], retry_rejected=retry_rejected)) recent_assets = [a for a in recent_assets if a["id"] in new_asset_ids] skipped = before_dedup - len(recent_assets) if skipped: @@ -196,8 +197,9 @@ def auto_configure(people: list[dict]) -> list[dict]: rprint(f" {name}: {len(all_assets)} total, {len(recent_assets)} recent") # Filter out assets already uploaded to Frigate + retry_rejected = os.environ.get("RETRY_REJECTED", "false").lower() in ("true", "1", "yes") before_dedup = len(recent_assets) - new_asset_ids = set(filter_already_uploaded([a["id"] for a in recent_assets])) + new_asset_ids = set(filter_already_uploaded([a["id"] for a in recent_assets], retry_rejected=retry_rejected)) recent_assets = [a for a in recent_assets if a["id"] in new_asset_ids] skipped = before_dedup - len(recent_assets) if skipped: @@ -334,7 +336,7 @@ def upload_to_frigate(jobs: list[dict]) -> None: # Mark this asset as uploaded so it's skipped on future runs asset_id = asset_map.get(fname) if asset_id: - mark_uploaded(asset_id) + mark_uploaded(asset_id, person_name=name) break else: @@ -350,7 +352,12 @@ def upload_to_frigate(jobs: list[dict]) -> None: error_detail = resp.json().get("message", resp.text[:100]) progress.console.print(f" [dim]{error_detail}[/dim]") except Exception: - progress.console.print(f" [dim]{resp.text[:100]}[/dim]") + error_detail = resp.text[:100] + progress.console.print(f" [dim]{error_detail}[/dim]") + if resp.status_code == 400 and "face" in error_detail.lower(): + asset_id = asset_map.get(fname) + if asset_id: + mark_rejected(asset_id, person_name=name) except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as exc: if attempt < max_retries: logger.warning(f"Upload attempt {attempt}/{max_retries} for {fname}: {type(exc).__name__}, retrying...") @@ -615,6 +622,19 @@ def main() -> None: rprint(f"Server: [dim]{Config.IMMICH_URL}[/dim]") rprint(f"Output: [dim]{Config.OUTPUT_DIR}[/dim]") + # Handle RESET_PERSON before anything else + reset_person_name = os.environ.get("RESET_PERSON", "").strip() + if reset_person_name: + reset_person(reset_person_name) + rprint(f"[bold yellow]Reset tracking data for: {reset_person_name}[/bold yellow]") + + # Show per-person tracker summary if data exists + summary = get_person_summary() + if summary: + rprint("\n[dim]Tracker summary:[/dim]") + for person_name, counts in summary.items(): + rprint(f" [dim]{person_name}: {counts['uploaded']} uploaded, {counts['rejected']} rejected[/dim]") + people = get_people() if not people: rprint("[bold red]Could not fetch people from Immich. Check URL/Key.[/bold red]") @@ -622,6 +642,10 @@ def main() -> None: # Check for non-interactive mode auto_mode = os.environ.get("AUTO_MODE", "false").lower() == "true" + dry_run = os.environ.get("DRY_RUN", "false").lower() in ("true", "1", "yes") + + if dry_run: + rprint("[bold yellow]DRY RUN — no images will be downloaded or uploaded[/bold yellow]") if auto_mode: rprint("[bold cyan]Running in AUTO mode (non-interactive)[/bold cyan]") @@ -631,7 +655,9 @@ def main() -> None: if jobs: _show_preview(jobs) - if auto_mode or Confirm.ask(f"Ready to process {sum(j['limit'] for j in jobs)} images?"): + if dry_run: + rprint("\n[bold yellow]Dry run complete — skipping execute and upload.[/bold yellow]") + elif auto_mode or Confirm.ask(f"Ready to process {sum(j['limit'] for j in jobs)} images?"): execute_jobs(jobs) upload_to_frigate(jobs) rprint("\n[bold green]Done! Happy Training.[/bold green]") diff --git a/if_curator/image_processing.py b/if_curator/image_processing.py index 6c85895..4b86f1a 100644 --- a/if_curator/image_processing.py +++ b/if_curator/image_processing.py @@ -14,6 +14,12 @@ logger = logging.getLogger(__name__) _yolo_model = None +def _save_jpeg(img: Image.Image, path: str) -> None: + if img.mode != "RGB": + img = img.convert("RGB") + img.save(path, format="JPEG") + + def get_yolo_model(): """Singleton for YOLO model.""" global _yolo_model @@ -110,7 +116,7 @@ def process_face_mode( scaled_landmarks = [[lm[0] * scale_x, lm[1] * scale_y] for lm in landmarks] aligned = align_face(img, scaled_landmarks) if aligned is not None: - aligned.save(os.path.join(output_dir, f"{count}.jpg"), format="JPEG") + _save_jpeg(aligned, os.path.join(output_dir, f"{count}.jpg")) return True # Fall back to bounding box crop with configurable margin @@ -124,7 +130,7 @@ def process_face_mode( ) face_crop = img.crop(crop_box) - face_crop.save(os.path.join(output_dir, f"{count}.jpg"), format="JPEG") + _save_jpeg(face_crop, os.path.join(output_dir, f"{count}.jpg")) return True @@ -149,9 +155,9 @@ def process_object_mode( 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() - img.crop((x1, y1, x2, y2)).save( + _save_jpeg( + img.crop((x1, y1, x2, y2)), os.path.join(output_dir, f"{count}_{class_idx}.jpg"), - format="JPEG", ) class_idx += 1 found = True @@ -164,6 +170,6 @@ def process_object_mode( 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") + _save_jpeg(img, os.path.join(output_dir, f"{count}.jpg")) return True diff --git a/if_curator/upload_tracker.py b/if_curator/upload_tracker.py index 8bca54b..957ee28 100644 --- a/if_curator/upload_tracker.py +++ b/if_curator/upload_tracker.py @@ -1,10 +1,13 @@ -"""Persistent tracker for Immich asset IDs already uploaded to Frigate. +"""Persistent tracker for Immich asset IDs already uploaded/rejected by Frigate. -Prevents duplicate uploads across runs by recording each successfully -uploaded asset ID in a JSON file within the configured CACHE_DIR. +Two separate JSON files in CACHE_DIR: + frigate_uploaded_ids.json — successfully uploaded assets + frigate_rejected_ids.json — assets Frigate rejected (e.g. no face detected) -To re-train from scratch, simply delete the tracker file -(frigate_uploaded_ids.json) from your cache directory. +Both are excluded from future candidate pools. To reset: + - All: delete both files + - One person: call reset_person("Name") or set RESET_PERSON=Name + - Rejects only: delete frigate_rejected_ids.json, or set RETRY_REJECTED=true """ import json @@ -14,62 +17,117 @@ from pathlib import Path logger = logging.getLogger(__name__) UPLOAD_TRACKER_FILE = "frigate_uploaded_ids.json" +REJECT_TRACKER_FILE = "frigate_rejected_ids.json" -def _tracker_path() -> Path: - """Return path to the tracker file, using Config.CACHE_DIR if available.""" +def _tracker_path(filename: str) -> Path: try: from .config import Config - - return Path(Config.CACHE_DIR) / UPLOAD_TRACKER_FILE + return Path(Config.CACHE_DIR) / filename except (ImportError, AttributeError): - return Path(UPLOAD_TRACKER_FILE) + return Path(filename) -def load_uploaded_ids() -> set[str]: - """Load the set of Immich asset IDs already uploaded to Frigate.""" - path = _tracker_path() +def _load(filename: str) -> dict: + path = _tracker_path(filename) if not path.exists(): - return set() + return {} try: with open(path) as f: - data = json.load(f) - return set(data.get("uploaded_asset_ids", [])) + return json.load(f) except (json.JSONDecodeError, OSError) as e: - logger.warning(f"Could not load upload tracker: {e}") - return set() + logger.warning(f"Could not load tracker {filename}: {e}") + return {} -def save_uploaded_ids(uploaded_ids: set[str]) -> None: - """Persist the set of uploaded Immich asset IDs to disk.""" - path = _tracker_path() +def _save(filename: str, data: dict) -> None: + path = _tracker_path(filename) path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w") as f: - json.dump({"uploaded_asset_ids": sorted(uploaded_ids)}, f, indent=2) + json.dump(data, f, indent=2) -def mark_uploaded(asset_id: str) -> None: - """Mark a single Immich asset ID as uploaded to Frigate.""" - ids = load_uploaded_ids() - ids.add(asset_id) - save_uploaded_ids(ids) - logger.debug(f"Marked asset {asset_id} as uploaded to Frigate") +def _flat_key(filename: str) -> str: + return "uploaded_asset_ids" if "uploaded" in filename else "rejected_asset_ids" -def is_uploaded(asset_id: str) -> bool: - """Check if an Immich asset ID has already been uploaded to Frigate.""" - return asset_id in load_uploaded_ids() +def _load_flat(filename: str) -> set[str]: + return set(_load(filename).get(_flat_key(filename), [])) -def filter_already_uploaded(asset_ids: list[str]) -> list[str]: - """Return only asset IDs that have NOT yet been uploaded to Frigate. +def _mark(filename: str, asset_id: str, person_name: str | None) -> None: + data = _load(filename) + flat_key = _flat_key(filename) + flat = set(data.get(flat_key, [])) + flat.add(asset_id) + data[flat_key] = sorted(flat) + if person_name: + by_person = data.setdefault("by_person", {}) + person_ids = set(by_person.get(person_name, [])) + person_ids.add(asset_id) + by_person[person_name] = sorted(person_ids) + _save(filename, data) - Logs how many were skipped so the user knows dedup is working. - """ - uploaded = load_uploaded_ids() - new_ids = [aid for aid in asset_ids if aid not in uploaded] + +# ── Public API ──────────────────────────────────────────────────────────────── + +def load_uploaded_ids() -> set[str]: + return _load_flat(UPLOAD_TRACKER_FILE) + + +def load_rejected_ids() -> set[str]: + return _load_flat(REJECT_TRACKER_FILE) + + +def mark_uploaded(asset_id: str, person_name: str | None = None) -> None: + _mark(UPLOAD_TRACKER_FILE, asset_id, person_name) + logger.debug(f"Marked {asset_id} as uploaded ({person_name})") + + +def mark_rejected(asset_id: str, person_name: str | None = None) -> None: + _mark(REJECT_TRACKER_FILE, asset_id, person_name) + logger.debug(f"Marked {asset_id} as rejected ({person_name})") + + +def reset_person(person_name: str) -> None: + """Remove all uploaded and rejected records for a given person.""" + for filename in (UPLOAD_TRACKER_FILE, REJECT_TRACKER_FILE): + data = _load(filename) + flat_key = _flat_key(filename) + by_person = data.get("by_person", {}) + person_ids = set(by_person.pop(person_name, [])) + if person_ids: + flat = set(data.get(flat_key, [])) - person_ids + data[flat_key] = sorted(flat) + data["by_person"] = by_person + _save(filename, data) + logger.info(f"Reset tracking data for {person_name}") + + +def get_person_summary() -> dict[str, dict[str, int]]: + """Return {person_name: {uploaded: N, rejected: N}} for display.""" + uploaded_by = _load(UPLOAD_TRACKER_FILE).get("by_person", {}) + rejected_by = _load(REJECT_TRACKER_FILE).get("by_person", {}) + names = set(uploaded_by) | set(rejected_by) + return { + name: { + "uploaded": len(uploaded_by.get(name, [])), + "rejected": len(rejected_by.get(name, [])), + } + for name in sorted(names) + } + + +def filter_already_uploaded( + asset_ids: list[str], + retry_rejected: bool = False, +) -> list[str]: + """Return asset IDs not yet uploaded (and not rejected, unless retry_rejected).""" + exclude = load_uploaded_ids() + if not retry_rejected: + exclude |= load_rejected_ids() + new_ids = [aid for aid in asset_ids if aid not in exclude] skipped = len(asset_ids) - len(new_ids) if skipped: - logger.info(f"Skipping {skipped} assets already uploaded to Frigate") + logger.info(f"Skipping {skipped} assets already uploaded or rejected by Frigate") return new_ids -