491 lines
16 KiB
Python
491 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""Postprocess Grok imagine_image_to_video clips into dense 2D sprites.
|
|
|
|
Pipeline steps (deterministic only — no creative generation):
|
|
extract → ffmpeg frames from mp4
|
|
clean → magenta flood-fill chroma + light despill
|
|
sample → even-index frame sets + feet/center normalize
|
|
process → extract + clean + sample in one shot
|
|
|
|
This skill is designed for Grok Build (imagine_text_to_image + imagine_image_to_video).
|
|
The script itself only needs ffmpeg, Pillow, and numpy.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import math
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from collections import deque
|
|
from pathlib import Path
|
|
from typing import Sequence
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
MAGENTA = np.array([255, 0, 255], dtype=np.float32)
|
|
|
|
|
|
def _ensure_dir(path: Path) -> Path:
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
return path
|
|
|
|
|
|
def _parse_counts(text: str) -> list[int]:
|
|
counts: list[int] = []
|
|
for part in text.split(","):
|
|
part = part.strip()
|
|
if not part:
|
|
continue
|
|
n = int(part)
|
|
if n < 1:
|
|
raise ValueError(f"frame count must be >= 1, got {n}")
|
|
counts.append(n)
|
|
if not counts:
|
|
raise ValueError("at least one frame count required")
|
|
return counts
|
|
|
|
|
|
def sample_indices(n_total: int, n_want: int) -> list[int]:
|
|
if n_total <= 0:
|
|
return []
|
|
if n_want >= n_total:
|
|
return list(range(n_total))
|
|
if n_want == 1:
|
|
return [0]
|
|
return [int(round(i * (n_total - 1) / (n_want - 1))) for i in range(n_want)]
|
|
|
|
|
|
def extract_frames(video: Path, out_dir: Path, fps: float = 0.0) -> list[Path]:
|
|
_ensure_dir(out_dir)
|
|
for old in out_dir.glob("frame_*.png"):
|
|
old.unlink()
|
|
ffmpeg = shutil.which("ffmpeg")
|
|
if not ffmpeg:
|
|
raise RuntimeError("ffmpeg not found on PATH. Install ffmpeg to extract video frames.")
|
|
pattern = str(out_dir / "frame_%04d.png")
|
|
cmd = [ffmpeg, "-y", "-i", str(video)]
|
|
if fps and fps > 0:
|
|
cmd += ["-vf", f"fps={fps}"]
|
|
else:
|
|
cmd += ["-vsync", "0"]
|
|
cmd.append(pattern)
|
|
proc = subprocess.run(cmd, capture_output=True, text=True)
|
|
if proc.returncode != 0:
|
|
raise RuntimeError("ffmpeg failed:\n" + (proc.stderr or proc.stdout or "unknown error"))
|
|
frames = sorted(out_dir.glob("frame_*.png"))
|
|
if not frames:
|
|
raise RuntimeError(f"no frames extracted into {out_dir}")
|
|
return frames
|
|
|
|
|
|
def _near_magenta_mask(rgb: np.ndarray, dist: float = 55.0) -> np.ndarray:
|
|
"""rgb: HxWx3 uint8 → bool mask of keyable magenta-ish pixels."""
|
|
f = rgb.astype(np.float32)
|
|
# Distance to pure magenta in RGB.
|
|
d = np.linalg.norm(f - MAGENTA, axis=2)
|
|
# Also catch bright pinks: high R+B, low G relative.
|
|
r, g, b = f[:, :, 0], f[:, :, 1], f[:, :, 2]
|
|
pinkish = (r > 160) & (b > 160) & (g < 140) & ((r + b) / 2 - g > 40)
|
|
return (d <= dist) | pinkish
|
|
|
|
|
|
def chroma_key_rgba(im: Image.Image, dist: float = 55.0) -> Image.Image:
|
|
"""Flood-fill magenta from corners, despill edges, return RGBA."""
|
|
rgba = im.convert("RGBA")
|
|
arr = np.array(rgba)
|
|
rgb = arr[:, :, :3]
|
|
h, w = rgb.shape[:2]
|
|
key = _near_magenta_mask(rgb, dist=dist)
|
|
|
|
visited = np.zeros((h, w), dtype=bool)
|
|
q: deque[tuple[int, int]] = deque()
|
|
for y, x in ((0, 0), (0, w - 1), (h - 1, 0), (h - 1, w - 1)):
|
|
if key[y, x]:
|
|
visited[y, x] = True
|
|
q.append((x, y))
|
|
# Also seed along edges where magenta is present.
|
|
for x in range(w):
|
|
for y in (0, h - 1):
|
|
if key[y, x] and not visited[y, x]:
|
|
visited[y, x] = True
|
|
q.append((x, y))
|
|
for y in range(h):
|
|
for x in (0, w - 1):
|
|
if key[y, x] and not visited[y, x]:
|
|
visited[y, x] = True
|
|
q.append((x, y))
|
|
|
|
while q:
|
|
x, y = q.popleft()
|
|
for nx, ny in ((x - 1, y), (x + 1, y), (x, y - 1), (x, y + 1)):
|
|
if 0 <= nx < w and 0 <= ny < h and not visited[ny, nx] and key[ny, nx]:
|
|
visited[ny, nx] = True
|
|
q.append((nx, ny))
|
|
|
|
out = arr.copy()
|
|
out[visited, 3] = 0
|
|
|
|
# Light despill on remaining near-magenta fringe (keep RGB, reduce alpha).
|
|
fringe = key & ~visited
|
|
if fringe.any():
|
|
# Pull toward less magenta and soften alpha.
|
|
fr = out[fringe].astype(np.float32)
|
|
r, g, b, a = fr[:, 0], fr[:, 1], fr[:, 2], fr[:, 3]
|
|
spill = np.maximum(0.0, (r + b) / 2.0 - g)
|
|
factor = np.clip(1.0 - spill / 180.0, 0.15, 1.0)
|
|
fr[:, 0] = np.clip(r - spill * 0.35, 0, 255)
|
|
fr[:, 2] = np.clip(b - spill * 0.35, 0, 255)
|
|
fr[:, 1] = np.clip(g + spill * 0.15, 0, 255)
|
|
fr[:, 3] = np.clip(a * factor, 0, 255)
|
|
out[fringe] = fr.astype(np.uint8)
|
|
|
|
# Fully transparent where alpha is 0.
|
|
out[out[:, :, 3] == 0, :3] = 0
|
|
return Image.fromarray(out, "RGBA")
|
|
|
|
|
|
def content_bbox(im: Image.Image, alpha_min: int = 32) -> tuple[int, int, int, int] | None:
|
|
arr = np.array(im.convert("RGBA"))
|
|
mask = arr[:, :, 3] > alpha_min
|
|
if not mask.any():
|
|
return None
|
|
ys, xs = np.where(mask)
|
|
return int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1
|
|
|
|
|
|
def normalize_sprite(
|
|
im: Image.Image,
|
|
cell: int = 128,
|
|
body_height: int = 100,
|
|
foot_y: int = 118,
|
|
anchor: str = "feet",
|
|
) -> Image.Image:
|
|
bb = content_bbox(im)
|
|
canvas = Image.new("RGBA", (cell, cell), (0, 0, 0, 0))
|
|
if not bb:
|
|
return canvas
|
|
crop = im.crop(bb)
|
|
cw, ch = crop.size
|
|
if ch <= 0 or cw <= 0:
|
|
return canvas
|
|
scale = body_height / float(ch)
|
|
nw = max(1, int(round(cw * scale)))
|
|
nh = max(1, int(round(ch * scale)))
|
|
if nw > cell - 4:
|
|
scale = (cell - 4) / float(cw)
|
|
nw = max(1, int(round(cw * scale)))
|
|
nh = max(1, int(round(ch * scale)))
|
|
if nh > cell - 4:
|
|
scale = (cell - 4) / float(ch)
|
|
nw = max(1, int(round(cw * scale)))
|
|
nh = max(1, int(round(ch * scale)))
|
|
resized = crop.resize((nw, nh), Image.Resampling.LANCZOS)
|
|
if anchor == "center":
|
|
x = (cell - nw) // 2
|
|
y = (cell - nh) // 2
|
|
else:
|
|
x = (cell - nw) // 2
|
|
y = foot_y - nh
|
|
if y < 0:
|
|
y = 0
|
|
if y + nh > cell:
|
|
y = max(0, cell - nh)
|
|
canvas.paste(resized, (x, y), resized)
|
|
return canvas
|
|
|
|
|
|
def clean_frames(
|
|
raw_dir: Path,
|
|
clean_dir: Path,
|
|
dist: float = 55.0,
|
|
) -> list[Path]:
|
|
_ensure_dir(clean_dir)
|
|
raws = sorted(raw_dir.glob("frame_*.png"))
|
|
if not raws:
|
|
raise RuntimeError(f"no raw frames in {raw_dir}")
|
|
outs: list[Path] = []
|
|
for i, path in enumerate(raws):
|
|
im = Image.open(path)
|
|
cleaned = chroma_key_rgba(im, dist=dist)
|
|
out = clean_dir / f"clean_{i:04d}.png"
|
|
cleaned.save(out)
|
|
outs.append(out)
|
|
if (i + 1) % 25 == 0 or i + 1 == len(raws):
|
|
print(f" cleaned {i + 1}/{len(raws)}")
|
|
return outs
|
|
|
|
|
|
def build_exports(
|
|
sprites: Sequence[Image.Image],
|
|
out_sprite_dir: Path,
|
|
tag: str,
|
|
n_frames: int,
|
|
gif_ms: int | None = None,
|
|
) -> dict:
|
|
_ensure_dir(out_sprite_dir)
|
|
sub = _ensure_dir(out_sprite_dir / tag) if tag else out_sprite_dir
|
|
paths = []
|
|
for i, sp in enumerate(sprites):
|
|
p = sub / f"sprite_{i + 1:02d}.png"
|
|
sp.save(p)
|
|
paths.append(str(p))
|
|
|
|
size = sprites[0].size[0]
|
|
strip = Image.new("RGBA", (size * len(sprites), size), (0, 0, 0, 0))
|
|
for i, sp in enumerate(sprites):
|
|
strip.paste(sp, (i * size, 0), sp)
|
|
strip_path = out_sprite_dir / f"run-strip-{n_frames}.png"
|
|
strip.save(strip_path)
|
|
|
|
cols = 8 if n_frames >= 16 else 4
|
|
rows = int(math.ceil(len(sprites) / cols))
|
|
grid = Image.new("RGBA", (size * cols, size * rows), (0, 0, 0, 0))
|
|
for i, sp in enumerate(sprites):
|
|
r, c = divmod(i, cols)
|
|
grid.paste(sp, (c * size, r * size), sp)
|
|
grid_path = out_sprite_dir / f"run-grid-{n_frames}.png"
|
|
grid.save(grid_path)
|
|
|
|
if gif_ms is None:
|
|
if n_frames >= 40:
|
|
gif_ms = 25
|
|
elif n_frames >= 20:
|
|
gif_ms = 40
|
|
elif n_frames >= 12:
|
|
gif_ms = 60
|
|
else:
|
|
gif_ms = 80
|
|
|
|
frames_gif = []
|
|
for sp in sprites:
|
|
bg = Image.new("RGBA", sp.size, (30, 30, 40, 255))
|
|
bg.paste(sp, (0, 0), sp)
|
|
frames_gif.append(bg.convert("P", palette=Image.ADAPTIVE, colors=255))
|
|
gif_path = out_sprite_dir / f"run-preview-{n_frames}.gif"
|
|
frames_gif[0].save(
|
|
gif_path,
|
|
save_all=True,
|
|
append_images=frames_gif[1:],
|
|
duration=gif_ms,
|
|
loop=0,
|
|
disposal=2,
|
|
)
|
|
|
|
# Legacy alias for 8-frame default
|
|
if n_frames == 8:
|
|
alias = out_sprite_dir / "run-preview.gif"
|
|
shutil.copy2(gif_path, alias)
|
|
|
|
return {
|
|
"count": n_frames,
|
|
"tag": tag,
|
|
"sprites": paths,
|
|
"strip": str(strip_path),
|
|
"grid": str(grid_path),
|
|
"gif": str(gif_path),
|
|
"gif_ms": gif_ms,
|
|
}
|
|
|
|
|
|
def sample_and_export(
|
|
clean_dir: Path,
|
|
out_dir: Path,
|
|
frame_counts: Sequence[int],
|
|
cell: int = 128,
|
|
body_height: int = 100,
|
|
foot_y: int = 118,
|
|
anchor: str = "feet",
|
|
) -> dict:
|
|
cleans = sorted(clean_dir.glob("clean_*.png"))
|
|
if not cleans:
|
|
raise RuntimeError(f"no cleaned frames in {clean_dir}")
|
|
sprite_dir = _ensure_dir(out_dir / "sprite")
|
|
n_total = len(cleans)
|
|
results = []
|
|
for n_want in frame_counts:
|
|
idxs = sample_indices(n_total, n_want)
|
|
sprites = []
|
|
for idx in idxs:
|
|
im = Image.open(cleans[idx]).convert("RGBA")
|
|
sprites.append(
|
|
normalize_sprite(
|
|
im,
|
|
cell=cell,
|
|
body_height=body_height,
|
|
foot_y=foot_y,
|
|
anchor=anchor,
|
|
)
|
|
)
|
|
tag = f"x{n_want}" if n_want != 8 else ""
|
|
# Always also write under xN for consistency when n!=8;
|
|
# for 8, write both root sprites and optional x8.
|
|
if n_want == 8:
|
|
# root-level sprite_01..08 for backwards compat
|
|
info = build_exports(sprites, sprite_dir, tag="", n_frames=n_want)
|
|
# also x8 folder
|
|
build_exports(sprites, sprite_dir, tag="x8", n_frames=n_want)
|
|
else:
|
|
info = build_exports(sprites, sprite_dir, tag=tag, n_frames=n_want)
|
|
info["indices"] = idxs
|
|
results.append(info)
|
|
print(f"exported {n_want} frames → {info['gif']}")
|
|
return {"total_clean": n_total, "sets": results}
|
|
|
|
|
|
def write_readme(out_dir: Path, meta: dict) -> None:
|
|
lines = [
|
|
"Video2dsprite output (Grok Build pipeline)",
|
|
"==========================================",
|
|
"base/ base still on #FF00FF",
|
|
"video/ imagine_image_to_video clip",
|
|
"frames-raw/ decoded frames",
|
|
"frames-clean/ chroma-keyed RGBA frames",
|
|
"sprite/ sampled normalized sprites + strips/grids/GIFs",
|
|
"pipeline-meta.json",
|
|
"",
|
|
"This folder was produced for Grok Build (imagine_text_to_image + imagine_image_to_video).",
|
|
"Codex/other agents cannot run the video step; they can still re-sample",
|
|
"existing frames with: python video2dsprite.py sample --clean-dir ...",
|
|
"",
|
|
json.dumps(meta, indent=2),
|
|
"",
|
|
]
|
|
(out_dir / "README.txt").write_text("\n".join(lines), encoding="utf-8")
|
|
|
|
|
|
def cmd_extract(args: argparse.Namespace) -> int:
|
|
frames = extract_frames(Path(args.video), Path(args.out_dir), fps=args.fps)
|
|
print(f"extracted {len(frames)} frames → {args.out_dir}")
|
|
return 0
|
|
|
|
|
|
def cmd_clean(args: argparse.Namespace) -> int:
|
|
outs = clean_frames(Path(args.raw_dir), Path(args.out_dir), dist=args.dist)
|
|
print(f"cleaned {len(outs)} frames → {args.out_dir}")
|
|
return 0
|
|
|
|
|
|
def cmd_sample(args: argparse.Namespace) -> int:
|
|
counts = _parse_counts(args.frame_counts)
|
|
meta = sample_and_export(
|
|
clean_dir=Path(args.clean_dir),
|
|
out_dir=Path(args.out_dir),
|
|
frame_counts=counts,
|
|
cell=args.cell_size,
|
|
body_height=args.body_height,
|
|
foot_y=args.foot_y,
|
|
anchor=args.anchor,
|
|
)
|
|
out = Path(args.out_dir)
|
|
full = {
|
|
"mode": "sample",
|
|
"clean_dir": str(Path(args.clean_dir).resolve()),
|
|
**meta,
|
|
}
|
|
(out / "pipeline-meta.json").write_text(json.dumps(full, indent=2), encoding="utf-8")
|
|
write_readme(out, full)
|
|
print("sample done")
|
|
return 0
|
|
|
|
|
|
def cmd_process(args: argparse.Namespace) -> int:
|
|
out = Path(args.out_dir)
|
|
raw_dir = out / "frames-raw"
|
|
clean_dir = out / "frames-clean"
|
|
video = Path(args.video)
|
|
if not video.is_file():
|
|
raise FileNotFoundError(video)
|
|
|
|
print(f"extract {video}")
|
|
frames = extract_frames(video, raw_dir, fps=args.fps)
|
|
print(f"clean {len(frames)} frames")
|
|
clean_frames(raw_dir, clean_dir, dist=args.dist)
|
|
counts = _parse_counts(args.frame_counts)
|
|
print(f"sample counts={counts}")
|
|
meta_sample = sample_and_export(
|
|
clean_dir=clean_dir,
|
|
out_dir=out,
|
|
frame_counts=counts,
|
|
cell=args.cell_size,
|
|
body_height=args.body_height,
|
|
foot_y=args.foot_y,
|
|
anchor=args.anchor,
|
|
)
|
|
meta = {
|
|
"skill": "video2dsprite",
|
|
"platform": "Grok Build (imagine_image_to_video required for generation step)",
|
|
"name": args.name,
|
|
"video": str(video.resolve()),
|
|
"out_dir": str(out.resolve()),
|
|
"raw_frames": len(frames),
|
|
"chroma_dist": args.dist,
|
|
"cell_size": args.cell_size,
|
|
"body_height": args.body_height,
|
|
"foot_y": args.foot_y,
|
|
"anchor": args.anchor,
|
|
**meta_sample,
|
|
}
|
|
(out / "pipeline-meta.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
|
|
write_readme(out, meta)
|
|
print("process done")
|
|
return 0
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
p = argparse.ArgumentParser(description="Video → dense 2D sprite postprocessor")
|
|
sub = p.add_subparsers(dest="command", required=True)
|
|
|
|
def add_common_sample(sp: argparse.ArgumentParser) -> None:
|
|
sp.add_argument("--frame-counts", default="8,16,24,48")
|
|
sp.add_argument("--cell-size", type=int, default=128)
|
|
sp.add_argument("--body-height", type=int, default=100)
|
|
sp.add_argument("--foot-y", type=int, default=118)
|
|
sp.add_argument("--anchor", choices=("feet", "center"), default="feet")
|
|
|
|
pe = sub.add_parser("extract", help="ffmpeg extract frames")
|
|
pe.add_argument("--video", required=True)
|
|
pe.add_argument("--out-dir", required=True)
|
|
pe.add_argument("--fps", type=float, default=0.0, help="0 = all frames")
|
|
pe.set_defaults(func=cmd_extract)
|
|
|
|
pc = sub.add_parser("clean", help="chroma-key raw frames")
|
|
pc.add_argument("--raw-dir", required=True)
|
|
pc.add_argument("--out-dir", required=True)
|
|
pc.add_argument("--dist", type=float, default=55.0)
|
|
pc.set_defaults(func=cmd_clean)
|
|
|
|
ps = sub.add_parser("sample", help="sample cleaned frames into sprite sets")
|
|
ps.add_argument("--clean-dir", required=True)
|
|
ps.add_argument("--out-dir", required=True)
|
|
add_common_sample(ps)
|
|
ps.set_defaults(func=cmd_sample)
|
|
|
|
pp = sub.add_parser("process", help="extract + clean + sample")
|
|
pp.add_argument("--video", required=True)
|
|
pp.add_argument("--out-dir", required=True)
|
|
pp.add_argument("--name", default="clip")
|
|
pp.add_argument("--fps", type=float, default=0.0)
|
|
pp.add_argument("--dist", type=float, default=55.0)
|
|
add_common_sample(pp)
|
|
pp.set_defaults(func=cmd_process)
|
|
|
|
return p
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
parser = build_parser()
|
|
args = parser.parse_args(argv)
|
|
try:
|
|
return int(args.func(args))
|
|
except Exception as exc: # noqa: BLE001 — CLI surface
|
|
print(f"error: {exc}", file=sys.stderr)
|
|
return 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|