app.py

import base64
import os
from io import BytesIO
from typing import Any, Dict, List, Tuple

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import requests
import streamlit as st
from PIL import Image


ENDPOINT_URL_DEFAULT = (
    "https://madsdoodle--audio-cnn-inference-audioclassifier-inference.modal.run"
)


def _split_layers(visualization: Dict[str, Dict[str, Any]]) -> Tuple[List[Tuple[str, Any]], Dict[str, List[Tuple[str, Any]]]]:
    main: List[Tuple[str, Any]] = []
    internals: Dict[str, List[Tuple[str, Any]]] = {}

    for name, data in visualization.items():
        if "." not in name:
            main.append((name, data))
        else:
            parent = name.split(".", 1)[0]
            internals.setdefault(parent, []).append((name, data))

    return main, internals


def _normalize_diverging(a: np.ndarray) -> np.ndarray:
    a = np.nan_to_num(a, nan=0.0, posinf=0.0, neginf=0.0)
    abs_max = float(np.max(np.abs(a))) if a.size else 0.0
    if abs_max == 0.0:
        return np.zeros_like(a, dtype=np.float32)
    return (a / abs_max).astype(np.float32)


def _heatmap_image(values_2d: List[List[float]], *, size: Tuple[int, int] = (720, 260), is_spectrogram: bool = False) -> Image.Image:
    arr = np.asarray(values_2d, dtype=np.float32)
    arr = _normalize_diverging(arr)

    cmap = matplotlib.colormaps.get_cmap("plasma")
    rgba = cmap((arr + 1.0) / 2.0)
    rgb = (rgba[:, :, :3] * 255).astype(np.uint8)
    img = Image.fromarray(rgb)

    if is_spectrogram:
        img = img.transpose(Image.Transpose.FLIP_TOP_BOTTOM)

    img = img.resize(size, Image.Resampling.BILINEAR)
    return img


def _waveform_plot(values: List[float], *, title: str) -> matplotlib.figure.Figure:
    y = np.asarray(values, dtype=np.float32)
    y = y[np.isfinite(y)]

    fig = plt.figure(figsize=(10, 3), dpi=100, facecolor='#0e1117')
    ax = fig.add_subplot(1, 1, 1)
    ax.set_facecolor('#0e1117')

    if y.size:
        x = np.linspace(0, 1, num=y.size)
        ax.plot(x, y, color='#ff4b4b', linewidth=1.5, alpha=0.9)
        ax.fill_between(x, y, alpha=0.3, color='#ff4b4b')
        ax.axhline(0, color='#262730', linewidth=1, linestyle='--', alpha=0.7)

    ax.set_title(title, fontsize=12, color='#fafafa', pad=12, weight='bold')
    ax.set_xlabel("Time (normalized)", fontsize=10, color='#a0a0a0')
    ax.set_ylabel("Amplitude", fontsize=10, color='#a0a0a0')
    ax.grid(True, alpha=0.15, color='#262730')
    ax.tick_params(colors='#a0a0a0', labelsize=9)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['bottom'].set_color('#262730')
    ax.spines['left'].set_color('#262730')
    fig.tight_layout()
    return fig


def _call_inference_endpoint(file_bytes: bytes, endpoint_url: str) -> Dict[str, Any]:
    b64 = base64.b64encode(file_bytes).decode("utf-8")

    r = requests.post(
        endpoint_url,
        headers={"Content-Type": "application/json"},
        json={"audio_data": b64},
        timeout=120,
    )
    r.raise_for_status()
    return r.json()


# Page configuration
st.set_page_config(
    page_title="CNN Audio Visualizer",
    page_icon="🎵",
    layout="wide",
    initial_sidebar_state="collapsed"
)

