Update app.py

app.py

@@ -5,28 +5,37 @@ import torch
 import matplotlib.pyplot as plt
 from transformers import AutoModelForAudioClassification, ASTFeatureExtractor
 import random
+import tempfile
 
 # Model and feature extractor loading
 model = AutoModelForAudioClassification.from_pretrained("./")
 feature_extractor = ASTFeatureExtractor.from_pretrained("./")
 
 def plot_waveform(waveform, sr):
-    plt.figure(figsize=(12, 4))  # Slightly larger plot for better visibility
+    plt.figure(figsize=(12, 4))
     plt.title('Waveform')
     plt.ylabel('Amplitude')
     plt.plot(np.linspace(0, len(waveform) / sr, len(waveform)), waveform)
     plt.xlabel('Time (s)')
-    return plt.gcf()
+    # Save plot to a temporary file
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png', dir='./')
+    plt.savefig(temp_file.name)
+    plt.close()  # Close the figure to free memory
+    return temp_file.name
 
 def plot_spectrogram(waveform, sr):
     S = librosa.feature.melspectrogram(y=waveform, sr=sr, n_mels=128)
     S_DB = librosa.power_to_db(S, ref=np.max)
-    plt.figure(figsize=(12, 6))  # Slightly larger plot for better visibility
+    plt.figure(figsize=(12, 6))
     librosa.display.specshow(S_DB, sr=sr, x_axis='time', y_axis='mel')
     plt.title('Mel Spectrogram')
     plt.colorbar(format='%+2.0f dB')
     plt.tight_layout()
-    return plt.gcf()
+    # Save plot to a temporary file
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png', dir='./')
+    plt.savefig(temp_file.name)
+    plt.close()  # Close the figure to free memory
+    return temp_file.name
 
 def custom_feature_extraction(audio, sr=16000, target_length=1024):
     features = feature_extractor(audio, sampling_rate=sr, return_tensors="pt", padding="max_length", max_length=target_length)