update code

.idea/.gitignore

@@ -0,0 +1,3 @@
+# Default ignored files
+/shelf/
+/workspace.xml

.idea/FRN.iml

@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+  <component name="PyDocumentationSettings">
+    <option name="format" value="PLAIN" />
+    <option name="myDocStringFormat" value="Plain" />
+  </component>
+</module>

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,23 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoredPackages">
+        <value>
+          <list size="10">
+            <item index="0" class="java.lang.String" itemvalue="tqdm" />
+            <item index="1" class="java.lang.String" itemvalue="scipy" />
+            <item index="2" class="java.lang.String" itemvalue="torchmetrics" />
+            <item index="3" class="java.lang.String" itemvalue="tensorboard" />
+            <item index="4" class="java.lang.String" itemvalue="scikit_learn" />
+            <item index="5" class="java.lang.String" itemvalue="matplotlib" />
+            <item index="6" class="java.lang.String" itemvalue="torch" />
+            <item index="7" class="java.lang.String" itemvalue="numpy" />
+            <item index="8" class="java.lang.String" itemvalue="einops" />
+            <item index="9" class="java.lang.String" itemvalue="pandas" />
+          </list>
+        </value>
+      </option>
+    </inspection_tool>
+  </profile>
+</component>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/FRN.iml" filepath="$PROJECT_DIR$/.idea/FRN.iml" />
+    </modules>
+  </component>
+</project>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
+</project>

main.py

@@ -104,21 +104,21 @@ if __name__ == '__main__':
         model.freeze()
         if args.mode == 'eval':
             model.cuda(device=0)
-            trainer = pl.Trainer(gpus=1)
+            trainer = pl.Trainer(accelerator='gpu', devices=1, enable_checkpointing=False, logger=False)
             testset = TestLoader()
             test_loader = DataLoader(testset, batch_size=1, num_workers=4)
             trainer.test(model, test_loader)
             print('Version', args.version)
             masking = CONFIG.DATA.EVAL.masking
             prob = CONFIG.DATA.EVAL.transition_probs[0]
-            loss_percent = (1 - prob[0]) / (2 - prob[0][prob[1]]) * 100
+            loss_percent = (1 - prob[0]) / (2 - prob[0] - prob[1]) * 100
             print('Evaluate with real trace' if masking == 'real' else
-                  'Evaluate with generated trace {}% packet loss'.format(str(prob)))
+                  'Evaluate with generated trace with {:.2f}% packet loss'.format(prob))
         elif args.mode == 'test':
             model.cuda(device=0)
             testset = BlindTestLoader(test_dir=CONFIG.TEST.in_dir)
             test_loader = DataLoader(testset, batch_size=1, num_workers=4)
-            trainer = pl.Trainer(gpus=1)
+            trainer = pl.Trainer(accelerator='gpu', devices=1, enable_checkpointing=False, logger=False)
             preds = trainer.predict(model, test_loader, return_predictions=True)
             mkdir_p(CONFIG.TEST.out_dir)
             for idx, path in enumerate(test_loader.dataset.data_list):

models/frn.py

@@ -160,8 +160,8 @@ class PLCModel(pl.LightningModule):
             sf.write(os.path.join(path, 'lossy_input.wav'), inp_wav, samplerate=CONFIG.DATA.sr, subtype='PCM_16')
             sf.write(os.path.join(path, 'target.wav'), tar_wav, samplerate=CONFIG.DATA.sr, subtype='PCM_16')
         if CONFIG.DATA.sr != 16000:
-            pred = librosa.resample(pred, 48000, 16000)
-            tar_wav = librosa.resample(tar_wav, 48000, 16000, res_type='kaiser_fast')
+            pred = librosa.resample(pred, orig_sr=48000, target_sr=16000)
+            tar_wav = librosa.resample(tar_wav, orig_sr=48000, target_sr=16000, res_type='kaiser_fast')
         ret = plcmos.run(pred, tar_wav)
         pesq = self.pesq(torch.tensor(pred), torch.tensor(tar_wav))
         metrics = {

utils/utils.py

@@ -52,7 +52,7 @@ def visualize(target, input, recon, path):
 
 
 def get_power(x, nfft):
-    S = librosa.stft(x, nfft)
+    S = librosa.stft(x, n_fft=nfft)
     S = np.log(np.abs(S) ** 2 + 1e-8)
     return S