| import time | |
| import streamlit as st | |
| # from transformers import pipeline | |
| import os | |
| import torch | |
| import datetime | |
| import numpy as np | |
| import soundfile | |
| from wavmark.utils import file_reader | |
| from audioseal import AudioSeal | |
| import torchaudio | |
| from pydub import AudioSegment | |
| import io | |
| import librosa | |
| import ffmpeg | |
| #from torchaudio.io import CodecConfig | |
| # import numpy | |
| # def my_read_file(audio_path, max_second): | |
| # signal, sr, audio_length_second = read_as_single_channel_16k(audio_path, default_sr) | |
| # if audio_length_second > max_second: | |
| # signal = signal[0:default_sr * max_second] | |
| # audio_length_second = max_second | |
| # return signal, sr, audio_length_second | |
| def create_default_value(): | |
| if "def_value" not in st.session_state: | |
| def_val_npy = np.random.choice([0, 1], size=32 - len_start_bit) | |
| def_val_str = "".join([str(i) for i in def_val_npy]) | |
| st.session_state.def_value = def_val_str | |
| def download_sample_audio(): | |
| url = "https://keithito.com/LJ-Speech-Dataset/LJ037-0171.wav" | |
| with open("test.wav", "wb") as f: | |
| resp = urllib.request.urlopen(url) | |
| f.write(resp.read()) | |
| wav, sample_rate = torchaudio.load("test.wav") | |
| return wav, sample_rate | |
| # Main web app | |
| def main(): | |
| create_default_value() | |
| # st.title("MDS07 Demo Presentation") | |
| # st.write("https://github.com/ravindi-r/audioseal") | |
| markdown_text = """ | |
| # MDS07 Demo Presentation | |
| [AudioSeal](https://github.com/ravindi-r/audioseal) is the next-generation watermarking tool driven by AI. | |
| You can upload an audio file and encode a custom 16-bit watermark or perform decoding from a watermarked audio. | |
| This page is for demonstration usage. | |
| If you have longer files for processing, we recommend using [our python toolkit](https://github.com/ravindi-r/audioseal). | |
| """ | |
| # 使用st.markdown渲染Markdown文本 | |
| 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()) | |
| 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", | |
| ) | |
| elif action == "Detect Watermark": | |
| detect_watermark_button = st.button("Detect Watermark", key="detect_watermark_btn") | |
| 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) | |
| if result2 < 0.5: | |
| st.markdown("This is likely an unwatermarked audio!") | |
| else: | |
| st.markdown("This is likely an watermarked audio!") | |
| except RuntimeError: | |
| 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)") | |
| # 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 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!") | |
| if __name__ == "__main__": | |
| default_sr = 16000 | |
| max_second_encode = 60 | |
| max_second_decode = 30 | |
| len_start_bit = 16 | |
| 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" | |
| # decoded_watermark, decode_cost = decode_watermark(audio_path) | |
| # print(decoded_watermark) |