feat: add wallpaper renaming tool and nix development environment
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.
This commit is contained in:
@@ -0,0 +1,217 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user