# Custom CSS
st.markdown("""
<style>
    .main {
        padding-top: 2rem;
    }
    
    .stButton>button {
        width: 100%;
        background: linear-gradient(135deg, #ff4b4b 0%, #ff6b6b 100%);
        color: white;
        border: none;
        padding: 0.75rem 2rem;
        font-size: 1.1rem;
        font-weight: 600;
        border-radius: 8px;
        transition: all 0.3s ease;
    }
    
    .stButton>button:hover {
        transform: translateY(-2px);
        box-shadow: 0 4px 12px rgba(255, 75, 75, 0.4);
    }
    
    .prediction-card {
        background: linear-gradient(135deg, #1e1e1e 0%, #262626 100%);
        padding: 1.5rem;
        border-radius: 12px;
        margin: 1rem 0;
        border: 1px solid #333;
    }
    
    .prediction-title {
        font-size: 1.1rem;
        font-weight: 600;
        color: #fafafa;
        margin-bottom: 0.5rem;
        text-transform: capitalize;
    }
    
    .prediction-score {
        background: #ff4b4b;
        color: white;
        padding: 0.25rem 0.75rem;
        border-radius: 20px;
        font-size: 0.9rem;
        font-weight: 600;
        display: inline-block;
    }
    
    .progress-bar {
        background: #1e1e1e;
        height: 12px;
        border-radius: 20px;
        overflow: hidden;
        margin-top: 0.75rem;
    }
    
    .progress-fill {
        background: linear-gradient(90deg, #ff4b4b 0%, #ff6b6b 100%);
        height: 100%;
        border-radius: 20px;
        transition: width 0.5s ease;
    }
    
    .layer-card {
        background: #1a1a1a;
        padding: 1rem;
        border-radius: 10px;
        margin: 0.75rem 0;
        border: 1px solid #2a2a2a;
    }
    
    .layer-title {
        font-size: 0.95rem;
        font-weight: 600;
        color: #ff4b4b;
        margin-bottom: 0.5rem;
    }
    
    .layer-shape {
        font-size: 0.8rem;
        color: #888;
        background: #0e1117;
        padding: 0.25rem 0.5rem;
        border-radius: 4px;
        display: inline-block;
        margin-top: 0.5rem;
    }
    
    h1 {
        background: linear-gradient(135deg, #ff4b4b 0%, #ff8e53 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        font-size: 3rem !important;
        font-weight: 800 !important;
        text-align: center;
        margin-bottom: 0.5rem !important;
    }
    
    .subtitle {
        text-align: center;
        color: #a0a0a0;
        font-size: 1.1rem;
        margin-bottom: 2rem;
    }
    
    .section-header {
        color: #fafafa;
        font-size: 1.5rem;
        font-weight: 700;
        margin-top: 2rem;
        margin-bottom: 1rem;
        padding-bottom: 0.5rem;
        border-bottom: 2px solid #ff4b4b;
    }
</style>
""", unsafe_allow_html=True)

# Header
st.markdown("# 🎵 CNN Audio Visualizer")
st.markdown('<p class="subtitle">Upload a WAV file to see real-time predictions and feature map visualizations</p>', unsafe_allow_html=True)

# Sidebar for settings
with st.sidebar:
    st.markdown("### ⚙️ Settings")
    endpoint_url = st.text_input(
        "Inference Endpoint",
        value=os.environ.get("AUDIO_CNN_ENDPOINT", ENDPOINT_URL_DEFAULT),
        help="API endpoint for audio classification"
    )

# File upload
col1, col2, col3 = st.columns([1, 2, 1])
with col2:
    uploaded_file = st.file_uploader(
        "📁 Upload WAV File",
        type=["wav"],
        help="Upload an audio file in WAV format"
    )
    
    analyze_button = st.button("🚀 Analyze Audio", use_container_width=True)

