app.py

import gradio as gr
import torch
import torchaudio
import tempfile
import logging
from audioseal import AudioSeal
import random
import string
from pathlib import Path

# Initialize logging
logging.basicConfig(level=logging.DEBUG, filename='app.log', filemode='w', format='%(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Helper function for generating a unique alphanumeric message
def generate_unique_message(length=16):
    characters = string.ascii_letters + string.digits
    return ''.join(random.choice(characters) for _ in range(length))

# Converts message to binary, ensuring it fits within the specified bit length
def message_to_binary(message, bit_length=16):
    binary_message = ''.join(format(ord(c), '08b') for c in message)
    return binary_message[:bit_length].ljust(bit_length, '0')

# Converts binary string back to ASCII message
def binary_to_message(binary_str):
    chars = [chr(int(binary_str[i:i+8], 2)) for i in range(0, len(binary_str), 8)]
    return ''.join(chars).rstrip('\x00')

# Converts binary string to hexadecimal
def binary_to_hex(binary_str):
    return hex(int(binary_str, 2))[2:].zfill(4)

# Load and resample audio file to match model's expected sample rate
def load_and_resample_audio(audio_file_path, target_sample_rate=16000):
    waveform, sample_rate = torchaudio.load(audio_file_path)
    if sample_rate != target_sample_rate:
        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_sample_rate)
        waveform = resampler(waveform)
    return waveform, target_sample_rate

# Main function for applying watermark to audio
def watermark_audio(audio_file_path, unique_message):
    waveform, sample_rate = load_and_resample_audio(audio_file_path, target_sample_rate=16000)
    waveform = torch.clamp(waveform, min=-1.0, max=1.0)
    generator = AudioSeal.load_generator("audioseal_wm_16bits")
    
    binary_message = message_to_binary(unique_message, bit_length=16)
    hex_message = binary_to_hex(binary_message)
    message_tensor = torch.tensor([int(bit) for bit in binary_message], dtype=torch.int32).unsqueeze(0)
    
    watermarked_audio = generator(waveform.unsqueeze(0), sample_rate=sample_rate, message=message_tensor).squeeze(0)
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.wav')
    torchaudio.save(temp_file.name, watermarked_audio, sample_rate)
    return temp_file.name, hex_message  # Display hex message to the user

# Function to detect watermark in audio
def detect_watermark(audio_file_path, original_hex_message=None):
    waveform, sample_rate = load_and_resample_audio(audio_file_path, target_sample_rate=16000)
    detector = AudioSeal.load_detector("audioseal_detector_16bits")
    
    result, message_tensor = detector.detect_watermark(waveform.unsqueeze(0), sample_rate=sample_rate)
    binary_message = ''.join(str(bit) for bit in message_tensor[0].tolist())
    detected_hex_message = binary_to_hex(binary_message)
    
    # Compare the detected message with the original, if provided
    match_result = "Not compared"
    if original_hex_message:
        match_result = "Match" if detected_hex_message == original_hex_message.upper() else "No Match"
    
    return result, detected_hex_message, match_result

# Setup for Gradio interface
style_path = Path("style.css")
style = style_path.read_text()

with gr.Blocks(css=style) as demo:
    with gr.Tab("Watermark Audio"):
        with gr.Column(scale=6):
            gr.Markdown("### How to Watermark Your Audio")
            gr.Markdown("""
This tool allows you to embed a unique, invisible watermark into your audio files. Here's how it works:

- **Upload Audio File**: Select the audio file you want to protect.
- **Generate Unique Message**: Click this button to create a unique code that identifies the audio as yours.
- **Unique Message**: This is the unique code generated for your audio. It's used to create the watermark.
- **Apply Watermark**: Embed the unique code into your audio file. This step converts the code into a format (hexadecimal) that's embedded into the audio without altering its quality.
- **Watermarked Audio**: After the watermark is applied, you can download and listen to your watermarked audio here. It will sound just like the original, but now it has your unique watermark.
- **Message Used for Watermarking**: Shows the code that was embedded into your audio. The actual embedded code is a hexadecimal version of this message, which is a more secure representation.
""")
            audio_input_watermark = gr.Audio(label="Upload Audio File for Watermarking", type="filepath")
            generate_message_button = gr.Button("Generate Unique Message")
            unique_message_output = gr.Textbox(label="Unique Message", value="Press Generate")
            watermark_button = gr.Button("Apply Watermark")
            watermarked_audio_output = gr.Audio(label="Watermarked Audio")
            message_output = gr.Textbox(label="Message Used for Watermarking")
            generate_message_button.click(fn=generate_unique_message, inputs=None, outputs=unique_message_output)
            watermark_button.click(fn=watermark_audio, inputs=[audio_input_watermark, unique_message_output], outputs=[watermarked_audio_output, message_output])

    with gr.Tab("Detect Watermark"):
        with gr.Column(scale=6):
            gr.Markdown("### How to Detect a Watermark in Your Audio")
            gr.Markdown("""
Use this tool to check if an audio file contains a specific watermark. Here's the process:

- **Upload Audio File**: Choose the audio file you suspect contains a watermark.
- **Original Hex Message for Comparison**: If you know the hexadecimal code of the watermark you're looking for, enter it here. This helps verify the specific watermark.
- **Detect Watermark**: Analyzes the audio to find any embedded watermarks.
- **Watermark Detection Result**: Indicates whether a watermark was found and its confidence level.
- **Detected Hex Message**: If a watermark is detected, this shows the found code in hexadecimal format.
- **Match Result**: Compares the detected hex code to the one you entered, indicating if they match or not. This confirms whether the detected watermark is the one you're looking for.
""")
            audio_input_detect_watermark = gr.Audio(label="Upload Audio File for Watermark Detection", type="filepath")
            original_hex_input = gr.Textbox(label="Original Hex Message for Comparison", placeholder="Enter the original hex message here")
            detect_watermark_button = gr.Button("Detect Watermark")
            watermark_detection_output = gr.Textbox(label="Watermark Detection Result")
            detected_message_output = gr.Textbox(label="Detected Hex Message")
            match_result_output = gr.Textbox(label="Match Result")
            detect_watermark_button.click(fn=detect_watermark, inputs=[audio_input_detect_watermark, original_hex_input], outputs=[watermark_detection_output, detected_message_output, match_result_output])

demo.launch()