Update app.py

app.py

@@ -7,11 +7,20 @@ from audioseal import AudioSeal
 import random
 import string
 from pathlib import Path
+from datetime import datetime
+import json
+import os
 
 # Initialize logging
 logging.basicConfig(level=logging.DEBUG, filename='app.log', filemode='w', format='%(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
+# File to store audio metadata
+metadata_file = 'audio_metadata.json'
+if not os.path.exists(metadata_file):
+    with open(metadata_file, 'w') as f:
+        json.dump({}, f)
+
 # Helper function for generating a unique alphanumeric message
 def generate_unique_message(length=16):
     characters = string.ascii_letters + string.digits
@@ -22,11 +31,6 @@ 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)
@@ -39,8 +43,36 @@ 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
+# Function to generate enhanced unique identifier with timestamp and sequential number
+def generate_enhanced_identifier():
+    timestamp = datetime.now().strftime('%Y%m%d%H%M%S%f')
+    sequential_number = str(get_next_sequential_number()).zfill(6)
+    return f"{timestamp}-{sequential_number}"
+
+# Function to increment and get the next sequential number from the metadata file
+def get_next_sequential_number():
+    with open(metadata_file, 'r+') as f:
+        data = json.load(f)
+        next_number = data.get('next_sequential_number', 1)
+        data['next_sequential_number'] = next_number + 1
+        f.seek(0)
+        json.dump(data, f, indent=4)
+        f.truncate()
+    return next_number
+
+# Function to save metadata for an audio file
+def save_audio_metadata(unique_id, original_hex, enhanced_id):
+    with open(metadata_file, 'r+') as f:
+        data = json.load(f)
+        data['audio_files'] = data.get('audio_files', {})
+        data['audio_files'][unique_id] = {'original_hex': original_hex, 'enhanced_id': enhanced_id}
+        f.seek(0)
+        json.dump(data, f, indent=4)
+        f.truncate()
+
+# Modify the watermark_audio function to include enhanced ID generation and saving metadata
 def watermark_audio(audio_file_path, unique_message):
+    # Original watermarking process
     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")
@@ -52,7 +84,12 @@ def watermark_audio(audio_file_path, unique_message):
     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
+    
+    # Enhanced ID generation and metadata saving
+    enhanced_id = generate_enhanced_identifier()
+    save_audio_metadata(unique_message, hex_message, enhanced_id)
+    
+    return temp_file.name, hex_message, enhanced_id  # Include enhanced ID in the return statement
 
 # Function to detect watermark in audio
 def detect_watermark(audio_file_path, original_hex_message=None):
@@ -70,45 +107,29 @@ def detect_watermark(audio_file_path, original_hex_message=None):
     
     return result, detected_hex_message, match_result
 
-# Setup for Gradio interface
 # Load the CSS styles
 style_path = Path("style.css")
-style = style_path.read_text()
+if style_path.exists():
+    style = style_path.read_text()
+else:
+    style = ""
 
+# Define Gradio interface
 with gr.Blocks(css=style) as demo:
     with gr.Tab("Watermark Audio"):
         with gr.Column(scale=6):
-            gr.Markdown("### How to Watermark Your Audio", elem_id="how_to_watermark")
-            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.
-""", elem_id="watermark_process_explanation")
             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")
+            unique_message_output = gr.Textbox(label="Unique Message")
             watermarked_audio_output = gr.Audio(label="Watermarked Audio")
             message_output = gr.Textbox(label="Message Used for Watermarking")
+            enhanced_id_output = gr.Textbox(label="Enhanced ID")
+            generate_message_button = gr.Button("Generate Unique Message")
+            watermark_button = gr.Button("Apply Watermark")
             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])
+            watermark_button.click(fn=watermark_audio, inputs=[audio_input_watermark, unique_message_output], outputs=[watermarked_audio_output, message_output, enhanced_id_output])
 
     with gr.Tab("Detect Watermark"):
         with gr.Column(scale=6):
-            gr.Markdown("### How to Detect a Watermark in Your Audio", elem_id="how_to_detect")
-            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.
-""", elem_id="detection_process_explanation")
             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")
@@ -117,4 +138,4 @@ Use this tool to check if an audio file contains a specific watermark. Here's th
             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()
+demo.launch()