| import gradio as gr | |
| import numpy as np | |
| import os | |
| import datetime | |
| import torch | |
| import soundfile | |
| from wavmark.utils import file_reader | |
| import wavmark | |
| def my_read_file(audio_path, max_second, default_sr=16000): | |
| signal, sr, audio_length_second = file_reader.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 add_watermark(audio_path, watermark_text, max_second_encode=60): | |
| assert len(watermark_text) == 16 | |
| watermark_npy = np.array([int(i) for i in watermark_text]) | |
| signal, sr, audio_length_second = my_read_file(audio_path, max_second_encode) | |
| watermarked_signal, _ = wavmark.encode_watermark(model, signal, watermark_npy, show_progress=False) | |
| tmp_file_name = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + "_" + watermark_text + ".wav" | |
| tmp_file_path = '/tmp/' + tmp_file_name | |
| soundfile.write(tmp_file_path, watermarked_signal, sr) | |
| return tmp_file_path | |
| def decode_watermark(audio_path, max_second_decode=30): | |
| assert os.path.exists(audio_path) | |
| signal, sr, audio_length_second = my_read_file(audio_path, max_second_decode) | |
| payload_decoded, _ = wavmark.decode_watermark(model, signal, show_progress=False) | |
| if payload_decoded is None: | |
| return "No Watermark" | |
| return "".join([str(i) for i in payload_decoded]) | |
| def create_default_value(len_start_bit=16): | |
| def_val_npy = np.random.choice([0, 1], size=32 - len_start_bit) | |
| return "".join([str(i) for i in def_val_npy]) | |
| def main(): | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| gr.Markdown("# Audio WaterMarking") | |
| with gr.Row(): | |
| gr.Markdown("You can upload an audio file and encode a custom 16-bit watermark or perform decoding from a watermarked audio. See [WaveMark toolkit](https://github.com/wavmark/wavmark) for further details.") | |
| with gr.Row(): | |
| audio_file = gr.Audio(label="Upload Audio", type="filepath") | |
| action = gr.Radio(["Add Watermark", "Decode Watermark"], label="Select Action") | |
| watermark_text = gr.Textbox(label="The watermark (0, 1 list of length-16):", value=create_default_value()) | |
| submit_button = gr.Button("Submit") | |
| with gr.Row(): | |
| output = gr.Audio(label="Processed Audio") | |
| decode_output = gr.Textbox(label="Decoded Watermark") | |
| def process_audio(audio_file, action, watermark_text): | |
| if action == "Add Watermark" and audio_file: | |
| return add_watermark(audio_file, watermark_text), None | |
| elif action == "Decode Watermark" and audio_file: | |
| return None, decode_watermark(audio_file) | |
| else: | |
| return None, None | |
| submit_button.click(process_audio, inputs=[audio_file, action, watermark_text], outputs=[output, decode_output]) | |
| demo.launch() | |
| 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) | |
| main() | |