Delete run_tflite.py

run_tflite.py

@@ -1,298 +0,0 @@
-import argparse
-from dataclasses import dataclass
-from pathlib import Path
-import sys
-
-import numpy as np
-import soundfile as sf
-import librosa
-from tflite_runtime.interpreter import Interpreter
-from tqdm import tqdm
-
-TFLITE_DIR = Path('./')
-
-# -----------------------------------------------------------------------------
-# Model registry
-# -----------------------------------------------------------------------------
-# 16 kHz models: WIN_LEN=320  (20 ms)
-# 48 kHz models: WIN_LEN=960  (20 ms)
-
-MODEL_CONFIG = {
-    # 16 kHz models
-    "baseline":  {"sr": 16000, "win_len": 320},
-    "dpdfnet2":  {"sr": 16000, "win_len": 320},
-    "dpdfnet4":  {"sr": 16000, "win_len": 320},
-    "dpdfnet8":  {"sr": 16000, "win_len": 320},
-
-    # 48 kHz models
-    "dpdfnet2_48khz_hr": {"sr": 48000, "win_len": 960},
-}
-
-
-def vorbis_window(window_len: int) -> np.ndarray:
-    window_size_h = window_len / 2
-    indices = np.arange(window_len)
-    sin = np.sin(0.5 * np.pi * (indices + 0.5) / window_size_h)
-    window = np.sin(0.5 * np.pi * sin * sin)
-    return window.astype(np.float32)
-
-
-def get_wnorm(window_len: int, frame_size: int) -> float:
-    # window_len - #samples of the window; frame_size - hop size
-    return 1.0 / (window_len ** 2 / (2 * frame_size))
-
-
-@dataclass(frozen=True)
-class STFTConfig:
-    sr: int
-    win_len: int
-    hop_size: int
-    win: np.ndarray
-    wnorm: float
-
-
-def make_stft_config(sr: int, win_len: int) -> STFTConfig:
-    hop_size = win_len // 2  # 50% hop
-    win = vorbis_window(win_len)
-    wnorm = get_wnorm(win_len, hop_size)
-    return STFTConfig(sr=sr, win_len=win_len, hop_size=hop_size, win=win, wnorm=wnorm)
-
-
-# -----------------------------------------------------------------------------
-# Pre/Post processing
-# -----------------------------------------------------------------------------
-
-def preprocessing(waveform: np.ndarray, cfg: STFTConfig) -> np.ndarray:
-    """
-    waveform: 1D float32 numpy array at cfg.sr, mono, range ~[-1,1]
-    Returns complex STFT as real/imag split: [B=1, T, F, 2] float32
-    """
-    # Librosa returns [F, T]; match original by using center=False here
-    spec = librosa.stft(
-        y=waveform.astype(np.float32, copy=False),
-        n_fft=cfg.win_len,
-        hop_length=cfg.hop_size,
-        win_length=cfg.win_len,
-        window=cfg.win,
-        center=True,
-        pad_mode="reflect",
-    )  # [F, T] complex64
-
-    spec = (spec.T * cfg.wnorm).astype(np.complex64)  # [T, F]
-    spec_ri = np.stack([spec.real, spec.imag], axis=-1).astype(np.float32)  # [T, F, 2]
-    return spec_ri[None, ...]  # [1, T, F, 2]
-
-
-def postprocessing(spec_e: np.ndarray, cfg: STFTConfig) -> np.ndarray:
-    """
-    spec_e: [1, T, F, 2] float32
-    Returns waveform (1D float32, cfg.sr)
-    """
-    # Recreate complex STFT with shape [F, T]
-    spec_c = spec_e[0].astype(np.float32)  # [T, F, 2]
-    spec = (spec_c[..., 0] + 1j * spec_c[..., 1]).T.astype(np.complex64)  # [F, T]
-
-    waveform_e = librosa.istft(
-        spec,
-        hop_length=cfg.hop_size,
-        win_length=cfg.win_len,
-        window=cfg.win,
-        center=True,
-        length=None,
-    ).astype(np.float32)
-
-    waveform_e = waveform_e / cfg.wnorm
-
-    # Keep the legacy alignment compensation behavior, scaled by win_len.
-    waveform_e = np.concatenate(
-        [waveform_e[cfg.win_len * 2 :], np.zeros(cfg.win_len * 2, dtype=np.float32)]
-    )
-    return waveform_e.astype(np.float32)
-
-
-# -----------------------------------------------------------------------------
-# Audio utilities
-# -----------------------------------------------------------------------------
-
-def to_mono(audio: np.ndarray) -> np.ndarray:
-    if audio.ndim == 1:
-        return audio
-    # Average channels to mono
-    return np.mean(audio, axis=1)
-
-
-def ensure_sr(waveform: np.ndarray, sr: int, target_sr: int) -> np.ndarray:
-    if sr == target_sr:
-        return waveform.astype(np.float32, copy=False)
-    return librosa.resample(
-        waveform.astype(np.float32, copy=False), orig_sr=sr, target_sr=target_sr
-    )
-
-
-def resample_back(waveform_model_sr: np.ndarray, model_sr: int, target_sr: int) -> np.ndarray:
-    if target_sr == model_sr:
-        return waveform_model_sr
-    return librosa.resample(
-        waveform_model_sr.astype(np.float32, copy=False),
-        orig_sr=model_sr,
-        target_sr=target_sr,
-    )
-
-
-def pcm16_safe(x: np.ndarray) -> np.ndarray:
-    x = np.clip(x, -1.0, 1.0)
-    return (x * 32767.0).astype(np.int16)
-
-
-# -----------------------------------------------------------------------------
-# Core processing
-# -----------------------------------------------------------------------------
-
-def _load_model_and_cfg(model_name: str) -> tuple[Interpreter, STFTConfig]:
-    """Create interpreter and return (interpreter, STFTConfig) for this model."""
-    if model_name not in MODEL_CONFIG:
-        raise ValueError(
-            f"Unknown model '{model_name}'. Add it to MODEL_CONFIG or pass a valid --model_name."
-        )
-
-    model_path = TFLITE_DIR / f"{model_name}.tflite"
-    if not model_path.exists():
-        raise FileNotFoundError(f"TFLite model not found: {model_path}")
-
-    interpreter = Interpreter(model_path=str(model_path))
-    interpreter.allocate_tensors()
-
-    cfg_dict = MODEL_CONFIG[model_name]
-    cfg = make_stft_config(sr=int(cfg_dict["sr"]), win_len=int(cfg_dict["win_len"]))
-
-    # Optional sanity-check: infer expected F from model input and compare
-    try:
-        input_details = interpreter.get_input_details()
-        shape = input_details[0].get("shape", None)
-        # Expect [1, 1, F, 2] (or [1, T, F, 2] for non-streaming)
-        if shape is not None and len(shape) >= 3:
-            F = int(shape[-2])  # ... F, 2
-            expected_F = cfg.win_len // 2 + 1
-            if F != expected_F:
-                raise ValueError(
-                    f"Model '{model_name}' input F={F} does not match win_len={cfg.win_len} "
-                    f"(expected F={expected_F}). Update MODEL_CONFIG for this model."
-                )
-    except Exception:
-        # Do not hard-fail on odd/unknown shapes; the runtime error will be informative.
-        pass
-
-    return interpreter, cfg
-
-
-def enhance_file(in_path: Path, out_path: Path, model_name: str) -> None:
-    # Load audio
-    audio, sr_in = sf.read(str(in_path), always_2d=False)
-    audio = to_mono(audio)
-    audio = audio.astype(np.float32, copy=False)
-
-    # Load model and its expected SR/STFT config
-    interpreter, cfg = _load_model_and_cfg(model_name)
-    input_details = interpreter.get_input_details()
-    output_details = interpreter.get_output_details()
-
-    # Resample to model SR
-    audio_model_sr = ensure_sr(audio, sr_in, cfg.sr)
-
-    # Alignment compensation #1
-    audio_pad = np.pad(audio_model_sr, (0, cfg.win_len), mode='constant', constant_values=0)
-
-    # STFT to frames (streaming)
-    spec = preprocessing(audio_pad, cfg)  # [1, T, F, 2]
-    num_frames = spec.shape[1]
-
-    # Frame-by-frame inference
-    outputs = []
-    for t in tqdm(range(num_frames), desc=f"{in_path.name}", unit="frm", leave=False):
-        frame = spec[:, t : t + 1]  # [1, 1, F, 2]
-        frame = np.ascontiguousarray(frame, dtype=np.float32)
-
-        interpreter.set_tensor(input_details[0]["index"], frame)
-        interpreter.invoke()
-        y = interpreter.get_tensor(output_details[0]["index"])  # expected [1,1,F,2]
-        outputs.append(np.ascontiguousarray(y, dtype=np.float32))
-
-    # Concatenate along time dimension
-    spec_e = np.concatenate(outputs, axis=1).astype(np.float32)  # [1, T, F, 2]
-
-    # iSTFT to waveform (model SR), then back to original SR for saving
-    enhanced_model_sr = postprocessing(spec_e, cfg)
-    enhanced = resample_back(enhanced_model_sr, cfg.sr, sr_in)
-
-    # Alignment compensation #2
-    enhanced = enhanced[: audio.size]
-
-    # Save as 16-bit PCM WAV, mono, original sample rate
-    out_path.parent.mkdir(parents=True, exist_ok=True)
-    sf.write(str(out_path), pcm16_safe(enhanced), sr_in, subtype="PCM_16")
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Enhance WAV files with a DPDFNet TFLite model (streaming)."
-    )
-    parser.add_argument(
-        "--noisy_dir",
-        type=str,
-        required=True,
-        help="Folder with noisy *.wav files (non-recursive).",
-    )
-    parser.add_argument(
-        "--enhanced_dir",
-        type=str,
-        required=True,
-        help="Output folder for enhanced WAVs.",
-    )
-    parser.add_argument(
-        "--model_name",
-        type=str,
-        default="dpdfnet8",
-        choices=sorted(MODEL_CONFIG.keys()),
-        help=(
-            "Name of the model to use. The script will automatically use the correct "
-            "sample-rate/STFT settings based on MODEL_CONFIG."
-        ),
-    )
-
-    args = parser.parse_args()
-    noisy_dir = Path(args.noisy_dir)
-    enhanced_dir = Path(args.enhanced_dir)
-    model_name = args.model_name
-
-    if not noisy_dir.is_dir():
-        print(
-            f"ERROR: --noisy_dir does not exist or is not a directory: {noisy_dir}",
-            file=sys.stderr,
-        )
-        sys.exit(1)
-
-    wavs = sorted(p for p in noisy_dir.glob("*.wav") if p.is_file())
-    if not wavs:
-        print(f"No .wav files found in {noisy_dir} (non-recursive).")
-        sys.exit(0)
-
-    cfg = MODEL_CONFIG.get(model_name, None)
-    print(f"Model: {model_name}")
-    if cfg is not None:
-        print(f"Model SR: {cfg['sr']} Hz | win_len: {cfg['win_len']} | hop: {cfg['win_len']//2}")
-    print(f"Input : {noisy_dir}")
-    print(f"Output: {enhanced_dir}")
-    print(f"Found {len(wavs)} file(s). Enhancing...\n")
-
-    for wav in wavs:
-        out_path = enhanced_dir / (wav.stem + f"_{model_name}.wav")
-        try:
-            enhance_file(wav, out_path, model_name)
-        except Exception as e:
-            print(f"[SKIP] {wav.name} due to error: {e}", file=sys.stderr)
-
-    print("\nProcessing complete. Outputs saved in:", enhanced_dir)
-
-
-if __name__ == "__main__":
-    main()