| |
| """GCRN 语音增强推理：4 秒 chunk STFT -> axmodel -> ISTFT -> 拼接输出。""" |
| from __future__ import annotations |
|
|
| import json |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from .audio import ( |
| SAMPLE_RATE, |
| chunks, |
| istft, |
| read_wav, |
| stft, |
| write_wav, |
| ) |
| from .runtime import InferenceSession, first_output |
|
|
| INPUT_NAME = "stft_input" |
| OUTPUT_NAME = "stft_output" |
| INPUT_SHAPE = (1, 2, 401, 161) |
|
|
|
|
| class GCRNDenoiser: |
| """AX650 GCRN 语音增强器。默认 .axmodel + axengine。""" |
|
|
| def __init__(self, model_path: str | Path, backend: str | None = None): |
| self.session = InferenceSession(model_path, backend) |
| if len(self.session.inputs) != 1 or self.session.inputs[0].shape != INPUT_SHAPE: |
| raise RuntimeError(f"unexpected GCRN input metadata: {self.session.inputs}") |
| if len(self.session.outputs) != 1 or self.session.outputs[0].shape != INPUT_SHAPE: |
| raise RuntimeError(f"unexpected GCRN output metadata: {self.session.outputs}") |
| self.input_name = self.session.inputs[0].name |
|
|
| def _warmup(self) -> None: |
| feed = stft(np.zeros(64000, dtype=np.float32)) |
| first_output(self.session, self.session.run({self.input_name: feed})) |
|
|
| def enhance_chunk(self, chunk: np.ndarray) -> np.ndarray: |
| feed = stft(chunk) |
| outputs = self.session.run({self.input_name: feed}) |
| output = np.asarray(first_output(self.session, outputs), dtype=np.float32) |
| return istft(output) |
|
|
| def enhance(self, pcm: np.ndarray) -> np.ndarray: |
| """int16 PCM -> 增强后 float32 波形（与输入等长）。""" |
| self._warmup() |
| enhanced = [ |
| self.enhance_chunk(chunk) |
| for chunk, _valid in chunks(pcm) |
| ] |
| return np.concatenate(enhanced)[: pcm.size] |
|
|
| def enhance_file( |
| self, input_path: str | Path, output_path: str | Path |
| ) -> dict: |
| """输入 16kHz 单声道 WAV，输出增强 WAV + 性能报告 dict。""" |
| samples, sample_rate = read_wav(input_path) |
| if sample_rate != SAMPLE_RATE: |
| raise ValueError( |
| f"GCRN requires {SAMPLE_RATE} Hz audio, got {sample_rate}" |
| ) |
| start = time.perf_counter() |
| enhanced = self.enhance(samples) |
| elapsed = time.perf_counter() - start |
| write_wav(output_path, enhanced, sample_rate) |
| audio_seconds = samples.size / sample_rate |
| report = { |
| "backend": self.session.backend, |
| "model": str(self.session.path), |
| "input": str(input_path), |
| "output": str(output_path), |
| "input_samples": int(samples.size), |
| "output_samples": int(enhanced.size), |
| "elapsed_seconds": round(elapsed, 4), |
| "audio_seconds": round(audio_seconds, 4), |
| "real_time_factor": round(elapsed / audio_seconds, 4), |
| } |
| return report |
|
|
| def save_report(self, report: dict, path: str | Path) -> None: |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") |
|
|
