Update src/utils.py

src/utils.py

@@ -1,83 +1,125 @@
-"""
-Utility functions for visualization and analysis
-"""
-
-import numpy as np
-import matplotlib.pyplot as plt
-import librosa
-import librosa.display
-
-def create_waveform_plot(y, sr, title="Audio Waveform"):
-    """
-    Create waveform visualization
-    
-    Args:
-        y (np.array): Audio time series
-        sr (int): Sample rate
-        title (str): Plot title
-    
-    Returns:
-        matplotlib.figure.Figure: Waveform plot
-    """
-    fig, ax = plt.subplots(figsize=(10, 3))
-    librosa.display.waveshow(y, sr=sr, ax=ax, color='#2E86DE')
-    ax.set_title(title, fontsize=14, fontweight='bold')
-    ax.set_xlabel('Time (seconds)', fontsize=11)
-    ax.set_ylabel('Amplitude', fontsize=11)
-    ax.grid(True, alpha=0.3)
-    plt.tight_layout()
-    return fig
-
-def create_spectrogram_plot(y, sr, title="Spectrogram"):
-    """
-    Create spectrogram visualization
-    
-    Args:
-        y (np.array): Audio time series
-        sr (int): Sample rate
-        title (str): Plot title
-    
-    Returns:
-        matplotlib.figure.Figure: Spectrogram plot
-    """
-    fig, ax = plt.subplots(figsize=(10, 4))
-    D = librosa.amplitude_to_db(np.abs(librosa.stft(y)), ref=np.max)
-    img = librosa.display.specshow(D, sr=sr, x_axis='time', y_axis='hz', 
-                                   ax=ax, cmap='viridis')
-    ax.set_title(title, fontsize=14, fontweight='bold')
-    ax.set_xlabel('Time (seconds)', fontsize=11)
-    ax.set_ylabel('Frequency (Hz)', fontsize=11)
-    fig.colorbar(img, ax=ax, format='%+2.0f dB')
-    plt.tight_layout()
-    return fig
-
-def create_mel_spectrogram_plot(y, sr, title="Mel Spectrogram"):
-    """
-    Create mel spectrogram visualization
-    
-    Args:
-        y (np.array): Audio time series
-        sr (int): Sample rate
-        title (str): Plot title
-    
-    Returns:
-        matplotlib.figure.Figure: Mel spectrogram plot
-    """
-    fig, ax = plt.subplots(figsize=(10, 4))
-    S = librosa.feature.melspectrogram(y=y, sr=sr)
-    S_dB = librosa.power_to_db(S, ref=np.max)
-    img = librosa.display.specshow(S_dB, sr=sr, x_axis='time', y_axis='mel',
-                                   ax=ax, cmap='magma')
-    ax.set_title(title, fontsize=14, fontweight='bold')
-    ax.set_xlabel('Time (seconds)', fontsize=11)
-    ax.set_ylabel('Mel Frequency', fontsize=11)
-    fig.colorbar(img, ax=ax, format='%+2.0f dB')
-    plt.tight_layout()
-    return fig
-
-def format_probability_text(prob_dict, top_k=None):
-    """
-    Format probability dictionary as text with progress bars
-    
-    Args:
-        prob_dict (dict): Dictionary of e
+"""
+Utility functions for visualization and analysis
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import librosa
+import librosa.display
+
+def create_waveform_plot(y, sr, title="Audio Waveform"):
+    """
+    Create waveform visualization
+    
+    Args:
+        y (np.array): Audio time series
+        sr (int): Sample rate
+        title (str): Plot title
+    
+    Returns:
+        matplotlib.figure.Figure: Waveform plot
+    """
+    fig, ax = plt.subplots(figsize=(10, 3))
+    librosa.display.waveshow(y, sr=sr, ax=ax, color='#2E86DE')
+    ax.set_title(title, fontsize=14, fontweight='bold')
+    ax.set_xlabel('Time (seconds)', fontsize=11)
+    ax.set_ylabel('Amplitude', fontsize=11)
+    ax.grid(True, alpha=0.3)
+    plt.tight_layout()
+    return fig
+
+def create_spectrogram_plot(y, sr, title="Spectrogram"):
+    """
+    Create spectrogram visualization
+    
+    Args:
+        y (np.array): Audio time series
+        sr (int): Sample rate
+        title (str): Plot title
+    
+    Returns:
+        matplotlib.figure.Figure: Spectrogram plot
+    """
+    fig, ax = plt.subplots(figsize=(10, 4))
+    D = librosa.amplitude_to_db(np.abs(librosa.stft(y)), ref=np.max)
+    img = librosa.display.specshow(D, sr=sr, x_axis='time', y_axis='hz', 
+                                   ax=ax, cmap='viridis')
+    ax.set_title(title, fontsize=14, fontweight='bold')
+    ax.set_xlabel('Time (seconds)', fontsize=11)
+    ax.set_ylabel('Frequency (Hz)', fontsize=11)
+    fig.colorbar(img, ax=ax, format='%+2.0f dB')
+    plt.tight_layout()
+    return fig
+
+def create_mel_spectrogram_plot(y, sr, title="Mel Spectrogram"):
+    """
+    Create mel spectrogram visualization
+    
+    Args:
+        y (np.array): Audio time series
+        sr (int): Sample rate
+        title (str): Plot title
+    
+    Returns:
+        matplotlib.figure.Figure: Mel spectrogram plot
+    """
+    fig, ax = plt.subplots(figsize=(10, 4))
+    S = librosa.feature.melspectrogram(y=y, sr=sr)
+    S_dB = librosa.power_to_db(S, ref=np.max)
+    img = librosa.display.specshow(S_dB, sr=sr, x_axis='time', y_axis='mel',
+                                   ax=ax, cmap='magma')
+    ax.set_title(title, fontsize=14, fontweight='bold')
+    ax.set_xlabel('Time (seconds)', fontsize=11)
+    ax.set_ylabel('Mel Frequency', fontsize=11)
+    fig.colorbar(img, ax=ax, format='%+2.0f dB')
+    plt.tight_layout()
+    return fig
+
+def format_probability_text(prob_dict, top_k=None):
+    """
+    Format probability dictionary as text with progress bars
+    
+    Args:
+        prob_dict (dict): Dictionary of emotion: probability
+        top_k (int): Show only top K emotions (None for all)
+    
+    Returns:
+        str: Formatted text
+    """
+    # Sort by probability
+    sorted_probs = sorted(prob_dict.items(), key=lambda x: x[1], reverse=True)
+    
+    if top_k:
+        sorted_probs = sorted_probs[:top_k]
+    
+    text = ""
+    for emotion, prob in sorted_probs:
+        # Create progress bar
+        bar_length = int(prob * 30)
+        bar = '█' * bar_length + '░' * (30 - bar_length)
+        
+        text += f"**{emotion.capitalize()}**: {bar} {prob*100:.2f}%\n"
+    
+    return text
+
+def get_emotion_emoji(emotion):
+    """
+    Get emoji for emotion
+    
+    Args:
+        emotion (str): Emotion name
+    
+    Returns:
+        str: Emoji character
+    """
+    emoji_map = {
+        'angry': '😠',
+        'calm': '😌',
+        'disgust': '🤢',
+        'fearful': '😨',
+        'happy': '😊',
+        'neutral': '😐',
+        'sad': '😢',
+        'surprised': '😲'
+    }
+    return emoji_map.get(emotion.lower(), '🎭')