mirror of
https://github.com/mr0xb/spritetagger.git
synced 2026-08-27 20:44:56 -04:00
423 lines
16 KiB
Python
Executable file
423 lines
16 KiB
Python
Executable file
#!/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()
|