initial commit

sprite_tagger.py

@@ -0,0 +1,423 @@
+#!/opt/homebrew/bin/python3.13
+"""Sprite frame tagger — navigate frames, assign labels, copy to labeled dirs.
+
+Usage:
+    sprite_tagger.py <frames_dir> <label1> [label2 ...] [--out <output_dir>]
+
+Example:
+    sprite_tagger.py ./sprites idle walk run jump
+
+Keys:
+    ← →         previous / next frame
+    1-9         assign label and advance
+    N           add a new label
+    Space       play/pause all frames
+    P           play/pause only the current frame's label (looping)
+    [ ]         slower / faster playback
+    X           clear label on current frame
+    Enter       export labeled frames to output dir
+"""
+
+import sys
+import shutil
+import tkinter as tk
+from tkinter import messagebox
+from pathlib import Path
+from PIL import Image, ImageDraw, ImageTk
+
+import labelstore
+import playback
+
+IMAGE_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp', '.tga'}
+
+BG        = '#1a1a2e'
+BG_PANEL  = '#16213e'
+ACCENT    = '#4ecca3'
+FG        = '#e0e0e0'
+FG_DIM    = '#556677'
+BTN_BG    = '#0f3460'
+BTN_ACT   = '#e94560'
+FONT_MONO = ('Courier', 11)
+FONT_TINY = ('Courier', 9)
+
+CANVAS_W  = 860
+CANVAS_H  = 620
+FPS_DEFAULT = 10
+FPS_MIN     = 1
+FPS_MAX     = 30
+
+
+def load_frames(directory: Path) -> list:
+    return sorted(
+        p for p in directory.iterdir()
+        if p.suffix.lower() in IMAGE_EXTENSIONS
+    )
+
+
+class SpriteTagger:
+    def __init__(self, root: tk.Tk, frames: list, labels: list, output_dir: Path,
+                 store_path: Path, assignments: dict = None):
+        self.root       = root
+        self.frames     = frames
+        self.labels     = labels
+        self.output_dir = output_dir
+        self.store_path = store_path
+        self.current    = 0
+        self.assignments: dict = assignments or {}   # frame index → label name
+        self._photo     = None        # ImageTk ref — must not be GC'd
+        self._playing   = False
+        self._fps       = FPS_DEFAULT
+        self._play_job  = None
+        self._preview_label = None    # None = play all; label str = preview that label
+
+        self._build_ui()
+        self._bind_keys()
+        self._show_frame()
+
+    # ── UI construction ──────────────────────────────────────────────────────
+
+    def _build_ui(self):
+        self.root.title('Sprite Tagger')
+        self.root.configure(bg=BG)
+        self.root.resizable(True, True)
+
+        # canvas
+        self.canvas = tk.Canvas(
+            self.root, bg=BG_PANEL,
+            width=CANVAS_W, height=CANVAS_H,
+            highlightthickness=1, highlightbackground='#2a2a4e',
+        )
+        self.canvas.pack(fill=tk.BOTH, expand=True, padx=10, pady=(10, 0))
+        self.canvas.bind('<Configure>', lambda _e: self._show_frame())
+
+        # status row
+        status = tk.Frame(self.root, bg=BG)
+        status.pack(fill=tk.X, padx=10, pady=4)
+
+        self._var_info  = tk.StringVar()
+        self._var_label = tk.StringVar()
+        self._var_fps   = tk.StringVar()
+
+        tk.Label(status, textvariable=self._var_info,  bg=BG, fg=FG,     font=FONT_MONO).pack(side=tk.LEFT)
+        tk.Label(status, textvariable=self._var_fps,   bg=BG, fg=FG_DIM, font=FONT_TINY).pack(side=tk.LEFT, padx=10)
+        tk.Label(status, textvariable=self._var_label, bg=BG, fg=ACCENT,  font=(*FONT_MONO, 'bold')).pack(side=tk.RIGHT)
+
+        # label buttons row
+        self._btn_row = tk.Frame(self.root, bg=BG)
+        self._btn_row.pack(fill=tk.X, padx=10, pady=(0, 4))
+
+        self._label_btns = {}
+        for label in self.labels:
+            self._add_label_button(label)
+
+        self._make_utility_buttons()
+
+        # hint row
+        tk.Label(
+            self.root,
+            text='← → navigate   1-9 label+advance   N new label   Space play all   P preview label   [ ] speed   X clear   ↵ export',
+            bg=BG, fg=FG_DIM, font=FONT_TINY,
+        ).pack(pady=(0, 6))
+
+    def _add_label_button(self, label: str):
+        idx = len(self._label_btns)
+        key = str(idx + 1) if idx < 9 else '-'
+        # tk.Label used instead of tk.Button — macOS ignores fg/bg on native buttons
+        lbl = tk.Label(
+            self._btn_row, text=f'[{key}] {label}',
+            bg=BTN_BG, fg='white',
+            padx=8, pady=4, font=FONT_MONO,
+            cursor='hand2', relief=tk.FLAT,
+        )
+        lbl.pack(side=tk.LEFT, padx=2)
+        lbl.bind('<Button-1>', lambda _e, l=label: self._assign(l))
+        lbl.bind('<Enter>',    lambda _e, w=lbl: w.configure(bg='#1a5a9a'))
+        lbl.bind('<Leave>',    lambda _e, w=lbl: self._restore_btn_color(w, label))
+        self._label_btns[label] = lbl
+
+    def _make_utility_buttons(self):
+        self._make_action_label(self._btn_row, '[X] clear',   '#5a1a3a', '#ff6688', self._clear,              side=tk.LEFT,  padx=(8, 2))
+        self._make_action_label(self._btn_row, '[N] + label', '#3a3a1a', '#dddd44', self._prompt_new_label,   side=tk.LEFT,  padx=2)
+        self._make_action_label(self._btn_row, '[↵] export',  '#0a3a1a', ACCENT,    self._export,             side=tk.RIGHT, padx=2, bold=True)
+
+    def _make_action_label(self, parent, text, bg, hover_bg, command, side, padx, bold=False):
+        font = (*FONT_MONO, 'bold') if bold else FONT_MONO
+        w = tk.Label(parent, text=text, bg=bg, fg='white', padx=8, pady=4, font=font, cursor='hand2', relief=tk.FLAT)
+        w.pack(side=side, padx=padx)
+        w.bind('<Button-1>', lambda _e: command())
+        w.bind('<Enter>',    lambda _e: w.configure(bg=hover_bg, fg='black'))
+        w.bind('<Leave>',    lambda _e: w.configure(bg=bg,       fg='white'))
+
+    def _restore_btn_color(self, widget: tk.Label, label: str):
+        current_label = self.assignments.get(self.current)
+        if label == current_label:
+            widget.configure(bg=ACCENT, fg='black')
+        else:
+            widget.configure(bg=BTN_BG, fg='white')
+
+    def _bind_keys(self):
+        self.root.bind('<Right>',  lambda _e: self._navigate(1))
+        self.root.bind('<Left>',   lambda _e: self._navigate(-1))
+        self.root.bind('<Return>', lambda _e: self._export())
+        self.root.bind('<space>',  lambda _e: self._toggle_play())
+        self.root.bind('p',        lambda _e: self._toggle_preview())
+        self.root.bind('P',        lambda _e: self._toggle_preview())
+        self.root.bind('[',        lambda _e: self._change_fps(-2))
+        self.root.bind(']',        lambda _e: self._change_fps(+2))
+        self.root.bind('x',        lambda _e: self._clear())
+        self.root.bind('X',        lambda _e: self._clear())
+        self.root.bind('n',        lambda _e: self._prompt_new_label())
+        self.root.bind('N',        lambda _e: self._prompt_new_label())
+        self._rebind_number_keys()
+
+    def _rebind_number_keys(self):
+        for i, label in enumerate(self.labels):
+            if i < 9:
+                self.root.bind(str(i + 1), lambda _e, l=label: self._assign(l))
+
+    # ── Frame display ────────────────────────────────────────────────────────
+
+    def _show_frame(self):
+        path  = self.frames[self.current]
+        img   = Image.open(path).convert('RGBA')
+
+        cw = self.canvas.winfo_width()  if self.canvas.winfo_width()  > 1 else CANVAS_W
+        ch = self.canvas.winfo_height() if self.canvas.winfo_height() > 1 else CANVAS_H
+
+        if img.width <= cw // 2 and img.height <= ch // 2:
+            # small sprite — zoom with nearest-neighbor to preserve pixels
+            zoom = min(cw // img.width, ch // img.height)
+            zoom = max(zoom, 1)
+            img = img.resize((img.width * zoom, img.height * zoom), Image.NEAREST)
+        elif img.width > cw or img.height > ch:
+            # large image — scale down with lanczos
+            scale = min(cw / img.width, ch / img.height)
+            img = img.resize((int(img.width * scale), int(img.height * scale)), Image.LANCZOS)
+
+        # checkerboard background for transparency
+        check_size = max(8, min(img.width, img.height) // 16)
+        bg = Image.new('RGBA', img.size, (40, 40, 60, 255))
+        draw = ImageDraw.Draw(bg)
+        for y in range(0, img.height, check_size):
+            for x in range(0, img.width, check_size):
+                if (x // check_size + y // check_size) % 2 == 0:
+                    draw.rectangle(
+                        [x, y, min(x + check_size, img.width) - 1, min(y + check_size, img.height) - 1],
+                        fill=(50, 50, 70, 255),
+                    )
+        bg.paste(img, mask=img.split()[3])
+
+        self._photo = ImageTk.PhotoImage(bg)
+        self.canvas.delete('all')
+        self.canvas.create_image(cw // 2, ch // 2, image=self._photo, anchor='center')
+
+        self._update_status()
+
+    def _update_status(self):
+        path    = self.frames[self.current]
+        labeled = len(self.assignments)
+        current_label = self.assignments.get(self.current)
+
+        self._var_info.set(f'{path.name}  [{self.current + 1}/{len(self.frames)}]  labeled: {labeled}')
+        self._var_label.set(f'▶ {current_label}' if current_label else '—')
+        play_sym = '⏸' if self._playing else '▶'
+        if self._preview_label is not None:
+            n = len(playback.frames_with_label(self.assignments, self._preview_label))
+            self._var_fps.set(f'{play_sym} preview: {self._preview_label} ({n}) ⟳  {self._fps}fps')
+        else:
+            self._var_fps.set(f'{play_sym} {self._fps}fps')
+
+        for label, btn in self._label_btns.items():
+            if label == current_label:
+                btn.configure(bg=ACCENT, fg='black')
+            else:
+                btn.configure(bg=BTN_BG, fg='white')
+
+    # ── Navigation ───────────────────────────────────────────────────────────
+
+    def _navigate(self, delta: int):
+        self.current = (self.current + delta) % len(self.frames)
+        self._show_frame()
+
+    # ── Labeling ─────────────────────────────────────────────────────────────
+
+    def _assign(self, label: str):
+        self.assignments[self.current] = label
+        self._save()
+        self._navigate(1)
+
+    def _clear(self):
+        self.assignments.pop(self.current, None)
+        self._save()
+        self._update_status()
+
+    def _save(self):
+        by_name = labelstore.to_name_keyed(self.assignments, self.frames)
+        labelstore.save_labels(self.store_path, by_name, self.labels)
+
+    # ── Playback ─────────────────────────────────────────────────────────────
+
+    def _toggle_play(self):
+        self._preview_label = None   # Space always plays all frames
+        if self._playing:
+            self._stop_play()
+        else:
+            self._start_play()
+
+    def _toggle_preview(self):
+        if self._preview_label is not None:
+            self._stop_play()
+            self._preview_label = None
+            self._update_status()
+            return
+        label = self.assignments.get(self.current)
+        if label is None:
+            self._flash('tag this frame first to preview its label')
+            return
+        self._preview_label = label
+        indices = playback.frames_with_label(self.assignments, label)
+        self.current = indices[0]
+        self._show_frame()
+        self._start_play()
+
+    def _start_play(self):
+        self._playing = True
+        delay = max(1, 1000 // self._fps)
+        self._play_job = self.root.after(delay, self._tick)
+        self._update_status()
+
+    def _stop_play(self):
+        self._playing = False
+        if self._play_job:
+            self.root.after_cancel(self._play_job)
+            self._play_job = None
+        self._update_status()
+
+    def _tick(self):
+        if not self._playing:
+            return
+        if self._preview_label is not None:
+            self._step_preview(1)
+        else:
+            self._navigate(1)
+        delay = max(1, 1000 // self._fps)
+        self._play_job = self.root.after(delay, self._tick)
+
+    def _step_preview(self, delta: int):
+        indices = playback.frames_with_label(self.assignments, self._preview_label)
+        if not indices:
+            self._stop_play()
+            self._preview_label = None
+            return
+        self.current = playback.next_in_cycle(indices, self.current, delta)
+        self._show_frame()
+
+    def _change_fps(self, delta: int):
+        self._fps = max(FPS_MIN, min(FPS_MAX, self._fps + delta))
+        self._update_status()
+
+    def _flash(self, msg: str):
+        self._var_label.set(msg)
+        self.root.after(1500, self._update_status)
+
+    # ── Label management ─────────────────────────────────────────────────────
+
+    def _prompt_new_label(self):
+        from tkinter import simpledialog
+        name = simpledialog.askstring('New label', 'Label name:', parent=self.root)
+        if not name:
+            return
+        name = name.strip().lower().replace(' ', '_')
+        if not name or name in self._label_btns:
+            return
+        self.labels.append(name)
+        self._add_label_button(name)
+        self._rebind_number_keys()
+        self._save()
+
+    # ── Export ───────────────────────────────────────────────────────────────
+
+    def _export(self):
+        if self._playing:
+            self._toggle_play()
+
+        if not self.assignments:
+            messagebox.showinfo('Nothing to export', 'No frames have been labeled yet.')
+            return
+
+        counts: dict = {}
+        for idx, label in self.assignments.items():
+            dest = self.output_dir / label
+            dest.mkdir(parents=True, exist_ok=True)
+            shutil.copy2(self.frames[idx], dest / self.frames[idx].name)
+            counts[label] = counts.get(label, 0) + 1
+
+        summary = '\n'.join(f'  {label}/: {n} frame{"s" if n != 1 else ""}' for label, n in sorted(counts.items()))
+        messagebox.showinfo('Export complete', f'Copied to:\n{self.output_dir}\n\n{summary}')
+
+
+# ── Entry point ───────────────────────────────────────────────────────────────
+
+def _parse_args(argv: list):
+    if len(argv) < 3:
+        print(__doc__)
+        sys.exit(1)
+
+    args = list(argv[1:])
+    output_dir = None
+
+    if '--out' in args:
+        i = args.index('--out')
+        if i + 1 >= len(args):
+            print('Error: --out requires a path argument')
+            sys.exit(1)
+        output_dir = Path(args[i + 1])
+        args = args[:i] + args[i + 2:]
+
+    frames_dir = Path(args[0])
+    labels = args[1:]
+
+    if not frames_dir.is_dir():
+        print(f'Error: {frames_dir} is not a directory')
+        sys.exit(1)
+    if not labels:
+        print('Error: provide at least one label')
+        sys.exit(1)
+
+    if output_dir is None:
+        output_dir = frames_dir / 'labeled'
+
+    return frames_dir, labels, output_dir
+
+
+def main():
+    frames_dir, labels, output_dir = _parse_args(sys.argv)
+
+    frames = load_frames(frames_dir)
+    if not frames:
+        print(f'No image files found in {frames_dir}')
+        print(f'Supported formats: {", ".join(sorted(IMAGE_EXTENSIONS))}')
+        sys.exit(1)
+
+    store_path = labelstore.sidecar_path(frames_dir)
+    saved_assignments, saved_labels = labelstore.load_labels(store_path)
+    for lbl in saved_labels:
+        if lbl not in labels:
+            labels.append(lbl)
+    assignments = labelstore.to_index_keyed(saved_assignments, frames)
+
+    print(f'Loaded {len(frames)} frames')
+    print(f'Labels : {", ".join(labels)}')
+    print(f'Output : {output_dir}')
+    if assignments:
+        print(f'Resumed {len(assignments)} labels from {store_path.name}')
+
+    root = tk.Tk()
+    root.geometry(f'{CANVAS_W + 20}x{CANVAS_H + 130}')
+    root.minsize(400, 400)
+    SpriteTagger(root, frames, labels, output_dir, store_path, assignments)
+    root.mainloop()
+
+
+if __name__ == '__main__':
+    main()