2455 lines
105 KiB
Python
2455 lines
105 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
photo_analyzer.py — AI photo analysis pipeline for Immich
|
||
|
||
Crawls a photo library, analyzes each image once via an OpenAI-compatible
|
||
vision model (any OpenAI-compatible provider — Google Gemini, OpenAI, …), stores results in
|
||
SQLite, and writes captions back into EXIF so Immich can index them.
|
||
Safe to interrupt and resume at any point.
|
||
|
||
Requirements:
|
||
pip install openai pillow rich
|
||
|
||
External tools:
|
||
exiftool (https://exiftool.org — must be on PATH)
|
||
|
||
Usage:
|
||
export LLM_API_KEY=... # or set it in photo_analyzer.env
|
||
python photo_analyzer.py --library /path/to/photos
|
||
|
||
# Dry run (no API calls, no EXIF writes):
|
||
python photo_analyzer.py --library /path/to/photos --dry-run
|
||
|
||
# Re-analyze already-processed files (e.g. after prompt change):
|
||
python photo_analyzer.py --library /path/to/photos --reanalyze
|
||
|
||
# Skip EXIF writing (index only):
|
||
python photo_analyzer.py --library /path/to/photos --no-exif
|
||
"""
|
||
|
||
import argparse
|
||
import base64
|
||
import hashlib
|
||
import json
|
||
import logging
|
||
import os
|
||
import re
|
||
import select
|
||
import shutil
|
||
import signal
|
||
import sqlite3
|
||
import subprocess
|
||
import sys
|
||
import termios
|
||
import threading
|
||
import time
|
||
import tty
|
||
from collections import deque
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from datetime import datetime
|
||
from itertools import groupby
|
||
from zoneinfo import ZoneInfo
|
||
from pathlib import Path
|
||
|
||
from openai import OpenAI, RateLimitError, APIError
|
||
from PIL import Image, ImageFile
|
||
|
||
# Decode truncated/partial JPEGs instead of erroring ("image file is truncated",
|
||
# "broken data stream") — renders the available portion, fine for captioning a
|
||
# damaged file. Without this, corrupt source images fail before the API call.
|
||
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
||
from rich.console import Console, Group
|
||
from rich.layout import Layout
|
||
from rich.live import Live
|
||
from rich.logging import RichHandler
|
||
from rich.panel import Panel
|
||
from rich.progress import (
|
||
BarColumn, MofNCompleteColumn, Progress, SpinnerColumn,
|
||
TaskProgressColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn,
|
||
)
|
||
from rich.table import Table
|
||
from rich.text import Text
|
||
|
||
# ──────────────────────────────────────────────
|
||
# Configuration
|
||
# ──────────────────────────────────────────────
|
||
|
||
SUPPORTED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".heic", ".heif", ".tiff", ".tif"}
|
||
|
||
# Config file for secrets (chiefly the API key). A real shell env var always
|
||
# wins; this is just a convenience so the key survives across sessions without
|
||
# re-exporting it. Searched in the current dir first, then next to this script.
|
||
# NB: deliberately NOT ".env" — that name is taken by Immich's docker-compose.
|
||
ENV_FILE = "photo_analyzer.env"
|
||
API_KEY_PLACEHOLDER = "sk-REPLACE_WITH_YOUR_KEY" # ignored until replaced
|
||
|
||
# OpenAI-compatible LLM endpoint. Defaults to Google Gemini Flash via the
|
||
# Generative Language OpenAI-compatibility endpoint, but ANY OpenAI-compatible
|
||
# provider works — override in photo_analyzer.env (or the shell) with
|
||
# LLM_BASE_URL / LLM_MODEL / LLM_API_KEY. e.g. OpenAI:
|
||
# LLM_BASE_URL=https://api.openai.com/v1
|
||
# LLM_MODEL=gpt-4o-mini
|
||
# LLM_API_KEY=sk-...
|
||
# (GEMINI_API_KEY / GOOGLE_API_KEY are honoured as fallbacks for the key.)
|
||
LLM_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
|
||
LLM_MODEL = "gemini-2.5-flash"
|
||
|
||
# Resize before sending — most vision APIs handle ≤ 4K well; 2048px long-edge is a good
|
||
# balance between quality and token cost (~30–60% fewer tokens than full-res)
|
||
MAX_LONG_EDGE = 2048
|
||
|
||
# Duplicate detection. Each photo gets a 64-bit DCT perceptual hash (phash); two
|
||
# photos are "the same picture" when their hashes differ by ≤ PHASH_THRESHOLD
|
||
# bits (Hamming distance). 0 = pixel-identical structure; ~5 tolerates
|
||
# recompression/resize/format-change across phone backups. Raise for looser
|
||
# matching (catches more, risks grouping similar-but-distinct shots), lower for
|
||
# stricter. Tunable via PHASH_THRESHOLD in photo_analyzer.env.
|
||
PHASH_THRESHOLD = 5
|
||
|
||
# Concurrency: tune to your provider's rate limits (free tiers are often low,
|
||
# paid tiers much higher). Keep conservative to avoid 429s; the throttle warning
|
||
# box and throttle_events.jsonl tell you when you've pushed too hard.
|
||
MAX_WORKERS = 3
|
||
RETRY_ATTEMPTS = 4
|
||
RETRY_BASE_DELAY = 10 # seconds; doubles on each retry
|
||
|
||
# Your provider's requests-per-day cap, shown as a live "RPD today: N / limit"
|
||
# gauge in the dashboard (Gemini Tier 1 flash = 10000). 0 = unknown → just count.
|
||
RPD_LIMIT = 10000
|
||
|
||
# Analysis prompt — structured JSON output for reliable parsing
|
||
ANALYSIS_PROMPT = """Analyze this photo and respond ONLY with a valid JSON object.
|
||
No markdown, no explanation, just the JSON.
|
||
|
||
{
|
||
"description": "One clear sentence describing what is happening in this photo.",
|
||
"tags": ["list", "of", "8-12", "descriptive", "keywords"],
|
||
"people_count": 0,
|
||
"setting": "indoor or outdoor",
|
||
"time_of_day": "morning | afternoon | evening | night | unknown",
|
||
"season": "spring | summer | autumn | winter | unknown",
|
||
"mood": "the specific emotional tone of THIS scene, 1-2 words",
|
||
"landmarks": ["recognizable landmarks, monuments, named buildings or place signs visible — [] if none"],
|
||
"location_hint": "best guess of WHERE this was taken, e.g. 'Cairo, Egypt' / 'Sardinia, Italy' / 'the Alps' — null if no basis",
|
||
"approx_year": null
|
||
}
|
||
|
||
Rules:
|
||
- people_count: count only clearly visible people (partial counts as 1)
|
||
- approx_year: integer only if strongly inferable from clothing/tech/decor or the album hint, otherwise null
|
||
- tags: specific and useful (prefer 'golden retriever' over 'dog', 'birthday cake' over 'food');
|
||
include any identified landmark or place name as a tag too
|
||
- mood: name the specific feeling of this particular scene — e.g. professional, theatrical,
|
||
serene, tense, celebratory, nostalgic, melancholic, candid, attentive, romantic. Do NOT
|
||
default to 'happy'/'joyful' just because people are smiling: a formal headshot is
|
||
'professional', a quiet landscape 'serene', a certificate 'accomplished'.
|
||
- landmarks: identify famous or named places from architecture, monuments, signage or
|
||
distinctive landscape (e.g. 'Eiffel Tower', 'Cologne Cathedral', 'Pyramids of Giza'). [] if none.
|
||
- location_hint: infer the place from landmarks, readable signs, architecture, vegetation/
|
||
landscape AND the album hint. Be as specific as the evidence allows (city > region > country).
|
||
Do NOT invent a precise place from weak evidence — when unsure, give the broader region or null.
|
||
- All values must be valid JSON types (string, integer, null, array)
|
||
"""
|
||
|
||
# ──────────────────────────────────────────────
|
||
# Console + Logging
|
||
# ──────────────────────────────────────────────
|
||
|
||
console = Console(highlight=False)
|
||
|
||
# ── Handlers ────────────────────────────────────────────────────────────────
|
||
|
||
_rich_handler = RichHandler(console=console, rich_tracebacks=True, show_path=False, markup=True)
|
||
_rich_handler.setLevel(logging.INFO) # bumped to DEBUG by --debug flag at runtime
|
||
|
||
_info_file_handler = logging.FileHandler("photo_analyzer.log", encoding="utf-8")
|
||
_info_file_handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)-7s %(message)s", datefmt="%H:%M:%S"))
|
||
_info_file_handler.setLevel(logging.INFO)
|
||
|
||
_debug_file_handler = logging.FileHandler("photo_analyzer_debug.log", encoding="utf-8")
|
||
_debug_file_handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)-7s %(name)s %(message)s", datefmt="%H:%M:%S"))
|
||
_debug_file_handler.setLevel(logging.DEBUG)
|
||
|
||
logging.basicConfig(
|
||
level=logging.DEBUG, # root at DEBUG so debug handler receives everything
|
||
format="%(message)s",
|
||
datefmt="[%H:%M:%S]",
|
||
handlers=[_rich_handler, _info_file_handler, _debug_file_handler],
|
||
)
|
||
log = logging.getLogger(__name__)
|
||
|
||
# Silence the per-request HTTP chatter from the OpenAI SDK / httpx — one
|
||
# "HTTP Request: POST … 200 OK" line per photo floods the console. Their
|
||
# warnings and errors still come through.
|
||
for _noisy in ("httpx", "httpcore", "openai"):
|
||
logging.getLogger(_noisy).setLevel(logging.WARNING)
|
||
|
||
# ── History logger (JSONL — one line per photo) ──────────────────────────────
|
||
|
||
_history_handler = logging.FileHandler("photo_analyzer_history.jsonl", encoding="utf-8")
|
||
_history_handler.setFormatter(logging.Formatter("%(message)s"))
|
||
|
||
history_log = logging.getLogger("history")
|
||
history_log.addHandler(_history_handler)
|
||
history_log.setLevel(logging.DEBUG)
|
||
history_log.propagate = False # never send to root / console
|
||
|
||
|
||
# ── Graceful shutdown ────────────────────────────────────────────────────────
|
||
|
||
_stop = threading.Event()
|
||
|
||
# Live-mode popup menu (hotkey 'm'). _menu = open/closed; _menu_state holds the
|
||
# highlighted row and which view is showing. Both threads touch it, so guard with
|
||
# the lock. Each item is (label, action) — action handled in _key_listener.
|
||
_menu = threading.Event()
|
||
_menu_lock = threading.Lock()
|
||
_menu_state = {"sel": 0, "view": "menu"}
|
||
_MENU_ITEMS = [
|
||
("Token usage", "token"), # switch overlay to session token stats
|
||
("Stop after current", "stop"), # graceful: finish in-flight, then quit
|
||
("Force quit now", "force"), # immediate exit
|
||
("Resume", "close"), # close menu, keep analyzing
|
||
]
|
||
|
||
# A single sqlite3.Connection is shared across worker threads (check_same_thread=
|
||
# False). SQLite serialises writes internally, but Python's connection object is
|
||
# not thread-safe — concurrent commits race ("cannot commit - no transaction is
|
||
# active"). This lock serialises every write; API latency dominates, so the cost
|
||
# is negligible.
|
||
_db_lock = threading.Lock()
|
||
|
||
def _handle_sigint(sig, frame):
|
||
# Audible bell on every press — instant confirmation the signal was received,
|
||
# even in the full-screen dashboard. The visual banner appears on the next
|
||
# render tick (see _Dashboard); plain mode shows the log line below.
|
||
try:
|
||
sys.__stderr__.write("\a")
|
||
sys.__stderr__.flush()
|
||
except Exception:
|
||
pass
|
||
if _stop.is_set():
|
||
log.warning("Force quit.")
|
||
sys.exit(1)
|
||
_stop.set()
|
||
log.warning("[yellow]Stopping after current batch… press Ctrl+C again to force quit.[/]")
|
||
|
||
signal.signal(signal.SIGINT, _handle_sigint)
|
||
|
||
|
||
def log_history(path: str, status: str, result: dict = None,
|
||
tokens: dict = None, error: str = None,
|
||
copied_from: str = None) -> None:
|
||
# status is typically "analyzed" | "error" | "copied" (variant inherited
|
||
# from its primary with no API call — copied_from names that primary).
|
||
entry = {"ts": time.strftime("%Y-%m-%dT%H:%M:%S"), "path": path, "status": status}
|
||
if copied_from:
|
||
entry["copied_from"] = copied_from
|
||
if tokens:
|
||
entry["tokens_total"] = tokens.get("total", 0)
|
||
entry["tokens_prompt"] = tokens.get("prompt", 0)
|
||
entry["tokens_completion"] = tokens.get("completion", 0)
|
||
if result:
|
||
entry["description"] = result.get("description")
|
||
entry["tags"] = result.get("tags", [])
|
||
entry["mood"] = result.get("mood")
|
||
entry["setting"] = result.get("setting")
|
||
entry["people_count"] = result.get("people_count", 0)
|
||
entry["location_hint"]= result.get("location_hint")
|
||
entry["approx_year"] = result.get("approx_year")
|
||
if error:
|
||
entry["error"] = error
|
||
history_log.info(json.dumps(entry, ensure_ascii=False))
|
||
|
||
|
||
# ──────────────────────────────────────────────
|
||
# Database
|
||
# ──────────────────────────────────────────────
|
||
|
||
DB_FILE = "photo_analysis.db"
|
||
|
||
SCHEMA = """
|
||
CREATE TABLE IF NOT EXISTS photos (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
path TEXT UNIQUE NOT NULL,
|
||
status TEXT NOT NULL DEFAULT 'pending',
|
||
-- pending | analyzed | exif_written | error | duplicate
|
||
phash TEXT, -- 64-bit perceptual hash (16 hex chars); content-based, survives resize/recompress/EXIF
|
||
file_sha1 TEXT, -- SHA-1 of the file bytes; exact identity, for rename/move reconciliation
|
||
dup_of TEXT, -- if status='duplicate', the canonical photo's path this duplicates
|
||
description TEXT,
|
||
tags TEXT, -- JSON array stored as string
|
||
people_count INTEGER,
|
||
setting TEXT,
|
||
time_of_day TEXT,
|
||
season TEXT,
|
||
mood TEXT,
|
||
location_hint TEXT,
|
||
approx_year INTEGER,
|
||
raw_response TEXT, -- full JSON from the model, for debugging
|
||
error_message TEXT,
|
||
analyzed_at TEXT,
|
||
exif_written_at TEXT
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_status ON photos(status);
|
||
CREATE INDEX IF NOT EXISTS idx_path ON photos(path);
|
||
-- idx_sha1 / idx_phash are created in _migrate_schema(), AFTER the columns are
|
||
-- guaranteed to exist (a pre-existing table won't have them until the ALTER).
|
||
|
||
CREATE VIRTUAL TABLE IF NOT EXISTS photos_fts USING fts5(
|
||
path, description, tags, mood, location_hint,
|
||
content=photos, content_rowid=id
|
||
);
|
||
|
||
CREATE TRIGGER IF NOT EXISTS photos_ai AFTER INSERT ON photos BEGIN
|
||
INSERT INTO photos_fts(rowid, path, description, tags, mood, location_hint)
|
||
VALUES (new.id, new.path, new.description, new.tags, new.mood, new.location_hint);
|
||
END;
|
||
|
||
CREATE TRIGGER IF NOT EXISTS photos_au AFTER UPDATE ON photos BEGIN
|
||
INSERT INTO photos_fts(photos_fts, rowid, path, description, tags, mood, location_hint)
|
||
VALUES ('delete', old.id, old.path, old.description, old.tags, old.mood, old.location_hint);
|
||
INSERT INTO photos_fts(rowid, path, description, tags, mood, location_hint)
|
||
VALUES (new.id, new.path, new.description, new.tags, new.mood, new.location_hint);
|
||
END;
|
||
|
||
CREATE TRIGGER IF NOT EXISTS photos_ad AFTER DELETE ON photos BEGIN
|
||
INSERT INTO photos_fts(photos_fts, rowid, path, description, tags, mood, location_hint)
|
||
VALUES ('delete', old.id, old.path, old.description, old.tags, old.mood, old.location_hint);
|
||
END;
|
||
"""
|
||
|
||
|
||
def get_db(db_path: str) -> sqlite3.Connection:
|
||
conn = sqlite3.connect(db_path, check_same_thread=False)
|
||
conn.row_factory = sqlite3.Row
|
||
conn.executescript(SCHEMA)
|
||
_migrate_schema(conn)
|
||
conn.commit()
|
||
return conn
|
||
|
||
|
||
def _migrate_schema(conn: sqlite3.Connection) -> None:
|
||
"""Add columns introduced after a DB was first created. CREATE TABLE IF NOT
|
||
EXISTS never alters an existing table, so pre-existing databases (e.g. the
|
||
original 25k-row one) need these added explicitly. Idempotent."""
|
||
have = {r["name"] for r in conn.execute("PRAGMA table_info(photos)")}
|
||
for col in ("phash", "file_sha1", "dup_of"):
|
||
if col not in have:
|
||
conn.execute(f"ALTER TABLE photos ADD COLUMN {col} TEXT")
|
||
conn.executescript(
|
||
"CREATE INDEX IF NOT EXISTS idx_sha1 ON photos(file_sha1);"
|
||
"CREATE INDEX IF NOT EXISTS idx_phash ON photos(phash);"
|
||
)
|
||
|
||
|
||
def purge_excluded(conn: sqlite3.Connection) -> int:
|
||
"""Remove DB entries for paths that should be excluded (thumbs, _IGNORE)."""
|
||
cur = conn.execute(
|
||
"DELETE FROM photos WHERE path LIKE '%/.@__thumb/%' OR path LIKE '%/_IGNORE/%'"
|
||
)
|
||
conn.commit()
|
||
return cur.rowcount
|
||
|
||
|
||
def prune_missing(conn: sqlite3.Connection, library: Path) -> int:
|
||
"""
|
||
Delete DB rows whose image file no longer exists on disk (moved or deleted),
|
||
so the stats/album breakdown don't keep listing folders that are gone.
|
||
|
||
Guarded by the library root existing: if the whole library is unreachable
|
||
(unmounted drive, wrong path) we skip pruning rather than wipe the database.
|
||
"""
|
||
if not library.exists():
|
||
log.warning(f"Library root not found ({library}); skipping stale-entry "
|
||
"prune so an unavailable filesystem can't wipe the DB.")
|
||
return 0
|
||
rows = conn.execute("SELECT path FROM photos").fetchall()
|
||
missing = [r["path"] for r in rows if not os.path.exists(r["path"])]
|
||
if missing:
|
||
with _db_lock:
|
||
conn.executemany("DELETE FROM photos WHERE path = ?",
|
||
[(p,) for p in missing])
|
||
conn.commit()
|
||
return len(missing)
|
||
|
||
|
||
def upsert_pending(conn: sqlite3.Connection, path: str):
|
||
with _db_lock:
|
||
conn.execute(
|
||
"INSERT OR IGNORE INTO photos (path, status) VALUES (?, 'pending')",
|
||
(path,)
|
||
)
|
||
conn.commit()
|
||
|
||
|
||
def mark_analyzed(conn: sqlite3.Connection, path: str, result: dict, raw: str):
|
||
with _db_lock:
|
||
conn.execute(
|
||
"""UPDATE photos SET
|
||
status = 'analyzed',
|
||
description = ?,
|
||
tags = ?,
|
||
people_count = ?,
|
||
setting = ?,
|
||
time_of_day = ?,
|
||
season = ?,
|
||
mood = ?,
|
||
location_hint = ?,
|
||
approx_year = ?,
|
||
raw_response = ?,
|
||
error_message = NULL,
|
||
analyzed_at = datetime('now')
|
||
WHERE path = ?""",
|
||
(
|
||
result.get("description"),
|
||
json.dumps(result.get("tags", []), ensure_ascii=False),
|
||
result.get("people_count"),
|
||
result.get("setting"),
|
||
result.get("time_of_day"),
|
||
result.get("season"),
|
||
result.get("mood"),
|
||
result.get("location_hint"),
|
||
result.get("approx_year"),
|
||
raw,
|
||
path,
|
||
),
|
||
)
|
||
conn.commit()
|
||
|
||
|
||
def mark_error(conn: sqlite3.Connection, path: str, message: str):
|
||
with _db_lock:
|
||
conn.execute(
|
||
"""UPDATE photos SET
|
||
status = 'error',
|
||
error_message = ?,
|
||
analyzed_at = datetime('now')
|
||
WHERE path = ?""",
|
||
(message, path),
|
||
)
|
||
conn.commit()
|
||
|
||
|
||
def mark_exif_written(conn: sqlite3.Connection, path: str):
|
||
with _db_lock:
|
||
conn.execute(
|
||
"""UPDATE photos SET
|
||
status = 'exif_written',
|
||
exif_written_at = datetime('now')
|
||
WHERE path = ?""",
|
||
(path,),
|
||
)
|
||
conn.commit()
|
||
|
||
|
||
def get_pending(conn: sqlite3.Connection, reanalyze: bool) -> list[str]:
|
||
# 'duplicate' rows never hit the API — they inherit from their canonical twin
|
||
# (or are simply skipped) so we don't spend a call, or an upload, on a copy.
|
||
if reanalyze:
|
||
rows = conn.execute(
|
||
"SELECT path FROM photos WHERE status != 'duplicate'"
|
||
).fetchall()
|
||
else:
|
||
rows = conn.execute(
|
||
"SELECT path FROM photos WHERE status IN ('pending', 'error')"
|
||
).fetchall()
|
||
return [r["path"] for r in rows]
|
||
|
||
|
||
def get_analyzed_no_exif(conn: sqlite3.Connection) -> list[sqlite3.Row]:
|
||
return conn.execute(
|
||
"SELECT * FROM photos WHERE status = 'analyzed'"
|
||
).fetchall()
|
||
|
||
|
||
# ──────────────────────────────────────────────
|
||
# Image preparation
|
||
# ──────────────────────────────────────────────
|
||
|
||
def prepare_image(path: Path) -> tuple[str, str]:
|
||
"""
|
||
Resize to MAX_LONG_EDGE if needed, return (base64_string, mime_type).
|
||
Handles HEIC by converting to JPEG in memory.
|
||
"""
|
||
suffix = path.suffix.lower()
|
||
|
||
with Image.open(path) as img:
|
||
img = img.convert("RGB") # normalise; drops alpha, converts HEIC
|
||
|
||
w, h = img.size
|
||
long_edge = max(w, h)
|
||
if long_edge > MAX_LONG_EDGE:
|
||
scale = MAX_LONG_EDGE / long_edge
|
||
img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
|
||
|
||
import io
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="JPEG", quality=85)
|
||
buf.seek(0)
|
||
b64 = base64.b64encode(buf.read()).decode("utf-8")
|
||
|
||
return b64, "image/jpeg"
|
||
|
||
|
||
# ──────────────────────────────────────────────
|
||
# API call with retry
|
||
# ──────────────────────────────────────────────
|
||
|
||
# Rolling record of recent rate-limit (429) / overload (503) hits so the dashboard
|
||
# can warn when the provider starts throttling — the cue to lower MAX_WORKERS.
|
||
# Workers append timestamps; the render thread reads a count over a short window.
|
||
# Every hit is ALSO appended to throttle_events.jsonl (separate, persistent backend)
|
||
# so the full throttling history survives across runs and the live UI.
|
||
_THROTTLE_LOG = Path(__file__).with_name("throttle_events.jsonl")
|
||
_ratelimit_hits: deque = deque(maxlen=500)
|
||
_ratelimit_total = 0
|
||
_ratelimit_lock = threading.Lock()
|
||
RATELIMIT_WINDOW = 30 # seconds to look back
|
||
RATELIMIT_WARN_AT = 5 # hits within the window → show the warning box
|
||
|
||
|
||
def _record_ratelimit(status: int, path, attempt: int, detail: str = ""):
|
||
"""Record one 429/503 hit: in-memory (for the warning box) + persistent JSONL.
|
||
`detail` is the provider's error body — for Gemini it names the exact quota
|
||
(PerMinute vs PerDay, TPM vs RPM) or a billing/balance message."""
|
||
global _ratelimit_total
|
||
with _ratelimit_lock:
|
||
_ratelimit_hits.append(time.monotonic())
|
||
_ratelimit_total += 1
|
||
try:
|
||
with _THROTTLE_LOG.open("a", encoding="utf-8") as fh:
|
||
fh.write(json.dumps({
|
||
"time": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||
"status": status, "model": LLM_MODEL,
|
||
"file": Path(path).name, "attempt": attempt,
|
||
"detail": (detail or "")[:500],
|
||
}, ensure_ascii=False) + "\n")
|
||
except Exception as e: # logging must never break the run
|
||
log.debug(f"throttle log write failed: {e}")
|
||
|
||
|
||
def _ratelimit_recent() -> int:
|
||
"""How many 429/503 hits in the last RATELIMIT_WINDOW seconds."""
|
||
cutoff = time.monotonic() - RATELIMIT_WINDOW
|
||
with _ratelimit_lock:
|
||
return sum(1 for t in _ratelimit_hits if t >= cutoff)
|
||
|
||
|
||
# Requests-per-day counter. Providers don't return remaining RPD, so count locally:
|
||
# every request fired = 1 RPD unit. Persisted per Pacific calendar day (RPD resets
|
||
# midnight PT) so the count is cumulative across multiple runs in the same day.
|
||
# ponytail: single-process counter — two concurrent runs would each undercount.
|
||
_RPD_FILE = Path(__file__).with_name("rpd_count.json")
|
||
_rpd_lock = threading.Lock()
|
||
_rpd = {"date": None, "count": 0}
|
||
|
||
|
||
def _pt_today() -> str:
|
||
return datetime.now(ZoneInfo("America/Los_Angeles")).strftime("%Y-%m-%d")
|
||
|
||
|
||
def _rpd_load():
|
||
"""Seed the counter from disk if it's still the same PT day, else start at 0."""
|
||
today = _pt_today()
|
||
try:
|
||
d = json.loads(_RPD_FILE.read_text(encoding="utf-8"))
|
||
if d.get("date") == today:
|
||
_rpd.update(date=today, count=int(d.get("count", 0)))
|
||
return
|
||
except Exception:
|
||
pass
|
||
_rpd.update(date=today, count=0)
|
||
|
||
|
||
def _rpd_increment():
|
||
with _rpd_lock:
|
||
today = _pt_today()
|
||
if _rpd["date"] != today: # rolled past midnight PT mid-run
|
||
_rpd["date"], _rpd["count"] = today, 0
|
||
_rpd["count"] += 1
|
||
|
||
|
||
def _rpd_persist():
|
||
with _rpd_lock:
|
||
snap = dict(_rpd)
|
||
try:
|
||
_RPD_FILE.write_text(json.dumps(snap), encoding="utf-8")
|
||
except Exception as e:
|
||
log.debug(f"rpd persist failed: {e}")
|
||
|
||
|
||
def _rpd_count() -> int:
|
||
with _rpd_lock:
|
||
return _rpd["count"]
|
||
|
||
|
||
def write_throttle_summary():
|
||
"""Append a run-summary line to the throttle log after the run completes."""
|
||
if _ratelimit_total == 0:
|
||
return
|
||
try:
|
||
with _THROTTLE_LOG.open("a", encoding="utf-8") as fh:
|
||
fh.write(json.dumps({
|
||
"time": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||
"type": "run_summary", "model": LLM_MODEL,
|
||
"workers": MAX_WORKERS, "total_throttle_hits": _ratelimit_total,
|
||
}, ensure_ascii=False) + "\n")
|
||
except Exception as e:
|
||
log.debug(f"throttle summary write failed: {e}")
|
||
log.warning(f"⚠ {_ratelimit_total} rate-limit/overload hits this run "
|
||
f"(logged to {_THROTTLE_LOG.name}). Consider lowering MAX_WORKERS.")
|
||
|
||
|
||
def analyze_image(client: OpenAI, path: Path, dry_run: bool) -> tuple[dict, str]:
|
||
"""
|
||
Call the configured OpenAI-compatible vision model (LLM_MODEL). Returns
|
||
(parsed_result, raw_json_string, tokens). Raises on unrecoverable error.
|
||
"""
|
||
if dry_run:
|
||
fake = {
|
||
"description": f"[DRY RUN] Photo at {path.name}",
|
||
"tags": ["dry-run", "test"],
|
||
"people_count": 0,
|
||
"setting": "unknown",
|
||
"time_of_day": "unknown",
|
||
"season": "unknown",
|
||
"mood": "neutral",
|
||
"location_hint": None,
|
||
"approx_year": None,
|
||
}
|
||
return fake, json.dumps(fake), {"prompt": 0, "completion": 0, "total": 0}
|
||
|
||
b64, mime = prepare_image(path)
|
||
image_url = f"data:{mime};base64,{b64}"
|
||
log.debug(f"Prepared image {path.name} as {mime} ({len(b64)} base64 chars)")
|
||
|
||
# Feed the album/folder name as context — folder names like "Ägypten 2010"
|
||
# carry place + year the model can't see from pixels alone. The image still
|
||
# wins on conflict; non-place folders (selfies, fun_pics) are ignored.
|
||
album = path.parent.name
|
||
prompt_text = ANALYSIS_PROMPT + (
|
||
f'\n\nAlbum hint: this photo is filed in a folder named "{album}". '
|
||
f'Folder names often contain the place and/or year — use it to inform '
|
||
f'"location_hint" and "approx_year", but trust the image if they conflict. '
|
||
f'Ignore the hint if it is clearly not a place or date (e.g. "selfies", "fun_pics").'
|
||
)
|
||
|
||
for attempt in range(1, RETRY_ATTEMPTS + 1):
|
||
log.debug(f"API attempt {attempt}/{RETRY_ATTEMPTS} for {path.name}")
|
||
try:
|
||
_rpd_increment() # every request fired counts toward the daily cap
|
||
response = client.chat.completions.create(
|
||
model=LLM_MODEL,
|
||
max_tokens=4096,
|
||
messages=[
|
||
{
|
||
"role": "user",
|
||
"content": [
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {"url": image_url},
|
||
},
|
||
{
|
||
"type": "text",
|
||
"text": prompt_text,
|
||
},
|
||
],
|
||
}
|
||
],
|
||
)
|
||
choice = response.choices[0] if response.choices else None
|
||
msg = choice.message if choice else None
|
||
raw = (msg.content if msg else None) or ""
|
||
|
||
# Safety/content-filter block: the model returns no message at all.
|
||
# Retrying the same model won't help — fail fast with a clear reason
|
||
# instead of crashing on None.content. Caption these on another model.
|
||
finish = getattr(choice, "finish_reason", "") or ""
|
||
if not raw and "content_filter" in finish:
|
||
raise RuntimeError(f"blocked by {LLM_MODEL} content filter ({finish})")
|
||
|
||
raw = raw.strip()
|
||
# Strip accidental markdown fences
|
||
if raw.startswith("```"):
|
||
raw = raw.split("```")[1]
|
||
if raw.startswith("json"):
|
||
raw = raw[4:]
|
||
raw = raw.strip()
|
||
|
||
# Empty response — retryable (model occasionally returns blank)
|
||
if not raw:
|
||
delay = RETRY_BASE_DELAY * attempt
|
||
log.warning(f"Empty response from {LLM_MODEL} for {path} (attempt {attempt}/{RETRY_ATTEMPTS}). Retry in {delay}s")
|
||
time.sleep(delay)
|
||
continue
|
||
|
||
result = json.loads(raw)
|
||
usage = response.usage
|
||
tokens = {
|
||
"prompt": usage.prompt_tokens if usage else 0,
|
||
"completion": usage.completion_tokens if usage else 0,
|
||
"total": usage.total_tokens if usage else 0,
|
||
}
|
||
log.debug(f"OK {path.name} — tokens: {tokens['total']} — raw: {raw[:200]}")
|
||
return result, raw, tokens
|
||
|
||
except RateLimitError as e:
|
||
detail = str(getattr(e, "message", "") or e)
|
||
_record_ratelimit(429, path, attempt, detail)
|
||
delay = RETRY_BASE_DELAY * (2 ** (attempt - 1))
|
||
log.warning(f"429 rate limited on {path.name}. Waiting {delay}s (attempt {attempt}/{RETRY_ATTEMPTS}) — {detail[:160]}")
|
||
time.sleep(delay)
|
||
|
||
except json.JSONDecodeError as e:
|
||
delay = RETRY_BASE_DELAY * attempt
|
||
log.warning(f"Bad JSON from {LLM_MODEL} for {path}: {e}. Raw: {raw[:300]!r} — retry in {delay}s (attempt {attempt}/{RETRY_ATTEMPTS})")
|
||
time.sleep(delay)
|
||
|
||
except APIError as e:
|
||
# 503 overload / 429 / any 5xx is a throttle — record + log EVERY one,
|
||
# even on the final attempt (so none go unlogged before the raise).
|
||
status = getattr(e, "status_code", None)
|
||
is_throttle = status in (429, 503) or (status or 0) >= 500
|
||
if is_throttle:
|
||
detail = str(getattr(e, "message", "") or e)
|
||
_record_ratelimit(status, path, attempt, detail)
|
||
log.warning(f"{status} throttle on {path.name} (attempt {attempt}/{RETRY_ATTEMPTS}) — {detail[:160]}")
|
||
if attempt < RETRY_ATTEMPTS:
|
||
delay = RETRY_BASE_DELAY * attempt
|
||
if not is_throttle:
|
||
log.warning(f"API error on {path.name}: {e}. Retry in {delay}s")
|
||
time.sleep(delay)
|
||
else:
|
||
raise
|
||
|
||
raise RuntimeError(f"Exhausted {RETRY_ATTEMPTS} retries for {path.name}")
|
||
|
||
|
||
# ──────────────────────────────────────────────
|
||
# EXIF writing via exiftool
|
||
# ──────────────────────────────────────────────
|
||
|
||
def build_exif_caption(row: sqlite3.Row) -> str:
|
||
"""Compose the string that will be written to EXIF ImageDescription."""
|
||
parts = []
|
||
if row["description"]:
|
||
parts.append(row["description"])
|
||
if row["tags"]:
|
||
try:
|
||
tags = json.loads(row["tags"])
|
||
parts.append("Tags: " + ", ".join(tags))
|
||
except Exception:
|
||
pass
|
||
if row["mood"]:
|
||
parts.append(f"Mood: {row['mood']}")
|
||
if row["location_hint"]:
|
||
parts.append(f"Location: {row['location_hint']}")
|
||
if row["approx_year"]:
|
||
parts.append(f"~{row['approx_year']}")
|
||
return " | ".join(parts)
|
||
|
||
|
||
def build_exif_caption_from_result(result: dict) -> str:
|
||
"""Same as build_exif_caption but from a raw API result dict."""
|
||
parts = []
|
||
if result.get("description"):
|
||
parts.append(result["description"])
|
||
tags = result.get("tags", [])
|
||
if tags:
|
||
parts.append("Tags: " + ", ".join(tags))
|
||
if result.get("mood"):
|
||
parts.append(f"Mood: {result['mood']}")
|
||
if result.get("location_hint"):
|
||
parts.append(f"Location: {result['location_hint']}")
|
||
if result.get("approx_year"):
|
||
parts.append(f"~{result['approx_year']}")
|
||
return " | ".join(parts)
|
||
|
||
|
||
def read_existing_exif(path: str) -> dict:
|
||
"""Read current ImageDescription, XPComment, Subject, Keywords from a file."""
|
||
try:
|
||
r = subprocess.run(
|
||
["exiftool", "-json", "-ImageDescription", "-XPComment", "-Subject", "-Keywords", path],
|
||
capture_output=True, text=True, timeout=30,
|
||
)
|
||
if r.returncode == 0 and r.stdout.strip():
|
||
return json.loads(r.stdout)[0]
|
||
except Exception:
|
||
pass
|
||
return {}
|
||
|
||
|
||
def _rec_has_nsfw(rec: dict) -> bool:
|
||
"""True if an exiftool JSON record's Keywords/Subject carries the 'nsfw' tag
|
||
written by the nsfwtag tool. Handles both list- and scalar-valued fields."""
|
||
vals = []
|
||
for field in ("Keywords", "Subject"):
|
||
v = rec.get(field)
|
||
if isinstance(v, list):
|
||
vals += v
|
||
elif v is not None:
|
||
vals.append(v)
|
||
return any(str(x).strip().lower() == "nsfw" for x in vals)
|
||
|
||
|
||
def filter_nsfw_tagged(paths: list[str]) -> tuple[list[str], list[str]]:
|
||
"""Split paths into (to_analyze, skipped), skipping any file the nsfwtag tool
|
||
marked with an 'nsfw' EXIF keyword. One batched exiftool call reads Keywords +
|
||
Subject for the whole list. Fails open (skips nothing) if exiftool can't run, so
|
||
a metadata hiccup never aborts an analysis run."""
|
||
if not paths:
|
||
return paths, []
|
||
skipped = set()
|
||
try:
|
||
r = subprocess.run(
|
||
["exiftool", "-j", "-Keywords", "-Subject", "-@", "-"],
|
||
input="\n".join(paths), capture_output=True, text=True,
|
||
timeout=max(120, len(paths) // 20), # metadata-only read is fast
|
||
)
|
||
if r.returncode == 0 and r.stdout.strip():
|
||
for rec in json.loads(r.stdout):
|
||
if _rec_has_nsfw(rec):
|
||
skipped.add(os.path.normpath(rec.get("SourceFile", "")))
|
||
except Exception as e:
|
||
log.warning(f"nsfw tag check failed ({e}); analyzing all files")
|
||
return paths, []
|
||
kept = [p for p in paths if os.path.normpath(p) not in skipped]
|
||
return kept, [p for p in paths if os.path.normpath(p) in skipped]
|
||
|
||
|
||
_REPAIR_LOG = Path(__file__).with_name("repairs.jsonl")
|
||
|
||
|
||
def _log_repair(record: dict) -> None:
|
||
"""Append one repair attempt to repairs.jsonl for later analysis."""
|
||
record["time"] = time.strftime("%Y-%m-%d %H:%M:%S")
|
||
try:
|
||
with _REPAIR_LOG.open("a", encoding="utf-8") as fh:
|
||
fh.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||
except Exception as e: # logging must never break the run
|
||
log.debug(f"could not write repair log: {e}")
|
||
|
||
|
||
# exiftool errors that a clean re-encode fixes: truncated/no-EOI JPEGs, corrupt
|
||
# EXIF directories (Bad format … IFD0), and extension/content mismatches (a JPEG
|
||
# named .png). Re-encoding rebuilds a valid file with sane metadata structure.
|
||
_REPAIRABLE_ERRORS = (
|
||
"EOI marker not found", # truncated JPEG
|
||
"Bad format", # corrupt IFD0 / EXIF directory
|
||
"looks more like", # wrong extension (e.g. "Not a valid PNG (looks more like a JPEG)")
|
||
"Bad IFD",
|
||
)
|
||
|
||
# Map file extension → the sips output format, so a re-encode produces a file that
|
||
# actually matches its extension (the fix for a JPEG mistakenly named .png).
|
||
_SIPS_FORMAT = {".png": "png", ".jpg": "jpeg", ".jpeg": "jpeg",
|
||
".tif": "tiff", ".tiff": "tiff", ".heic": "jpeg", ".heif": "jpeg"}
|
||
|
||
|
||
def repair_image(path: str) -> bool:
|
||
"""
|
||
Re-encode a structurally broken image into a clean one so exiftool can write
|
||
to it. Handles truncated JPEGs (no EOI), corrupt EXIF directories, and
|
||
extension/content mismatches. Backs up the original as <path>.orig and
|
||
re-encodes via sips to the format matching the file's extension.
|
||
Logs the outcome to repairs.jsonl. Returns True on success.
|
||
"""
|
||
p = Path(path)
|
||
fmt = _SIPS_FORMAT.get(p.suffix.lower(), "jpeg")
|
||
backup = p.with_suffix(p.suffix + ".orig")
|
||
tmp = p.with_suffix(p.suffix + ".repair_tmp")
|
||
rec = {"path": path, "action": "reencode", "format": fmt, "backup": str(backup)}
|
||
try:
|
||
size_before = p.stat().st_size
|
||
if not backup.exists():
|
||
shutil.copy2(p, backup)
|
||
r = subprocess.run(
|
||
["sips", "-s", "format", fmt, "-s", "formatOptions", "best",
|
||
str(backup if backup.exists() else p), "--out", str(tmp)],
|
||
capture_output=True, text=True, timeout=120,
|
||
)
|
||
if r.returncode != 0 or not tmp.exists():
|
||
rec.update(status="failed", error=(r.stderr or "sips produced no output").strip())
|
||
_log_repair(rec)
|
||
log.error(f"repair failed (sips) for {path}: {rec['error']}")
|
||
tmp.unlink(missing_ok=True)
|
||
return False
|
||
os.replace(tmp, p) # same-dir, atomic; avoids cross-volume mv issues
|
||
rec.update(status="ok", size_before=size_before, size_after=p.stat().st_size)
|
||
_log_repair(rec)
|
||
log.warning(f"repaired image (re-encoded to {fmt}, backup at {backup.name}): {path}")
|
||
return True
|
||
except Exception as e:
|
||
rec.update(status="failed", error=str(e))
|
||
_log_repair(rec)
|
||
log.error(f"repair failed for {path}: {e}")
|
||
tmp.unlink(missing_ok=True)
|
||
return False
|
||
|
||
|
||
def write_exif(path: str, caption: str, keywords: list[str]) -> bool:
|
||
"""
|
||
Add AI caption and keywords to EXIF without overwriting existing values.
|
||
- ImageDescription / XPComment: appended with ' | AI: ...' if already populated
|
||
- Subject / Keywords: merged with existing, deduplicated
|
||
exiftool must be on PATH.
|
||
"""
|
||
existing = read_existing_exif(path)
|
||
|
||
# ── Text fields: append if already set, write if empty ──────────────────
|
||
def merge_text(field: str) -> str:
|
||
cur = (existing.get(field) or "").strip()
|
||
if cur:
|
||
if caption in cur: # idempotency: don't append twice
|
||
return cur
|
||
return f"{caption} | {cur}"
|
||
return caption
|
||
|
||
final_desc = merge_text("ImageDescription")
|
||
final_comment = merge_text("XPComment")
|
||
|
||
# ── Keyword fields: merge + deduplicate ──────────────────────────────────
|
||
def to_list(val) -> list[str]:
|
||
if not val:
|
||
return []
|
||
return [val] if isinstance(val, str) else list(val)
|
||
|
||
existing_kws = to_list(existing.get("Keywords")) + to_list(existing.get("Subject"))
|
||
merged_kws = list(dict.fromkeys(existing_kws + keywords)) # preserves order, deduplicates
|
||
|
||
# ── Build exiftool command ───────────────────────────────────────────────
|
||
cmd = [
|
||
"exiftool",
|
||
"-m", # ignore minor errors/warnings (e.g. bad MakerNotes offsets) so writes still apply
|
||
"-overwrite_original",
|
||
f"-ImageDescription={final_desc}",
|
||
f"-XPComment={final_comment}",
|
||
]
|
||
for kw in merged_kws:
|
||
cmd.append(f"-Subject={kw}")
|
||
cmd.append(f"-Keywords={kw}")
|
||
|
||
cmd.append(path)
|
||
|
||
try:
|
||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
|
||
if result.returncode != 0:
|
||
err = result.stderr.strip()
|
||
# Structurally broken image exiftool can't write to — re-encode and retry once.
|
||
if any(sig in err for sig in _REPAIRABLE_ERRORS) and repair_image(path):
|
||
retry = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
|
||
if retry.returncode == 0:
|
||
return True
|
||
log.error(f"exiftool error after repair for {path}: {retry.stderr.strip()}")
|
||
return False
|
||
log.error(f"exiftool error for {path}: {err}")
|
||
return False
|
||
return True
|
||
except FileNotFoundError:
|
||
log.error("exiftool not found on PATH. Install from https://exiftool.org")
|
||
sys.exit(1)
|
||
except subprocess.TimeoutExpired:
|
||
log.error(f"exiftool timed out for {path}")
|
||
return False
|
||
|
||
|
||
# ──────────────────────────────────────────────
|
||
# Discovery
|
||
# ──────────────────────────────────────────────
|
||
|
||
def discover_photos(library: Path) -> list[Path]:
|
||
"""Recursively find all supported image files, sorted for deterministic order."""
|
||
photos = []
|
||
for ext in SUPPORTED_EXTENSIONS:
|
||
photos.extend(library.rglob(f"*{ext}"))
|
||
photos.extend(library.rglob(f"*{ext.upper()}"))
|
||
# Never process files inside _IGNORE/ folders (anywhere in the tree)
|
||
photos = [p for p in photos if "_IGNORE" not in p.parts and ".@__thumb" not in p.parts]
|
||
return sorted(set(photos))
|
||
|
||
|
||
# ──────────────────────────────────────────────
|
||
# Variant grouping
|
||
#
|
||
# Many photos exist as several near-identical files that differ only by an edit
|
||
# marker or a downscaled copy, e.g.
|
||
# _MG_1432.JPG
|
||
# _MG_1432-bearbeitet.jpg (edited)
|
||
# _MG_1432 (640x427).jpg (resized)
|
||
# _MG_1432 (640x427)-bearbeitet.jpg
|
||
# With --group-variants we analyze only ONE per group (the largest file with no
|
||
# markers) via the API, then copy that result into the siblings' DB rows and
|
||
# EXIF — no extra API calls.
|
||
# ──────────────────────────────────────────────
|
||
|
||
# " (640x427)", " (1920 x 1080)" etc. — a trailing pixel-dimensions suffix.
|
||
_PIXEL_RE = re.compile(r"\s*\(\s*\d+\s*[x×]\s*\d+\s*\)", re.IGNORECASE)
|
||
# Edit markers to strip. Extend this tuple if your library uses others.
|
||
_EDIT_RE = re.compile(r"[-_ ]?bearbeitet", re.IGNORECASE)
|
||
|
||
|
||
def _strip_variant_markers(stem: str) -> str:
|
||
"""Reduce a filename stem to its base identity (drop pixel size + edit tag)."""
|
||
s = _PIXEL_RE.sub("", stem)
|
||
s = _EDIT_RE.sub("", s)
|
||
return s.strip()
|
||
|
||
|
||
def _is_variant(path: Path) -> bool:
|
||
"""True if the name carries a pixel-size suffix or an edit marker."""
|
||
stem = path.stem
|
||
return bool(_PIXEL_RE.search(stem)) or bool(_EDIT_RE.search(stem))
|
||
|
||
|
||
def _variant_key(path: Path) -> tuple[str, str]:
|
||
"""Group key: same folder + same base name = same picture."""
|
||
return (str(path.parent), _strip_variant_markers(path.stem).lower())
|
||
|
||
|
||
def _safe_size(path: Path) -> int:
|
||
try:
|
||
return path.stat().st_size
|
||
except OSError:
|
||
return 0
|
||
|
||
|
||
def build_variant_groups(paths: list[Path]) -> dict[str, list[str]]:
|
||
"""
|
||
Group files that are variants of the same picture. Returns a mapping
|
||
{primary_path: [secondary_path, ...]} for groups of 2+. The primary is the
|
||
largest file with no markers (falling back to the largest file overall).
|
||
Standalone pictures are not included.
|
||
"""
|
||
from collections import defaultdict
|
||
buckets: dict[tuple, list[Path]] = defaultdict(list)
|
||
for p in paths:
|
||
buckets[_variant_key(p)].append(p)
|
||
|
||
members_of: dict[str, list[str]] = {}
|
||
for members in buckets.values():
|
||
if len(members) < 2:
|
||
continue
|
||
clean = [p for p in members if not _is_variant(p)]
|
||
primary = max(clean or members, key=_safe_size)
|
||
secs = sorted(str(p) for p in members if p != primary)
|
||
members_of[str(primary)] = secs
|
||
return members_of
|
||
|
||
|
||
def _row_to_result(row: sqlite3.Row) -> dict:
|
||
"""Reconstruct an analysis result dict from a stored primary row."""
|
||
try:
|
||
tags = json.loads(row["tags"] or "[]")
|
||
except Exception:
|
||
tags = []
|
||
return {
|
||
"description": row["description"],
|
||
"tags": tags,
|
||
"people_count": row["people_count"],
|
||
"setting": row["setting"],
|
||
"time_of_day": row["time_of_day"],
|
||
"season": row["season"],
|
||
"mood": row["mood"],
|
||
"location_hint": row["location_hint"],
|
||
"approx_year": row["approx_year"],
|
||
}
|
||
|
||
|
||
def propagate_variants(conn: sqlite3.Connection, members_of: dict[str, list[str]],
|
||
args) -> None:
|
||
"""
|
||
Copy each analyzed primary's data into its sibling variants — DB row first
|
||
(durable), then EXIF — without calling the API. Skips siblings already
|
||
finished, and groups whose primary isn't analyzed yet.
|
||
"""
|
||
# Find work: primaries that are analyzed, with siblings not yet done.
|
||
todo = [] # (secondary_path, result_dict, primary_path)
|
||
for primary, secs in members_of.items():
|
||
prow = conn.execute(
|
||
"SELECT * FROM photos WHERE path = ?", (primary,)
|
||
).fetchone()
|
||
if not prow or prow["status"] not in ("analyzed", "exif_written"):
|
||
continue # primary not ready — its siblings wait for a later run
|
||
result = _row_to_result(prow)
|
||
for sec in secs:
|
||
srow = conn.execute(
|
||
"SELECT status FROM photos WHERE path = ?", (sec,)
|
||
).fetchone()
|
||
if srow and srow["status"] == "exif_written":
|
||
continue # already fully done
|
||
if args.no_exif and srow and srow["status"] == "analyzed":
|
||
continue # DB already carries the data; nothing more to do
|
||
todo.append((sec, result, primary))
|
||
|
||
if not todo:
|
||
return
|
||
|
||
log.info(f"Propagating analysis to {len(todo):,} variant copies (no API calls)")
|
||
|
||
ok = err = 0
|
||
with Progress(
|
||
SpinnerColumn(spinner_name="dots", style="magenta"),
|
||
TextColumn("[bold magenta] Variants[/]"),
|
||
BarColumn(bar_width=None, style="dim magenta", complete_style="magenta", finished_style="green"),
|
||
MofNCompleteColumn(),
|
||
TaskProgressColumn(style="bold"),
|
||
TimeElapsedColumn(),
|
||
console=console,
|
||
refresh_per_second=4,
|
||
) as progress:
|
||
task = progress.add_task("", total=len(todo))
|
||
for sec, result, primary in todo:
|
||
# 1) DB first — durable. Record that this row was copied, not analyzed,
|
||
# and from which primary (so it's traceable in raw_response).
|
||
raw = json.dumps({"copied_variant": True, "copied_from": primary},
|
||
ensure_ascii=False)
|
||
mark_analyzed(conn, sec, result, raw)
|
||
# History line is marked "copied" (not "analyzed") with its source.
|
||
log_history(sec, "copied", result=result, copied_from=primary)
|
||
progress.console.print(
|
||
f"[magenta][/] [bold]{Path(sec).name}[/] "
|
||
f"[dim]copied from[/] {Path(primary).name}"
|
||
)
|
||
# 2) EXIF into the sibling image (unless disabled).
|
||
if not args.no_exif and not args.dry_run:
|
||
caption = build_exif_caption_from_result(result)
|
||
keywords = result.get("tags", []) or []
|
||
if write_exif(sec, caption, keywords):
|
||
mark_exif_written(conn, sec)
|
||
ok += 1
|
||
else:
|
||
err += 1
|
||
else:
|
||
ok += 1
|
||
progress.update(task, advance=1)
|
||
if _stop.is_set():
|
||
progress.console.print("[yellow]Stopped. Progress saved.[/]")
|
||
break
|
||
|
||
log.info(f"Variant propagation complete: {ok:,} OK, {err:,} errors")
|
||
|
||
|
||
# ──────────────────────────────────────────────
|
||
# Content hashing, duplicate detection, move reconciliation
|
||
#
|
||
# Two hashes per photo, both stored in the DB (the permanent, resume-safe
|
||
# ledger):
|
||
# • file_sha1 — SHA-1 of the raw bytes. Exact identity. Changes if EXIF is
|
||
# rewritten, so it tracks a moved/renamed file only while its bytes are
|
||
# untouched. Used by reconcile_moved to follow files across a reorg.
|
||
# • phash — 64-bit DCT perceptual hash. Content identity that SURVIVES resize,
|
||
# recompression, format conversion and EXIF edits. Used to spot the same
|
||
# picture arriving from different phone backups. Compared by Hamming
|
||
# distance (see PHASH_THRESHOLD).
|
||
# ──────────────────────────────────────────────
|
||
|
||
def _sha1_file(path: Path) -> str | None:
|
||
"""SHA-1 of a file's bytes, streamed. None if unreadable."""
|
||
try:
|
||
h = hashlib.sha1()
|
||
with open(path, "rb") as f:
|
||
for chunk in iter(lambda: f.read(1 << 20), b""):
|
||
h.update(chunk)
|
||
return h.hexdigest()
|
||
except OSError as e:
|
||
log.debug(f"sha1 failed for {path}: {e}")
|
||
return None
|
||
|
||
|
||
def _phash_image(path: Path) -> str | None:
|
||
"""DCT perceptual hash → 16 hex chars (64 bits). Matches the standard
|
||
imagehash.phash recipe: grayscale → 32×32 → 2D DCT → keep the top-left 8×8
|
||
low-frequency block → bit = coefficient > median. None if unreadable."""
|
||
try:
|
||
from scipy.fftpack import dct
|
||
import numpy as np
|
||
with Image.open(path) as img:
|
||
small = img.convert("L").resize((32, 32), Image.LANCZOS)
|
||
a = np.asarray(small, dtype=np.float64)
|
||
d = dct(dct(a, axis=0), axis=1) # 2D DCT-II
|
||
low = d[:8, :8] # 64 lowest frequencies
|
||
bits = (low > np.median(low)).flatten()
|
||
val = 0
|
||
for b in bits:
|
||
val = (val << 1) | int(b)
|
||
return f"{val:016x}"
|
||
except Exception as e:
|
||
log.debug(f"phash failed for {path}: {e}")
|
||
return None
|
||
|
||
|
||
def ensure_hashes(conn: sqlite3.Connection, paths, force: bool = False) -> int:
|
||
"""Compute + store phash and file_sha1 for photos that lack them (or all, if
|
||
force). One-time cost per file; resume-safe (stored, committed periodically,
|
||
interruptible). Rows must already exist (upsert_pending)."""
|
||
if force:
|
||
todo = [str(p) for p in paths]
|
||
else:
|
||
have = {r["path"] for r in conn.execute(
|
||
"SELECT path FROM photos WHERE phash IS NOT NULL AND file_sha1 IS NOT NULL"
|
||
)}
|
||
todo = [str(p) for p in paths if str(p) not in have]
|
||
if not todo:
|
||
return 0
|
||
|
||
log.info(f"Hashing {len(todo):,} photo(s) (perceptual + SHA-1)…")
|
||
done = 0
|
||
for i, p in enumerate(todo, 1):
|
||
if _stop.is_set():
|
||
break
|
||
ph = _phash_image(Path(p))
|
||
s1 = _sha1_file(Path(p))
|
||
with _db_lock:
|
||
# COALESCE(new, old): a failed hash (None) leaves any existing value be.
|
||
conn.execute(
|
||
"UPDATE photos SET phash = COALESCE(?, phash), "
|
||
"file_sha1 = COALESCE(?, file_sha1) WHERE path = ?",
|
||
(ph, s1, p),
|
||
)
|
||
if i % 200 == 0:
|
||
conn.commit()
|
||
done += 1
|
||
if i % 1000 == 0:
|
||
log.info(f" hashed {i:,}/{len(todo):,}")
|
||
with _db_lock:
|
||
conn.commit()
|
||
log.info(f"Hashing complete: {done:,} updated")
|
||
return done
|
||
|
||
|
||
def cluster_duplicates(conn: sqlite3.Connection, threshold: int) -> list[list[dict]]:
|
||
"""Group photos whose perceptual hashes are within `threshold` bits into
|
||
duplicate clusters (union-find over the Hamming graph). Each returned cluster
|
||
is a list of {path, size, status} dicts sorted largest-file-first — so
|
||
cluster[0] is the canonical (highest-quality) copy. Only clusters of ≥2.
|
||
|
||
ponytail: O(n²) vectorised Hamming scan — a few seconds at ~25k photos. If the
|
||
library grows past ~100k, swap the inner scan for a BK-tree.
|
||
"""
|
||
import numpy as np
|
||
rows = conn.execute(
|
||
"SELECT path, phash, status FROM photos WHERE phash IS NOT NULL"
|
||
).fetchall()
|
||
n = len(rows)
|
||
if n < 2:
|
||
return []
|
||
|
||
arr = np.array([int(r["phash"], 16) for r in rows], dtype=np.uint64)
|
||
popcount8 = np.array([bin(i).count("1") for i in range(256)], dtype=np.uint16)
|
||
|
||
parent = list(range(n))
|
||
def find(i):
|
||
while parent[i] != i:
|
||
parent[i] = parent[parent[i]]
|
||
i = parent[i]
|
||
return i
|
||
|
||
for i in range(n):
|
||
tail = arr[i + 1:]
|
||
if not len(tail):
|
||
break
|
||
x = tail ^ arr[i] # uint64 XOR
|
||
dist = popcount8[x.view(np.uint8).reshape(-1, 8)].sum(axis=1)
|
||
for off in np.nonzero(dist <= threshold)[0]:
|
||
ri, rj = find(i), find(i + 1 + int(off))
|
||
if ri != rj:
|
||
parent[rj] = ri
|
||
|
||
groups: dict[int, list[int]] = {}
|
||
for i in range(n):
|
||
groups.setdefault(find(i), []).append(i)
|
||
|
||
clusters = []
|
||
for members in groups.values():
|
||
if len(members) < 2:
|
||
continue
|
||
items = [{"path": rows[m]["path"],
|
||
"size": _safe_size(Path(rows[m]["path"])),
|
||
"status": rows[m]["status"]} for m in members]
|
||
items.sort(key=lambda it: it["size"], reverse=True)
|
||
clusters.append(items)
|
||
return clusters
|
||
|
||
|
||
def list_duplicates(conn: sqlite3.Connection, threshold: int) -> None:
|
||
"""Report perceptual-duplicate clusters. Read-only — changes nothing."""
|
||
clusters = cluster_duplicates(conn, threshold)
|
||
if not clusters:
|
||
log.info(f"No perceptual duplicates found (threshold {threshold}).")
|
||
return
|
||
total_dupes = sum(len(c) - 1 for c in clusters)
|
||
console.print(f"\n[bold]{len(clusters):,} duplicate cluster(s), "
|
||
f"{total_dupes:,} redundant copy/copies[/] "
|
||
f"[dim](Hamming ≤ {threshold})[/]\n")
|
||
for c in clusters:
|
||
canonical = c[0]
|
||
console.print(f"[green] KEEP[/] [bold]{canonical['path']}[/] "
|
||
f"[dim]({canonical['size']:,} B)[/]")
|
||
for dup in c[1:]:
|
||
console.print(f" [yellow] dup[/] {dup['path']} "
|
||
f"[dim]({dup['size']:,} B, {dup['status']})[/]")
|
||
console.print()
|
||
console.print("[dim]Run with --dedupe to mark the copies as 'duplicate' "
|
||
"(skipped by analysis + excluded from upload).[/]\n")
|
||
|
||
|
||
def mark_duplicates(conn: sqlite3.Connection, threshold: int) -> tuple[int, int]:
|
||
"""Mark every non-canonical member of each cluster status='duplicate' with
|
||
dup_of pointing at the canonical (largest) copy. Returns (marked, clusters)."""
|
||
clusters = cluster_duplicates(conn, threshold)
|
||
marked = 0
|
||
with _db_lock:
|
||
for c in clusters:
|
||
canonical = c[0]["path"]
|
||
for dup in c[1:]:
|
||
conn.execute(
|
||
"UPDATE photos SET status = 'duplicate', dup_of = ? WHERE path = ?",
|
||
(canonical, dup["path"]),
|
||
)
|
||
marked += 1
|
||
conn.commit()
|
||
return marked, len(clusters)
|
||
|
||
|
||
def reconcile_moved(conn: sqlite3.Connection, library: Path) -> int:
|
||
"""Follow files that moved/renamed since the last run: match on-disk files
|
||
that aren't in the DB against DB rows whose path is gone, keyed by file_sha1,
|
||
and update the row's path in place — preserving its analysis, EXIF state and
|
||
hashes. Must run BEFORE prune_missing so moved files aren't deleted and
|
||
re-analyzed. Cheap when nothing moved (no missing rows → early return)."""
|
||
if not library.exists():
|
||
return 0
|
||
db_rows = conn.execute("SELECT path, file_sha1 FROM photos").fetchall()
|
||
db_paths = {r["path"] for r in db_rows}
|
||
# sha1 → old path, only for rows whose file is gone and that carry a sha1.
|
||
missing = {r["file_sha1"]: r["path"] for r in db_rows
|
||
if r["file_sha1"] and not os.path.exists(r["path"])}
|
||
if not missing:
|
||
return 0
|
||
new_files = [p for p in discover_photos(library) if str(p) not in db_paths]
|
||
if not new_files:
|
||
return 0
|
||
|
||
moved = 0
|
||
for p in new_files:
|
||
s = _sha1_file(p)
|
||
old = missing.pop(s, None) if s else None
|
||
if old:
|
||
with _db_lock:
|
||
conn.execute("UPDATE photos SET path = ? WHERE path = ?", (str(p), old))
|
||
conn.commit()
|
||
moved += 1
|
||
if not missing:
|
||
break
|
||
if moved:
|
||
log.info(f"Reconciled {moved:,} moved/renamed file(s) to their new paths "
|
||
f"(analysis + EXIF preserved, no re-upload).")
|
||
return moved
|
||
|
||
|
||
# ──────────────────────────────────────────────
|
||
# Main pipeline
|
||
# ──────────────────────────────────────────────
|
||
|
||
def run_analysis(args, conn: sqlite3.Connection, client: OpenAI):
|
||
"""
|
||
Analyze images via the configured vision model. Each worker writes the DB record first, then
|
||
immediately writes EXIF into that image (unless --no-exif/--dry-run). The
|
||
trailing run_exif_write() pass then only has to mop up any inline failures.
|
||
"""
|
||
|
||
library = Path(args.library).expanduser().resolve()
|
||
if not library.exists():
|
||
log.error(f"Library path does not exist: {library}")
|
||
sys.exit(1)
|
||
|
||
log.info(f"Scanning {library} for photos...")
|
||
all_photos = discover_photos(library)
|
||
log.info(f"Found {len(all_photos):,} photos")
|
||
|
||
# Remove any previously-registered excluded paths
|
||
purged = purge_excluded(conn)
|
||
if purged:
|
||
log.info(f"Purged {purged:,} excluded entries from DB (.@__thumb, _IGNORE)")
|
||
|
||
# Register all files as pending (idempotent)
|
||
for p in all_photos:
|
||
upsert_pending(conn, str(p))
|
||
|
||
# Content hashes (phash + sha1) for any file that lacks them — the ledger
|
||
# duplicate detection and move-reconciliation both read from. Only-missing,
|
||
# so it's a one-time cost per file and skipped on resume.
|
||
if not args.dry_run:
|
||
ensure_hashes(conn, all_photos)
|
||
|
||
# Variant grouping: analyze only the primary of each group via the API; its
|
||
# siblings inherit the result afterwards (propagate_variants), no API call.
|
||
members_of: dict[str, list[str]] = {}
|
||
secondaries: set[str] = set()
|
||
if getattr(args, "group_variants", False):
|
||
members_of = build_variant_groups(all_photos)
|
||
secondaries = {s for secs in members_of.values() for s in secs}
|
||
if members_of:
|
||
log.info(f"Variant grouping: {len(secondaries):,} copies will inherit from "
|
||
f"{len(members_of):,} primaries (no API call for the copies)")
|
||
|
||
pending = get_pending(conn, args.reanalyze)
|
||
if secondaries:
|
||
pending = [p for p in pending if p not in secondaries]
|
||
# Skip anything the nsfwtag tool flagged 'nsfw' — those stay local, never hit the API.
|
||
pending, nsfw_skipped = filter_nsfw_tagged(pending)
|
||
if nsfw_skipped:
|
||
log.info(f"Skipping {len(nsfw_skipped):,} photo(s) tagged 'nsfw' by nsfwtag")
|
||
log.info(f"{len(pending):,} photos to analyze (skipping already processed)")
|
||
|
||
# Flush variant copies whose primary is ALREADY analyzed (e.g. from earlier
|
||
# runs) up front — before the dashboard seeds its album counts — so those
|
||
# copies are counted. Copies whose primary is still pending are skipped here
|
||
# and handled inline when that primary is analyzed below.
|
||
if members_of:
|
||
propagate_variants(conn, members_of, args)
|
||
|
||
if not pending:
|
||
log.info("No primaries left to analyze.")
|
||
else:
|
||
# Overall progress across all runs. Count truly-done rows directly —
|
||
# not total minus pending — because variant copies are excluded from
|
||
# `pending` yet aren't done, which would otherwise inflate the figure.
|
||
total_in_db = conn.execute("SELECT COUNT(*) FROM photos").fetchone()[0]
|
||
already_done = conn.execute(
|
||
"SELECT COUNT(*) FROM photos WHERE status IN ('analyzed', 'exif_written')"
|
||
).fetchone()[0]
|
||
log.info(f"Overall progress: {already_done:,} / {total_in_db:,} done "
|
||
f"({already_done / total_in_db * 100:.1f}%)" if total_in_db else "")
|
||
|
||
# Sort by album (leaf folder) — finish one album fully before the next.
|
||
pending_sorted = sorted(pending, key=lambda p: (str(Path(p).parent), p))
|
||
folder_groups = [
|
||
(folder, list(group))
|
||
for folder, group in groupby(pending_sorted, key=lambda p: str(Path(p).parent))
|
||
]
|
||
|
||
# Live dashboard only when attached to a real terminal; piped/redirected
|
||
# output (logs, CI, nohup) falls back to plain line-by-line printing.
|
||
if console.is_terminal:
|
||
try:
|
||
ok, err, copied = _run_analysis_live(
|
||
args, conn, client, library, pending, folder_groups,
|
||
total_in_db, already_done, members_of=members_of,
|
||
)
|
||
except Exception as e:
|
||
log.warning(f"Live dashboard failed ({e}); falling back to plain mode")
|
||
log.debug("Live dashboard traceback", exc_info=True)
|
||
ok, err, copied = _run_analysis_plain(
|
||
args, conn, client, library, pending, folder_groups,
|
||
total_in_db, already_done, members_of=members_of,
|
||
)
|
||
else:
|
||
ok, err, copied = _run_analysis_plain(
|
||
args, conn, client, library, pending, folder_groups,
|
||
total_in_db, already_done, members_of=members_of,
|
||
)
|
||
|
||
log.info(f"Analysis complete: {ok:,} OK, {err:,} errors, {copied:,} copied to variants")
|
||
if total_in_db:
|
||
done_total = already_done + ok + copied
|
||
log.info(f"Overall: {done_total:,} / {total_in_db:,} done "
|
||
f"({done_total / total_in_db * 100:.1f}%)")
|
||
|
||
|
||
# ──────────────────────────────────────────────
|
||
# Analysis engine + display (shared by live + plain modes)
|
||
# ──────────────────────────────────────────────
|
||
|
||
def _make_progress() -> Progress:
|
||
"""The progress bar, identical in both display modes."""
|
||
return Progress(
|
||
SpinnerColumn(spinner_name="dots", style="cyan"),
|
||
TextColumn("[bold cyan]\uf03e Analyzing[/]"),
|
||
BarColumn(bar_width=None, style="dim blue", complete_style="bright_blue", finished_style="green"),
|
||
MofNCompleteColumn(),
|
||
TaskProgressColumn(style="bold"),
|
||
TimeElapsedColumn(),
|
||
TextColumn("[dim]·[/]"),
|
||
TimeRemainingColumn(),
|
||
TextColumn(" [dim]ok=[/][green]{task.fields[ok]}[/] [dim]err=[/][red]{task.fields[err]}[/] [dim]overall=[/][yellow]{task.fields[overall]}[/]"),
|
||
console=console,
|
||
refresh_per_second=4,
|
||
)
|
||
|
||
|
||
def _run_folder_loop(folder_groups, library, pending, conn, client, args,
|
||
progress, task_id, total_in_db, already_done,
|
||
renderer, stats_lock, folder_done=None, members_of=None,
|
||
token_box=None):
|
||
"""
|
||
Shared engine: walk albums in order, run workers, drive the progress bar,
|
||
and hand each result to the renderer. DB writes happen inside worker threads
|
||
here — entirely outside any render path — so a render crash can never corrupt
|
||
the database.
|
||
|
||
When members_of maps a primary to its variant copies, each copy is propagated
|
||
(DB + EXIF) right after its primary is analyzed — so copies complete, and show
|
||
up in the feed, incrementally instead of in a separate pass at the end.
|
||
Returns (ok, err, copied).
|
||
"""
|
||
ok = err = copied = 0
|
||
total_tokens = 0
|
||
members_of = members_of or {}
|
||
|
||
write_exif_inline = not args.no_exif and not args.dry_run
|
||
|
||
def process(path_str: str):
|
||
# Bail out before any network call if a stop was requested. This makes a
|
||
# draining thread-pool queue (e.g. after Ctrl+C → shutdown) collapse into
|
||
# instant no-ops instead of firing off the rest of the folder's API calls.
|
||
if _stop.is_set():
|
||
return "skip", path_str, None, None, []
|
||
path = Path(path_str)
|
||
try:
|
||
result, raw, tokens = analyze_image(client, path, args.dry_run)
|
||
# 1) DB first — the durable record. If anything below fails, the
|
||
# photo is already safely 'analyzed' and can be retried.
|
||
mark_analyzed(conn, path_str, result, raw)
|
||
log_history(path_str, "analyzed", result=result, tokens=tokens)
|
||
# 2) Write EXIF into the image right away, in this same worker. On
|
||
# failure we leave status 'analyzed' — the trailing EXIF pass (or
|
||
# --exif-only) retries it, so nothing is lost.
|
||
caption = build_exif_caption_from_result(result)
|
||
keywords = result.get("tags", []) or []
|
||
if write_exif_inline:
|
||
if write_exif(path_str, caption, keywords):
|
||
mark_exif_written(conn, path_str)
|
||
# 3) Propagate to this primary's variant copies — no API call. Same
|
||
# DB-first ordering, marked "copied" with the source primary.
|
||
copied_secs = []
|
||
for sec in members_of.get(path_str, ()):
|
||
raw_c = json.dumps({"copied_variant": True, "copied_from": path_str},
|
||
ensure_ascii=False)
|
||
mark_analyzed(conn, sec, result, raw_c)
|
||
log_history(sec, "copied", result=result, copied_from=path_str)
|
||
if write_exif_inline and write_exif(sec, caption, keywords):
|
||
mark_exif_written(conn, sec)
|
||
copied_secs.append(sec)
|
||
return "ok", path_str, result, tokens, copied_secs
|
||
except Exception as e:
|
||
mark_error(conn, path_str, str(e))
|
||
log_history(path_str, "error", error=str(e))
|
||
return "err", path_str, str(e), {}, []
|
||
|
||
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
|
||
for folder, folder_paths in folder_groups:
|
||
if _stop.is_set():
|
||
break
|
||
|
||
album = album_label(Path(folder_paths[0]), library)
|
||
renderer.folder_start(album, len(folder_paths))
|
||
|
||
futures = {executor.submit(process, p): p for p in folder_paths}
|
||
for future in as_completed(futures):
|
||
outcome = future.result()
|
||
if outcome[0] == "skip": # worker no-op'd because we're stopping
|
||
continue
|
||
est = 0
|
||
copied_secs = []
|
||
if outcome[0] == "ok":
|
||
ok += 1
|
||
tokens = outcome[3]
|
||
copied_secs = outcome[4]
|
||
with stats_lock:
|
||
total_tokens += tokens["total"]
|
||
est = int((total_tokens / ok) * (len(pending) - (ok + err)))
|
||
if folder_done is not None:
|
||
folder_done[album] = folder_done.get(album, 0) + 1 + len(copied_secs)
|
||
if token_box is not None:
|
||
t = tokens["total"]
|
||
token_box["n"] += 1
|
||
token_box["total"] += t
|
||
token_box["prompt"] += tokens.get("prompt", 0)
|
||
token_box["completion"] += tokens.get("completion", 0)
|
||
token_box["min"] = t if token_box["n"] == 1 else min(token_box["min"], t)
|
||
token_box["max"] = max(token_box["max"], t)
|
||
copied += len(copied_secs)
|
||
else:
|
||
err += 1
|
||
|
||
ctx = {"album": album, "ok": ok, "err": err,
|
||
"total_tokens": total_tokens, "est": est}
|
||
renderer.photo(outcome, ctx)
|
||
# Show each inherited copy in the feed, marked as copied.
|
||
for sec in copied_secs:
|
||
renderer.copied(sec, outcome[1], outcome[2])
|
||
|
||
done_now = already_done + ok + err + copied
|
||
overall = f"{done_now / total_in_db * 100:.1f}%" if total_in_db else "?"
|
||
progress.update(task_id, advance=1, ok=ok, err=err, overall=overall)
|
||
|
||
if _stop.is_set():
|
||
for f in futures:
|
||
f.cancel()
|
||
renderer.stopped()
|
||
break
|
||
|
||
return ok, err, copied
|
||
|
||
|
||
# ── Plain renderer (non-TTY / fallback) ───────────────────────────────────────
|
||
|
||
class _PlainRenderer:
|
||
"""Prints one block per photo — the original behaviour, unchanged."""
|
||
|
||
def __init__(self, progress: Progress):
|
||
self.progress = progress
|
||
|
||
def folder_start(self, album: str, n: int):
|
||
self.progress.console.print(
|
||
f"\n[bold blue]\uf07b {album}[/] [dim]({n} photos)[/]"
|
||
)
|
||
|
||
def photo(self, outcome, ctx):
|
||
if outcome[0] == "ok":
|
||
path_str, result, tokens = outcome[1], outcome[2], outcome[3]
|
||
tags = result.get("tags", [])
|
||
desc = result.get("description", "")
|
||
caption = build_exif_caption_from_result(result)
|
||
mood = result.get("mood", "")
|
||
setting = result.get("setting", "")
|
||
people = result.get("people_count", 0)
|
||
self.progress.console.print(
|
||
f"\n[green]\uf058[/] [bold]{path_str}[/]\n"
|
||
f" [cyan]\uf03e[/] [dim]Desc [/] {desc}\n"
|
||
f" [yellow]\uf02b[/] [dim]Tags [/] [dim]{', '.join(tags)}[/]\n"
|
||
f" [magenta]\uf005[/] [dim]Mood [/] {mood} [dim]·[/] {setting} [dim]·[/] {people} [dim]people[/]\n"
|
||
f" [blue]\uf044[/] [dim]EXIF [/] [dim]{caption[:120]}{'…' if len(caption) > 120 else ''}[/]\n"
|
||
f" [yellow]⚡[/] [dim]Tokens [/] [bold]{tokens['total']:,}[/] "
|
||
f"[dim](prompt {tokens['prompt']:,} + compl {tokens['completion']:,}) · "
|
||
f"session [/][bold]{ctx['total_tokens']:,}[/][dim] · est. remaining [/][bold]~{ctx['est']:,}[/]"
|
||
)
|
||
else:
|
||
self.progress.console.print(
|
||
f"\n[red]\uf057[/] [bold red]{outcome[1]}[/]\n"
|
||
f" [dim]{outcome[2]}[/]"
|
||
)
|
||
|
||
def copied(self, sec, primary, result):
|
||
self.progress.console.print(
|
||
f"\uf0c5 [bold]{Path(sec).name}[/] "
|
||
f"[dim]\u2190 copied from {Path(primary).name} (no API call)[/]"
|
||
)
|
||
|
||
def stopped(self):
|
||
self.progress.console.print(
|
||
"[yellow]Stopped. In-flight requests completed, rest skipped.[/]"
|
||
)
|
||
|
||
|
||
def _run_analysis_plain(args, conn, client, library, pending, folder_groups,
|
||
total_in_db, already_done, members_of=None):
|
||
progress = _make_progress()
|
||
stats_lock = threading.Lock()
|
||
with progress:
|
||
task_id = progress.add_task("", total=len(pending), ok=0, err=0, overall="0.0%")
|
||
renderer = _PlainRenderer(progress)
|
||
return _run_folder_loop(
|
||
folder_groups, library, pending, conn, client, args,
|
||
progress, task_id, total_in_db, already_done,
|
||
renderer, stats_lock, folder_done=None, members_of=members_of,
|
||
)
|
||
|
||
|
||
# ── Live dashboard ────────────────────────────────────────────────────────────
|
||
|
||
_STATS_ROWS = 14 # max album rows shown; the rest are summarised
|
||
|
||
|
||
def _text_bar(done: int, total: int, width: int = 18) -> Text:
|
||
pct = (done / total) if total else 0.0
|
||
filled = max(0, min(width, int(round(pct * width))))
|
||
complete = total > 0 and done >= total
|
||
bar = Text("█" * filled, style="green" if complete else "bright_blue")
|
||
bar.append("░" * (width - filled), style="grey30")
|
||
return bar
|
||
|
||
|
||
def _stats_panel(folder_total, folder_done, stats_lock) -> Panel:
|
||
"""Album progress — leaf folder = album. Reads in-memory dicts only."""
|
||
try:
|
||
with stats_lock:
|
||
items = [(a, folder_done.get(a, 0), t) for a, t in folder_total.items()]
|
||
|
||
def rank(it):
|
||
a, d, t = it
|
||
if t and 0 < d < t: # active first
|
||
return (0, a)
|
||
if d == 0: # not started
|
||
return (1, a)
|
||
return (2, a) # complete last
|
||
items.sort(key=rank)
|
||
|
||
done_albums = sum(1 for _, d, t in items if t and d >= t)
|
||
title = (f"[b]Album Progress[/] [dim]\u2014 leaf folder = album "
|
||
f"({done_albums:,}/{len(items):,} albums)[/]")
|
||
|
||
table = Table.grid(padding=(0, 2))
|
||
table.add_column(justify="left", no_wrap=True)
|
||
table.add_column(justify="right", no_wrap=True)
|
||
table.add_column(no_wrap=True)
|
||
table.add_column(justify="right", no_wrap=True)
|
||
|
||
for album, done, total in items[:_STATS_ROWS]:
|
||
pct = (done / total * 100) if total else 0
|
||
if total and done >= total:
|
||
flag = Text("✓", style="green")
|
||
elif done == 0:
|
||
flag = Text("0%", style="grey50")
|
||
else:
|
||
flag = Text(f"{pct:3.0f}%", style="yellow")
|
||
table.add_row(
|
||
Text(fit_label(album, 30), style="white"),
|
||
Text(f"{done:,}/{total:,}", style="grey62"),
|
||
_text_bar(done, total),
|
||
flag,
|
||
)
|
||
|
||
hidden = items[_STATS_ROWS:]
|
||
if hidden:
|
||
hdone = sum(1 for _, d, t in hidden if t and d >= t)
|
||
table.add_row(
|
||
Text(f"\u2026 +{len(hidden):,} more albums", style="grey50"),
|
||
Text(f"{hdone:,} done", style="grey50"), Text(""), Text(""),
|
||
)
|
||
|
||
return Panel(table, title=title, title_align="left",
|
||
border_style="grey37", padding=(0, 1))
|
||
except Exception as e:
|
||
return Panel(Text(f"stats render error: {e}", style="red"), border_style="red")
|
||
|
||
|
||
def _feed_panel(log_buffer, log_lock) -> Panel:
|
||
"""Recent per-photo results. Reads the shared deque under its lock."""
|
||
try:
|
||
with log_lock:
|
||
rows = list(log_buffer)
|
||
|
||
grid = Table.grid(padding=(0, 1))
|
||
grid.add_column(overflow="ellipsis", no_wrap=True)
|
||
|
||
if not rows:
|
||
grid.add_row(Text("Waiting for first result\u2026", style="grey50"))
|
||
|
||
for e in rows:
|
||
kind = e.get("kind")
|
||
if kind == "folder":
|
||
grid.add_row(Text(f"\uf07b {e['album']} ({e['n']})", style="bold blue"))
|
||
elif kind == "stop":
|
||
grid.add_row(Text("Stopped \u2014 in-flight done, rest skipped.", style="yellow"))
|
||
elif kind == "copied":
|
||
line = Text("\uf0c5 ", style="magenta")
|
||
line.append(e["name"], style="bold magenta")
|
||
line.append(f" \u2190 {e.get('from', '')} (copied)", style="grey50")
|
||
grid.add_row(line)
|
||
if e.get("desc"):
|
||
grid.add_row(Text(f" {e['desc']}", style="grey62"))
|
||
elif e.get("ok"):
|
||
line = Text("✓ ", style="green")
|
||
line.append(e["name"], style="bold white")
|
||
grid.add_row(line)
|
||
if e.get("desc"):
|
||
grid.add_row(Text(f" {e['desc']}", style="grey62"))
|
||
else:
|
||
line = Text("✗ ", style="red")
|
||
line.append(e["name"], style="bold red")
|
||
line.append(f" \u2014 {e.get('err', '')}", style="grey50")
|
||
grid.add_row(line)
|
||
|
||
return Panel(grid, title="[b]Recent[/]", title_align="left",
|
||
border_style="grey37", padding=(0, 1))
|
||
except Exception as e:
|
||
return Panel(Text(f"feed render error: {e}", style="red"), border_style="red")
|
||
|
||
|
||
def _balance_text(bal: dict | None) -> Text:
|
||
"""One-line API balance, coloured by how much is left. Reads cached data."""
|
||
if not bal or bal.get("available_balance") is None:
|
||
return Text(" ⚡ balance: unavailable", style="grey50")
|
||
avail = bal.get("available_balance")
|
||
try:
|
||
amount = float(avail)
|
||
colour = "red" if amount < 1 else "yellow" if amount < 5 else "green"
|
||
shown = f"${amount:,.2f}"
|
||
except (TypeError, ValueError):
|
||
colour, shown = "grey62", str(avail)
|
||
line = Text(" ⚡ balance: ", style="grey62")
|
||
line.append(shown, style=f"bold {colour}")
|
||
return line
|
||
|
||
|
||
def _rpd_text() -> Text:
|
||
"""One-line requests-per-day gauge (local count vs RPD_LIMIT), coloured by
|
||
how close to the daily cap. Resets at midnight Pacific."""
|
||
n = _rpd_count()
|
||
line = Text(" 📅 RPD today: ", style="grey62")
|
||
if RPD_LIMIT > 0:
|
||
frac = n / RPD_LIMIT
|
||
colour = "red" if frac >= 0.95 else "yellow" if frac >= 0.8 else "green"
|
||
line.append(f"{n:,} / {RPD_LIMIT:,}", style=f"bold {colour}")
|
||
else:
|
||
line.append(f"{n:,}", style="bold grey70")
|
||
return line
|
||
|
||
|
||
def _tokens_panel(box, lock) -> Panel:
|
||
"""Session token usage — overlay toggled with 't'. Reads the shared box."""
|
||
with lock:
|
||
n, tot = box["n"], box["total"]
|
||
pr, co = box["prompt"], box["completion"]
|
||
mn, mx = box["min"], box["max"]
|
||
if not n:
|
||
body = Text("No analyses yet this session.", style="grey50")
|
||
else:
|
||
t = Table.grid(padding=(0, 3))
|
||
t.add_column(style="grey62", no_wrap=True)
|
||
t.add_column(justify="right", style="bold white", no_wrap=True)
|
||
t.add_row("Photos this session", f"{n:,}")
|
||
t.add_row("Avg total / photo", f"{tot / n:,.0f}")
|
||
t.add_row("Avg prompt (image)", f"{pr / n:,.0f}")
|
||
t.add_row("Avg completion", f"{co / n:,.0f}")
|
||
t.add_row("Min / max total", f"{mn:,} / {mx:,}")
|
||
t.add_row("Session total", f"{tot:,}")
|
||
body = t
|
||
return Panel(body, title="[b]Token usage[/] [dim]— m or Esc to go back[/]",
|
||
title_align="left", border_style="cyan", padding=(0, 1))
|
||
|
||
|
||
def _menu_panel(box, lock) -> Panel:
|
||
"""The popup overlay: either the navigable menu list or the token view."""
|
||
with _menu_lock:
|
||
sel, view = _menu_state["sel"], _menu_state["view"]
|
||
if view == "token":
|
||
return _tokens_panel(box, lock)
|
||
grid = Table.grid(padding=(0, 1))
|
||
grid.add_column(no_wrap=True)
|
||
for i, (label, _) in enumerate(_MENU_ITEMS):
|
||
if i == sel:
|
||
grid.add_row(Text(f" ▸ {label} ", style="bold black on cyan"))
|
||
else:
|
||
grid.add_row(Text(f" {label} ", style="white"))
|
||
hint = Text("↑/↓ or j/k · Enter select · m/Esc close", style="grey42")
|
||
return Panel(Group(grid, Text(""), hint), title="[b]Menu[/]",
|
||
title_align="left", border_style="cyan", padding=(0, 1))
|
||
|
||
|
||
def _read_key(timeout):
|
||
"""One keypress from cbreak stdin, normalised. Arrow keys arrive as ESC [ A/B;
|
||
a lone ESC (no follow-up byte) is reported as 'esc'."""
|
||
r, _, _ = select.select([sys.stdin], [], [], timeout)
|
||
if not r:
|
||
return None
|
||
ch = sys.stdin.read(1)
|
||
if ch == "\x1b":
|
||
r2, _, _ = select.select([sys.stdin], [], [], 0.001)
|
||
if not r2:
|
||
return "esc"
|
||
return {"[A": "up", "[B": "down"}.get(sys.stdin.read(2))
|
||
if ch in ("\r", "\n"):
|
||
return "enter"
|
||
return ch
|
||
|
||
|
||
def _key_listener(stop_evt):
|
||
"""Drive the 'm' popup menu. cbreak (not raw) keeps ISIG on, so Ctrl+C still
|
||
raises SIGINT exactly as before. No-op without a real TTY."""
|
||
if not sys.stdin.isatty():
|
||
return
|
||
fd = sys.stdin.fileno()
|
||
old = termios.tcgetattr(fd)
|
||
try:
|
||
tty.setcbreak(fd)
|
||
while not stop_evt.is_set():
|
||
k = _read_key(0.2)
|
||
if k is None:
|
||
continue
|
||
if not _menu.is_set():
|
||
if k in ("m", "M"):
|
||
with _menu_lock:
|
||
_menu_state["sel"], _menu_state["view"] = 0, "menu"
|
||
_menu.set()
|
||
continue
|
||
with _menu_lock:
|
||
view = _menu_state["view"]
|
||
if view == "token": # token panel: only way out is back
|
||
if k in ("m", "M", "esc", "enter"):
|
||
with _menu_lock:
|
||
_menu_state["view"] = "menu"
|
||
continue
|
||
if k in ("up", "k"):
|
||
with _menu_lock:
|
||
_menu_state["sel"] = (_menu_state["sel"] - 1) % len(_MENU_ITEMS)
|
||
elif k in ("down", "j"):
|
||
with _menu_lock:
|
||
_menu_state["sel"] = (_menu_state["sel"] + 1) % len(_MENU_ITEMS)
|
||
elif k in ("m", "M", "esc"):
|
||
_menu.clear()
|
||
elif k == "enter":
|
||
with _menu_lock:
|
||
action = _MENU_ITEMS[_menu_state["sel"]][1]
|
||
if action == "token":
|
||
with _menu_lock:
|
||
_menu_state["view"] = "token"
|
||
elif action == "close":
|
||
_menu.clear()
|
||
elif action == "stop":
|
||
_stop.set()
|
||
_menu.clear()
|
||
elif action == "force":
|
||
# Reuse the tested SIGINT path: _stop already set means the
|
||
# handler force-quits on the main thread, which unwinds Live
|
||
# (restores the screen) and our finally restores termios.
|
||
_stop.set()
|
||
os.kill(os.getpid(), signal.SIGINT)
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
termios.tcsetattr(fd, termios.TCSADRAIN, old)
|
||
|
||
|
||
class _Dashboard:
|
||
"""Re-rendered by Live every tick; rebuilds the layout from live state."""
|
||
|
||
def __init__(self, progress, folder_total, folder_done, stats_lock,
|
||
log_buffer, log_lock, bal_holder=None, token_box=None):
|
||
self.progress = progress
|
||
self.folder_total = folder_total
|
||
self.folder_done = folder_done
|
||
self.stats_lock = stats_lock
|
||
self.log_buffer = log_buffer
|
||
self.log_lock = log_lock
|
||
self.bal_holder = bal_holder or {}
|
||
self.token_box = token_box
|
||
|
||
def __rich__(self):
|
||
try:
|
||
# Cached balance (refreshed off-thread); never a network call here.
|
||
prog_block = Group(
|
||
self.progress,
|
||
_balance_text(self.bal_holder.get("value")),
|
||
_rpd_text(),
|
||
Text(" press m for menu", style="grey42"),
|
||
)
|
||
sections = []
|
||
# Popup menu overlay, toggled with 'm' — same pattern as the stop
|
||
# banner: state is set off-thread, drawn on the next render tick.
|
||
if _menu.is_set() and self.token_box is not None:
|
||
sections.append(Layout(_menu_panel(self.token_box, self.stats_lock),
|
||
name="menu", size=10))
|
||
# Immediate, unmissable feedback that Ctrl+C was received. Shown on
|
||
# the next refresh tick (~250 ms) — the log line would be hidden by
|
||
# the full-screen Live, so we render it inside the dashboard instead.
|
||
if _stop.is_set():
|
||
banner = Text(
|
||
" STOPPING — finishing current work. "
|
||
"Ctrl+C again to force-quit. ",
|
||
style="bold black on yellow", justify="center",
|
||
)
|
||
sections.append(Layout(Panel(banner, border_style="bold yellow",
|
||
padding=(0, 1)), name="banner", size=3))
|
||
# Throttling warning — provider returning 429/503 in clusters. Auto-
|
||
# clears once hits age out of the window. Cue to lower MAX_WORKERS.
|
||
rl = _ratelimit_recent()
|
||
if rl >= RATELIMIT_WARN_AT:
|
||
warn = Text(
|
||
f" ⚠ {rl} rate-limit/overload hits in {RATELIMIT_WINDOW}s — "
|
||
f"provider is throttling. Lower MAX_WORKERS if this persists. ",
|
||
style="bold white on red", justify="center",
|
||
)
|
||
sections.append(Layout(Panel(warn, border_style="bold red",
|
||
padding=(0, 1)), name="throttle", size=3))
|
||
sections.append(Layout(_stats_panel(self.folder_total, self.folder_done, self.stats_lock),
|
||
name="stats", ratio=2))
|
||
sections.append(Layout(_feed_panel(self.log_buffer, self.log_lock),
|
||
name="feed", ratio=2))
|
||
sections.append(Layout(Panel(prog_block, title="[b]Analyzing[/]", title_align="left",
|
||
border_style="grey37", padding=(0, 1)),
|
||
name="progress", size=5))
|
||
layout = Layout()
|
||
layout.split_column(*sections)
|
||
return layout
|
||
except Exception as e:
|
||
return Text(f"dashboard render error: {e}", style="red")
|
||
|
||
|
||
class _LiveRenderer:
|
||
"""Workers never touch the terminal — they only append to the shared deque."""
|
||
|
||
def __init__(self, log_buffer, log_lock):
|
||
self.log_buffer = log_buffer
|
||
self.log_lock = log_lock
|
||
|
||
def folder_start(self, album: str, n: int):
|
||
with self.log_lock:
|
||
self.log_buffer.appendleft({"kind": "folder", "album": album, "n": n})
|
||
|
||
def photo(self, outcome, ctx):
|
||
if outcome[0] == "ok":
|
||
path_str, result = outcome[1], outcome[2]
|
||
entry = {"kind": "photo", "ok": True, "name": Path(path_str).name,
|
||
"desc": (result.get("description") or "")}
|
||
else:
|
||
entry = {"kind": "photo", "ok": False, "name": Path(outcome[1]).name,
|
||
"err": outcome[2]}
|
||
with self.log_lock:
|
||
self.log_buffer.appendleft(entry)
|
||
|
||
def copied(self, sec, primary, result):
|
||
entry = {"kind": "copied", "name": Path(sec).name,
|
||
"from": Path(primary).name,
|
||
"desc": (result.get("description") or "")}
|
||
with self.log_lock:
|
||
self.log_buffer.appendleft(entry)
|
||
|
||
def stopped(self):
|
||
with self.log_lock:
|
||
self.log_buffer.appendleft({"kind": "stop"})
|
||
|
||
|
||
def _run_analysis_live(args, conn, client, library, pending, folder_groups,
|
||
total_in_db, already_done, members_of=None):
|
||
progress = _make_progress()
|
||
task_id = progress.add_task("", total=len(pending), ok=0, err=0, overall="0.0%")
|
||
|
||
stats_lock = threading.Lock()
|
||
log_lock = threading.Lock()
|
||
log_buffer = deque(maxlen=10)
|
||
|
||
# Seed album totals + done counts from the DB once, up front. No queries
|
||
# run inside the 4 Hz render loop — workers mutate folder_done in memory.
|
||
folder_total: dict = {}
|
||
folder_done: dict = {}
|
||
for row in conn.execute("SELECT path, status FROM photos"):
|
||
album = album_label(Path(row["path"]), library)
|
||
folder_total[album] = folder_total.get(album, 0) + 1
|
||
if row["status"] in ("analyzed", "exif_written"):
|
||
folder_done[album] = folder_done.get(album, 0) + 1
|
||
|
||
# API balance shown in the Analyzing panel. Fetched once now, then refreshed
|
||
# on a slow daemon thread — never on the render path, never per photo (would
|
||
# waste the 3 RPM budget). The render only reads the cached value.
|
||
bal_holder = {"value": fetch_balance(client.api_key) if not args.dry_run else None}
|
||
bal_stop = threading.Event()
|
||
|
||
_rpd_load() # seed today's request count (cumulative across runs this PT day)
|
||
|
||
def _refresh_balance():
|
||
while not bal_stop.wait(90): # every 90s until told to stop
|
||
v = fetch_balance(client.api_key)
|
||
if v is not None:
|
||
bal_holder["value"] = v
|
||
_rpd_persist() # checkpoint the RPD count alongside the balance refresh
|
||
|
||
bal_thread = None
|
||
if not args.dry_run:
|
||
bal_thread = threading.Thread(target=_refresh_balance, daemon=True)
|
||
bal_thread.start()
|
||
|
||
token_box = {"n": 0, "total": 0, "prompt": 0, "completion": 0, "min": 0, "max": 0}
|
||
|
||
dashboard = _Dashboard(progress, folder_total, folder_done, stats_lock,
|
||
log_buffer, log_lock, bal_holder=bal_holder,
|
||
token_box=token_box)
|
||
renderer = _LiveRenderer(log_buffer, log_lock)
|
||
|
||
key_stop = threading.Event()
|
||
key_thread = threading.Thread(target=_key_listener, args=(key_stop,), daemon=True)
|
||
key_thread.start()
|
||
|
||
try:
|
||
with Live(dashboard, console=console, refresh_per_second=4, screen=True):
|
||
return _run_folder_loop(
|
||
folder_groups, library, pending, conn, client, args,
|
||
progress, task_id, total_in_db, already_done,
|
||
renderer, stats_lock, folder_done=folder_done, members_of=members_of,
|
||
token_box=token_box,
|
||
)
|
||
finally:
|
||
key_stop.set() # restore terminal mode promptly
|
||
bal_stop.set() # stop the refresher promptly
|
||
_rpd_persist() # final RPD checkpoint so the count survives this run
|
||
|
||
|
||
def run_exif_write(args, conn: sqlite3.Connection):
|
||
"""Phase 2: write analysis results into EXIF of analyzed files."""
|
||
|
||
if args.no_exif:
|
||
log.info("--no-exif set, skipping EXIF write phase.")
|
||
return
|
||
|
||
rows = get_analyzed_no_exif(conn)
|
||
log.info(f"{len(rows):,} photos need EXIF written")
|
||
|
||
if not rows:
|
||
return
|
||
|
||
ok = err = 0
|
||
with Progress(
|
||
SpinnerColumn(spinner_name="dots", style="yellow"),
|
||
TextColumn("[bold yellow]\uf044 Writing EXIF[/]"),
|
||
BarColumn(bar_width=None, style="dim yellow", complete_style="yellow", finished_style="green"),
|
||
MofNCompleteColumn(),
|
||
TaskProgressColumn(style="bold"),
|
||
TimeElapsedColumn(),
|
||
TextColumn("[dim]·[/]"),
|
||
TimeRemainingColumn(),
|
||
TextColumn(" [dim]ok=[/][green]{task.fields[ok]}[/] [dim]err=[/][red]{task.fields[err]}[/]"),
|
||
console=console,
|
||
refresh_per_second=4,
|
||
) as progress:
|
||
task_id = progress.add_task("", total=len(rows), ok=0, err=0)
|
||
|
||
for row in rows:
|
||
if args.dry_run:
|
||
mark_exif_written(conn, row["path"])
|
||
ok += 1
|
||
progress.update(task_id, advance=1, ok=ok, err=err)
|
||
continue
|
||
|
||
caption = build_exif_caption(row)
|
||
try:
|
||
keywords = json.loads(row["tags"] or "[]")
|
||
except Exception:
|
||
keywords = []
|
||
|
||
if write_exif(row["path"], caption, keywords):
|
||
mark_exif_written(conn, row["path"])
|
||
ok += 1
|
||
else:
|
||
err += 1
|
||
|
||
progress.update(task_id, advance=1, ok=ok, err=err)
|
||
|
||
if _stop.is_set():
|
||
progress.console.print("[yellow]Stopped. Progress saved.[/]")
|
||
break
|
||
|
||
log.info(f"EXIF write complete: {ok:,} OK, {err:,} errors")
|
||
|
||
|
||
def album_label(photo: Path, library: Path = None) -> str:
|
||
"""
|
||
Album = the leaf folder directly containing the photo, shown as a path
|
||
relative to the library root so nested albums stay distinct
|
||
(Urlaub/Rom vs Urlaub/Venedig). Falls back to the absolute parent if the
|
||
photo lies outside the library.
|
||
"""
|
||
parent = photo.parent
|
||
if library:
|
||
try:
|
||
rel = parent.relative_to(library)
|
||
return "(root)" if str(rel) == "." else str(rel)
|
||
except ValueError:
|
||
pass
|
||
return parent.name or str(parent)
|
||
|
||
|
||
def fit_label(label: str, width: int) -> str:
|
||
"""Left-truncate so the album name (path tail) stays visible: Urlaub/…/Rom."""
|
||
if len(label) <= width:
|
||
return label
|
||
return "…" + label[-(width - 1):]
|
||
|
||
|
||
def print_stats(conn: sqlite3.Connection, library: Path = None):
|
||
DONE = {"analyzed", "exif_written"}
|
||
|
||
rows = conn.execute(
|
||
"SELECT status, COUNT(*) as n FROM photos GROUP BY status"
|
||
).fetchall()
|
||
log.info("── Database summary ──────────────────────")
|
||
for r in rows:
|
||
log.info(f" {r['status']:20s} {r['n']:>6,}")
|
||
|
||
# Album breakdown — an album is the leaf folder directly containing the
|
||
# photos (e.g. Urlaub/Rom and Urlaub/Venedig are two albums, not one).
|
||
all_rows = conn.execute("SELECT path, status FROM photos").fetchall()
|
||
if not all_rows:
|
||
log.info("──────────────────────────────────────────")
|
||
return
|
||
|
||
folder_total: dict[str, int] = {}
|
||
folder_done: dict[str, int] = {}
|
||
|
||
for row in all_rows:
|
||
folder = album_label(Path(row["path"]), library)
|
||
folder_total[folder] = folder_total.get(folder, 0) + 1
|
||
if row["status"] in DONE:
|
||
folder_done[folder] = folder_done.get(folder, 0) + 1
|
||
|
||
log.info("── Album breakdown ───────────────────────")
|
||
for folder in sorted(folder_total):
|
||
done = folder_done.get(folder, 0)
|
||
total = folder_total[folder]
|
||
pct = done / total * 100
|
||
flag = "[green]✓[/]" if done == total else f"[yellow]{pct:3.0f}%[/]"
|
||
log.info(f" {fit_label(folder, 40):<40s} {done:>5,}/{total:<5,} {flag}")
|
||
|
||
log.info("──────────────────────────────────────────")
|
||
|
||
|
||
# ──────────────────────────────────────────────
|
||
# CLI
|
||
# ──────────────────────────────────────────────
|
||
|
||
def load_env_file() -> Path | None:
|
||
"""
|
||
Load KEY=VALUE pairs from a .env file into os.environ. Existing shell env
|
||
vars are never overwritten, so `export LLM_API_KEY=...` still wins.
|
||
Looks in the current directory first, then beside this script. Tiny manual
|
||
parser — no python-dotenv dependency. Returns the file used, or None.
|
||
"""
|
||
candidates = [Path.cwd() / ENV_FILE, Path(__file__).resolve().parent / ENV_FILE]
|
||
for path in candidates:
|
||
if not path.is_file():
|
||
continue
|
||
try:
|
||
for raw in path.read_text(encoding="utf-8").splitlines():
|
||
line = raw.strip()
|
||
if not line or line.startswith("#"):
|
||
continue
|
||
if line.startswith("export "):
|
||
line = line[len("export "):]
|
||
if "=" not in line:
|
||
continue
|
||
key, _, value = line.partition("=")
|
||
key = key.strip()
|
||
value = value.strip()
|
||
# Strip an inline comment: a '#' at the start of the value or
|
||
# preceded by whitespace. Skip if the value is quoted, and keep
|
||
# '#' that's part of a token (e.g. a URL fragment, pass#word).
|
||
if value[:1] not in ("'", '"'):
|
||
for i, ch in enumerate(value):
|
||
if ch == "#" and (i == 0 or value[i - 1].isspace()):
|
||
value = value[:i]
|
||
break
|
||
value = value.strip().strip('"').strip("'")
|
||
if key and key not in os.environ:
|
||
os.environ[key] = value
|
||
except Exception as e:
|
||
log.warning(f"Could not read {path}: {e}")
|
||
return path
|
||
return None
|
||
|
||
|
||
# Read config from the env file (or real shell env). These let photo_analyzer.env
|
||
# drive every setting, so the script can run with no flags at all.
|
||
def _env_str(name: str, fallback=None):
|
||
v = os.environ.get(name)
|
||
return v if v not in (None, "") else fallback
|
||
|
||
|
||
def _env_int(name: str, fallback: int) -> int:
|
||
v = os.environ.get(name)
|
||
if v in (None, ""):
|
||
return fallback
|
||
try:
|
||
return int(v)
|
||
except ValueError:
|
||
log.warning(f"{name}={v!r} is not an integer; using {fallback}")
|
||
return fallback
|
||
|
||
|
||
def _env_bool(name: str) -> bool:
|
||
return os.environ.get(name, "").strip().lower() in ("1", "true", "yes", "on")
|
||
|
||
|
||
def fetch_balance(api_key: str) -> dict | None:
|
||
"""GET /v1/users/me/balance. Returns the data dict, or None on any failure."""
|
||
import urllib.request
|
||
import urllib.error
|
||
|
||
url = LLM_BASE_URL.rstrip("/") + "/users/me/balance"
|
||
req = urllib.request.Request(url, headers={"Authorization": f"Bearer {api_key}"})
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||
data = json.loads(resp.read().decode("utf-8"))
|
||
return data.get("data", data) or {}
|
||
except urllib.error.HTTPError as e:
|
||
body = e.read().decode("utf-8", "replace")[:300]
|
||
log.debug(f"Balance query failed: HTTP {e.code} {e.reason} {body}")
|
||
except Exception as e:
|
||
log.debug(f"Balance query failed: {e}")
|
||
return None
|
||
|
||
|
||
def query_balance(api_key: str) -> None:
|
||
"""Print remaining API balance (works on providers that expose a balance endpoint)."""
|
||
d = fetch_balance(api_key)
|
||
if d is None:
|
||
log.error("Balance query failed or unsupported by this provider (see debug log).")
|
||
return
|
||
log.info("── API balance ───────────────────────────")
|
||
log.info(f" available {d.get('available_balance')}")
|
||
log.info(f" cash {d.get('cash_balance')}")
|
||
log.info(f" voucher {d.get('voucher_balance')}")
|
||
log.info("──────────────────────────────────────────")
|
||
|
||
|
||
def check_quota(client: OpenAI) -> bool:
|
||
"""Fire ONE minimal request to test whether today's per-model quota is
|
||
available, instead of launching the whole pipeline to find out. Returns
|
||
True if the model accepted the request. A 429 here isn't billed and (for a
|
||
daily cap) doesn't consume quota."""
|
||
log.info(f"Probing {LLM_MODEL} with one minimal request…")
|
||
try:
|
||
client.chat.completions.create(
|
||
model=LLM_MODEL, max_tokens=1,
|
||
messages=[{"role": "user", "content": "ping"}],
|
||
)
|
||
log.info(f"✓ {LLM_MODEL} is accepting requests — safe to run.")
|
||
return True
|
||
except RateLimitError as e:
|
||
detail = str(getattr(e, "message", "") or e)
|
||
low = detail.lower()
|
||
if "per_day" in low or "perday" in low:
|
||
log.error(f"✗ Daily quota still exhausted for {LLM_MODEL}. "
|
||
f"Wait for reset (~midnight Pacific) or switch model.")
|
||
elif "per_minute" in low or "perminute" in low:
|
||
log.warning(f"⚠ Per-minute limit only (transient) — daily quota looks OK, "
|
||
f"just lower MAX_WORKERS. {detail[:140]}")
|
||
else:
|
||
log.error(f"✗ Rate limited: {detail[:200]}")
|
||
return False
|
||
except APIError as e:
|
||
log.error(f"Quota check failed: {e}")
|
||
return False
|
||
|
||
|
||
def main():
|
||
# Load photo_analyzer.env first so every setting below can default from it.
|
||
# Real shell env vars still win (load_env_file never overwrites them), and a
|
||
# CLI flag still overrides the env value for one-off runs.
|
||
env_file = load_env_file()
|
||
|
||
# Tuning knobs — env-driven, no flags. See photo_analyzer.env for docs.
|
||
global MAX_WORKERS, MAX_LONG_EDGE, RETRY_ATTEMPTS, RETRY_BASE_DELAY
|
||
global LLM_BASE_URL, LLM_MODEL
|
||
MAX_WORKERS = _env_int("MAX_WORKERS", MAX_WORKERS)
|
||
MAX_LONG_EDGE = _env_int("MAX_LONG_EDGE", MAX_LONG_EDGE)
|
||
RETRY_ATTEMPTS = _env_int("RETRY_ATTEMPTS", RETRY_ATTEMPTS)
|
||
RETRY_BASE_DELAY = _env_int("RETRY_BASE_DELAY", RETRY_BASE_DELAY)
|
||
LLM_BASE_URL = _env_str("LLM_BASE_URL", LLM_BASE_URL)
|
||
LLM_MODEL = _env_str("LLM_MODEL", LLM_MODEL)
|
||
global RPD_LIMIT, PHASH_THRESHOLD
|
||
RPD_LIMIT = _env_int("RPD_LIMIT", RPD_LIMIT)
|
||
PHASH_THRESHOLD = _env_int("PHASH_THRESHOLD", PHASH_THRESHOLD)
|
||
|
||
parser = argparse.ArgumentParser(
|
||
description="Analyze photos with an OpenAI-compatible vision model and write results to EXIF for Immich. "
|
||
"Every option can be set in photo_analyzer.env instead of on the CLI."
|
||
)
|
||
parser.add_argument("--library", default=_env_str("LIBRARY"), help="Root path of your photo library (env: LIBRARY; required except for --balance)")
|
||
parser.add_argument("--db", default=_env_str("DB", DB_FILE), help=f"SQLite DB path (env: DB; default: {DB_FILE})")
|
||
parser.add_argument("--dry-run", action="store_true", default=_env_bool("DRY_RUN"), help="Discover + register files, skip API and EXIF (env: DRY_RUN)")
|
||
parser.add_argument("--reanalyze", action="store_true", default=_env_bool("REANALYZE"), help="Re-analyze already-processed files (env: REANALYZE)")
|
||
parser.add_argument("--no-exif", action="store_true", default=_env_bool("NO_EXIF"), help="Skip EXIF write phase (index only) (env: NO_EXIF)")
|
||
parser.add_argument("--exif-only", action="store_true", default=_env_bool("EXIF_ONLY"), help="Skip analysis, only write EXIF for already-analyzed photos (env: EXIF_ONLY)")
|
||
parser.add_argument("--group-variants", action="store_true", default=_env_bool("GROUP_VARIANTS"), help="Analyze one photo per group of size/edit variants and copy the result to the rest — saves API calls (env: GROUP_VARIANTS)")
|
||
parser.add_argument("--backfill-phash", action="store_true", default=_env_bool("BACKFILL_PHASH"), help="Compute + store perceptual hash and SHA-1 for all photos, then exit — seeds the dedup ledger for an existing DB (env: BACKFILL_PHASH)")
|
||
parser.add_argument("--list-dupes", action="store_true", default=_env_bool("LIST_DUPES"), help="Report perceptual-duplicate clusters and exit — read-only, changes nothing (env: LIST_DUPES)")
|
||
parser.add_argument("--dedupe", action="store_true", default=_env_bool("DEDUPE"), help="Mark non-canonical duplicates status='duplicate' (skipped by analysis + excluded from upload), then exit (env: DEDUPE)")
|
||
parser.add_argument("--stats", action="store_true", default=_env_bool("STATS"), help="Print DB summary and exit (env: STATS)")
|
||
parser.add_argument("--balance", action="store_true", default=_env_bool("BALANCE"), help="Print remaining API balance and exit, if the provider exposes one (env: BALANCE)")
|
||
parser.add_argument("--quota-check", action="store_true", default=_env_bool("QUOTA_CHECK"), help="Fire ONE tiny request to test if today's quota is available, then exit (env: QUOTA_CHECK)")
|
||
parser.add_argument("--debug", action="store_true", default=_env_bool("DEBUG"), help="Show DEBUG messages on console (env: DEBUG; always written to photo_analyzer_debug.log)")
|
||
args = parser.parse_args()
|
||
|
||
if not args.library and not args.balance and not args.quota_check:
|
||
parser.error("--library is required (set LIBRARY in photo_analyzer.env, or pass --library)")
|
||
|
||
if args.debug:
|
||
_rich_handler.setLevel(logging.DEBUG)
|
||
|
||
# LLM_API_KEY is the provider-agnostic name; GEMINI_API_KEY / GOOGLE_API_KEY
|
||
# are honoured as fallbacks for the default Google provider.
|
||
api_key = _env_str("LLM_API_KEY") or _env_str("GEMINI_API_KEY") or _env_str("GOOGLE_API_KEY")
|
||
if api_key == API_KEY_PLACEHOLDER:
|
||
api_key = None # they created .env but haven't pasted a real key yet
|
||
|
||
no_api_mode = args.dry_run or args.stats or args.backfill_phash or args.list_dupes or args.dedupe
|
||
if not api_key and not no_api_mode:
|
||
log.error("No API key set (LLM_API_KEY or GEMINI_API_KEY).")
|
||
if env_file:
|
||
log.error(f"Edit {env_file} and set: LLM_API_KEY=your-key")
|
||
else:
|
||
log.error(f"Create a {ENV_FILE} file with: LLM_API_KEY=your-key")
|
||
log.error("Or export it in your shell. Get a key from your LLM provider (default: Google AI Studio, https://aistudio.google.com/apikey).")
|
||
sys.exit(1)
|
||
|
||
if args.balance:
|
||
query_balance(api_key)
|
||
return
|
||
|
||
conn = get_db(args.db)
|
||
client = OpenAI(api_key=api_key or "dry-run", base_url=LLM_BASE_URL)
|
||
|
||
if args.quota_check:
|
||
sys.exit(0 if check_quota(client) else 1)
|
||
|
||
library = Path(args.library).expanduser().resolve()
|
||
|
||
# Follow files that moved/renamed since the last run (by SHA-1) BEFORE pruning,
|
||
# so a folder reorg updates paths in place instead of deleting + re-analyzing.
|
||
reconcile_moved(conn, library)
|
||
|
||
# Drop DB entries for files moved/deleted since the last run, so the stats
|
||
# and album breakdown reflect what's actually on disk now.
|
||
pruned = prune_missing(conn, library)
|
||
if pruned:
|
||
log.info(f"Pruned {pruned:,} stale entries (files moved or deleted)")
|
||
|
||
if args.stats:
|
||
print_stats(conn, library)
|
||
return
|
||
|
||
# Content-hash / duplicate tools — operate on the DB ledger, no API needed.
|
||
if args.backfill_phash or args.list_dupes or args.dedupe:
|
||
photos = discover_photos(library)
|
||
for p in photos:
|
||
upsert_pending(conn, str(p))
|
||
ensure_hashes(conn, photos) # only-missing; seeds the ledger
|
||
if args.backfill_phash:
|
||
return
|
||
if args.list_dupes:
|
||
list_duplicates(conn, PHASH_THRESHOLD)
|
||
return
|
||
marked, clusters = mark_duplicates(conn, PHASH_THRESHOLD)
|
||
log.info(f"Marked {marked:,} duplicate(s) across {clusters:,} cluster(s) "
|
||
f"— status='duplicate', excluded from analysis + upload.")
|
||
return
|
||
|
||
# EXIF writing needs exiftool — verify once, up front, cleanly (instead of
|
||
# failing deep inside a worker thread or mid-backlog).
|
||
will_write_exif = not args.no_exif and not args.dry_run
|
||
if will_write_exif and shutil.which("exiftool") is None:
|
||
log.error("exiftool not found on PATH. Install it (brew install exiftool),")
|
||
log.error("or run with --no-exif to analyze into the DB only.")
|
||
sys.exit(1)
|
||
|
||
# On every start, first flush photos that were already analyzed in a previous
|
||
# run but never had their EXIF written (e.g. analyzed before exiftool was
|
||
# installed). This writes the stored DB data into those images up front,
|
||
# independent of — and before — the long analysis queue.
|
||
run_exif_write(args, conn)
|
||
|
||
if not args.exif_only:
|
||
run_analysis(args, conn, client) # analyzes pending; writes EXIF inline
|
||
run_exif_write(args, conn) # mop up any inline write failures
|
||
|
||
write_throttle_summary() # append run-summary line to throttle_events.jsonl
|
||
print_stats(conn, library)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|