Update app.py

app.py

@@ -1,33 +1,37 @@
 import gradio as gr
-import librosa
-import numpy as np
 import torch
 import torchaudio
-import torchaudio.transforms as T
+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):
-    """Generates a unique alphanumeric message of the given length."""
     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):
-    """Converts a message to binary format, truncating or padding to `bit_length`."""
     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):
-    """Converts a binary string back to its ASCII message representation."""
     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:
@@ -35,30 +39,33 @@ def load_and_resample_audio(audio_file_path, target_sample_rate=16000):
         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")
     
-    # Convert the unique message to binary format expected by the model
-    binary_message = message_to_binary(unique_message, bit_length=128)  # Adjust bit_length as needed
+    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, unique_message
+    return temp_file.name, hex_message  # Display hex message to the user
 
+# Function to detect watermark in audio
 def detect_watermark(audio_file_path):
     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_message = binary_to_message(binary_message)
+    detected_message = binary_to_hex(binary_message)  # Convert detected binary message to hex for display
     
     return result, detected_message
 
+# Setup for Gradio interface
 style_path = Path("style.css")
 style = style_path.read_text()
 
@@ -82,4 +89,4 @@ with gr.Blocks(css=style) as demo:
             detected_message_output = gr.Textbox(label="Detected Unique Message")
             detect_watermark_button.click(fn=detect_watermark, inputs=audio_input_detect_watermark, outputs=[watermark_detection_output, detected_message_output])
 
-demo.launch()
+demo.launch()