Update app.py

app.py

@@ -50,18 +50,30 @@ def plot_spectrogram(waveform, sample_rate):
     buf.seek(0)
     return Image.open(buf)
 
-def detect_watermark(waveform, sample_rate):
-    if not 'AudioSeal' in globals():
-        return "AudioSeal not available", plot_spectrogram(waveform, sample_rate)
-
+def detect_watermark(audio_data, sample_rate):
+    # Ensure AudioSeal is available
+    if 'AudioSeal' not in globals():
+        spectrogram_image = plot_spectrogram(audio_data, sample_rate)
+        return "AudioSeal not available", spectrogram_image
+    
+    # Load audio data correctly
+    waveform, sr = load_and_resample_audio(audio_data, target_sample_rate=16000)
+    
+    # Ensure waveform is a tensor before passing to the detector
+    if not isinstance(waveform, torch.Tensor):
+        return "Error: waveform is not a tensor.", plot_spectrogram(waveform, sr)
+    
+    # Load the detector
     detector = AudioSeal.load_detector("audioseal_detector_16bits")
-    # Assuming `forward` method exists and accepts sample_rate; adjust accordingly
-    results, messages = detector.forward(waveform, sample_rate=sample_rate)  # Adjusted to pass sample_rate explicitly
+    
+    # Process waveform with the detector
+    results, messages = detector.forward(waveform.unsqueeze(0), sample_rate=sample_rate)  # Ensure waveform is in batch form
     detect_probs = results[:, 1, :]
     result = detect_probs.mean().cpu().item()
     message = f"Detection result: {'Watermarked or AI-Generated audio' if result > 0.5 else 'Not watermarked or AI-Generated'}\nProbability of watermark: {result}"
-    spectrogram_image = plot_spectrogram(waveform, sample_rate)
-    return message, spectrogram_image  # Return both message and image
+    spectrogram_image = plot_spectrogram(waveform, sr)
+    
+    return message, spectrogram_image
 
 
 def main(audio_file_path):