Update DPTNet_eval/DPTNet_quant_sep.py

DPTNet_eval/DPTNet_quant_sep.py

@@ -61,37 +61,54 @@ def load_dpt_model():
     return model
 
 
-def dpt_sep_process(wav_path, model=None, outfilename=None):
-    if model is None:
-        model = load_dpt_model()
-
-    x, sr = torchaudio.load(wav_path)
-    x = x.cpu()
-
-    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
-
-    # 正規化
-    max_abs = x[0].abs().max().item()
-    sep_1 = sep_1 * max_abs / sep_1.abs().max().item()
-    sep_2 = sep_2 * max_abs / sep_2.abs().max().item()
-
-    # 增加 channel 維度，變為 (1, T)
-    sep_1 = sep_1.unsqueeze(0)
-    sep_2 = sep_2.unsqueeze(0)
+import torchaudio
+import tempfile
 
-    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)
-        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)
+def dpt_sep_process(wav_path, model=None, outfilename=None):
+    try:
+        if model is None:
+            model = load_dpt_model()
+
+        # 使用 torchaudio 的通用加載方法
+        x, sr = torchaudio.load(wav_path, format=wav_path.split('.')[-1])
+        x = x.mean(dim=0, keepdim=True)  # 強制轉單聲道
+
+        # 自動重採樣處理
+        if sr != 16000:
+            resampler = torchaudio.transforms.Resample(sr, 16000)
+            x = resampler(x)
+            sr = 16000
+
+        with torch.no_grad():
+            est_sources = model(x)
+
+        # 後處理修正
+        est_sources = est_sources.squeeze(0)
+        sep_1, sep_2 = est_sources[0], est_sources[1]
+
+        # 正規化增強
+        peak = 0.9 * torch.max(torch.abs(x))
+        sep_1 = peak * sep_1 / torch.max(torch.abs(sep_1))
+        sep_2 = peak * sep_2 / torch.max(torch.abs(sep_2))
+
+        # 使用臨時輸出目錄
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            sep1_path = os.path.join(tmp_dir, "sep1.wav")
+            sep2_path = os.path.join(tmp_dir, "sep2.wav")
+            
+            torchaudio.save(sep1_path, sep_1.unsqueeze(0), sr)
+            torchaudio.save(sep2_path, sep_2.unsqueeze(0), sr)
+
+            # 移動檔案到最終位置
+            final_sep1 = outfilename.replace('.wav', '_sep1.wav')
+            final_sep2 = outfilename.replace('.wav', '_sep2.wav')
+            os.replace(sep1_path, final_sep1)
+            os.replace(sep2_path, final_sep2)
+
+        return final_sep1, final_sep2
+
+    except Exception as e:
+        raise RuntimeError(f"分離過程錯誤: {str(e)}") from e
 
 
 if __name__ == '__main__':