Updating to gradio

#1
by derek-thomas HF staff - opened
Files changed (1) hide show
  1. app.py +34 -71
app.py CHANGED
@@ -1,104 +1,69 @@
1
- import time
2
- import wavmark
3
- import streamlit as st
4
  import os
5
- import torch
6
  import datetime
7
- import numpy as np
8
  import soundfile
9
  from wavmark.utils import file_reader
10
- import subprocess
11
- import sys
12
- import time
13
-
14
-
15
 
16
- def my_read_file(audio_path, max_second):
17
  signal, sr, audio_length_second = file_reader.read_as_single_channel_16k(audio_path, default_sr)
18
  if audio_length_second > max_second:
19
  signal = signal[0:default_sr * max_second]
20
  audio_length_second = max_second
21
-
22
  return signal, sr, audio_length_second
23
 
24
-
25
- def add_watermark(audio_path, watermark_text):
26
- #t1 = time.time()
27
  assert len(watermark_text) == 16
28
  watermark_npy = np.array([int(i) for i in watermark_text])
29
  signal, sr, audio_length_second = my_read_file(audio_path, max_second_encode)
30
  watermarked_signal, _ = wavmark.encode_watermark(model, signal, watermark_npy, show_progress=False)
31
-
32
  tmp_file_name = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + "_" + watermark_text + ".wav"
33
  tmp_file_path = '/tmp/' + tmp_file_name
34
  soundfile.write(tmp_file_path, watermarked_signal, sr)
35
- #encode_time_cost = time.time() - t1
36
  return tmp_file_path
37
 
38
- #def encode_water()
39
-
40
- def decode_watermark(audio_path):
41
  assert os.path.exists(audio_path)
42
-
43
  signal, sr, audio_length_second = my_read_file(audio_path, max_second_decode)
44
  payload_decoded, _ = wavmark.decode_watermark(model, signal, show_progress=False)
45
-
46
  if payload_decoded is None:
47
- return "No Watermark"
48
-
49
- payload_decoded_str = "".join([str(i) for i in payload_decoded])
50
- st.write("Result:", payload_decoded_str)
51
-
52
-
53
- def create_default_value():
54
- if "def_value" not in st.session_state:
55
- def_val_npy = np.random.choice([0, 1], size=32 - len_start_bit)
56
- def_val_str = "".join([str(i) for i in def_val_npy])
57
- st.session_state.def_value = def_val_str
58
-
59
 
 
 
 
60
 
61
  def main():
62
- create_default_value()
63
-
64
-
65
- markdown_text = """
66
- # Audio WaterMarking
67
- You can upload an audio file and encode a custom 16-bit watermark or perform decoding from a watermarked audio.
68
-
69
- See [WaveMarktoolkit](https://github.com/wavmark/wavmark) for further details.
70
- """
71
-
72
- st.markdown(markdown_text)
73
-
74
- audio_file = st.file_uploader("Upload Audio", type=["wav", "mp3"], accept_multiple_files=False)
75
-
76
- if audio_file:
77
 
78
- tmp_input_audio_file = os.path.join("/tmp/", audio_file.name)
79
- with open(tmp_input_audio_file, "wb") as f:
80
- f.write(audio_file.getbuffer())
 
 
81
 
 
 
 
82
 
 
 
 
 
 
 
 
83
 
84
- action = st.selectbox("Select Action", ["Add Watermark", "Decode Watermark"])
85
-
86
- if action == "Add Watermark":
87
- watermark_text = st.text_input("The watermark (0, 1 list of length-16):", value=st.session_state.def_value)
88
- add_watermark_button = st.button("Add Watermark", key="add_watermark_btn")
89
- if add_watermark_button:
90
- if audio_file and watermark_text:
91
- with st.spinner("Adding Watermark..."):
92
- watermarked_audio = add_watermark(tmp_input_audio_file, watermark_text)
93
- st.write("Watermarked Audio:")
94
- print("watermarked_audio:", watermarked_audio)
95
- st.audio(watermarked_audio, format="audio/wav")
96
-
97
- elif action == "Decode Watermark":
98
- if st.button("Decode"):
99
- with st.spinner("Decoding..."):
100
- decode_watermark(tmp_input_audio_file)
101
 
 
102
 
