Update app.py

app.py

@@ -101,8 +101,8 @@ def main():
         #Load the WAV file using torchaudio
         if file_extension in [".wav", ".flac"]: 
             wav, sample_rate = torchaudio.load("test.wav")
-            st.markdown("Before unsquueze wav")
-            st.markdown(wav)
+            # st.markdown("Before unsquueze wav")
+            # st.markdown(wav)
             file_extension_ori =".wav"
             #Unsqueeze for line 176
             wav= wav.unsqueeze(0)
@@ -194,49 +194,52 @@ def main():
         elif action == "Detect Watermark":
             detect_watermark_button = st.button("Detect Watermark", key="detect_watermark_btn")
 
-            if audio_file:
-                #1st attempt
-                watermark = model.get_watermark(wav, default_sr)
-                watermarked_audio = wav + watermark
-                print(watermarked_audio.size())
-                size = watermarked_audio.size()
-                #st.markdown(size)
-
-                #print(watermarked_audio.squeeze())
-                # squeeze = watermarked_audio.squeeze(1)
-                # shape = squeeze.size()
+            # if audio_file:
+            #     #1st attempt
+            #     watermark = model.get_watermark(wav, default_sr)
+            #     watermarked_audio = wav + watermark
+            #     print(watermarked_audio.size())
+            #     size = watermarked_audio.size()
+            #     #st.markdown(size)
+
                 
             if detect_watermark_button:
                 with st.spinner("Detecting..."):
                     result, message = detector.detect_watermark(watermarked_audio, sample_rate=default_sr, message_threshold=0.5)
+                    st.markdown("This is likely a watermarked audio:")
                     st.markdown(result)
                     print(f"\nThis is likely a watermarked audio: {result}")
 
                     # Run on an unwatermarked audio
-                    tmp_input_audio_file = os.path.join("/tmp/", audio_file.name)
-                    file_extension = os.path.splitext(tmp_input_audio_file)[1].lower()
-                    if file_extension in [".wav", ".flac"]:
-                        with open("test.wav", "wb") as f:
-                            f.write(audio_file.getbuffer())
 
-                    elif file_extension == ".mp3":
-                        with open("test.mp3", "wb") as f:
-                            f.write(audio_file.getbuffer())
+                    audio_file = st.file_uploader("Upload Audio", type=["wav", "mp3"], accept_multiple_files=False)
+
+                    if audio_file:
+                        tmp_input_audio_file = os.path.join("/tmp/", audio_file.name)
+                        file_extension = os.path.splitext(tmp_input_audio_file)[1].lower()
+                        if file_extension in [".wav", ".flac"]:
+                            with open("test.wav", "wb") as f:
+                                f.write(audio_file.getbuffer())
+
+                        elif file_extension == ".mp3":
+                            with open("test.mp3", "wb") as f:
+                                f.write(audio_file.getbuffer())
 
                     #Load the WAV file using torchaudio
-                    if file_extension in [".wav", ".flac"]: 
-                        wav, sample_rate = torchaudio.load("test.wav")
+                        if file_extension in [".wav", ".flac"]: 
+                            wav, sample_rate = torchaudio.load("test.wav")
 
-                    elif file_extension == ".mp3":
-                        # Load an MP3 file
-                        audio = AudioSegment.from_mp3("test.mp3")
+                        elif file_extension == ".mp3":
+                            # Load an MP3 file
+                            audio = AudioSegment.from_mp3("test.mp3")
 
-                        # Export it as a WAV file
-                        audio.export("test.wav", format="wav")
-                        wav3, sample_rate = torchaudio.load("test.wav")
-                        wav3= wav3.unsqueeze(0)
+                            # Export it as a WAV file
+                            audio.export("test.wav", format="wav")
+                            wav3, sample_rate = torchaudio.load("test.wav")
+                            wav3= wav3.unsqueeze(0)
                         
                     result2, message2 = detector.detect_watermark(wav3, sample_rate=default_sr, message_threshold=0.5)
+                    st.markdown("This is likely an unwatermarked audio:")
                     print(f"This is likely an unwatermarked audio: {result2}")
                     st.markdown(result2)