добавлен поиск порогов

tuning/search_thresholds.py

@@ -0,0 +1,36 @@
+from utils import init_jit_model, predict, calculate_best_thresholds, SileroVadDataset, SileroVadPadder
+from omegaconf import OmegaConf
+import torch
+torch.set_num_threads(1)
+
+if __name__ == '__main__':
+    config = OmegaConf.load('config.yml')
+
+    loader = torch.utils.data.DataLoader(SileroVadDataset(config, mode='val'),
+                                         batch_size=config.batch_size,
+                                         collate_fn=SileroVadPadder,
+                                         num_workers=config.num_workers)
+
+    if config.jit_model_path:
+        print(f'Loading model from the local folder: {config.jit_model_path}')
+        model = init_jit_model(config.jit_model_path, device=config.device)
+    else:
+        if config.use_torchhub:
+            print('Loading model using torch.hub')
+            model, _ = torch.hub.load(repo_or_dir='snakers4/silero-vad',
+                                      model='silero_vad',
+                                      onnx=False,
+                                      force_reload=True)
+        else:
+            print('Loading model using silero-vad library')
+            from silero_vad import load_silero_vad
+            model = load_silero_vad(onnx=False)
+
+    print('Model loaded')
+    model.to(config.device)
+
+    print('Making predicts...')
+    all_predicts, all_gts = predict(model, loader, config.device, sr=8000 if config.tune_8k else 16000)
+    print('Calculating thresholds...')
+    best_ths_enter, best_ths_exit, best_acc = calculate_best_thresholds(all_predicts, all_gts)
+    print(f'Best threshold: {best_ths_enter}\nBest exit threshold: {best_ths_exit}\nBest accuracy: {best_acc}')