Update DPTNet_eval/DPTNet_quant_sep.py

DPTNet_eval/DPTNet_quant_sep.py

@@ -1,25 +1,12 @@
-# DPTNet_quant_sep.py
-
-import warnings
-warnings.filterwarnings("ignore", message="Failed to initialize NumPy: _ARRAY_API not found")
-
 import os
 import torch
 import numpy as np
 import torchaudio
-from huggingface_hub import hf_hub_download
-
-# 動態導入 asteroid_test 中的 DPTNet
-try:
-    from . import asteroid_test
-except ImportError as e:
-    raise ImportError("無法載入 asteroid_test 模組，請確認該模組與訓練時相同") from e
-
-torchaudio.set_audio_backend("sox_io")
+import yaml
+from . import asteroid_test
 
 
 def get_conf():
-    """取得模型參數設定"""
     conf_filterbank = {
         'n_filters': 64,
         'kernel_size': 16,
@@ -45,61 +32,19 @@ def get_conf():
 
 def load_dpt_model():
     print('Load Separation Model...')
-
-    speech_sep_token = os.getenv("SpeechSeparation")
-    if not speech_sep_token:
-        raise EnvironmentError("環境變數 SpeechSeparation 未設定！")
-
-    model_path = hf_hub_download(
-        repo_id="DeepLearning101/speech-separation",
-        filename="train_dptnet_aishell_partOverlap_B2_300epoch_quan-int8.p",
-        token=speech_sep_token
-    )
-
+    now_path = os.path.split(os.path.realpath(__file__))[0]
     conf_filterbank, conf_masknet = get_conf()
-
-    try:
-        model_class = getattr(asteroid_test, "DPTNet")
-        model = model_class(**conf_filterbank, **conf_masknet)
-    except Exception as e:
-        raise RuntimeError("模型結構錯誤：請確認 asteroid_test.py 是否與訓練時相同") from e
-
-    model = torch.quantization.quantize_dynamic(
-        model,
-        {torch.nn.LSTM, torch.nn.Linear},
-        dtype=torch.qint8
-    )
-
+    model_path = os.path.join(now_path, "trained_model/train_dptnet_aishell_partOverlap_B2_300epoch_quan-int8.p")
+    model = getattr(asteroid_test, "DPTNet")(**conf_filterbank, **conf_masknet)
+    model = torch.quantization.quantize_dynamic(model, {torch.nn.LSTM, torch.nn.Linear}, dtype=torch.qint8)
     state_dict = torch.load(model_path, map_location="cpu")
-    own_state = model.state_dict()
-
-    # 只保留是 torch.Tensor 的 key-value pairs
-    filtered_state_dict = {}
-    for k, v in state_dict.items():
-        if k in own_state:
-            if isinstance(v, torch.Tensor) and isinstance(own_state[k], torch.Tensor):
-                if v.shape == own_state[k].shape:
-                    filtered_state_dict[k] = v
-                else:
-                    print(f"Skip '{k}': shape mismatch")
-            else:
-                print(f"Skip '{k}': not a tensor")
-
-    missing_keys, unexpected_keys = model.load_state_dict(filtered_state_dict, strict=False)
-
-    if missing_keys:
-        print("⚠️ Missing keys:", missing_keys)
-    if unexpected_keys:
-        print("ℹ️ Unexpected keys:", unexpected_keys)
-
+    model.load_state_dict(state_dict)
     model.eval()
     return model
 
-
 def dpt_sep_process(wav_path, model=None, outfilename=None):
-    """進行語音分離處理"""
     if model is None:
-        model = load_dpt_model()
+        model = load_model()
 
     x, sr = torchaudio.load(wav_path)
     x = x.cpu()
@@ -107,8 +52,10 @@ def dpt_sep_process(wav_path, model=None, outfilename=None):
     with torch.no_grad():
         est_sources = model(x)  # shape: (1, 2, T)
 
+    # 確保 est_sources 是 (1, 2, T)，再拆分
     est_sources = est_sources.squeeze(0)  # shape: (2, T)
-    sep_1, sep_2 = est_sources  # 拆成兩個 (T,) 的 tensor
+
+    sep_1, sep_2 = est_sources  # 拆成兩個 (T, ) 的 tensor
 
     # 正規化
     max_abs = x[0].abs().max().item()
@@ -119,7 +66,6 @@ def dpt_sep_process(wav_path, model=None, outfilename=None):
     sep_1 = sep_1.unsqueeze(0)
     sep_2 = sep_2.unsqueeze(0)
 
-    # 儲存結果
     if outfilename is not None:
         torchaudio.save(outfilename.replace('.wav', '_sep1.wav'), sep_1, sr)
         torchaudio.save(outfilename.replace('.wav', '_sep2.wav'), sep_2, sr)
@@ -127,7 +73,28 @@ def dpt_sep_process(wav_path, model=None, outfilename=None):
     else:
         torchaudio.save(wav_path.replace('.wav', '_sep1.wav'), sep_1, sr)
         torchaudio.save(wav_path.replace('.wav', '_sep2.wav'), sep_2, sr)
-
+        
+# def dpt_sep_process(wav_path, model=None, outfilename=None):
+#     if model == None:
+#         model = load_model()
+#     x, sr = torchaudio.load(wav_path)
+#     x = x.cpu()
+#     with torch.no_grad():
+#         est_sources = model(x)
+
+#     est_sources_np = est_sources.squeeze(0)
+
+#     sep_1, sep_2 = est_sources_np
+#     sep_1 = sep_1 * x[0].abs().max().item() / sep_1.abs().max().item()
+#     sep_2 = sep_2 * x[0].abs().max().item() / sep_2.abs().max().item()
+
+#     if outfilename != None:
+#         torchaudio.save(outfilename.replace('.wav', '_sep1.wav'), sep_1, sr)
+#         torchaudio.save(outfilename.replace('.wav', '_sep2.wav'), sep_2, sr)
+#         torchaudio.save(outfilename.replace('.wav', '_mix.wav'), x, sr)
+#     else:
+#         torchaudio.save(wav_path.replace('.wav', '_sep1.wav'), sep_1, sr)
+#         torchaudio.save(wav_path.replace('.wav', '_sep2.wav'), sep_2, sr)
 
 if __name__ == '__main__':
     print("This module should be used via Flask or Gradio.")