# Process when button is clicked
if analyze_button and uploaded_file is not None:
    with st.spinner("🔄 Processing audio..."):
        try:
            # Read file bytes
            file_bytes = uploaded_file.read()
            
            # Call API
            data = _call_inference_endpoint(file_bytes, endpoint_url)
            
            preds = data.get("predictions", [])
            spect = data.get("input_spectrogram", {})
            wave = data.get("waveform", {})
            viz = data.get("visualization", {})
            
            # Success message
            st.success("✅ Analysis complete!")
            
            # Top Predictions
            st.markdown('<div class="section-header">🎯 Top Predictions</div>', unsafe_allow_html=True)
            
            top_preds = preds[:3]
            for i, pred in enumerate(top_preds):
                cls = str(pred.get("class", "")).replace("_", " ")
                conf = float(pred.get("confidence", 0.0))
                pct = max(0.0, min(100.0, conf * 100.0))
                
                st.markdown(f"""
                <div class="prediction-card">
                    <div style="display: flex; justify-content: space-between; align-items: center;">
                        <div class="prediction-title">#{i+1} {cls}</div>
                        <div class="prediction-score">{pct:.1f}%</div>
                    </div>
                    <div class="progress-bar">
                        <div class="progress-fill" style="width: {pct}%;"></div>
                    </div>
                </div>
                """, unsafe_allow_html=True)
            
            # Spectrogram and Waveform
            st.markdown('<div class="section-header">📊 Audio Analysis</div>', unsafe_allow_html=True)
            
            col_a, col_b = st.columns(2)
            
            with col_a:
                st.markdown("#### 🌈 Input Spectrogram")
                spect_values = spect.get("values", [])
                spect_shape = spect.get("shape", [])
                if spect_values:
                    spect_img = _heatmap_image(spect_values, size=(800, 300), is_spectrogram=True)
                    st.image(spect_img, use_container_width=True)
                    shape_str = " × ".join(str(x) for x in spect_shape) if isinstance(spect_shape, list) else str(spect_shape)
                    st.caption(f"Shape: {shape_str}")
            
            with col_b:
                st.markdown("#### 〰️ Audio Waveform")
                wave_values = wave.get("values", [])
                sr = int(wave.get("sample_rate", 0) or 0)
                dur = float(wave.get("duration", 0.0) or 0.0)
                wave_title = f"{dur:.2f}s @ {sr}Hz" if sr else f"{dur:.2f}s"
                if wave_values:
                    wave_fig = _waveform_plot(wave_values, title=wave_title)
                    st.pyplot(wave_fig)
                    plt.close(wave_fig)
            
            # Convolutional Layers
            st.markdown('<div class="section-header">🧠 Convolutional Layer Outputs</div>', unsafe_allow_html=True)
            
            main, internals = _split_layers(viz)
            main = sorted(main, key=lambda t: t[0])
            
            # Display all main layers in a uniform horizontal row
            if main:
                num_layers = len(main)
                cols = st.columns(num_layers)
                
                for idx, (main_name, main_data) in enumerate(main):
                    with cols[idx]:
                        main_values = main_data.get("values", [])
                        main_shape = main_data.get("shape", [])
                        
                        st.markdown(f'<div class="layer-title">{main_name}</div>', unsafe_allow_html=True)
                        
                        if main_values:
                            # Fixed height for uniform appearance
                            main_img = _heatmap_image(main_values, size=(300, 200), is_spectrogram=False)
                            st.image(main_img, use_container_width=True)
                        
                        shape_str = " × ".join(str(x) for x in main_shape) if isinstance(main_shape, list) else str(main_shape)
                        st.markdown(f'<div class="layer-shape">Shape: {shape_str}</div>', unsafe_allow_html=True)
                        
                        # Show internal layers in expander
                        internal_items = internals.get(main_name, [])
                        if internal_items:
                            with st.expander(f"Show {len(internal_items)} internal layers", expanded=False):
                                for layer_name, layer_data in sorted(internal_items, key=lambda t: t[0]):
                                    vals = layer_data.get("values", [])
                                    title = layer_name.replace(f"{main_name}.", "")
                                    st.markdown(f"**{title}**")
                                    if vals:
                                        # Uniform size for internal layers
                                        img = _heatmap_image(vals, size=(280, 180), is_spectrogram=False)
                                        st.image(img, use_container_width=True)
                                    st.divider()
            
        except Exception as e:
            st.error(f"❌ Error: {str(e)}")

elif analyze_button and uploaded_file is None:
    st.warning("⚠️ Please upload a WAV file first!")

# Footer
st.markdown("---")
st.markdown(
    '<p style="text-align: center; color: #666; font-size: 0.9rem;">Built with Streamlit • CNN Audio Classification</p>',
    unsafe_allow_html=True
)