noise_pipeline.py

import os
import glob
from pathlib import Path

import librosa
import librosa.display
import matplotlib.pyplot as plt
import numpy as np
import soundfile as sf
from scipy import ndimage


class AcousticPurifier:

    HZ = 16000
    FFT_SZ = 512
    HOP_SZ = 160
    WIN_SZ = 400
    TIME_SMOOTH_SPAN = 5

    def __init__(self, over_sub: float = 1.5, min_floor: float = 0.02, quiet_pct: float = 10.0):
        self.over_sub = over_sub
        self.min_floor = min_floor
        self.quiet_pct = quiet_pct

    def process_and_report(self, in_audio_path: Path, out_wav_path: Path, out_plot_dir: Path) -> None:
        wave_data, _ = librosa.load(in_audio_path, sr=self.HZ, mono=True)

        z_matrix = librosa.stft(wave_data, n_fft=self.FFT_SZ, hop_length=self.HOP_SZ, win_length=self.WIN_SZ)
        amplitude_matrix = np.abs(z_matrix)

        blurred_amp = ndimage.uniform_filter1d(amplitude_matrix, size=self.TIME_SMOOTH_SPAN, axis=1, mode='nearest')

        frame_power = np.mean(blurred_amp ** 2, axis=0)
        power_db = 10.0 * np.log10(frame_power + 1e-12) 
        
        threshold_val = float(np.percentile(power_db, self.quiet_pct))
        is_background = power_db <= threshold_val
        
        if not np.any(is_background):
            is_background[np.argmin(power_db)] = True

        avg_bg_spectrum = np.mean(blurred_amp[:, is_background], axis=1, keepdims=True)

        target_reduction = blurred_amp - (self.over_sub * avg_bg_spectrum)
        safety_floor = self.min_floor * blurred_amp
        clean_amplitude = np.where(target_reduction > safety_floor, target_reduction, safety_floor)

        phase_angles = np.divide(z_matrix, amplitude_matrix + 1e-12)
        reconstructed_z = clean_amplitude * phase_angles
        
        clean_wave = librosa.istft(reconstructed_z, hop_length=self.HOP_SZ, win_length=self.WIN_SZ)

        sf.write(out_wav_path, clean_wave, self.HZ)

        self._render_figures(
            in_audio_path.stem, out_plot_dir, 
            amplitude_matrix, clean_amplitude, 
            power_db, threshold_val, is_background, avg_bg_spectrum
        )

    def _render_figures(self, track_id: str, plot_dir: Path, raw_amp: np.ndarray, 
                        clean_amp: np.ndarray, pwr_db: np.ndarray, thresh: float, 
                        bg_mask: np.ndarray, bg_spec: np.ndarray) -> None:
        
        plot_dir.mkdir(parents=True, exist_ok=True)
        bg_hex = "#eef2f3"
        grid_hex = "#d3d3d3"

        fig_a, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5), sharey=True)
        fig_a.patch.set_facecolor(bg_hex)
        
        raw_heatmap = librosa.amplitude_to_db(raw_amp, ref=np.max)
        clean_heatmap = librosa.amplitude_to_db(clean_amp, ref=np.max)

        hm1 = librosa.display.specshow(raw_heatmap, sr=self.HZ, hop_length=self.HOP_SZ, 
                                       x_axis='time', y_axis='hz', ax=ax1, cmap='inferno')
        ax1.set_title("Input (Set C)", size=10, weight='semibold')
        ax1.set_xlabel("Progression [s]")
        ax1.set_ylabel("Frequency Range [Hz]")
        fig_a.colorbar(hm1, ax=ax1, format="%+2.0f dB")

        hm2 = librosa.display.specshow(clean_heatmap, sr=self.HZ, hop_length=self.HOP_SZ, 
                                       x_axis='time', y_axis='hz', ax=ax2, cmap='inferno')
        ax2.set_title("Output (Purified)", size=10, weight='semibold')
        ax2.set_xlabel("Progression [s]")
        fig_a.colorbar(hm2, ax=ax2, format="%+2.0f dB")

        fig_a.suptitle(f"Spectrographic Shift | ID: {track_id}", size=12)
        plt.tight_layout()
        plt.savefig(plot_dir / f"{track_id}_heatmaps.png", dpi=110)
        plt.close(fig_a)

        time_vector = librosa.frames_to_time(np.arange(len(pwr_db)), sr=self.HZ, hop_length=self.HOP_SZ)
        
        fig_b, ax_b = plt.subplots(figsize=(10, 4))
        fig_b.patch.set_facecolor(bg_hex)
        ax_b.set_facecolor("white")
        
        ax_b.plot(time_vector, pwr_db, color="teal", lw=1.0, label="Frame Power")
        ax_b.axhline(thresh, color="crimson", linestyle=":", lw=1.5, label=f"Cutoff ({thresh:.1f} dB)")
        
        bg_times = time_vector[bg_mask]
        bg_vals = pwr_db[bg_mask]
        ax_b.scatter(bg_times, bg_vals, color="black", marker="s", s=15, label="Sampled Background")
        
        ax_b.set_title(f"Temporal Power Dynamics | ID: {track_id}", size=10, weight='semibold')
        ax_b.set_xlabel("Progression [s]")
        ax_b.set_ylabel("Frame Power [dB]")
        ax_b.grid(color=grid_hex, linestyle='-')
        ax_b.legend(loc="upper right")
        
        plt.tight_layout()
        plt.savefig(plot_dir / f"{track_id}_power_dynamics.png", dpi=110)
        plt.close(fig_b)

        freq_vector = librosa.fft_frequencies(sr=self.HZ, n_fft=self.FFT_SZ)
        bg_spec_1d = bg_spec.flatten() 
        bg_spec_db = 10.0 * np.log10((bg_spec_1d ** 2) + 1e-12)
        
        fig_c, ax_c = plt.subplots(figsize=(10, 4))
        fig_c.patch.set_facecolor(bg_hex)
        ax_c.set_facecolor("white")

        ax_c.plot(freq_vector, bg_spec_db, color="navy", lw=1.2)
        ax_c.fill_between(freq_vector, bg_spec_db, bg_spec_db.min(), color="steelblue", alpha=0.3)
        
        ax_c.set_title(f"Extracted Noise Manifold | ID: {track_id}", size=10, weight='semibold')
        ax_c.set_xlabel("Frequency Bins [Hz]")
        ax_c.set_ylabel("Power Amplitude [dB]")
        ax_c.set_xlim(0, self.HZ / 2)
        ax_c.grid(color=grid_hex, linestyle='-')

        plt.tight_layout()
        plt.savefig(plot_dir / f"{track_id}_noise_manifold.png", dpi=110)
        plt.close(fig_c)


