File size: 2,654 Bytes
8c867d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | """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
|