import os, math, random, wave, struct, zipfile, io, shutil from pathlib import Path import numpy as np ROOT = Path("./data/synth_dataset") SR = 16000 FPS = 30 BS_DIM = 24 # Change this to match --bs_dim in your train.py (e.g., 52) TRAIN_N = 12 # number of training clips VAL_N = 4 # number of validation clips DUR_RANGE = (2.5, 6.0) # seconds def write_wav_mono16(path, audio, sr=SR): """Write mono float32 [-1,1] to 16-bit PCM WAV.""" audio = np.clip(audio, -0.999, 0.999) pcm = (audio * 32767.0).astype(np.int16) with wave.open(str(path), 'wb') as wf: wf.setnchannels(1) wf.setsampwidth(2) wf.setframerate(sr) wf.writeframes(pcm.tobytes()) def smooth(x, k=5): if k <= 1: return x k = int(k) pad = (k-1)//2 xpad = np.pad(x, (pad, pad), mode='edge') kernel = np.ones(k)/k return np.convolve(xpad, kernel, mode='valid') def make_am_env(num_samples, bumps=4): """Create a smooth amplitude envelope (0..1).""" t = np.linspace(0, 1, num_samples, endpoint=False) env = np.zeros_like(t) for _ in range(bumps): c = random.uniform(0.1, 0.9) w = random.uniform(0.05, 0.25) env += np.exp(-0.5*((t-c)/w)**2) env = env / (env.max() + 1e-8) env = smooth(env, k=201) env = env / (env.max() + 1e-8) return env def make_f0(num_samples, fmin=110, fmax=220): """A gentle varying f0 contour in Hz.""" t = np.linspace(0, 1, num_samples, endpoint=False) base = random.uniform(fmin, fmax) vibrato = 2.0 * np.sin(2*np.pi*(2.0)*t) # small vibrato trend = 10.0 * np.sin(2*np.pi*random.uniform(0.3,0.6)*t + random.uniform(0, 2*np.pi)) f0 = base + vibrato + trend f0 = np.clip(f0, fmin, fmax) return f0 def synth_audio(duration_sec): """Return audio float32 [-1,1].""" n = int(SR*duration_sec) env = make_am_env(n) f0 = make_f0(n) phase = 0.0 audio = np.zeros(n, dtype=np.float32) for i in range(n): phase += 2*math.pi*f0[i]/SR s = 0.6*np.sin(phase) + 0.2*np.sin(2*phase) + 0.1*np.sin(3*phase) audio[i] = env[i]*s # add small noise audio += 0.01*np.random.randn(n).astype(np.float32) # normalize safely audio /= max(0.8*np.max(np.abs(audio)), 1e-6) return audio def audio_energy_frames(audio, fps=FPS, sr=SR): """Frame-level energy [0,1] aligned to fps.""" T = int(len(audio)/sr*fps) if T <= 0: return np.zeros(1, dtype=np.float32) # compute RMS per frame win = int(sr/fps) rms = [] for i in range(T): s = i*win e = min((i+1)*win, len(audio)) seg = audio[s:e] if len(seg)==0: rms.append(0.0) else: rms.append(np.sqrt(np.mean(seg**2))) rms = np.array(rms, dtype=np.float32) # normalize to 0..1 if rms.max() > 0: rms = rms / rms.max() return smooth(rms, k=5) def make_blinks(T, length=7, interval=(50,160)): """Return a blink curve of length T with several spikes shaped like eye blinks.""" eye1 = np.array([ 0.36537236,0.95023572,0.95593375,0.91671562,0.36725610,0.11911326,0.02535750], dtype=np.float32) eye2 = np.array([0.23477617,0.90995198,0.94475806,0.77786213,0.19107167,0.23543704,0.08916393], dtype=np.float32) eye3 = np.array([0.87004077,0.94983369,0.94941854,0.69591165,0.19107167,0.07257628,0.00710890], dtype=np.float32) eye4 = np.array([0.00030799,0.55670142,0.95265675,0.94234562,0.42585719,0.14833522,0.01765949], dtype=np.float32) templates = [eye1, eye2, eye3, eye4] curve = np.zeros(T, dtype=np.float32) i = random.randint(0, max(1, interval[0]-1)) while i < T - length: tpl = random.choice(templates) curve[i:i+length] = np.maximum(curve[i:i+length], tpl[:min(length, T-i)]) i += random.randint(interval[0], interval[1]) return np.clip(curve, 0, 1) def synth_bs(audio, bs_dim=BS_DIM, fps=FPS, sr=SR): """Create (T, bs_dim) with plausible correlations to audio energy + random smooth motion.""" energy = audio_energy_frames(audio, fps=fps, sr=sr) # (T,) T = len(energy) bs = np.zeros((T, bs_dim), dtype=np.float32) # Correlate "jawOpen" like dim with energy jaw_idx = min(24, bs_dim-1) bs[:, jaw_idx] = np.clip(energy * 1.2, 0, 1) # Add smiles correlated with smoothed energy trend if bs_dim > 45: smile_l = 43 # mouthSmileLeft in the reference list smile_r = 44 # mouthSmileRight trend = smooth(energy, k=31) bs[:, smile_l] = np.clip(0.6*trend, 0, 1) bs[:, smile_r] = np.clip(0.6*trend, 0, 1) # Add blinks for left/right eyes if available if bs_dim > 9: blink = make_blinks(T) bs[:, 8] = np.maximum(bs[:, 8], blink) # eyeBlinkLeft bs[:, 9] = np.maximum(bs[:, 9], blink) # eyeBlinkRight # Fill the rest with smooth low-amplitude motions for d in range(bs_dim): if d in (jaw_idx, 8, 9, 43, 44): continue noise = smooth(np.random.randn(T).astype(np.float32), k=21) noise = (noise - noise.min()) / (noise.max() - noise.min() + 1e-8) # 0..1 bs[:, d] = np.clip(0.15*noise, 0, 1) return bs def build_split(split_name, n_samples): audio_dir = ROOT / split_name / "audio" bs_dir = ROOT / split_name / "bs" audio_dir.mkdir(parents=True, exist_ok=True) bs_dir.mkdir(parents=True, exist_ok=True) for i in range(n_samples): dur = random.uniform(*DUR_RANGE) audio = synth_audio(dur) bs = synth_bs(audio, bs_dim=BS_DIM) stem = f"{split_name}_{i:03d}" wav_path = audio_dir / f"{stem}.wav" bs_path = bs_dir / f"{stem}.npy" write_wav_mono16(wav_path, audio, sr=SR) np.save(bs_path, bs, allow_pickle=False) return audio_dir, bs_dir # Clean previous and generate if ROOT.exists(): shutil.rmtree(ROOT) ROOT.mkdir(parents=True, exist_ok=True) train_paths = build_split("train", TRAIN_N) val_paths = build_split("val", VAL_N) # Zip the dataset for easy download zip_path = Path("./data/synth_dataset.zip") with zipfile.ZipFile(zip_path, 'w', compression=zipfile.ZIP_DEFLATED) as zf: for p in ROOT.rglob("*"): zf.write(p, p.relative_to(ROOT.parent)) print("Created dataset at:", ROOT) print("Train:", train_paths) print("Val:", val_paths) print("ZIP:", zip_path)