Spaces:
Sleeping
Sleeping
| """Procedural captioned video clips: coloured shapes moving on simple backgrounds. | |
| Everything is generated on the fly from a seeded RNG, so there is no dataset on | |
| disk and every clip comes with ground-truth metadata (per-frame trajectory) that | |
| the evaluator can check generated videos against. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| from dataclasses import dataclass | |
| import numpy as np | |
| import torch | |
| # ---------------------------------------------------------------- vocabulary | |
| SHAPES = ["circle", "square", "triangle", "star", "ring", "cross"] | |
| COLORS = { | |
| "red": (230, 40, 40), | |
| "green": (40, 200, 60), | |
| "blue": (50, 90, 240), | |
| "yellow": (245, 220, 40), | |
| "cyan": (40, 220, 230), | |
| "magenta": (220, 50, 220), | |
| "white": (245, 245, 245), | |
| "orange": (250, 140, 20), | |
| } | |
| SIZES = {"small": 0.11, "medium": 0.16, "big": 0.23} # radius as a fraction of frame size | |
| BACKGROUNDS = ["black", "white", "gray", "navy", "gradient"] | |
| # motion name -> caption words | |
| MOTIONS = { | |
| "left": "moves left", | |
| "right": "moves right", | |
| "up": "moves up", | |
| "down": "moves down", | |
| "up-left": "moves up left", | |
| "up-right": "moves up right", | |
| "down-left": "moves down left", | |
| "down-right": "moves down right", | |
| "bounces": "bounces", | |
| "grows": "grows", | |
| "shrinks": "shrinks", | |
| "spins": "spins", | |
| "still": "stays still", | |
| } | |
| DIRECTIONS = { | |
| "left": (-1, 0), "right": (1, 0), "up": (0, -1), "down": (0, 1), | |
| "up-left": (-1, -1), "up-right": (1, -1), "down-left": (-1, 1), "down-right": (1, 1), | |
| } | |
| ROTATIONALLY_SYMMETRIC = {"circle", "ring"} # spinning these is invisible, so never caption it | |
| # Combos never rendered during training; used to test compositional generalisation. | |
| HELDOUT_SHAPE_COLOR = {("star", "red"), ("circle", "yellow"), ("triangle", "cyan"), ("square", "magenta")} | |
| HELDOUT_SHAPE_MOTION = {("ring", "down"), ("cross", "grows"), ("triangle", "bounces"), ("star", "left")} | |
| SPECIALS = ["<pad>", "<unk>", "<null>"] | |
| _WORDS = ["a", "and", "on", "background", "moves", "stays", "still", "left", "right", "up", "down", | |
| "bounces", "grows", "shrinks", "spins", *SIZES, *COLORS, *SHAPES, *BACKGROUNDS] | |
| VOCAB = SPECIALS + sorted(set(_WORDS)) | |
| PAD, UNK, NULL = 0, 1, 2 | |
| MAX_TOKENS = 24 | |
| class Tokenizer: | |
| def __init__(self, max_len: int = MAX_TOKENS): | |
| self.max_len = max_len | |
| self.index = {w: i for i, w in enumerate(VOCAB)} | |
| def encode(self, caption: str) -> list[int]: | |
| ids = [self.index.get(w, UNK) for w in caption.lower().split()] | |
| return (ids or [NULL])[: self.max_len] | |
| def __call__(self, captions: list[str]) -> torch.Tensor: | |
| out = torch.full((len(captions), self.max_len), PAD, dtype=torch.long) | |
| for i, c in enumerate(captions): | |
| ids = self.encode(c) | |
| out[i, : len(ids)] = torch.tensor(ids) | |
| return out | |
| def null(self, n: int) -> torch.Tensor: | |
| """Token rows for the unconditional (classifier-free guidance) branch.""" | |
| out = torch.full((n, self.max_len), PAD, dtype=torch.long) | |
| out[:, 0] = NULL | |
| return out | |
| # ---------------------------------------------------------------- sampling captions | |
| class ClipConfig: | |
| frames: int = 16 | |
| size: int = 48 | |
| p_two_objects: float = 0.3 | |
| p_mention_bg: float = 0.85 | |
| supersample: int = 2 | |
| # Curriculum: restrict what gets sampled (None = everything). | |
| shapes: tuple[str, ...] | None = None | |
| colors: tuple[str, ...] | None = None | |
| sizes: tuple[str, ...] | None = None | |
| motions: tuple[str, ...] | None = None | |
| backgrounds: tuple[str, ...] | None = None | |
| def simple(cls, **kw) -> "ClipConfig": | |
| """Stage-1 curriculum: one object, black background, 3 shapes, 4 colours, 5 motions.""" | |
| base = dict(p_two_objects=0.0, p_mention_bg=0.0, shapes=("circle", "square", "triangle"), | |
| colors=("red", "green", "blue", "yellow"), motions=("left", "right", "up", "down", "bounces"), | |
| backgrounds=("black",)) | |
| return cls(**{**base, **kw}) | |
| def _pick(rng: np.random.Generator, options, allowed): | |
| opts = [o for o in options if allowed is None or o in allowed] | |
| return opts[rng.integers(len(opts))] | |
| def is_heldout(meta: dict) -> bool: | |
| return any((o["shape"], o["color"]) in HELDOUT_SHAPE_COLOR or (o["shape"], o["motion"]) in HELDOUT_SHAPE_MOTION | |
| for o in meta["objects"]) | |
| def _clashes(color: str, bg: str) -> bool: | |
| return (color, bg) in {("white", "white"), ("blue", "navy")} | |
| def _random_object(rng: np.random.Generator, bg: str, cfg: ClipConfig) -> dict: | |
| while True: | |
| shape = _pick(rng, SHAPES, cfg.shapes) | |
| color = _pick(rng, COLORS, cfg.colors) | |
| motion = _pick(rng, MOTIONS, cfg.motions) | |
| size = _pick(rng, SIZES, cfg.sizes) | |
| if not _clashes(color, bg) and not (motion == "spins" and shape in ROTATIONALLY_SYMMETRIC): | |
| return {"shape": shape, "color": color, "size": size, "motion": motion} | |
| def _random_meta(rng: np.random.Generator, cfg: ClipConfig, split: str) -> dict: | |
| """split: 'train' (no held-out combos), 'heldout' (one object, held-out combo), 'any'.""" | |
| while True: | |
| bg = _pick(rng, BACKGROUNDS, cfg.backgrounds) | |
| if split == "heldout": | |
| obj = _random_object(rng, bg, cfg) | |
| if rng.random() < 0.5: | |
| obj["shape"], obj["color"] = sorted(HELDOUT_SHAPE_COLOR)[rng.integers(len(HELDOUT_SHAPE_COLOR))] | |
| else: | |
| obj["shape"], obj["motion"] = sorted(HELDOUT_SHAPE_MOTION)[rng.integers(len(HELDOUT_SHAPE_MOTION))] | |
| objects = [obj] | |
| else: | |
| n = 2 if rng.random() < cfg.p_two_objects else 1 | |
| objects = [_random_object(rng, bg, cfg) for _ in range(n)] | |
| meta = {"objects": objects, "bg": bg, "mention_bg": bool(rng.random() < cfg.p_mention_bg)} | |
| if any(_clashes(o["color"], bg) for o in objects): | |
| continue | |
| if split == "train" and is_heldout(meta): | |
| continue | |
| return meta | |
| def caption_of(meta: dict) -> str: | |
| parts = [] | |
| for o in meta["objects"]: | |
| size = "" if o["size"] == "medium" else o["size"] + " " | |
| parts.append(f"a {size}{o['color']} {o['shape']} {MOTIONS[o['motion']]}") | |
| cap = " and ".join(parts) | |
| if meta.get("mention_bg", True): | |
| cap += f" on a {meta['bg']} background" | |
| return cap | |
| def parse_caption(caption: str) -> dict: | |
| """Inverse of caption_of (without trajectories). Unmentioned size -> 'medium', bg -> None.""" | |
| words = caption.lower().split() | |
| bg = None | |
| if "on" in words: | |
| i = words.index("on") | |
| bg = words[i + 2] if i + 2 < len(words) else None | |
| words = words[:i] | |
| phrases, cur = [], [] | |
| for w in words: | |
| if w == "and": | |
| phrases.append(cur) | |
| cur = [] | |
| else: | |
| cur.append(w) | |
| phrases.append(cur) | |
| motion_by_words = {v: k for k, v in MOTIONS.items()} | |
| objects = [] | |
| for p in phrases: | |
| p = [w for w in p if w != "a"] | |
| size = next((w for w in p if w in SIZES), "medium") | |
| color = next((w for w in p if w in COLORS), None) | |
| shape = next((w for w in p if w in SHAPES), None) | |
| rest = " ".join(p[p.index(shape) + 1:]) if shape in p else "" | |
| objects.append({"shape": shape, "color": color, "size": size, "motion": motion_by_words.get(rest)}) | |
| return {"objects": objects, "bg": bg, "mention_bg": bg is not None} | |
| def sample_prompts(n: int, split: str = "train", seed: int = 0, cfg: ClipConfig | None = None) -> list[tuple[str, dict]]: | |
| rng = np.random.default_rng(seed) | |
| cfg = cfg or ClipConfig() | |
| out = [] | |
| for _ in range(n): | |
| meta = _random_meta(rng, cfg, split) | |
| out.append((caption_of(meta), meta)) | |
| return out | |
| # ---------------------------------------------------------------- rendering | |
| def _background(bg: str, S: int) -> np.ndarray: | |
| img = np.empty((S, S, 3), np.float32) | |
| if bg == "gradient": | |
| ramp = np.linspace(20, 120, S, dtype=np.float32)[:, None, None] | |
| img[:] = ramp * np.array([0.6, 0.8, 1.0], np.float32) | |
| else: | |
| img[:] = {"black": (0, 0, 0), "white": (250, 250, 250), "gray": (128, 128, 128), "navy": (15, 20, 80)}[bg] | |
| return img | |
| def _polygon_mask(px, py, verts) -> np.ndarray: | |
| """Even-odd point-in-polygon test, vectorised over pixels.""" | |
| inside = np.zeros(px.shape, bool) | |
| n = len(verts) | |
| for i in range(n): | |
| x1, y1 = verts[i] | |
| x2, y2 = verts[(i + 1) % n] | |
| crosses = (y1 > py) != (y2 > py) | |
| xint = (x2 - x1) * (py - y1) / (y2 - y1 + 1e-12) + x1 | |
| inside ^= crosses & (px < xint) | |
| return inside | |
| def _regular(n: int, r: float, angle: float, cx: float, cy: float, inner: float | None = None): | |
| pts = [] | |
| k = 2 * n if inner else n | |
| for i in range(k): | |
| a = angle - math.pi / 2 + 2 * math.pi * i / k | |
| rr = (inner if i % 2 else r) if inner else r | |
| pts.append((cx + rr * math.cos(a), cy + rr * math.sin(a))) | |
| return pts | |
| def _shape_mask(shape: str, px, py, cx, cy, r, angle) -> np.ndarray: | |
| dx, dy = px - cx, py - cy | |
| if shape == "circle": | |
| return dx * dx + dy * dy <= r * r | |
| if shape == "ring": | |
| d2 = dx * dx + dy * dy | |
| return (d2 <= r * r) & (d2 >= (0.55 * r) ** 2) | |
| if shape == "square": | |
| return _polygon_mask(px, py, _regular(4, r * 1.15, angle + math.pi / 4, cx, cy)) | |
| if shape == "triangle": | |
| return _polygon_mask(px, py, _regular(3, r * 1.2, angle, cx, cy)) | |
| if shape == "star": | |
| return _polygon_mask(px, py, _regular(5, r * 1.25, angle, cx, cy, inner=r * 0.5)) | |
| if shape == "cross": | |
| c, s = math.cos(angle), math.sin(angle) | |
| u, v = c * dx + s * dy, -s * dx + c * dy | |
| arm = 0.33 * r | |
| return ((np.abs(u) <= r) & (np.abs(v) <= arm)) | ((np.abs(v) <= r) & (np.abs(u) <= arm)) | |
| raise ValueError(shape) | |
| def _trajectory(rng: np.random.Generator, obj: dict, T: int) -> list[tuple[float, float, float, float]]: | |
| """Per-frame (cx, cy, r, angle) in normalised [0,1] frame coordinates.""" | |
| r0 = SIZES[obj["size"]] * rng.uniform(0.9, 1.1) | |
| angle0 = rng.uniform(0, 2 * math.pi) | |
| m = obj["motion"] | |
| ts = np.linspace(0, 1, T) | |
| if m in DIRECTIONS: | |
| dx, dy = DIRECTIONS[m] | |
| dist = rng.uniform(0.35, 0.6) / (math.sqrt(2) if dx and dy else 1) | |
| margin = r0 + 0.02 | |
| def start(d): | |
| if d > 0: | |
| return rng.uniform(margin, max(margin, 1 - margin - dist)) | |
| if d < 0: | |
| return rng.uniform(min(1 - margin, margin + dist), 1 - margin) | |
| return rng.uniform(margin, 1 - margin) | |
| x0, y0 = start(dx), start(dy) | |
| return [(x0 + dx * dist * t, y0 + dy * dist * t, r0, angle0) for t in ts] | |
| if m == "bounces": | |
| theta = rng.uniform(0, 2 * math.pi) | |
| speed = rng.uniform(1.4, 2.2) | |
| lo, hi = r0, 1 - r0 | |
| x, y = rng.uniform(lo, hi, size=2) | |
| vx, vy = speed * math.cos(theta), speed * math.sin(theta) | |
| out, dt = [], 1 / (T - 1) | |
| for _ in range(T): | |
| out.append((x, y, r0, angle0)) | |
| x, y = x + vx * dt, y + vy * dt | |
| if x < lo or x > hi: | |
| vx, x = -vx, float(np.clip(2 * np.clip(x, lo, hi) - x, lo, hi)) | |
| if y < lo or y > hi: | |
| vy, y = -vy, float(np.clip(2 * np.clip(y, lo, hi) - y, lo, hi)) | |
| return out | |
| margin = r0 * (1.6 if m in ("grows", "shrinks") else 1.0) + 0.02 | |
| cx, cy = rng.uniform(margin, 1 - margin, size=2) | |
| if m == "grows": | |
| return [(cx, cy, r0 * (0.5 + 1.1 * t), angle0) for t in ts] | |
| if m == "shrinks": | |
| return [(cx, cy, r0 * (1.6 - 1.1 * t), angle0) for t in ts] | |
| if m == "spins": | |
| spin = rng.choice([-1, 1]) * rng.uniform(0.6, 1.0) * math.pi | |
| return [(cx, cy, r0, angle0 + spin * t) for t in ts] | |
| return [(cx, cy, r0, angle0) for _ in ts] # still | |
| def render(meta: dict, rng: np.random.Generator, cfg: ClipConfig) -> np.ndarray: | |
| """Render meta (adds meta['objects'][i]['traj']) to uint8 [T, H, W, 3].""" | |
| T, S, ss = cfg.frames, cfg.size, cfg.supersample | |
| big = S * ss | |
| coords = (np.arange(big, dtype=np.float32) + 0.5) / big | |
| px, py = np.meshgrid(coords, coords) | |
| bg = _background(meta["bg"], big) | |
| for o in meta["objects"]: | |
| o["traj"] = _trajectory(rng, o, T) | |
| frames = np.empty((T, S, S, 3), np.uint8) | |
| for f in range(T): | |
| img = bg.copy() | |
| for o in meta["objects"]: | |
| cx, cy, r, ang = o["traj"][f] | |
| mask = _shape_mask(o["shape"], px, py, cx, cy, r, ang) | |
| img[mask] = COLORS[o["color"]] | |
| img = img.reshape(S, ss, S, ss, 3).mean(axis=(1, 3)) | |
| frames[f] = np.clip(img + 0.5, 0, 255).astype(np.uint8) | |
| return frames | |
| def make_clip(rng: np.random.Generator, cfg: ClipConfig | None = None, split: str = "train"): | |
| cfg = cfg or ClipConfig() | |
| meta = _random_meta(rng, cfg, split) | |
| video = render(meta, rng, cfg) | |
| return video, caption_of(meta), meta | |
| # ---------------------------------------------------------------- torch dataset | |
| def to_tensor(video: np.ndarray) -> torch.Tensor: | |
| """uint8 [T,H,W,3] -> float [3,T,H,W] in [-1,1].""" | |
| return torch.from_numpy(video).permute(3, 0, 1, 2).float().div(127.5).sub(1) | |
| class SynthVideoDataset(torch.utils.data.IterableDataset): | |
| """Infinite stream of (video float [3,T,H,W] in [-1,1], caption).""" | |
| def __init__(self, cfg: ClipConfig | None = None, seed: int = 0, include_heldout: bool = False): | |
| self.cfg = cfg or ClipConfig() | |
| self.seed = seed | |
| self.split = "any" if include_heldout else "train" | |
| def __iter__(self): | |
| info = torch.utils.data.get_worker_info() | |
| wid = info.id if info else 0 | |
| rng = np.random.default_rng([self.seed, wid]) | |
| while True: | |
| video, caption, _ = make_clip(rng, self.cfg, self.split) | |
| yield to_tensor(video), caption | |
| def collate(batch): | |
| videos, captions = zip(*batch) | |
| return torch.stack(videos), list(captions) | |