Update app.py

app.py

@@ -10,8 +10,8 @@ from dataset import MaskGenerator
 import onnxruntime, onnx
 import matplotlib.pyplot as plt
 from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
-from torchmetrics.audio.stoi import ShortTimeObjectiveIntelligibility as STOI
-from torchmetrics.audio.pesq import PerceptualEvaluationSpeechQuality as PESQ
+from pystoi import stoi
+from pesq import pesq
 import pandas as pd
 import torchaudio
 
@@ -125,38 +125,34 @@ if st.button('Сгенерировать потери'):
     st.audio('lossy.wav')
     st.text('Улучшенное аудио')
     st.audio('enhanced.wav')
-    data_clean, samplerate = torchaudio.load('target.wav')
 
-    data_lossy, samplerate = torchaudio.load('lossy.wav')
-    data_enhanced, samplerate = torchaudio.load('enhanced.wav')
-    min_len = min(data_clean.shape[1], data_lossy.shape[1], data_enhanced.shape[1])
-    data_clean = data_clean[:, :min_len]
-    data_lossy = data_lossy[:, :min_len]
-    data_enhanced = data_enhanced[:, :min_len]
-
-
-    stoi = STOI(samplerate)
-    stoi_orig = round(float(stoi(data_clean, data_clean)),3)
-    stoi_lossy = round(float(stoi(data_clean, data_lossy)),5)
-    stoi_enhanced = round(float(stoi(data_clean, data_enhanced)),5)
+    
+    data_clean, samplerate = sf.read('target.wav')
+    data_lossy, samplerate = sf.read('lossy.wav')
+    data_enhanced, samplerate = sf.read('enhanced.wav')
+    min_len = min(data_clean.shape[0], data_lossy.shape[0], data_enhanced.shape[0])
+    data_clean = data_clean[:min_len]
+    data_lossy = data_lossy[:min_len]
+    data_enhanced = data_enhanced[:min_len]
+
+
+    stoi_orig = round(stoi(data_clean, data_clean, samplerate, extended=False),5)
+    stoi_lossy  = round(stoi(data_clean, data_lossy , samplerate, extended=False),5)
+    stoi_enhanced = round(stoi(data_clean, data_enhanced, samplerate, extended=False),5)
+    
     stoi_mass=[stoi_orig, stoi_lossy, stoi_enhanced]
 
 
-    #pesq = PESQ(16000, 'nb')
-    #data_clean = data_clean.cpu().numpy()
-    #data_lossy = data_lossy.detach().cpu().numpy()
-    #data_enhanced = data_enhanced.cpu().numpy()
-
-    #if samplerate != 16000:
-    #    data_lossy = librosa.resample(data_lossy, orig_sr=48000, target_sr=16000)
-    #    data_clean = librosa.resample(data_clean, orig_sr=48000, target_sr=16000)
-    #    data_enhanced = librosa.resample(data_enhanced, orig_sr=48000, target_sr=16000)
+    if samplerate != 16000:
+        data_lossy = librosa.resample(data_lossy, orig_sr=48000, target_sr=16000)
+        data_clean = librosa.resample(data_clean, orig_sr=48000, target_sr=16000)
+        data_enhanced = librosa.resample(data_enhanced, orig_sr=48000, target_sr=16000)
 
-    #    pesq_orig = np.array(pesq(torch.tensor(data_clean), torch.tensor(data_clean)))
-    #    pesq_lossy = np.array(pesq(torch.tensor(data_lossy), torch.tensor(data_clean)))
-    #    pesq_enhanced = np.array(pesq(torch.tensor(data_enhanced), torch.tensor(data_clean)))
+    pesq_orig = pesq(fs = 16000, ref = data_clean, deg = data_clean, mode='nb')
+    pesq_lossy = pesq(fs = 16000, ref = data_clean, deg = data_lossy, mode='nb')
+    pesq_enhanced = pesq(fs = 16000, ref = data_clean, deg = data_enhanced, mode='nb')
 
-    #psq_mas=[pesq_orig, pesq_lossy, pesq_enhanced]
+    psq_mas=[pesq_orig, pesq_lossy, pesq_enhanced]
 
 
 
@@ -164,7 +160,7 @@ if st.button('Сгенерировать потери'):
 
     df['Audio'] = ['Clean', 'Lossy', 'Enhanced']
 
-    #df['PESQ'] = psq_mas
+    df['PESQ'] = psq_mas
 
     df['STOI'] = stoi_mass