deleted files that are not needed for app

.github/workflows/python-app.yml

@@ -1,39 +0,0 @@
-# This workflow will install Python dependencies, run tests and lint with a single version of Python
-# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
-
-name: Python application
-
-on:
-  push:
-    branches: [ "main" ]
-  pull_request:
-    branches: [ "main" ]
-
-permissions:
-  contents: read
-
-jobs:
-  build:
-
-    runs-on: ubuntu-latest
-
-    steps:
-    - uses: actions/checkout@v3
-    - name: Set up Python 3.8
-      uses: actions/setup-python@v3
-      with:
-        python-version: "3.8"
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install flake8 pytest
-        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-    - name: Lint with flake8
-      run: |
-        # stop the build if there are Python syntax errors or undefined names
-        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
-    - name: Test with pytest
-      run: |
-        pytest

app.py

@@ -10,6 +10,20 @@ import torchaudio
 import yaml
 
 
+def denoising_transform(audio, model):
+    src_path = Path("cache_wav/original/{}.wav".format(str(uuid.uuid4())))
+    tgt_path = Path("cache_wav/denoised/{}.wav".format(str(uuid.uuid4())))
+    src_path.parent.mkdir(exist_ok=True, parents=True)
+    tgt_path.parent.mkdir(exist_ok=True, parents=True)
+    (ffmpeg.input(audio)
+     .output(src_path.as_posix(), acodec='pcm_s16le', ac=1, ar=22050)
+     .run()
+     )
+    wav, rate = torchaudio.load(audio)
+    reduced_noise = model.predict(wav)
+    torchaudio.save(tgt_path, reduced_noise, rate)
+    return tgt_path
+
 
 def run_app(model_filename, config_filename):
     model_path = hf_hub_download(repo_id="BorisovMaksim/demucs", filename=model_filename)
@@ -20,37 +34,27 @@ def run_app(model_filename, config_filename):
     checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
     model.load_state_dict(checkpoint['model_state_dict'])
 
