Update app.py

app.py

@@ -55,15 +55,20 @@ def watermark_audio(audio_file_path, 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):
+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_message = binary_to_hex(binary_message)  # Convert detected binary message to hex for display
+    detected_hex_message = binary_to_hex(binary_message)
     
-    return result, detected_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")
@@ -84,9 +89,11 @@ with gr.Blocks(css=style) as demo:
     with gr.Tab("Detect Watermark"):
         with gr.Column(scale=6):
             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 Unique Message")
-            detect_watermark_button.click(fn=detect_watermark, inputs=audio_input_detect_watermark, outputs=[watermark_detection_output, detected_message_output])
+            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()
+demo.launch()