update desc

.gitignore

@@ -1,5 +1,5 @@
 .DS_Store
-
+说明.txt
 #idea
 .idea
 wandb/

app.py

@@ -1,27 +1,31 @@
-import pdb
 import time
 import wavmark
 import streamlit as st
 import os
 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
 
 
+def my_read_file(audio_path, max_second):
+    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):
     t1 = time.time()
     assert len(watermark_text) == 16
     watermark_npy = np.array([int(i) for i in watermark_text])
-    # todo: 控制时间
-
-    signal, sr, audio_length_second = file_reader.read_as_single_channel_16k(audio_path, 16000)
+    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') + "_" + str(uuid.uuid4()) + ".wav"
+    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)
     encode_time_cost = time.time() - t1
@@ -32,12 +36,8 @@ def decode_watermark(audio_path):
     assert os.path.exists(audio_path)
 
     t1 = time.time()
-    watermarked_signal, _, audio_length_second = file_reader.read_as_single_channel_16k(audio_path, 16000)
-    max_second = 5
-    if audio_length_second > max_second:
-        watermarked_signal = watermarked_signal[0:16000 * max_second]
-
-    payload_decoded, _ = wavmark.decode_watermark(model, watermarked_signal, show_progress=False)
+    signal, sr, audio_length_second = my_read_file(audio_path, max_second_decode)
+    payload_decoded, _ = wavmark.decode_watermark(model, signal, show_progress=False)
     decode_cost = time.time() - t1
 
     if payload_decoded is None:
@@ -59,8 +59,19 @@ def create_default_value():
 def main():
     create_default_value()
 
-    st.title("WavMark")
-    st.write("https://github.com/wavmark/wavmark")
+    # st.title("WavMark")
+    # st.write("https://github.com/wavmark/wavmark")
+    markdown_text = """
+    # WavMark
+    [WavMark](https://github.com/wavmark/wavmark) 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 and only process **the first minute** of the audio. 
+    If you have longer files for processing, we recommend using [our python toolkit](https://github.com/wavmark/wavmark).
+    """
+
+    # 使用st.markdown渲染Markdown文本
+    st.markdown(markdown_text)
 
     audio_file = st.file_uploader("Upload Audio", type=["wav", "mp3"], accept_multiple_files=False)
 
@@ -95,8 +106,10 @@ def main():
 
 
 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()