d

app.py

@@ -1,60 +1,51 @@
 import pdb
 import time
-
-
+import wavmark
 import streamlit as st
 import os
-from utils import wm_add_v2, file_reader, model_util, wm_decode_v2, bin_util
-from models import my_model_v7_recover
 import torch
 import uuid
 import datetime
 import numpy as np
 import soundfile
 from huggingface_hub import hf_hub_download, HfApi
+from wavmark.utils import file_reader
 
 
 # Function to add watermark to audio
 def add_watermark(audio_path, watermark_text):
-    assert len(watermark_text) == 5
-
-    start_bit, msg_bit, watermark = wm_add_v2.create_parcel_message(len_start_bit, 32, watermark_text)
+    assert len(watermark_text) == 16
+    watermark_npy = np.array([int(i) for i in watermark_text])
+    # todo: 控制时间
 
-    data, sr, audio_length_second = file_reader.read_as_single_channel_16k(audio_path, 16000)
-
-    _, signal_wmd, time_cost = wm_add_v2.add_watermark(watermark, data, 16000, 0.1, device, model)
+    signal, sr, audio_length_second = file_reader.read_as_single_channel_16k(audio_path, 16000)
+    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') + "_" + str(uuid.uuid4()) + ".wav"
     tmp_file_path = '/tmp/' + tmp_file_name
-    soundfile.write(tmp_file_path, signal_wmd, sr)
+    soundfile.write(tmp_file_path, watermarked_signal, sr)
     return tmp_file_path
 
 
 # Function to decode watermark from audio
 def decode_watermark(audio_path):
-    data, sr, audio_length_second = file_reader.read_as_single_channel_16k(audio_path, 16000)
-    data = data[0:5 * sr]
-    start_bit = wm_add_v2.fix_pattern[0:len_start_bit]
-    support_count, mean_result, results = wm_decode_v2.extract_watermark_v2(
-        data,
-        start_bit,
-        0.1,
-        16000,
-        0.3,
-        model,
-        device, "best")
-
-    if mean_result is None:
+    watermarked_signal, sr, audio_length_second = file_reader.read_as_single_channel_16k(audio_path, 16000)
+    payload_decoded, _ = wavmark.decode_watermark(model, watermarked_signal, show_progress=False)
+
+    if payload_decoded is None:
         return "No Watermark"
 
-    payload = mean_result[len_start_bit:]
-    return bin_util.binArray2HexStr(payload)
+    return payload_decoded
 
 
 # Main web app
 def main():
+    max_upload_size = 20 * 1024 * 1024  # 20 MB in bytes
+
     if "def_value" not in st.session_state:
-        st.session_state.def_value = bin_util.binArray2HexStr(np.random.choice([0, 1], size=32 - len_start_bit))
+        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
 
     st.title("Neural Audio Watermark")
     st.write("Choose the action you want to perform:")
@@ -62,14 +53,15 @@ def main():
     action = st.selectbox("Select Action", ["Add Watermark", "Decode Watermark"])
 
     if action == "Add Watermark":
-        audio_file = st.file_uploader("Upload Audio File (WAV)", type=["wav"], accept_multiple_files=False)
+        audio_file = st.file_uploader("Upload Audio File (WAV)", type=["wav"], accept_multiple_files=False,
+                                      max_upload_size=max_upload_size)
         if audio_file:
             tmp_input_audio_file = os.path.join("/tmp/", audio_file.name)
             with open(tmp_input_audio_file, "wb") as f:
                 f.write(audio_file.getbuffer())
             st.audio(tmp_input_audio_file, format="audio/wav")
 
-            watermark_text = st.text_input("Enter Watermark Text (5 English letters)", value=st.session_state.def_value)
+            watermark_text = st.text_input("Enter Watermark", value=st.session_state.def_value)
 
             add_watermark_button = st.button("Add Watermark", key="add_watermark_btn")
             if add_watermark_button:  # 点击按钮后执行的
@@ -90,7 +82,8 @@ def main():
                         # st.button("Add Watermark", disabled=False)
 
     elif action == "Decode Watermark":
-        audio_file = st.file_uploader("Upload Audio File (WAV/MP3)", type=["wav", "mp3"], accept_multiple_files=False)
+        audio_file = st.file_uploader("Upload Audio File (WAV/MP3)", type=["wav", "mp3"], accept_multiple_files=False,
+                                      max_upload_size=max_upload_size)
         if audio_file:
             if st.button("Decode Watermark"):
                 # 1.保存
@@ -110,36 +103,10 @@ def main():
                 st.write("Time Cost:%d seconds" % (decode_cost))
 
 
-def load_model(resume_path):
-    n_fft = 1000
-    hop_length = 400
-    # https://huggingface.co/M4869/InvertibleWM/blob/main/step59000_snr39.99_pesq4.35_BERP_none0.30_mean1.81_std1.81.pkl
-    # api_key = st.secrets["api_key"]
-    # print(api_key, api_key)
-    api_key = "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"
-    model_ckpt_path = hf_hub_download(repo_id="M4869/InvertibleWM",
-                                      filename="step59000_snr39.99_pesq4.35_BERP_none0.30_mean1.81_std1.81.pkl",
-                                      token=api_key
-                                      )
-    # print("model_ckpt_path", model_ckpt_path)
-    resume_path = model_ckpt_path
-    # return
-
-    model = my_model_v7_recover.Model(16000, 32, n_fft, hop_length,
-                                      use_recover_layer=False, num_layers=8).to(device)
-    checkpoint = torch.load(resume_path, map_location=torch.device('cpu'))
-    state_dict = model_util.map_state_dict(checkpoint['model'])
-    model.load_state_dict(state_dict, strict=True)
-    model.eval()
-    return model
-
-
 if __name__ == "__main__":
-    len_start_bit = 12
+    len_start_bit = 16
 
     device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-
-    model = load_model("./data/step59000_snr39.99_pesq4.35_BERP_none0.30_mean1.81_std1.81.pkl")
+    model = wavmark.load_model().to(device)
 
     main()
-    # decode_watermark("/Users/my/Downloads/7a95b353a46893903e9f946c24170b210ce14e8c52c63bb2ab3d144e.wav")

requirements.txt

@@ -1,6 +1,4 @@
+wavmark
 soundfile
 librosa
-resampy
-torch==1.13.1
-torchvision==0.14.1
-torchaudio==0.13.1
+resampy