| """Synthetic media generators shared by tests, data building and benchmarks.""" |
| import numpy as np |
| from PIL import Image, ImageDraw |
|
|
| COLORS = {"red": (220, 40, 40), "blue": (40, 80, 220), "green": (40, 170, 60), "yellow": (230, 200, 30), |
| "purple": (140, 60, 180), "orange": (240, 130, 30)} |
| SHAPES = ["circle", "square", "triangle"] |
|
|
|
|
| def draw_shape(d, shape, color, cx, cy, r): |
| c = COLORS[color] |
| if shape == "circle": |
| d.ellipse([cx - r, cy - r, cx + r, cy + r], fill=c) |
| elif shape == "square": |
| d.rectangle([cx - r, cy - r, cx + r, cy + r], fill=c) |
| else: |
| d.polygon([(cx, cy - r), (cx - r, cy + r), (cx + r, cy + r)], fill=c) |
|
|
|
|
| def shapes_image(objs, size=384, bg=(245, 245, 240)): |
| """objs: list of (shape, color, cx, cy, r) in [0,1] coords for cx, cy and fraction r.""" |
| im = Image.new("RGB", (size, size), bg) |
| d = ImageDraw.Draw(im) |
| for shape, color, cx, cy, r in objs: |
| draw_shape(d, shape, color, cx * size, cy * size, r * size) |
| return im |
|
|
|
|
| def beeps(n, sr=16000, dur=4.0, freq=880.0, rng=None): |
| """n short tone beeps spread over `dur` seconds, plus light noise.""" |
| rng = rng or np.random.default_rng(0) |
| t = np.zeros(int(sr * dur), np.float32) |
| if n: |
| starts = np.linspace(0.3, dur - 0.5, n) + rng.uniform(-0.05, 0.05, n) |
| for s in starts: |
| a = int(s * sr); L = int(0.18 * sr) |
| tt = np.arange(L) / sr |
| t[a:a + L] += 0.5 * np.sin(2 * np.pi * freq * tt) * np.hanning(L) |
| t += rng.normal(0, 0.01, t.shape).astype(np.float32) |
| return t |
|
|
|
|
| def moving_video(shape, color, direction, n_frames=8, size=320, flashes=0, change_to=None, rng=None): |
| """A shape moving across the frame. flashes: number of frames where the background flashes white-yellow. |
| change_to: the colour the shape switches to halfway through (None = no change).""" |
| rng = rng or np.random.default_rng(0) |
| frames = [] |
| x0, y0 = rng.uniform(0.35, 0.65), rng.uniform(0.35, 0.65) |
| dx = {"left": (-1, 0), "right": (1, 0), "up": (0, -1), "down": (0, 1)}[direction] |
| span = 0.6 |
| flash_idx = set(rng.choice(np.arange(1, n_frames - 1), size=flashes, replace=False).tolist()) if flashes else set() |
| for i in range(n_frames): |
| f = i / (n_frames - 1) - 0.5 |
| cx = x0 + dx[0] * f * span; cy = y0 + dx[1] * f * span |
| cx = min(max(cx, 0.1), 0.9); cy = min(max(cy, 0.1), 0.9) |
| bg = (255, 250, 150) if i in flash_idx else (40, 40, 48) |
| col = change_to if (change_to and i >= n_frames // 2) else color |
| frames.append(shapes_image([(shape, col, cx, cy, 0.09)], size=size, bg=bg)) |
| return frames |
|
|