def execute_batch_job(src_dir: str, dest_dir: str) -> None:
    source_path = Path(src_dir)
    dest_path = Path(dest_dir)
    dest_path.mkdir(parents=True, exist_ok=True)

    tracks = glob.glob(str(source_path / "*.wav")) + glob.glob(str(source_path / "*.WAV"))
    
    if not tracks:
        print(f"[ERR] No valid audio found at {source_path}")
        return

    purifier = AcousticPurifier(over_sub=1.5, min_floor=0.02, quiet_pct=10.0)

    print("\n--- Start ---")
    print(f"Total targets: {len(tracks)}")
    
    for idx, t_path in enumerate(tracks, 1):
        target = Path(t_path)
        out_wav = dest_path / f"{target.stem}.wav"
        
        print(f"[{idx:03d}/{len(tracks):03d}] Cleaning: {target.name} ... ", end="")
        
        try:
            purifier.process_and_report(target, out_wav, dest_path)
            print("OK")
        except Exception as e:
            print(f"FAIL ({e})")
            
    print("--- End---\n")


if __name__ == "__main__":
    DIR_IN = os.environ.get("SET_C_INPUT_DIR", "./dataset/set_C_raw")
    DIR_OUT = os.environ.get("SET_C_OUTPUT_DIR", "./dataset/set_C_cleaneds")
    
    execute_batch_job(DIR_IN, DIR_OUT)