103
  if __name__ == "__main__":
104
  default_sr = 16000
@@ -108,5 +73,3 @@ if __name__ == "__main__":
108
  device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
109
  model = wavmark.load_model().to(device)
110
  main()
111
-
112
-
 
1
+ import gradio as gr
2
+ import numpy as np
 
3
  import os
 
4
  import datetime
5
+ import torch
6
  import soundfile
7
  from wavmark.utils import file_reader
8
+ import wavmark
 
 
 
 
9
 
10
+ def my_read_file(audio_path, max_second, default_sr=16000):
11
  signal, sr, audio_length_second = file_reader.read_as_single_channel_16k(audio_path, default_sr)
12
  if audio_length_second > max_second:
13
  signal = signal[0:default_sr * max_second]
14
  audio_length_second = max_second
 
15
  return signal, sr, audio_length_second
16
 
17
+ def add_watermark(audio_path, watermark_text, max_second_encode=60):
 
 
18
  assert len(watermark_text) == 16
19
  watermark_npy = np.array([int(i) for i in watermark_text])
20
  signal, sr, audio_length_second = my_read_file(audio_path, max_second_encode)
21
  watermarked_signal, _ = wavmark.encode_watermark(model, signal, watermark_npy, show_progress=False)
 
22
  tmp_file_name = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + "_" + watermark_text + ".wav"
23
  tmp_file_path = '/tmp/' + tmp_file_name
24
  soundfile.write(tmp_file_path, watermarked_signal, sr)
 
25
  return tmp_file_path
26
 
27
+ def decode_watermark(audio_path, max_second_decode=30):
 
 
28
  assert os.path.exists(audio_path)
 
29
  signal, sr, audio_length_second = my_read_file(audio_path, max_second_decode)
30
  payload_decoded, _ = wavmark.decode_watermark(model, signal, show_progress=False)
 
31
  if payload_decoded is None:
32
+ return "No Watermark"
33
+ return "".join([str(i) for i in payload_decoded])
 
 
 
 
 
 
 
 
 
 
34
 
35
+ def create_default_value(len_start_bit=16):
36
+ def_val_npy = np.random.choice([0, 1], size=32 - len_start_bit)
37
+ return "".join([str(i) for i in def_val_npy])
38
 
39
  def main():
40
+ with gr.Blocks() as demo:
41
+ with gr.Row():
42
+ gr.Markdown("# Audio WaterMarking")
43
+ with gr.Row():
44
+ 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.")
 
 
 
 
 
 
 
 
 
 
45
 
46
+ with gr.Row():
47
+ audio_file = gr.Audio(label="Upload Audio", type="filepath")
48
+ action = gr.Radio(["Add Watermark", "Decode Watermark"], label="Select Action")
49
+ watermark_text = gr.Textbox(label="The watermark (0, 1 list of length-16):", value=create_default_value())
50
+ submit_button = gr.Button("Submit")
51
 
52
+ with gr.Row():
53
+ output = gr.Audio(label="Processed Audio")
54
+ decode_output = gr.Textbox(label="Decoded Watermark")
55
 
56
+ def process_audio(audio_file, action, watermark_text):
57
+ if action == "Add Watermark" and audio_file:
58
+ return add_watermark(audio_file, watermark_text), None
59
+ elif action == "Decode Watermark" and audio_file:
60
+ return None, decode_watermark(audio_file)
61
+ else:
62
+ return None, None
63
 
64
+ submit_button.click(process_audio, inputs=[audio_file, action, watermark_text], outputs=[output, decode_output])
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
65
 
66
+ demo.launch()
67
 
68
  if __name__ == "__main__":
69
  default_sr = 16000
 
73
  device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
74
  model = wavmark.load_model().to(device)
75
  main()