Update app.py

app.py

@@ -197,10 +197,16 @@ def main():
                         #st.write("Time Cost: %d seconds" % encode_time_cost)
 
 #                         # st.button("Add Watermark", disabled=False)
-        # elif action == "Decode Watermark":
-        #     if st.button("Decode"):
-        #         with st.spinner("Decoding..."):
-        #             decode_watermark(tmp_input_audio_file)
+        elif action == "Detect Watermark":
+            if st.button("Detect"):
+                with st.spinner("Detecting..."):
+                    result, message = detector.detect_watermark(watermarked_audio, sample_rate=default_sr, message_threshold=0.5)
+
+                    print(f"\nThis is likely a watermarked audio: {result}")
+
+                    # Run on an unwatermarked audio
+                    result2, message2 = detector.detect_watermark(wav, sample_rate=default_sr, message_threshold=0.5)
+                    print(f"This is likely an unwatermarked audio: {result2}")
 
 
 if __name__ == "__main__":
@@ -211,6 +217,7 @@ if __name__ == "__main__":
     device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
     # model = wavmark.load_model().to(device)
     model = AudioSeal.load_generator("audioseal_wm_16bits")
+    detector = AudioSeal.load_detector(("audioseal_detector_16bits"))
     main()
 
     # audio_path = "/Users/my/Library/Mobile Documents/com~apple~CloudDocs/CODE/PycharmProjects/4_语音水印/419_huggingface水印/WavMark/example.wav"