Kiuyha/Audio-Separation-Inspector

1import streamlit as st2import os3import torch4import numpy as np5import pandas as pd6import plotly.express as px7import scipy.io.wavfile8import io9 10BASE_PATH = 'Models'11 12st.set_page_config(layout="wide", page_title="Audio Source Separation Inspector")13 14def process_audio(file_path, gain_factor):15    try:16        # 1. FIX: Check if file is actually a Git LFS pointer (text file)17        with open(file_path, 'rb') as f:18            header = f.read(50)19            if header.startswith(b'version https://git-lfs'):20                st.error(f"❌ **LFS Error:** `{os.path.basename(file_path)}` is a Git LFS pointer, not a WAV file. Run `git lfs pull` in your terminal.")21                return None22 23        sample_rate, data = scipy.io.wavfile.read(file_path)24        25        if data.dtype == np.int16:26            data = data.astype(np.float32) / 32768.027        elif data.dtype == np.int32:28            data = data.astype(np.float32) / 2147483648.029        30        data = data * gain_factor31        32        data = np.clip(data, -1.0, 1.0)33        34        data = (data * 32767).astype(np.int16)35        36        virtual_file = io.BytesIO()37        scipy.io.wavfile.write(virtual_file, sample_rate, data)38        return virtual_file39    except Exception as e:40        st.error(f"Error processing audio: {e}")41        return file_path42 43def get_subdirs(path):44    if not os.path.exists(path):45        return []46    return [d for d in os.listdir(path) if os.path.isdir(os.path.join(path, d))]47 48def load_spectrogram_interactive(pt_path, title="Spectrogram"):49    try:50        # 2. FIX: Added weights_only=False to fix PyTorch 2.6+ error51        spec_tensor = torch.load(pt_path, map_location='cpu', weights_only=False)52 53        if spec_tensor.dim() == 4: 54            spec_tensor = spec_tensor[0]55        if spec_tensor.dim() == 3: 56            spec_data = spec_tensor.mean(dim=0).numpy()57        else: 58            spec_data = spec_tensor.numpy()59 60        if spec_data.min() >= 0:61            spec_data = np.log1p(spec_data)62 63        fig = px.imshow(64            spec_data,65            origin='lower',66            aspect='auto',67            color_continuous_scale='Viridis',68            labels=dict(x="Time Frame", y="Frequency Bin", color="Log Magnitude"),69            title=title70        )71        fig.update_layout(margin=dict(l=0, r=0, t=30, b=0), height=300)72        return fig73    except Exception as e:74        st.error(f"Error loading spectrogram: {e}")75        return None76 77def load_feature_map_interactive(pt_path):78    try:79        # 3. FIX: Added weights_only=False here as well80        feat_tensor = torch.load(pt_path, map_location='cpu', weights_only=False)81 82        if feat_tensor.dim() == 4:83            feat_tensor = feat_tensor[0]84 85        mean_activation = feat_tensor.mean(dim=0).numpy()86 87        fig = px.imshow(88            mean_activation,89            origin='lower',90            aspect='auto',91            color_continuous_scale='Viridis',92            labels=dict(x="Time", y="Freq/Feature", color="Activation"),93            title=f"Mean Activation (Shape: {list(feat_tensor.shape)})"94        )95        fig.update_layout(margin=dict(l=0, r=0, t=40, b=0))96        return fig97    except Exception as e:98        return None99 100st.title("🎵 Audio Source Separation Inspector")101 102st.markdown("""103### Model Interpretation Guide104This tool helps you evaluate how well the model separates audio sources.105* **Audio Quality:** Listen for "artifacts" (robotic sounds or clicking) in the Prediction compared to the Target.106* **Spectrogram Clarity:** In the visuals below, distinct horizontal lines represent clear tones. Vertical smear usually indicates percussion or noise. 107* **Error Analysis:** If the Prediction looks "blurry" compared to the Target, the model is losing high-frequency details.108""")109 110if not os.path.exists(BASE_PATH):111    st.error(f"Models directory not found at {BASE_PATH}. Please ensure your data was uploaded correctly.")112    st.stop()113 114models = get_subdirs(BASE_PATH)115selected_model = st.sidebar.selectbox("Select Model", models)116 117st.sidebar.markdown("### Audio Settings")118volume_boost = st.sidebar.slider(119    "Volume Boost (Gain)", 120    min_value=1.0, 121    max_value=20.0, 122    value=1.0, 123    step=0.5,124    help="Digitally increases the amplitude of the audio signal."125)126 127if selected_model:128    model_path = os.path.join(BASE_PATH, selected_model)129    artifacts_path = os.path.join(model_path, "test_artifacts")130 131    if os.path.exists(artifacts_path):132        samples = get_subdirs(artifacts_path)133        samples.sort(key=lambda x: int(x.split('_')[-1]) if '_' in x else 0)134 135        selected_sample = st.sidebar.selectbox("Select Sample ID", samples)136 137        if selected_sample:138            sample_path = os.path.join(artifacts_path, selected_sample)139            audio_dir = os.path.join(sample_path, "audio")140            specs_dir = os.path.join(sample_path, "specs")141            feats_dir = os.path.join(sample_path, "feats")142 143            all_files = os.listdir(audio_dir)144            target_files = [f for f in all_files if f.startswith("target_") and f.endswith(".wav")]145            classes = [f.replace("target_", "").replace(".wav", "") for f in target_files]146 147            selected_class = st.sidebar.selectbox("Focus Class", classes)148 149            tab1, tab2, tab3 = st.tabs(["🎧 Audio & Spectrograms", "🧠 Internal Activations", "📊 Model Metadata"])150 151            with tab1:152                st.header(f"Sample {selected_sample} | Focus: {selected_class.capitalize()}")153                154                st.subheader("1. Mixture (Input)")155                st.markdown("The raw input containing all sound sources mixed together.")156                mix_audio = os.path.join(audio_dir, "mixture.wav")157                mix_spec = os.path.join(specs_dir, "mixture.pt")158 159                c1, c2 = st.columns([1, 3])160                with c1:161                    if os.path.exists(mix_audio):162                        st.markdown("**Audio:**")163                        processed_mix = process_audio(mix_audio, volume_boost)164                        if processed_mix:165                            st.audio(processed_mix, format='audio/wav')166                with c2:167                    if os.path.exists(mix_spec):168                        fig = load_spectrogram_interactive(mix_spec, title="Mixture Mel-Spectrogram")169                        if fig: st.plotly_chart(fig, width='stretch')170 171                st.divider()172 173                st.subheader(f"2. Target: {selected_class}")174                st.markdown(f"**Interpretation:** This is the 'Ground Truth'. Look at the spectrogram structure here—this is the ideal output.")175                tgt_audio = os.path.join(audio_dir, f"target_{selected_class}.wav")176                tgt_spec = os.path.join(specs_dir, f"target_{selected_class}.pt")177 178                c1, c2 = st.columns([1, 3])179                with c1:180                    if os.path.exists(tgt_audio):181                        st.markdown("**Audio:**")182                        processed_tgt = process_audio(tgt_audio, volume_boost)183                        if processed_tgt:184                            st.audio(processed_tgt, format='audio/wav')185                with c2:186                    if os.path.exists(tgt_spec):187                        fig = load_spectrogram_interactive(tgt_spec, title=f"Target Mel-Spectrogram ({selected_class})")188                        if fig: st.plotly_chart(fig, width='stretch')189 190                st.divider()191 192                st.subheader(f"3. Prediction: {selected_class}")193                st.markdown(f"**Interpretation:** Compare this to the Target above. If you see 'fuzziness' in the dark areas, the model is not silencing background noise correctly.")194                pred_audio = os.path.join(audio_dir, f"pred_{selected_class}.wav")195                pred_spec = os.path.join(specs_dir, f"pred_{selected_class}.pt")196 197                c1, c2 = st.columns([1, 3])198                with c1:199                    if os.path.exists(pred_audio):200                        st.markdown("**Audio:**")201                        processed_pred = process_audio(pred_audio, volume_boost)202                        if processed_pred:203                            st.audio(processed_pred, format='audio/wav')204                with c2:205                    if os.path.exists(pred_spec):206                        fig = load_spectrogram_interactive(pred_spec, title=f"Predicted Mel-Spectrogram ({selected_class})")207                        if fig: st.plotly_chart(fig, width='stretch')208 209            with tab2:210                st.header("Internal Feature Maps")211                212                st.markdown("These heatmaps visualize the neural network's internal state. Bright spots indicate features the model considers important for separation.")213 214                if os.path.exists(feats_dir):215                    feat_files = sorted(os.listdir(feats_dir))216 217                    if feat_files:218                        selected_layer = st.selectbox("Select Probed Layer", feat_files)219                        if selected_layer:220                            st.write(f"Layer: **{selected_layer.replace('.pt', '')}**")221                            fig = load_feature_map_interactive(os.path.join(feats_dir, selected_layer))222                            if fig:223                                st.plotly_chart(fig, width='stretch')224                    else:225                        st.warning("No feature maps found for this sample.")226                else:227                    st.error("Features directory not found.")228 229            with tab3:230                st.header("Training and Testing Logs")231                232                st.markdown("Use these graphs to check for **Overfitting**. If Training Loss decreases but Test Metrics stagnate or drop, the model is memorizing data rather than learning general features.")233 234                c1, c2 = st.columns(2)235                with c1:236                    results_csv = os.path.join(model_path, "test_results.csv")237                    if os.path.exists(results_csv):238                        st.subheader("Test Metrics")239                        df = pd.read_csv(results_csv)240                        x_axis = 'Batch_Index' if 'Batch_Index' in df.columns else df.index241                        numeric_cols = df.select_dtypes(include=np.number).columns242                        fig = px.line(df, title="Test Metrics", x=x_axis, y=numeric_cols)243                        st.plotly_chart(fig, width='stretch')244                        st.dataframe(df, width='stretch')245                    else:246                        st.info("No `test_results.csv` found.")247 248                with c2:249                    loss_csv = os.path.join(model_path, "loss.csv")250                    if os.path.exists(loss_csv):251                        st.subheader("Training Loss")252                        try:253                            df_loss = pd.read_csv(loss_csv)254                            x_axis = 'epoch' if 'epoch' in df_loss.columns else df_loss.index255 256                            numeric_cols = df_loss.select_dtypes(include=np.number).columns257                            fig = px.line(df_loss, x=x_axis, y=numeric_cols, title="Loss Curves")258                            st.plotly_chart(fig, width='stretch')259                            st.dataframe(df_loss, width='stretch')260                        except Exception as e:261                            st.write("Could not parse `loss.csv`.", e)262                    else:263                        st.info("No `loss.csv` found.")264 265    else:266        st.warning(f"No 'test_artifacts' folder found in {selected_model}")