update

examples/dtln/run.sh

@@ -2,18 +2,14 @@
 
 : <<'END'
 
-sh run.sh --stage 2 --stop_stage 2 --system_version windows --file_folder_name file_dir --final_model_name dfnet-nx-speech \
---noise_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/noise" \
---speech_dir "E:/Users/tianx/HuggingDatasets/nx_noise/data/speech"
-
 sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir-256 --final_model_name dtln-256-nx-dns3 \
 --config_file "yaml/config-256.yaml" \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/dns3-noise" \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
 
-sh run.sh --stage 2 --stop_stage 2 --system_version centos --file_folder_name file_dir-512 --final_model_name dtln-512-nx-dns3 \
+sh run.sh --stage 1 --stop_stage 2 --system_version centos --file_folder_name file_dir-512 --final_model_name dtln-512-nx-dns3 \
 --config_file "yaml/config-512.yaml" \
---noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise/dns3-noise" \
+--noise_dir "/data/tianxing/HuggingDatasets/nx_noise/data/noise" \
 --speech_dir "/data/tianxing/HuggingDatasets/nx_noise/data/speech/dns3-speech"
 
 

main.py

@@ -7,12 +7,16 @@ import logging
 from pathlib import Path
 import platform
 import shutil
+import tempfile
+import time
 from typing import Tuple
 import zipfile
-import time
 
 import gradio as gr
 from huggingface_hub import snapshot_download
+import librosa
+import librosa.display
+import matplotlib.pyplot as plt
 import numpy as np
 
 import log
@@ -108,6 +112,23 @@ def load_denoise_model(infer_cls, **kwargs):
     return infer_engine
 
 
+def generate_spectrogram(signal: np.ndarray, sample_rate: int = 8000, title: str = "Spectrogram"):
+    mag = np.abs(librosa.stft(signal))
+    mag_db = librosa.amplitude_to_db(mag, ref=np.max)
+
+    # 保存为临时图片文件
+    plt.figure(figsize=(10, 3))
+    librosa.display.specshow(mag_db, sr=sample_rate)
+    # librosa.display.specshow(mag_db, sr=sample_rate, x_axis='time', y_axis='log')
+    # plt.colorbar(format='%+2.0f dB')
+    plt.title(title)
+
+    temp_file = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+    plt.savefig(temp_file.name, bbox_inches="tight")
+    plt.close()
+    return temp_file.name
+
+
 def when_click_denoise_button(noisy_audio_file_t = None, noisy_audio_microphone_t = None, engine: str = None):
     if noisy_audio_file_t is None and noisy_audio_microphone_t is None:
         raise gr.Error(f"audio file and microphone is null.")
@@ -137,6 +158,9 @@ def when_click_denoise_button(noisy_audio_file_t = None, noisy_audio_microphone_
         enhanced_audio = infer_engine.enhancement_by_ndarray(noisy_audio)
         time_cost = time.time() - begin
 
+        noisy_mag_db = generate_spectrogram(noisy_audio, title="noisy")
+        denoise_mag_db = generate_spectrogram(enhanced_audio, title="denoise")
+
         fpr = time_cost / audio_duration
 
         info = {
@@ -151,7 +175,7 @@ def when_click_denoise_button(noisy_audio_file_t = None, noisy_audio_microphone_
         raise gr.Error(f"enhancement failed, error type: {type(e)}, error text: {str(e)}.")
 
     enhanced_audio_t = (sample_rate, enhanced_audio)
-    return enhanced_audio_t, message
+    return enhanced_audio_t, message, noisy_mag_db, denoise_mag_db
 
 
 def main():
@@ -207,18 +231,23 @@ def main():
                         dn_engine = gr.Dropdown(choices=denoise_engine_choices, value=denoise_engine_choices[0], label="engine")
                         dn_button = gr.Button(variant="primary")
                     with gr.Column(variant="panel", scale=5):
-                        dn_enhanced_audio = gr.Audio(label="enhanced_audio")
-                        dn_message = gr.Textbox(lines=1, max_lines=20, label="message")
+                        with gr.Tabs():
+                            with gr.TabItem("audio"):
+                                dn_enhanced_audio = gr.Audio(label="enhanced_audio")
+                                dn_message = gr.Textbox(lines=1, max_lines=20, label="message")
+                            with gr.TabItem("mag_db"):
+                                dn_noisy_mag_db = gr.Image(label="noisy_mag_db")
+                                dn_denoise_mag_db = gr.Image(label="denoise_mag_db")
 
                 dn_button.click(
                     when_click_denoise_button,
                     inputs=[dn_noisy_audio_file, dn_noisy_audio_microphone, dn_engine],
-                    outputs=[dn_enhanced_audio, dn_message]
+                    outputs=[dn_enhanced_audio, dn_message, dn_noisy_mag_db, dn_denoise_mag_db]
                 )
                 gr.Examples(
                     examples=examples,
                     inputs=[dn_noisy_audio_file, dn_noisy_audio_microphone, dn_engine],
-                    outputs=[dn_enhanced_audio, dn_message],
+                    outputs=[dn_enhanced_audio, dn_message, dn_noisy_mag_db, dn_denoise_mag_db],
                     fn=when_click_denoise_button,
                     # cache_examples=True,
                     # cache_mode="lazy",

toolbox/torchaudio/models/dtln/modeling_dtln.py

@@ -344,7 +344,9 @@ class DTLNPretrainedModel(DTLNModel):
 
 
 def main():
-    config = DTLNConfig()
+    config = DTLNConfig(fft_size=512,
+                        hop_size=128,
+                        )
     model = DTLNPretrainedModel(config)
     model.eval()
 
@@ -354,6 +356,8 @@ def main():
         denoise = model.forward(noisy)
     print(f"denoise.shape: {denoise.shape}")
     print(denoise[:, :, 300: 302])
+    print(denoise[:, :, 8000: 8002])
+    print(denoise[:, :, 15600: 15602])
     print(denoise[:, :, 15680: 15682])
     print(denoise[:, :, 15760: 15762])
     print(denoise[:, :, 15840: 15842])
@@ -362,6 +366,8 @@ def main():
     print(f"denoise.shape: {denoise.shape}")
     # denoise = denoise[:, :, (config.fft_size - config.hop_size):]
     print(denoise[:, :, 300: 302])
+    print(denoise[:, :, 8000: 8002])
+    print(denoise[:, :, 15600: 15602])
     print(denoise[:, :, 15680: 15682])
     print(denoise[:, :, 15760: 15762])
     print(denoise[:, :, 15840: 15842])