| """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
|
|
|
|
|
|
|
| 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}
|
| BACKGROUNDS = ["black", "white", "gray", "navy", "gradient"]
|
|
|
| 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"}
|
|
|
|
|
| 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
|
|
|
|
|
|
|
|
|
| @dataclass
|
| class ClipConfig:
|
| frames: int = 16
|
| size: int = 48
|
| p_two_objects: float = 0.3
|
| p_mention_bg: float = 0.85
|
| supersample: int = 2
|
|
|
| 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
|
|
|
| @classmethod
|
| 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
|
|
|
|
|
|
|
|
|
| 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]
|
|
|
|
|
| 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
|
|
|
|
|
|
|
|
|
| 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)
|
|
|