-    def denoising_transform(audio):
-        # Path(__file__).parent.resolve()
-        src_path = Path("cache_wav/original/{}.wav".format(str(uuid.uuid4())))
-        tgt_path = Path("cache_wav/denoised/{}.wav".format(str(uuid.uuid4())))
-        src_path.parent.mkdir(exist_ok=True, parents=True)
-        tgt_path.parent.mkdir(exist_ok=True, parents=True)
-        (ffmpeg.input(audio)
-         .output(src_path.as_posix(), acodec='pcm_s16le', ac=1, ar=22050)
-         .run()
-         )
-        wav, rate = torchaudio.load(audio)
-        reduced_noise = model.predict(wav)
-        torchaudio.save(tgt_path, reduced_noise, rate)
-        return tgt_path
-
-    demo = gr.Interface(
-        fn=denoising_transform,
+    interface_demucs = gr.Interface(
+        fn=lambda x: denoising_transform(x, model),
         inputs=gr.Audio(label="Source Audio", source="microphone", type='filepath'),
-        outputs=gr.Audio(label="Target Audio", type='filepath'),
-        examples=[
-            ["testing/wavs/p232_071.wav"],
-            ["testing/wavs/p232_284.wav"],
-        ],
-        title="Denoising"
+        outputs=gr.Audio(label="Demucs", type='filepath'),
+        allow_flagging='never'
     )
+    interface_spectral_gating = gr.Interface(
+        fn=lambda x: denoising_transform(x, SpectralGating()),
+        inputs=gr.Audio(label="Source Audio", source="microphone", type='filepath'),
+        outputs=gr.Audio(label="Spectral Gating", type='filepath'),
+        allow_flagging='never'
+    )
+    gr.Parallel(interface_demucs, interface_spectral_gating,
+                title="Denoising",
+                examples=[
+                    ["testing/wavs/p232_071.wav"],
+                    ["testing/wavs/p232_284.wav"],
+                ]).launch()
 
 
-    demo.launch()
-
 if __name__ == "__main__":
     model_filename = "original_sr/Demucs_original_sr_epoch3.pt"
     config_filename = "original_sr/config.yaml"
     run_app(model_filename, config_filename)
-

checkpoing_saver.py

@@ -1,53 +0,0 @@
-import os
-import numpy as np
-import logging
-import torch
-import wandb
-
-class CheckpointSaver:
-    def __init__(self, dirpath, run_name='', decreasing=True, top_n=5):
-        """
-        dirpath: Directory path where to store all model weights
-        decreasing: If decreasing is `True`, then lower metric is better
-        top_n: Total number of models to track based on validation metric value
-        """
-        if not os.path.exists(dirpath): os.makedirs(dirpath)
-        self.dirpath = dirpath
-        self.top_n = top_n
-        self.decreasing = decreasing
-        self.top_model_paths = []
-        self.best_metric_val = np.Inf if decreasing else -np.Inf
-        self.run_name = run_name
-
-
-    def __call__(self, model, epoch, metric_val, optimizer, loss):
-        model_path = os.path.join(self.dirpath, model.__class__.__name__ + f'_{self.run_name}_epoch{epoch}.pt')
-        save = metric_val < self.best_metric_val if self.decreasing else metric_val > self.best_metric_val
-        if save:
-            logging.info(
-                f"Current metric value better than {metric_val} better than best {self.best_metric_val}, saving model at {model_path}, & logging model weights to W&B.")
-            self.best_metric_val = metric_val
-            torch.save(
-                {  # Save our checkpoint loc
-                    'epoch': epoch,
-                    'model_state_dict': model.state_dict(),
-                    'optimizer_state_dict': optimizer.state_dict(),
-                    'loss': loss,
-                }, model_path)
-            self.log_artifact(f'model-ckpt-epoch-{epoch}.pt', model_path, metric_val)
-            self.top_model_paths.append({'path': model_path, 'score': metric_val})
-            self.top_model_paths = sorted(self.top_model_paths, key=lambda o: o['score'], reverse=not self.decreasing)
-        if len(self.top_model_paths) > self.top_n:
-            self.cleanup()
-
-    def log_artifact(self, filename, model_path, metric_val):
-        artifact = wandb.Artifact(filename, type='model', metadata={'Validation score': metric_val})
-        artifact.add_file(model_path)
-        wandb.run.log_artifact(artifact)
-
-    def cleanup(self):
-        to_remove = self.top_model_paths[self.top_n:]
-        logging.info(f"Removing extra models.. {to_remove}")
-        for o in to_remove:
-            os.remove(o['path'])
-        self.top_model_paths = self.top_model_paths[:self.top_n]

conf/config.yaml

@@ -1,34 +0,0 @@
-defaults:
-  - dataset: valentini
-  - model: demucs
-  - loss: L1_Multi_STFT
-  - optimizer: sgd
-
-training:
-  num_epochs: 10
-  model_save_path: /media/public/checkpoints
-
-dataloader:
-  max_seconds: 2
-  sample_rate: 16000
-  train_batch_size: 12
-  valid_batch_size: 12
-  normalize: False
-  num_workers: 12
-
-validation:
-  path: /media/public/datasets/denoising/DS_10283_2791/noisy_testset_wav
-  sample_rate: 48000
-
-
-wandb:
-  entity: borisovmaksim
-  run_name: default
-  project: denoising
-  log_interval: 100
-  api_key: local-e23d01ece807cb31e69b2cf4137e4998e4b9856f
-  host: http://localhost:8080/
-  notes: "Experiment note"
-  resume: False
-
-gpu: 1

conf/dataset/valentini.yaml

@@ -1,4 +0,0 @@
-valentini:
- dataset_path: /media/public/datasets/denoising/DS_10283_2791/
- val_fraction: 0.2
- sample_rate:  48000

conf/loss/L1_Multi_STFT.yaml

@@ -1 +0,0 @@
-name: L1_Multi_STFT

conf/loss/mse.yaml

@@ -1 +0,0 @@
-name: mse

conf/model/demucs.yaml

@@ -1,19 +0,0 @@
-demucs:
- H: 64
- L: 3
- encoder:
-   conv1:
-     kernel_size: 8
-     stride: 2
-   conv2:
-     kernel_size: 1
-     stride: 1
-
- decoder:
-   conv1:
-     kernel_size: 1
-     stride: 1
-   conv2:
-     kernel_size: 8
-     stride: 2
-

conf/optimizer/sgd.yaml

@@ -1,3 +0,0 @@
-sgd:
-  lr: 0.001
-  momentum: 0.9

datasets/__init__.py

@@ -1,17 +0,0 @@
-from typing import Tuple
-from torch.utils.data import Dataset
-
-from datasets.valentini import Valentini
-from transforms import Transform
-
-DATASETS_POOL = {
-    'valentini': Valentini
-}
-
-
-def get_datasets(cfg) -> Tuple[Dataset, Dataset]:
-    name, dataset_params = list(cfg['dataset'].items())[0]
-    transform = Transform(input_sample_rate=dataset_params['sample_rate'], **cfg['dataloader'])
-    train_dataset = DATASETS_POOL[name](valid=False, transform=transform, **dataset_params)
-    valid_dataset = DATASETS_POOL[name](valid=True, transform=transform, **dataset_params)
-    return train_dataset, valid_dataset

datasets/minimal.py

@@ -1,23 +0,0 @@
-import torch
-from torch.utils.data import Dataset
-from pathlib import Path
-import torchaudio
-import numpy as np
-from torchaudio.transforms import Resample
-
-
-class Minimal(Dataset):
-    def __init__(self, cfg):
-        self.wavs = ['p232_284.wav', 'p232_071.wav', 'p257_171.wav']
-        self.dataset_path = cfg['validation']['path']
-        self.target_rate = cfg['dataloader']['sample_rate']
-        self.resampler = Resample(orig_freq=cfg['validation']['sample_rate'],
-                                  new_freq=cfg['dataloader']['sample_rate'])
-
-    def __len__(self):
-        return len(self.wavs)
-
-    def __getitem__(self, idx):
-        wav, rate = torchaudio.load(Path(self.dataset_path) / self.wavs[idx])
-        wav = self.resampler(wav)
-        return wav, self.target_rate

datasets/valentini.py

@@ -1,44 +0,0 @@
-import torch
-from torch.utils.data import Dataset
-from pathlib import Path
-import torchaudio
-
-MAX_RANDOM_SEED = 1000
-
-class Valentini(Dataset):
-    def __init__(self, dataset_path, val_fraction, transform=None, valid=False, *args, **kwargs):
-        clean_path = Path(dataset_path) / 'clean_trainset_56spk_wav'
-        noisy_path = Path(dataset_path) / 'noisy_trainset_56spk_wav'
-        clean_wavs = list(clean_path.glob("*"))
-        noisy_wavs = list(noisy_path.glob("*"))
-        valid_threshold = int(len(clean_wavs) * val_fraction)
-        if valid:
-            self.clean_wavs = clean_wavs[:valid_threshold]
-            self.noisy_wavs = noisy_wavs[:valid_threshold]
-        else:
-            self.clean_wavs = clean_wavs[valid_threshold:]
-            self.noisy_wavs = noisy_wavs[valid_threshold:]
-
-        assert len(self.clean_wavs) == len(self.noisy_wavs)
-
-        self.transform = transform
-        self.valid = valid
-
-    def __len__(self):
-        return len(self.clean_wavs)
-
-    def __getitem__(self, idx):
-        noisy_wav, noisy_sr = torchaudio.load(self.noisy_wavs[idx])
-        clean_wav, clean_sr = torchaudio.load(self.clean_wavs[idx])
-
-        if self.transform:
-            random_seed = 0 if self.valid else torch.randint(MAX_RANDOM_SEED, (1,))[0]
-            torch.manual_seed(random_seed)
-            noisy_wav = self.transform(noisy_wav)
-            torch.manual_seed(random_seed)
-            clean_wav = self.transform(clean_wav)
-        return noisy_wav, clean_wav
-
-
-
-

denoisers/SpectralGating.py

@@ -4,7 +4,7 @@ import torchaudio
 
 
 class SpectralGating(torch.nn.Module):
-    def __init__(self, rate=16000):
+    def __init__(self, rate=48000):
         super(SpectralGating, self).__init__()
         self.rate = rate
 
@@ -12,11 +12,8 @@ class SpectralGating(torch.nn.Module):
         reduced_noise = torch.Tensor(nr.reduce_noise(y=wav, sr=self.rate))
         return reduced_noise
     
-    def predict(self, wav_path, out_path):
-        data, rate = torchaudio.load(wav_path)
-        reduced_noise = torch.Tensor(nr.reduce_noise(y=data, sr=rate))
-        torchaudio.save(out_path, reduced_noise, rate)
-        return out_path
+    def predict(self, wav):
+        return self.forward(wav)
 
 
 

losses.py

@@ -1,185 +0,0 @@
-# -*- coding: utf-8 -*-
-# Copyright (c) Facebook, Inc. and its affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-
-# Original copyright 2019 Tomoki Hayashi
-#  MIT License (https://opensource.org/licenses/MIT)
-
-
-import torch
-import torch.nn.functional as F
-
-# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-"""STFT-based Loss modules."""
-
-
-def stft(x, fft_size, hop_size, win_length, window):
-    """Perform STFT and convert to magnitude spectrogram.
-    Args:
-        x (Tensor): Input signal tensor (B, T).
-        fft_size (int): FFT size.
-        hop_size (int): Hop size.
-        win_length (int): Window length.
-        window (str): Window function type.
-    Returns:
-        Tensor: Magnitude spectrogram (B, #frames, fft_size // 2 + 1).
-    """
-    x_stft = torch.stft(x[:, 0, :], fft_size, hop_size, win_length, window, return_complex=True)
-    x_stft = torch.view_as_real(x_stft)
-    real = x_stft[..., 0]
-    imag = x_stft[..., 1]
-
-    # NOTE(kan-bayashi): clamp is needed to avoid nan or inf
-    return torch.sqrt(torch.clamp(real ** 2 + imag ** 2, min=1e-7)).transpose(2, 1)
-
-
-class SpectralConvergengeLoss(torch.nn.Module):
-    """Spectral convergence loss module."""
-
-    def __init__(self):
-        """Initilize spectral convergence loss module."""
-        super(SpectralConvergengeLoss, self).__init__()
-
-    def forward(self, x_mag, y_mag):
-        """Calculate forward propagation.
-        Args:
-            x_mag (Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
-            y_mag (Tensor): Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).
-        Returns:
-            Tensor: Spectral convergence loss value.
-        """
-        return torch.norm(y_mag - x_mag, p="fro") / torch.norm(y_mag, p="fro")
-
-
-class LogSTFTMagnitudeLoss(torch.nn.Module):
-    """Log STFT magnitude loss module."""
-
-    def __init__(self):
-        """Initilize los STFT magnitude loss module."""
-        super(LogSTFTMagnitudeLoss, self).__init__()
-
-    def forward(self, x_mag, y_mag):
-        """Calculate forward propagation.
-        Args:
-            x_mag (Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
-            y_mag (Tensor): Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).
-        Returns:
-            Tensor: Log STFT magnitude loss value.
-        """
-        return F.l1_loss(torch.log(y_mag), torch.log(x_mag))
-
-
-class STFTLoss(torch.nn.Module):
-    """STFT loss module."""
-
-    def __init__(self, fft_size=1024, shift_size=120, win_length=600, window="hann_window"):
-        """Initialize STFT loss module."""
-        super(STFTLoss, self).__init__()
-        self.fft_size = fft_size
-        self.shift_size = shift_size
-        self.win_length = win_length
-        self.register_buffer("window", getattr(torch, window)(win_length))
-        self.spectral_convergenge_loss = SpectralConvergengeLoss()
-        self.log_stft_magnitude_loss = LogSTFTMagnitudeLoss()
-
-    def forward(self, x, y):
-        """Calculate forward propagation.
-        Args:
-            x (Tensor): Predicted signal (B, T).
-            y (Tensor): Groundtruth signal (B, T).
-        Returns:
-            Tensor: Spectral convergence loss value.
-            Tensor: Log STFT magnitude loss value.
-        """
-        x_mag = stft(x, self.fft_size, self.shift_size, self.win_length, self.window)
-        y_mag = stft(y, self.fft_size, self.shift_size, self.win_length, self.window)
-        sc_loss = self.spectral_convergenge_loss(x_mag, y_mag)
-        mag_loss = self.log_stft_magnitude_loss(x_mag, y_mag)
-
-        return sc_loss, mag_loss
-
-
-class MultiResolutionSTFTLoss(torch.nn.Module):
-    """Multi resolution STFT loss module."""
-
-    def __init__(self,
-                 fft_sizes=(1024, 2048, 512),
-                 hop_sizes=(120, 240, 50),
-                 win_lengths=(600, 1200, 240),
-                 window="hann_window", factor_sc=0.1, factor_mag=0.1):
-        """Initialize Multi resolution STFT loss module.
-        Args:
-            fft_sizes (list): List of FFT sizes.
-            hop_sizes (list): List of hop sizes.
-            win_lengths (list): List of window lengths.
-            window (str): Window function type.
-            factor (float): a balancing factor across different losses.
-        """
-        super(MultiResolutionSTFTLoss, self).__init__()
-        assert len(fft_sizes) == len(hop_sizes) == len(win_lengths)
-        self.stft_losses = torch.nn.ModuleList()
-        for fs, ss, wl in zip(fft_sizes, hop_sizes, win_lengths):
-            self.stft_losses += [STFTLoss(fs, ss, wl, window)]
-        self.factor_sc = factor_sc
-        self.factor_mag = factor_mag
-
-    def forward(self, x, y):
-        """Calculate forward propagation.
-        Args:
-            x (Tensor): Predicted signal (B, T).
-            y (Tensor): Groundtruth signal (B, T).
-        Returns:
-            Tensor: Multi resolution spectral convergence loss value.
-            Tensor: Multi resolution log STFT magnitude loss value.
-        """
-        sc_loss = 0.0
-        mag_loss = 0.0
-        for f in self.stft_losses:
-            sc_l, mag_l = f(x, y)
-            sc_loss += sc_l
-            mag_loss += mag_l
-        sc_loss /= len(self.stft_losses)
-        mag_loss /= len(self.stft_losses)
-
-        return self.factor_sc*sc_loss, self.factor_mag*mag_loss
-
-
-
-
-class L1_Multi_STFT(torch.nn.Module):
-    """STFT loss module."""
-
-    def __init__(self):
-        """Initialize STFT loss module."""
-        super(L1_Multi_STFT, self).__init__()
-        self.multi_STFT_loss = MultiResolutionSTFTLoss()
-        self.l1_loss = torch.nn.L1Loss()
-
-    def forward(self, x, y):
-        """Calculate forward propagation.
-        Args:
-            x (Tensor): Predicted signal (B, T).
-            y (Tensor): Groundtruth signal (B, T).
-        Returns:
-            Tensor: Spectral convergence loss value.
-            Tensor: Log STFT magnitude loss value.
-        """
-        sc_loss, mag_loss = self.multi_STFT_loss(x, y)
-        l1_loss = self.l1_loss(x, y)
-        return sc_loss + mag_loss + l1_loss
-
-
-LOSSES = {
-    'mse': torch.nn.MSELoss(),
-    'L1': torch.nn.L1Loss(),
-    'Multi_STFT': MultiResolutionSTFTLoss(),
-    'L1_Multi_STFT': L1_Multi_STFT()
-}
-
-
-def get_loss(loss_config, device):
-    return LOSSES[loss_config['name']].to(device)

main.py

@@ -1,12 +0,0 @@
-import hydra
-from omegaconf import DictConfig
-from train import train
-
-
-@hydra.main(version_base=None, config_path="conf", config_name="config")
-def main(cfg: DictConfig):
-    train(cfg)
-
-
-if __name__ == '__main__':
-    main()

optimizers.py

@@ -1,10 +0,0 @@
-import torch
-
-OPTIMIZERS_POOL = {
-    'sgd': torch.optim.SGD,
-}
-
-def get_optimizer(model_params, optimizer_config):
-    name, params = list(optimizer_config.items())[0]
-    optimizer = OPTIMIZERS_POOL[name](model_params, **params)
-    return optimizer

testing/evaluation.py

@@ -1,62 +0,0 @@
-import argparse
-from tqdm import tqdm
-
-from utils import load_wav, collect_valentini_paths
-from metrics import Metrics
-from denoisers.SpectralGating import SpectralGating
-
-
-PARSERS = {
-    'valentini': collect_valentini_paths
-}
-MODELS = {
-    'baseline': SpectralGating
-}
-
-
-
-def evaluate_on_dataset(model_name, dataset_path, dataset_type):
-    if model_name is not None:
-        model = MODELS[model_name]()
-    parser = PARSERS[dataset_type]
-    clean_wavs, noisy_wavs = parser(dataset_path)
-
-    metrics = Metrics()
-    mean_scores = {'PESQ': 0, 'STOI': 0}
-    for clean_path, noisy_path in tqdm(zip(clean_wavs, noisy_wavs), total=len(clean_wavs)):
-        clean_wav = load_wav(clean_path)
-        noisy_wav = load_wav(noisy_path)
-
-        if model_name is None:
-            scores = metrics.calculate(denoised=noisy_wav, clean=clean_wav)
-        else:
-            denoised_wav = model(noisy_wav)
-            scores = metrics.calculate(denoised=denoised_wav, clean=clean_wav)
-
-        mean_scores['PESQ'] += scores['PESQ']
-        mean_scores['STOI'] += scores['STOI']
-
-    mean_scores['PESQ'] = mean_scores['PESQ'].numpy() / len(clean_wavs)
-    mean_scores['STOI'] = mean_scores['STOI'].numpy() / len(clean_wavs)
-
-    return mean_scores
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser(prog='Program to evaluate denoising')
-    parser.add_argument('--dataset_path', type=str,
-                        default='/media/public/dataset/denoising/DS_10283_2791/',
-                        help='Path to dataset folder')
-    parser.add_argument('--dataset_type', type=str, required=True,
-                        choices=['valentini'])
-    parser.add_argument('--model_name', type=str,
-                        choices=['baseline'])
-
-
-    args = parser.parse_args()
-
-    mean_scores = evaluate_on_dataset(model_name=args.model_name,
-                        dataset_path=args.dataset_path,
-                        dataset_type=args.dataset_type)
-    print(f"Metrics on {args.dataset_type} dataset with "
-          f"{args.model_name if args.model_name is not None else 'ideal denoising'} = {mean_scores}")

testing/metrics.py

@@ -1,36 +0,0 @@
-import pesq
-from torchmetrics.audio.pesq import PerceptualEvaluationSpeechQuality
-from torchmetrics.audio.stoi import ShortTimeObjectiveIntelligibility
-from torchaudio.transforms import Resample
-import torch
-import torchaudio
-from torchmetrics import SignalNoiseRatio
-
-class Metrics(torch.nn.Module):
-    def __init__(self, source_rate, target_rate=16000, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.source_rate = source_rate
-        self.target_rate = target_rate
-        self.resampler = Resample(orig_freq=source_rate, new_freq=target_rate)
-        self.nb_pesq = PerceptualEvaluationSpeechQuality(target_rate, 'wb')
-        self.stoi = ShortTimeObjectiveIntelligibility(target_rate, False)
-        self.snr = SignalNoiseRatio()
-        
-    def forward(self, denoised, clean):
-        pesq_scores, stoi_scores = 0, 0
-        for denoised_wav, clean_wav in zip(denoised, clean):
-            if self.source_rate != self.target_rate:
-                denoised_wav = self.resampler(denoised_wav)
-                clean_wav = self.resampler(clean_wav)
-            try:
-                pesq_scores += self.nb_pesq(denoised_wav, clean_wav).item()
-                stoi_scores += self.stoi(denoised_wav, clean_wav).item()
-            except pesq.NoUtterancesError as e:
-                print(e)
-            except ValueError as e:
-                print(e)
-
-        return {'PESQ': pesq_scores,
-                'STOI': stoi_scores}
-    
-

testing/test_inference.py

@@ -1,6 +0,0 @@
-def func(x):
-    return x + 1
-
-
-def test_answer():
-    assert func(3) == 4

train.py

@@ -1,123 +0,0 @@
-import os
-import torch
-from torch.utils.data import DataLoader
-import omegaconf
-from omegaconf import DictConfig
-import wandb
-
-from checkpoing_saver import CheckpointSaver
-from denoisers import get_model
-from optimizers import get_optimizer
-from losses import get_loss
-from datasets import get_datasets
-from testing.metrics import Metrics
-from datasets.minimal import Minimal
-from tqdm import tqdm
-
-def init_wandb(cfg):
-    wandb.login(key=cfg['wandb']['api_key'], host=cfg['wandb']['host'])
-    wandb.init(project=cfg['wandb']['project'],
-               notes=cfg['wandb']['notes'],
-               config=omegaconf.OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True),
-               resume=cfg['wandb']['resume'],
-               name=cfg['wandb']['run_name'])
-    if wandb.run.resumed:
-        api = wandb.Api()
-        runs = api.runs(f"{cfg['wandb']['entity']}/{cfg['wandb']['project']}",
-                        order='train_pesq')
-        run = [run for run in runs if run.name == cfg['wandb']['run_name'] and run.state != 'running'][0]
-        artifacts = run.logged_artifacts()
-        best_model = [artifact for artifact in artifacts if artifact.type == 'model'][0]
-
-        best_model.download()
-
-def train(cfg: DictConfig):
-    device = torch.device(f'cuda:{cfg.gpu}' if torch.cuda.is_available() else 'cpu')
-    init_wandb(cfg)
-    checkpoint_saver = CheckpointSaver(dirpath=cfg['training']['model_save_path'], run_name=wandb.run.name,
-                                       decreasing=False)
-    metrics = Metrics(source_rate=cfg['dataloader']['sample_rate']).to(device)
-
-    model = get_model(cfg['model']).to(device)
-    optimizer = get_optimizer(model.parameters(), cfg['optimizer'])
-    loss_fn = get_loss(cfg['loss'], device)
-    train_dataset, valid_dataset = get_datasets(cfg)
-    minimal_dataset = Minimal(cfg)
-
-    dataloaders = {
-        'train':  DataLoader(train_dataset, batch_size=cfg['dataloader']['train_batch_size'], shuffle=True,
-                             num_workers=cfg['dataloader']['num_workers']),
-        'val': DataLoader(valid_dataset, batch_size=cfg['dataloader']['valid_batch_size'], shuffle=False,
-                          num_workers=cfg['dataloader']['num_workers']),
-        'minimal': DataLoader(minimal_dataset)
-    }
-
-    wandb.watch(model, log_freq=cfg['wandb']['log_interval'])
-    epoch = 0
-    while epoch < cfg['training']['num_epochs']:
-        for phase in ['train', 'val']:
-            if phase == 'train':
-                model.train()
-            else:
-                model.eval()
-
-            running_loss, running_pesq, running_stoi = 0.0, 0.0, 0.0
-            loop = tqdm(dataloaders[phase])
-            for i, (inputs, labels) in enumerate(loop):
-                inputs = inputs.to(device)
-                labels = labels.to(device)
-
-                optimizer.zero_grad()
-
-                with torch.set_grad_enabled(phase == 'train'):
-                    outputs = model(inputs)
-                    loss = loss_fn(outputs, labels)
-
-                    if phase == 'train':
-                        loss.backward()
-                        optimizer.step()
-
-                running_metrics = metrics(denoised=outputs, clean=labels)
-                running_loss += loss.item() * inputs.size(0)
-                running_pesq += running_metrics['PESQ']
-                running_stoi += running_metrics['STOI']
-
-                loop.set_description(f"Epoch [{epoch}/{cfg['training']['num_epochs']}][{phase}]")
-                loop.set_postfix(loss=running_loss / (i + 1) / inputs.size(0),
-                                 pesq=running_pesq / (i + 1) / inputs.size(0),
-                                 stoi=running_stoi / (i + 1) / inputs.size(0))
-
-                if phase == 'train' and i % cfg['wandb']['log_interval'] == 0:
-                    wandb.log({"train_loss": running_loss / (i + 1) / inputs.size(0),
-                               "train_pesq": running_pesq / (i + 1) / inputs.size(0),
-                               "train_stoi": running_stoi / (i + 1) / inputs.size(0)})
-
-            epoch_loss = running_loss / len(dataloaders[phase].dataset)
-            eposh_pesq = running_pesq / len(dataloaders[phase].dataset)
-            eposh_stoi = running_stoi / len(dataloaders[phase].dataset)
-
-            wandb.log({f"{phase}_loss": epoch_loss,
-                       f"{phase}_pesq": eposh_pesq,
-                       f"{phase}_stoi": eposh_stoi})
-
-            if phase == 'val':
-                for i, (wav, rate) in enumerate(dataloaders['minimal']):
-                    if cfg['dataloader']['normalize']:
-                        std = torch.std(wav)
-                        wav = wav / std
-                        prediction = model(wav.to(device))
-                        prediction = prediction * std
-                    else:
-                        prediction = model(wav.to(device))
-                    wandb.log({
-                        f"{i}_example": wandb.Audio(
-                            prediction.detach().cpu().numpy()[0][0],
-                            sample_rate=rate)})
-
-                checkpoint_saver(model, epoch, metric_val=eposh_pesq,
-                                 optimizer=optimizer, loss=epoch_loss)
-        epoch += 1
-
-
-if __name__ == "__main__":
-    pass

transforms.py

@@ -1,31 +0,0 @@
-import torch
-
-from torchaudio.transforms import Resample
-from torchvision.transforms import RandomCrop
-
-
-
-class Transform(torch.nn.Module):
-    def __init__(
-            self,
-            input_sample_rate,
-            sample_rate,
-            max_seconds,
-            normalize,
-            *args,
-            **kwargs
-    ):
-        super().__init__()
-        self.input_sample_rate = input_sample_rate
-        self.sample_rate = sample_rate
-        self.resample = Resample(orig_freq=input_sample_rate, new_freq=sample_rate)
-        self.random_crop = RandomCrop((1, int(max_seconds * sample_rate)), pad_if_needed=True)
-        self.normalize = normalize
-
-    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
-        if self.input_sample_rate != self.sample_rate:
-            waveform = self.resample(waveform)
-        if self.normalize:
-            waveform = waveform / torch.std(waveform)
-        cropped = self.random_crop(waveform)
-        return cropped

upload_to_HF.py

@@ -1,56 +0,0 @@
-import os
-import wandb
-from huggingface_hub import HfApi
-from pathlib import Path
-import huggingface_hub
-import ssl
-import os
-os.environ['CURL_CA_BUNDLE'] = ''
-
-ssl._create_default_https_context = ssl._create_unverified_context
-
-class Uploader:
-    def __init__(self, entity, project, run_name, repo_id, username):
-        self.entity = entity
-        self.project = project
-        self.run_name = run_name
-        self.hf_api = HfApi()
-        self.wandb_api = wandb.Api()
-        self.repo_id = repo_id
-        self.username = username
-        huggingface_hub.login(os.environ.get('HUGGINGFACE_TOKEN'))
-
-    def get_model_from_wandb_run(self):
-        runs = self.wandb_api.runs(f"{self.entity}/{self.project}",
-                        # order='+summary_metrics.train_pesq'
-                        )
-        run = [run for run in runs if run.name == self.run_name][0]
-        artifacts = run.logged_artifacts()
-        best_model = [artifact for artifact in artifacts if artifact.type == 'model'][0]
-        artifact_dir = best_model.download()
-        model_path = list(Path(artifact_dir).glob("*.pt"))[0].absolute().as_posix()
-        print(f"Model validation score = {best_model.metadata['Validation score']}")
-        return model_path
-
-    def upload_to_HF(self):
-        model_path = self.get_model_from_wandb_run()
-        self.hf_api.upload_file(
-            path_or_fileobj=model_path,
-            path_in_repo=Path(model_path).name,
-            repo_id=f'{self.username}/{self.repo_id}',
-        )
-
-    def create_repo(self):
-        self.hf_api.create_repo(repo_id=self.repo_id, exist_ok=True)
-
-
-
-if __name__ == '__main__':
-    uploader = Uploader(entity='borisovmaksim',
-                        project='denoising',
-                        run_name='wav_normalization',
-                        repo_id='demucs',
-                        username='BorisovMaksim')
-    uploader.create_repo()
-    uploader.upload_to_HF()
-

utils.py

@@ -1,7 +1,3 @@
-import torchaudio
-import torch
-import matplotlib.pyplot as plt
-from pathlib import Path
 from torch.nn.functional import pad
 
 
@@ -17,45 +13,3 @@ def pad_cut_batch_audio(wavs, new_shape):
     return wavs
 
 
-def collect_valentini_paths(dataset_path):
-    clean_path = Path(dataset_path) / 'clean_testset_wav'
-    noisy_path = Path(dataset_path) / 'noisy_testset_wav'
-
-    clean_wavs = list(clean_path.glob("*"))
-    noisy_wavs = list(noisy_path.glob("*"))
-
-    return clean_wavs, noisy_wavs
-
-
-
-
-def plot_spectrogram(stft, title="Spectrogram", xlim=None):
-    magnitude = stft.abs()
-    spectrogram = 20 * torch.log10(magnitude + 1e-8).numpy()
-    figure, axis = plt.subplots(1, 1)
-    img = axis.imshow(spectrogram, cmap="viridis", vmin=-100, vmax=0, origin="lower", aspect="auto")
-    figure.suptitle(title)
-    plt.colorbar(img, ax=axis)
-    plt.show()
-
-
-def plot_mask(mask, title="Mask", xlim=None):
-    mask = mask.numpy()
-    figure, axis = plt.subplots(1, 1)
-    img = axis.imshow(mask, cmap="viridis", origin="lower", aspect="auto")
-    figure.suptitle(title)
-    plt.colorbar(img, ax=axis)
-    plt.show()
-
-
-    
-
-def generate_mixture(waveform_clean, waveform_noise, target_snr):
-    
-    power_clean_signal = waveform_clean.pow(2).mean()
-    power_noise_signal = waveform_noise.pow(2).mean()
-    current_snr = 10 * torch.log10(power_clean_signal / power_noise_signal)
-    
-    waveform_noise *= 10 ** (-(target_snr - current_snr) / 20)
-    return waveform_clean + waveform_noise
-