Olivier
3a45ba0f20
Implement a Python script that uses a local vision model (via Ollama) to automatically rename image files to descriptive snake_case names. - Add `rename-wallpapers.py` to handle image processing, resizing with ffmpeg, and querying the Ollama API. - Add `flake.nix` to provide a reproducible development environment including Python, Ollama, ffmpeg, and other necessary dependencies.
218 lines
7.5 KiB
Python
218 lines
7.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Rename image files to descriptive snake_case names using a local vision model.
|
|
|
|
The model (default: llava) is queried through a running Ollama server. This
|
|
script ensures the server is up and the model is pulled before processing.
|
|
"""
|
|
import argparse
|
|
import base64
|
|
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
|
|
IMG_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp"}
|
|
PROMPT = (
|
|
"You are naming an image file. Look at the image and respond with ONLY a "
|
|
"filename, nothing else. Rules: all lowercase; snake_case (underscores only, "
|
|
"no spaces, hyphens, or special characters); start with 'a_' or 'an_'; describe "
|
|
"the main subject, its state, and the background; 5 to 10 words; end with the "
|
|
"correct extension (.jpg/.png/.jpeg). Example: a_snowy_mountain_with_grey_sky.jpg"
|
|
)
|
|
|
|
|
|
def run(cmd, **kw):
|
|
return subprocess.run(cmd, capture_output=True, text=True, **kw)
|
|
|
|
|
|
def ollama_ready():
|
|
r = run(["curl", "-s", "-o", "/dev/null", "-w", "%{http_code}",
|
|
"--max-time", "3", "http://localhost:11434/api/tags"])
|
|
return r.returncode == 0 and r.stdout.strip() == "200"
|
|
|
|
|
|
def ensure_server(timeout=90):
|
|
if ollama_ready():
|
|
return
|
|
print("Starting ollama server...", flush=True)
|
|
proc = subprocess.Popen(["ollama", "serve"],
|
|
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
|
deadline = time.time() + timeout
|
|
while time.time() < deadline:
|
|
if ollama_ready():
|
|
return
|
|
time.sleep(1)
|
|
raise SystemExit("ollama server did not become ready in time")
|
|
|
|
|
|
def ensure_model(model, timeout=600):
|
|
listing = run(["ollama", "list"], timeout=30)
|
|
have = set()
|
|
for line in listing.stdout.splitlines()[1:]:
|
|
have.add(line.split()[0].replace(":latest", ""))
|
|
if model.replace(":latest", "") in have:
|
|
return
|
|
print(f"Pulling model '{model}' (one-time download)...", flush=True)
|
|
r = run(["ollama", "pull", model], timeout=timeout)
|
|
if r.returncode != 0:
|
|
raise SystemExit("failed to pull model:\n" + r.stderr)
|
|
|
|
|
|
def resize(path):
|
|
out = tempfile.mktemp(suffix=".jpg")
|
|
run(["ffmpeg", "-y", "-loglevel", "error", "-i", path,
|
|
"-vf", "scale='min(768,iw)':'min(768,ih)'", "-q:v", "4", out],
|
|
timeout=120)
|
|
if not os.path.exists(out):
|
|
# fall back to the original file if ffmpeg chokes
|
|
return path
|
|
return out
|
|
|
|
|
|
def describe(path, model):
|
|
r = resize(path)
|
|
with open(r, "rb") as f:
|
|
b64 = base64.b64encode(f.read()).decode()
|
|
payload = {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": PROMPT, "images": [b64]}],
|
|
"stream": False,
|
|
"options": {"temperature": 0.2},
|
|
}
|
|
pf = tempfile.NamedTemporaryFile("w", suffix=".json", delete=False)
|
|
json.dump(payload, pf)
|
|
pf.close()
|
|
last = None
|
|
for _ in range(3):
|
|
try:
|
|
p = run(["curl", "-s", "http://localhost:11434/api/chat",
|
|
"-H", "Content-Type: application/json", "-d", f"@{pf.name}"],
|
|
timeout=240)
|
|
return json.loads(p.stdout)["message"]["content"].strip()
|
|
except Exception as e: # noqa: BLE001
|
|
last = e
|
|
raise last
|
|
|
|
|
|
def sanitize(model_out, ext):
|
|
m = re.search(r"[a-z0-9_]+\.(?:jpg|jpeg|png|gif|webp)", model_out.lower())
|
|
name = m.group(0) if m else re.sub(r"[^a-z0-9_]", "_", model_out.lower())
|
|
stem, _ = os.path.splitext(name)
|
|
stem = re.sub(r"[^a-z0-9_]", "_", stem)
|
|
stem = re.sub(r"_+", "_", stem).strip("_")
|
|
if not (stem.startswith("a_") or stem.startswith("an_")):
|
|
stem = "a_" + stem
|
|
return stem + ext.lower()
|
|
|
|
|
|
def already_named(fn):
|
|
s = fn.lower()
|
|
return s.startswith("a_") or s.startswith("an_") or s == "retro_mountain_red.jpg"
|
|
|
|
|
|
def find_images(directory, recursive):
|
|
found = []
|
|
walker = os.walk(directory) if recursive else [(directory, [], os.listdir(directory))]
|
|
for root, _, files in walker:
|
|
for fn in files:
|
|
if os.path.splitext(fn)[1].lower() in IMG_EXTS:
|
|
found.append(os.path.join(root, fn))
|
|
return sorted(found)
|
|
|
|
|
|
def pick_with_fzf(paths):
|
|
# Feed full paths but display only the basename (--with-nth); fzf still
|
|
# outputs the full path, and the preview renders an image thumbnail.
|
|
by_full = {p: p for p in paths}
|
|
by_base = {os.path.basename(p): p for p in paths}
|
|
preview_cmd = os.environ.get(
|
|
"RENAME_PREVIEW_CMD",
|
|
'chafa --size 50x32 --color-space rgb {}')
|
|
r = run(["fzf", "--multi", "--prompt", "pick images to rename> ",
|
|
"--delimiter", "/", "--with-nth", "-1",
|
|
"--preview", preview_cmd,
|
|
"--preview-window", "right:50%:wrap",
|
|
"--border", "--height", "100%"],
|
|
input="\n".join(sorted(paths)), timeout=None)
|
|
if r.returncode != 0:
|
|
return []
|
|
chosen = [line for line in r.stdout.splitlines() if line]
|
|
return [by_full.get(c, by_base.get(os.path.basename(c), c)) for c in chosen]
|
|
|
|
|
|
def process(path, model, force, dry_run):
|
|
fn = os.path.basename(path)
|
|
ext = os.path.splitext(fn)[1].lower()
|
|
if not force and already_named(fn):
|
|
print(f"SKIP {fn}", flush=True)
|
|
return None
|
|
try:
|
|
raw = describe(path, model)
|
|
new = sanitize(raw, ext)
|
|
except Exception as e: # noqa: BLE001
|
|
print(f"FAIL {fn}: {e}", flush=True)
|
|
return None
|
|
target = path if force else os.path.join(os.path.dirname(path), new)
|
|
i = 1
|
|
while os.path.exists(target) and os.path.abspath(target) != os.path.abspath(path):
|
|
target = os.path.join(os.path.dirname(path),
|
|
f"{os.path.splitext(new)[0]}_{i}{ext}")
|
|
i += 1
|
|
base = os.path.basename(target)
|
|
if dry_run:
|
|
print(f"{fn} => {base} (raw: {raw})", flush=True)
|
|
else:
|
|
shutil.move(path, target)
|
|
print(f"{fn} => {base}", flush=True)
|
|
return base
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(
|
|
description="Rename images to descriptive snake_case names via a local vision model.")
|
|
ap.add_argument("files", nargs="*", help="image files to rename")
|
|
ap.add_argument("--dir", default=None, help="directory to scan (default: current dir)")
|
|
ap.add_argument("--recursive", action="store_true", help="scan directories recursively")
|
|
ap.add_argument("--pick", action="store_true", help="interactively pick with fzf")
|
|
ap.add_argument("--model", default=os.environ.get("RENAME_MODEL", "llava"),
|
|
help="ollama vision model (default: llava)")
|
|
ap.add_argument("--force", action="store_true", help="rename even already-named files")
|
|
ap.add_argument("--dry-run", action="store_true", help="print names without renaming")
|
|
args = ap.parse_args()
|
|
|
|
ensure_server()
|
|
ensure_model(args.model)
|
|
|
|
paths = list(args.files)
|
|
if args.dir:
|
|
paths += find_images(args.dir, args.recursive)
|
|
if not paths:
|
|
paths = find_images(args.dir or ".", args.recursive)
|
|
# normalise and dedupe
|
|
paths = sorted({os.path.abspath(p) for p in paths})
|
|
|
|
if args.pick:
|
|
paths = pick_with_fzf(paths)
|
|
if not paths:
|
|
print("Nothing selected.", flush=True)
|
|
return
|
|
|
|
if not paths:
|
|
print("No images found.", flush=True)
|
|
return
|
|
|
|
done = 0
|
|
for p in paths:
|
|
if process(p, args.model, args.force, args.dry_run):
|
|
done += 1
|
|
print(f"\nDone. {done} image(s) renamed"
|
|
+ (" (dry-run)" if args.dry_run else "") + ".", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|