Update app.py

app.py

@@ -59,64 +59,44 @@ def main():
     st.markdown(markdown_text)
 
     audio_file = st.file_uploader("Upload Audio", type=["wav", "mp3"], accept_multiple_files=False)
+    try:
+        if audio_file:
+            #2nd attempt
+            # Save file to local storage
+            tmp_input_audio_file = os.path.join("/tmp/", audio_file.name)
+            file_extension = os.path.splitext(tmp_input_audio_file)[1].lower()
+            #st.markdown(file_extension)
+            if file_extension in [".wav", ".flac"]:
+                with open("test.wav", "wb") as f:
+                    f.write(audio_file.getbuffer())
 
-    if audio_file:
-        # 保存文件到本地：
-        # tmp_input_audio_file = os.path.join("/tmp/", audio_file.name)
-        # st.markdown(tmp_input_audio_file)
-        # with open(tmp_input_audio_file, "wb") as f:
-        #     f.write(audio_file.getbuffer())
-        # st.audio(tmp_input_audio_file, format="mp3/wav")
+                st.audio("test.wav", format="audio/wav")
 
+            elif file_extension == ".mp3":
+                with open("test.mp3", "wb") as f:
+                    f.write(audio_file.getbuffer())
 
-        #1st attempt
-        #audio_path = " audio_file.name"
-
-        # audio, sr = torchaudio.load(audio_file)
-        # st.audio(audio_file, format="audio/mpeg")
-        # audio= audio.unsqueeze(0)
-
-        # st.markdown("SR")
-        # st.markdown(sr)
-        # st.markdown("after unsqueeze wav or mp3")
-        # st.markdown(audio)
-        
-        #2nd attempt
-        # Save file to local storage
-        tmp_input_audio_file = os.path.join("/tmp/", audio_file.name)
-        file_extension = os.path.splitext(tmp_input_audio_file)[1].lower()
-        #st.markdown(file_extension)
-        if file_extension in [".wav", ".flac"]:
-            with open("test.wav", "wb") as f:
-                f.write(audio_file.getbuffer())
-
-            st.audio("test.wav", format="audio/wav")
-
-        elif file_extension == ".mp3":
-            with open("test.mp3", "wb") as f:
-                f.write(audio_file.getbuffer())
-
-            st.audio("test.mp3", format="audio/mpeg")
+                st.audio("test.mp3", format="audio/mpeg")
 
         #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")
             # st.markdown("Before unsquueze wav")
             # st.markdown(wav)
-            file_extension_ori =".wav"
+                file_extension_ori =".wav"
             #Unsqueeze for line 176
-            wav= wav.unsqueeze(0)
+                wav= wav.unsqueeze(0)
                 
-        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")
-            wav= wav3.unsqueeze(0)
-            file_extension_ori =".mp3"
-            file_extension =".wav"
+                audio.export("test.wav", format="wav")
+                wav3, sample_rate = torchaudio.load("test.wav")
+                wav= wav3.unsqueeze(0)
+                file_extension_ori =".mp3"
+                file_extension =".wav"
 
             #RuntimeError: Could not infer dtype of numpy.float32
             #wav = torch.tensor(wav3).float() / 32768.0
@@ -130,15 +110,15 @@ def main():
         #Unsqueeze for line 176
         # wav= wav.unsqueeze(0)
 
-        action = st.selectbox("Select Action", ["Add Watermark", "Detect Watermark"])
+            action = st.selectbox("Select Action", ["Add Watermark", "Detect Watermark"])
 
-        if action == "Add Watermark":
-            #watermark_text = st.text_input("The watermark (0, 1 list of length-16):", value=st.session_state.def_value)
-            add_watermark_button = st.button("Add Watermark", key="add_watermark_btn")
-            if add_watermark_button:  # 点击按钮后执行的
-                #if audio_file and watermark_text:
-                if audio_file:
-                    with st.spinner("Adding Watermark..."):
+            if action == "Add Watermark":
+                #watermark_text = st.text_input("The watermark (0, 1 list of length-16):", value=st.session_state.def_value)
+                add_watermark_button = st.button("Add Watermark", key="add_watermark_btn")
+                if add_watermark_button:  # 点击按钮后执行的
+                    #if audio_file and watermark_text:
+                    if audio_file:
+                        with st.spinner("Adding Watermark..."):
                         #wav = my_read_file(wav,max_second_encode)
                         
                         #1st attempt
@@ -154,84 +134,216 @@ def main():
                         #st.markdown(shape)
 
                         #st.markdown(squeeze)
