app.py

import gradio as gr
from audioseal import AudioSeal
import torch
import numpy as np
import traceback

# Function to handle audio data as NumPy arrays
def detect_watermark(audio_data, sample_rate):
    try:
        # Extract the audio array from the tuple (audio_data, sample_rate)
        audio_array, _ = audio_data

        # Ensure audio_array is 2D (channels, samples). If it's mono, add an axis.
        if audio_array.ndim == 1:
            audio_array = np.expand_dims(audio_array, axis=0)
        
        # Convert NumPy array to tensor
        waveform = torch.tensor(audio_array, dtype=torch.float32)
        
        # Ensure waveform is 2D (batch, channels, samples) for AudioSeal
        if waveform.ndim == 2:
            waveform = waveform.unsqueeze(0)

        # Initialize and use the AudioSeal detector
        detector = AudioSeal.load_detector("audioseal_detector_16bits")
        result, message = detector.detect_watermark(waveform, message_threshold=0.5)
        
        # Interpret and return the detection result
        detection_result = "AI-generated" if result else "genuine"
        return f"This audio is likely {detection_result} based on watermark detection."
    except Exception as e:
        error_message = f"An error occurred: {str(e)}"
        traceback_str = ''.join(traceback.format_tb(e.__traceback__))
        full_error_message = f"{error_message}\n{traceback_str}"
        return full_error_message

interface = gr.Interface(fn=detect_watermark,
                         inputs=[gr.Audio(label="Upload your audio", type="numpy"),
                                 gr.Number(label="Sample Rate", value=44100, visible=False)],
                         outputs="text",
                         title="Deep Fake Defender: AI Voice Cloning Detection",
                         description="Upload an audio file to check if it's AI-generated or genuine.")

if __name__ == "__main__":
    interface.launch(debug=True)