Update

app.py

@@ -11,22 +11,29 @@ from PIL import Image
 
 def plot_spectrogram(waveform, sample_rate):
     """Plot and return a spectrogram."""
+    # Ensure waveform is 2D (channels, time) after squeeze
+    if waveform.ndim == 1:
+        waveform = waveform.unsqueeze(0)  # Add a channel dimension if it's missing
+
     spectrogram_transform = T.Spectrogram()
     spectrogram = spectrogram_transform(waveform)
     spectrogram_db = torchaudio.transforms.AmplitudeToDB()(spectrogram)
 
     plt.figure(figsize=(10, 4))
-    plt.imshow(spectrogram_db[0].numpy(), cmap='hot', aspect='auto', origin='lower')
+    # Ensure we're plotting the first channel for 2D data
+    if spectrogram_db.ndim == 3:
+        spectrogram_db = spectrogram_db[0]
+    plt.imshow(spectrogram_db.numpy(), cmap='hot', aspect='auto', origin='lower')
     plt.colorbar(format='%+2.0f dB')
     plt.title('Spectrogram')
     plt.xlabel('Time Frame')
     plt.ylabel('Frequency')
-    
+
     buf = io.BytesIO()
     plt.savefig(buf, format='png')
     plt.close()
     buf.seek(0)
-    
+
     return Image.open(buf)
 
 def detect_watermark(audio_file_path, threshold=0.99):