-                    if file_extension_ori in [".wav", ".flac"]:
-                        torchaudio.save("output.wav", squeeze, default_sr, bits_per_sample=16)
-                        watermarked_wav = torchaudio.save("output.wav", squeeze, default_sr, bits_per_sample=16)
-
-                        st.audio("output.wav", format="audio/wav")
-
-                        with open("output.wav", "rb") as file:
-                            #file.read()
-                            #file.write(watermarked_wav.getbuffer())
-                            binary_data = file.read()
-                            btn = st.download_button(
-                                label="Download watermarked audio",
-                                data=binary_data,
-                                file_name="output.wav",
-                                mime="audio/wav",
-                            )
-
-
-                    elif file_extension_ori == ".mp3":
-                        torchaudio.save("output.wav", squeeze, default_sr)
-                        watermarked_mp3 = torchaudio.save("output.wav", squeeze, default_sr)
-                        audio = AudioSegment.from_wav("output.wav")
-
-                        # Export as MP3
-                        audio.export("output.mp3", format="mp3")
-                        st.audio("output.mp3", format="audio/mpeg")
+                        if file_extension_ori in [".wav", ".flac"]:
+                            torchaudio.save("output.wav", squeeze, default_sr, bits_per_sample=16)
+                            watermarked_wav = torchaudio.save("output.wav", squeeze, default_sr, bits_per_sample=16)
+
+                            st.audio("output.wav", format="audio/wav")
+
+                            with open("output.wav", "rb") as file:
+                                #file.read()
+                                #file.write(watermarked_wav.getbuffer())
+                                binary_data = file.read()
+                                btn = st.download_button(
+                                    label="Download watermarked audio",
+                                    data=binary_data,
+                                    file_name="output.wav",
+                                    mime="audio/wav",
+                                )
+
+
+                        elif file_extension_ori == ".mp3":
+                            torchaudio.save("output.wav", squeeze, default_sr)
+                            watermarked_mp3 = torchaudio.save("output.wav", squeeze, default_sr)
+                            audio = AudioSegment.from_wav("output.wav")
+
+                            # Export as MP3
+                            audio.export("output.mp3", format="mp3")
+                            st.audio("output.mp3", format="audio/mpeg")
+                        
+                            with open("output.mp3", "rb") as file:
+                                #file.write(watermarked_wav.getbuffer())
+                                binary_data = file.read()
+                                st.download_button(
+                                    label="Download watermarked audio",
+                                    data=binary_data,
+                                    file_name="output.mp3",
+                                    mime="audio/mpeg",
+    except error:
+        st.error("Please input audio with one channel (mono-channel)")
+        
+    # if audio_file:
+    #     # 保存文件到本地：
+    #     # tmp_input_audio_file = os.path.join("/tmp/", audio_file.name)
+    #     # st.markdown(tmp_input_audio_file)
+    #     # with open(tmp_input_audio_file, "wb") as f:
+    #     #     f.write(audio_file.getbuffer())
+    #     # st.audio(tmp_input_audio_file, format="mp3/wav")
+
+
+    #     #1st attempt
+    #     #audio_path = " audio_file.name"
+
+    #     # audio, sr = torchaudio.load(audio_file)
+    #     # st.audio(audio_file, format="audio/mpeg")
+    #     # audio= audio.unsqueeze(0)
+
+    #     # st.markdown("SR")
+    #     # st.markdown(sr)
+    #     # st.markdown("after unsqueeze wav or mp3")
+    #     # st.markdown(audio)
+        
+    #     #2nd attempt
+    #     # Save file to local storage
+    #     tmp_input_audio_file = os.path.join("/tmp/", audio_file.name)
+    #     file_extension = os.path.splitext(tmp_input_audio_file)[1].lower()
+    #     #st.markdown(file_extension)
+    #     if file_extension in [".wav", ".flac"]:
+    #         with open("test.wav", "wb") as f:
+    #             f.write(audio_file.getbuffer())
+
+    #         st.audio("test.wav", format="audio/wav")
+
+    #     elif file_extension == ".mp3":
+    #         with open("test.mp3", "wb") as f:
+    #             f.write(audio_file.getbuffer())
+
+    #         st.audio("test.mp3", format="audio/mpeg")
+
+    #     #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)
+    #         file_extension_ori =".wav"
+    #         #Unsqueeze for line 176
+    #         wav= wav.unsqueeze(0)
+                
+    #     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")
+    #         wav= wav3.unsqueeze(0)
+    #         file_extension_ori =".mp3"
+    #         file_extension =".wav"
+
+    #         #RuntimeError: Could not infer dtype of numpy.float32
+    #         #wav = torch.tensor(wav3).float() / 32768.0
+            
+    #         #RuntimeError: Numpy is not available
+    #         # wav = torch.from_numpy(wav3) #/32768.0
+    #         # wav = wav.unsqueeze(0).unsqueeze(0)
+    #         # st.markdown("Before unsqueeze mp3")
+    #         # st.markdown(wav)
+        
+    #     #Unsqueeze for line 176
+    #     # wav= wav.unsqueeze(0)
+
+    #     action = st.selectbox("Select Action", ["Add Watermark", "Detect Watermark"])
+
+    #     if action == "Add Watermark":
+    #         #watermark_text = st.text_input("The watermark (0, 1 list of length-16):", value=st.session_state.def_value)
+    #         add_watermark_button = st.button("Add Watermark", key="add_watermark_btn")
+    #         if add_watermark_button:  # 点击按钮后执行的
+    #             #if audio_file and watermark_text:
+    #             if audio_file:
+    #                 with st.spinner("Adding Watermark..."):
+    #                     #wav = my_read_file(wav,max_second_encode)
+                        
+    #                     #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()
+    #                     #st.markdown(shape)
+
+    #                     #st.markdown(squeeze)
+    #                 if file_extension_ori in [".wav", ".flac"]:
+    #                     torchaudio.save("output.wav", squeeze, default_sr, bits_per_sample=16)
+    #                     watermarked_wav = torchaudio.save("output.wav", squeeze, default_sr, bits_per_sample=16)
+
+    #                     st.audio("output.wav", format="audio/wav")
+
+    #                     with open("output.wav", "rb") as file:
+    #                         #file.read()
+    #                         #file.write(watermarked_wav.getbuffer())
+    #                         binary_data = file.read()
+    #                         btn = st.download_button(
+    #                             label="Download watermarked audio",
+    #                             data=binary_data,
+    #                             file_name="output.wav",
+    #                             mime="audio/wav",
+    #                         )
+
+
+    #                 elif file_extension_ori == ".mp3":
+    #                     torchaudio.save("output.wav", squeeze, default_sr)
+    #                     watermarked_mp3 = torchaudio.save("output.wav", squeeze, default_sr)
+    #                     audio = AudioSegment.from_wav("output.wav")
+
+    #                     # Export as MP3
+    #                     audio.export("output.mp3", format="mp3")
+    #                     st.audio("output.mp3", format="audio/mpeg")
                         
-                        with open("output.mp3", "rb") as file:
-                            #file.write(watermarked_wav.getbuffer())
-                            binary_data = file.read()
-                            st.download_button(
-                                label="Download watermarked audio",
-                                data=binary_data,
-                                file_name="output.mp3",
-                                mime="audio/mpeg",
-                            )
-                    # except RuntimeError:
-                    #     st.error("Please input audio with one channel (mono-channel)")
+    #                     with open("output.mp3", "rb") as file:
+    #                         #file.write(watermarked_wav.getbuffer())
+    #                         binary_data = file.read()
+    #                         st.download_button(
+    #                             label="Download watermarked audio",
+    #                             data=binary_data,
+    #                             file_name="output.mp3",
+    #                             mime="audio/mpeg",
+    #                         )
+    #                 # except RuntimeError:
+    #                 #     st.error("Please input audio with one channel (mono-channel)")
                                  
-        elif action == "Detect Watermark":
-            detect_watermark_button = st.button("Detect Watermark", key="detect_watermark_btn")
+    #     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)
+    #         # 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("Probability of audio being watermarked: ")
-                    # st.markdown(result)
-                    # st.markdown("This is likely a watermarked audio!")
-                    # print(f"\nThis is likely a watermarked audio: {result}")
-
-                    #Run on an unwatermarked audio
-
-                    if file_extension in [".wav", ".flac"]: 
-                        wav, sample_rate = torchaudio.load("test.wav")
-                        wav= wav.unsqueeze(0)
-
-                    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")
-                        wav, sample_rate = torchaudio.load("test.wav")
-                        wav= wav.unsqueeze(0)
+    #         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("Probability of audio being watermarked: ")
+    #                 # st.markdown(result)
+    #                 # st.markdown("This is likely a watermarked audio!")
+    #                 # print(f"\nThis is likely a watermarked audio: {result}")
+
+    #                 #Run on an unwatermarked audio
+
+    #                 if file_extension in [".wav", ".flac"]: 
+    #                     wav, sample_rate = torchaudio.load("test.wav")
+    #                     wav= wav.unsqueeze(0)
+
+    #                 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")
+    #                     wav, sample_rate = torchaudio.load("test.wav")
+    #                     wav= wav.unsqueeze(0)
                         
-                    result2, message2 = detector.detect_watermark(wav, sample_rate=default_sr, message_threshold=0.5)
-                    print(f"This is likely an unwatermarked audio: {result2}")
-                    st.markdown("Probability of audio being watermarked: ")
-                    st.markdown(result2)
-                    st.markdown("This is likely 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}")
+    #                 st.markdown("Probability of audio being watermarked: ")
+    #                 st.markdown(result2)
+    #                 st.markdown("This is likely an unwatermarked audio!")
 
 
 if __name__ == "__main__":