update

main.py

@@ -1,6 +1,7 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
 import argparse
+import json
 from functools import lru_cache
 import logging
 from pathlib import Path
@@ -8,6 +9,7 @@ import platform
 import shutil
 from typing import Tuple
 import zipfile
+import time
 
 import gradio as gr
 from huggingface_hub import snapshot_download
@@ -18,6 +20,8 @@ from project_settings import environment, project_path, log_directory
 from toolbox.os.command import Command
 from toolbox.torchaudio.models.mpnet.inference_mpnet import InferenceMPNet
 from toolbox.torchaudio.models.frcrn.inference_frcrn import InferenceFRCRN
+from toolbox.torchaudio.models.dfnet.inference_dfnet import InferenceDfNet
+
 
 log.setup_size_rotating(log_directory=log_directory)
 
@@ -63,7 +67,7 @@ def shell(cmd: str):
 
 denoise_engines = {
     "dfnet-nx-dns3": {
-        "infer_cls": InferenceFRCRN,
+        "infer_cls": InferenceDfNet,
         "kwargs": {
             "pretrained_model_path_or_zip_file": (project_path / "trained_models/dfnet-nx-dns3.zip").as_posix()
         }
@@ -99,6 +103,7 @@ def when_click_denoise_button(noisy_audio_file_t = None, noisy_audio_microphone_
     noisy_audio_t: Tuple = noisy_audio_file_t or noisy_audio_microphone_t
 
     sample_rate, signal = noisy_audio_t
+    audio_duration = signal.shape[-1] // 8000
 
     # Test: 使用 microphone 时，显示采样率是 44100，但 signal 实际是按 8000 的采样率的。
     logger.info(f"run denoise; engine: {engine}, sample_rate: {sample_rate}, signal dtype: {signal.dtype}, signal shape: {signal.shape}")
@@ -114,13 +119,25 @@ def when_click_denoise_button(noisy_audio_file_t = None, noisy_audio_microphone_
         kwargs = infer_engine_param["kwargs"]
         infer_engine = load_denoise_model(infer_cls=infer_cls, **kwargs)
 
+        begin = time.time()
         enhanced_audio = infer_engine.enhancement_by_ndarray(noisy_audio)
+        time_cost = time.time() - begin
+
+        fpr = time_cost / audio_duration
+
+        info = {
+            "time_cost": round(time_cost, 4),
+            "audio_duration": round(audio_duration, 4),
+            "fpr": round(fpr, 4)
+        }
+        message = json.dumps(info, ensure_ascii=False, indent=4)
+
         enhanced_audio = np.array(enhanced_audio * (1 << 15), dtype=np.int16)
     except Exception as e:
         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
+    return enhanced_audio_t, message
 
 
 def main():
@@ -177,16 +194,17 @@ def main():
                         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")
 
                 dn_button.click(
                     when_click_denoise_button,
                     inputs=[dn_noisy_audio_file, dn_noisy_audio_microphone, dn_engine],
-                    outputs=[dn_enhanced_audio]
+                    outputs=[dn_enhanced_audio, dn_message]
                 )
                 gr.Examples(
                     examples=examples,
                     inputs=[dn_noisy_audio_file, dn_noisy_audio_microphone, dn_engine],
-                    outputs=[dn_enhanced_audio],
+                    outputs=[dn_enhanced_audio, dn_message],
                     fn=when_click_denoise_button,
                     # cache_examples=True,
                     # cache_mode="lazy",