1"""Spectrogram features of the dataset clips, cached on disk.
3Each clip becomes a log-magnitude spectrogram in dB relative to its own
4maximum, floored at -80 dB and scaled to [0, 1], over 0 to FMAX Hz. The level
5of the clip therefore does not enter (the WAVs are peak-normalized anyway).
6The frequency resolution matters more than the time resolution here: the
7sheet's low partials are tens of hertz apart and ring for seconds, so the
8window is long (N_FFT / 48000 = 171 ms) and the hop coarse (21 ms).
9"""
11import json
12import os
13import sys
14import wave
15from concurrent.futures import ProcessPoolExecutor
17import numpy as np
19FS = 48000
20N_FFT = 8192
21HOP = 1024
22FMAX = 6000.0
23DB_FLOOR = 80.0
24N_BINS = int(FMAX / FS * N_FFT) # 1024 bins of 5.86 Hz
25# Log-frequency variant: 48 bins per octave over 8 octaves from F_LO. On this
26# axis a change of the sheet's frequency scale (kappa / L^2) translates the
27# whole pattern of partials, which a convolutional network detects naturally.
28F_LO = 20.0
29BINS_PER_OCTAVE = 48
30N_LOGF = 8 * BINS_PER_OCTAVE
33def read_wav(path):
34 with wave.open(path) as w:
35 assert w.getnchannels() == 1 and w.getsampwidth() == 2 and w.getframerate() == FS, path
36 x = np.frombuffer(w.readframes(w.getnframes()), dtype="<i2").astype(np.float32) / 32768.0
37 return x
40def spectrogram(x):
41 """(N_BINS, n_frames) float16 in [0, 1]; 1 at the clip's loudest bin."""
42 win = np.hanning(N_FFT + 1)[:-1].astype(np.float32)
43 n_frames = 1 + (len(x) - N_FFT) // HOP
44 frames = np.lib.stride_tricks.as_strided(x, shape=(n_frames, N_FFT), strides=(x.strides[0] * HOP, x.strides[0]))
45 spec = np.abs(np.fft.rfft(frames * win, axis=1))[:, :N_BINS].T
46 db = 20 * np.log10(spec + 1e-9)
47 db -= db.max()
48 return np.clip(db / DB_FLOOR + 1, 0, 1).astype(np.float16)
51_FB = None
54def logf_filterbank():
55 """(N_LOGF, N_BINS) triangular filters centered on the log-spaced frequencies, at least one linear bin wide."""
56 global _FB
57 if _FB is None:
58 f_lin = np.arange(N_BINS) * FS / N_FFT
59 centers = F_LO * 2 ** (np.arange(N_LOGF) / BINS_PER_OCTAVE)
60 fb = np.zeros((N_LOGF, N_BINS), np.float32)
61 for k, fc in enumerate(centers):
62 half = max(fc * (2 ** (1 / BINS_PER_OCTAVE) - 1), FS / N_FFT)
63 w = np.clip(1 - np.abs(f_lin - fc) / half, 0, None)
64 fb[k] = w / max(w.sum(), 1e-9)
65 _FB = fb
66 return _FB
69def spectrogram_logf(x):
70 """(N_LOGF, n_frames) float16 in [0, 1] on a log-frequency axis."""
71 win = np.hanning(N_FFT + 1)[:-1].astype(np.float32)
72 n_frames = 1 + (len(x) - N_FFT) // HOP
73 frames = np.lib.stride_tricks.as_strided(x, shape=(n_frames, N_FFT), strides=(x.strides[0] * HOP, x.strides[0]))
74 power = np.abs(np.fft.rfft(frames * win, axis=1))[:, :N_BINS] ** 2
75 db = 10 * np.log10(logf_filterbank() @ power.T + 1e-18)
76 db -= db.max()
77 return np.clip(db / DB_FLOOR + 1, 0, 1).astype(np.float16)
80def _one(args):
81 clip_path, out_path, kind = args
82 np.save(out_path, (spectrogram_logf if kind == "logf" else spectrogram)(read_wav(clip_path)))
83 return out_path
86def load_index(data_dir):
87 rows = []
88 with open(os.path.join(data_dir, "index.jsonl")) as f:
89 for line in f:
90 if line.strip():
91 rows.append(json.loads(line))
92 # Keep only clips whose WAV is present (the index may run ahead of a copy).
93 return [r for r in rows if os.path.exists(os.path.join(data_dir, "clips", r["id"] + ".wav"))]
96def build_cache(data_dir, cache_dir, workers=8, kind="lin"):
97 """Computes the spectrograms (kind "lin" or "logf") of every indexed clip not yet cached; returns the index rows and the stacked features."""
98 os.makedirs(cache_dir, exist_ok=True)
99 rows = load_index(data_dir)
100 todo = []
101 for r in rows:
102 out = os.path.join(cache_dir, r["id"] + ".npy")
103 if not os.path.exists(out):
104 todo.append((os.path.join(data_dir, "clips", r["id"] + ".wav"), out, kind))
105 if todo:
106 with ProcessPoolExecutor(workers) as ex:
107 for _ in ex.map(_one, todo, chunksize=4):
108 pass
109 feats = np.stack([np.load(os.path.join(cache_dir, r["id"] + ".npy")) for r in rows])
110 return rows, feats
113if __name__ == "__main__":
114 data_dir, cache_dir = sys.argv[1], sys.argv[2]
115 rows, feats = build_cache(data_dir, cache_dir)
116 print(len(rows), "clips, features", feats.shape, feats.dtype)