180 lines
6.2 KiB
Python
180 lines
6.2 KiB
Python
|
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)
|