| |
| """GCRN 依赖极简的音频预处理：16kHz 单声道 WAV <-> STFT [1,2,401,161]。 |
| |
| 与 GCRN.AXERA 的 board_inference/audio.py 保持一致（hamming sym 窗 + numpy rfft）。 |
| """ |
| from __future__ import annotations |
|
|
| import wave |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| SAMPLE_RATE = 16000 |
| CHUNK_SAMPLES = 4 * SAMPLE_RATE |
| N_FFT = 320 |
| HOP_LENGTH = 160 |
| WIN_LENGTH = 320 |
|
|
|
|
| def read_wav(path: str | Path) -> tuple[np.ndarray, int]: |
| with wave.open(str(path), "rb") as source: |
| if source.getnchannels() != 1 or source.getsampwidth() != 2: |
| raise ValueError(f"expected mono 16-bit PCM WAV: {path}") |
| sample_rate = source.getframerate() |
| samples = np.frombuffer( |
| source.readframes(source.getnframes()), dtype=np.int16 |
| ).copy() |
| return samples, sample_rate |
|
|
|
|
| def write_wav(path: str | Path, samples: np.ndarray, sample_rate: int) -> None: |
| values = np.asarray(samples, dtype=np.float32).reshape(-1) |
| pcm = np.clip(values, -1.0, 1.0) |
| pcm = np.rint(pcm * 32767.0).astype(np.int16) |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with wave.open(str(path), "wb") as output: |
| output.setnchannels(1) |
| output.setsampwidth(2) |
| output.setframerate(sample_rate) |
| output.writeframes(pcm.tobytes()) |
|
|
|
|
| def _symmetric_hamming(size: int) -> np.ndarray: |
| index = np.arange(size, dtype=np.float32) |
| return (0.54 - 0.46 * np.cos(2.0 * np.pi * index / (size - 1))).astype( |
| np.float32 |
| ) |
|
|
|
|
| def chunks(samples: np.ndarray) -> list[tuple[np.ndarray, int]]: |
| """int16 PCM -> 4 秒 float32 chunk 序列 [(chunk, valid_samples)]。""" |
| values = np.asarray(samples, dtype=np.float32).reshape(-1) / 32768.0 |
| if values.size == 0: |
| raise ValueError("input WAV is empty") |
| result = [] |
| for start in range(0, values.size, CHUNK_SAMPLES): |
| end = min(start + CHUNK_SAMPLES, values.size) |
| valid = end - start |
| chunk = np.zeros(CHUNK_SAMPLES, dtype=np.float32) |
| chunk[:valid] = values[start:end] |
| result.append((chunk, valid)) |
| return result |
|
|
|
|
| def stft(chunk: np.ndarray) -> np.ndarray: |
| """4 秒波形 -> [1, 2, 401, 161]（batch, real/imag, time, freq）。""" |
| signal = np.asarray(chunk, dtype=np.float32).reshape(-1) |
| if signal.size != CHUNK_SAMPLES: |
| raise ValueError(f"expected {CHUNK_SAMPLES} samples, got {signal.size}") |
| pad = N_FFT // 2 |
| signal = np.pad(signal, (pad, pad), mode="constant") |
| frame_count = 1 + (signal.size - WIN_LENGTH) // HOP_LENGTH |
| frames = np.lib.stride_tricks.sliding_window_view(signal, WIN_LENGTH)[ |
| ::HOP_LENGTH |
| ][:frame_count] |
| spectrum = np.fft.rfft(frames * _symmetric_hamming(WIN_LENGTH)[None, :], axis=1) |
| value = np.stack((spectrum.real, spectrum.imag), axis=0) |
| return value.astype(np.float32)[None, ...] |
|
|
|
|
| def istft(value: np.ndarray) -> np.ndarray: |
| """[1,2,401,161] 增强谱 -> 64000 个 float32 样本。""" |
| output = np.asarray(value, dtype=np.float32) |
| if output.shape != (1, 2, 401, 161): |
| raise ValueError(f"unexpected GCRN output shape: {output.shape}") |
| value = output[0] |
| spectrum = value[0].astype(np.float64) + 1j * value[1].astype(np.float64) |
| frames = np.fft.irfft(spectrum, n=N_FFT, axis=1) |
| window = _symmetric_hamming(WIN_LENGTH).astype(np.float64) |
| frames *= window[None, :] |
|
|
| output_length = N_FFT + HOP_LENGTH * (frames.shape[0] - 1) |
| waveform = np.zeros(output_length, dtype=np.float64) |
| window_sum = np.zeros(output_length, dtype=np.float64) |
| for index, frame in enumerate(frames): |
| start = index * HOP_LENGTH |
| waveform[start : start + N_FFT] += frame |
| window_sum[start : start + N_FFT] += window * window |
| valid = window_sum > np.finfo(np.float64).eps |
| waveform[valid] /= window_sum[valid] |
| pad = N_FFT // 2 |
| return waveform[pad:-pad].astype(np.float32) |
|
|
