Dual-Path TFC-TDF UNet for Music Source Separation
import torch
from src.dp_tdf.dp_tdf_net import DPTDFNet
from src.layers.chunk_size import wave_to_batches

import dotenv
from omegaconf import OmegaConf
dotenv.load_dotenv(override=True)

def run_once(src_name):
    # create dummy data
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # create dummy config
    yaml = rf"G:\sdx2023fork\configs\model\{src_name}.yaml"
    # yaml = r"G:\sdx2023fork\configs\model\TDF_mono.yaml"
    # load yaml into dict
    confg = OmegaConf.load(yaml)

    stft = torch.rand(1, 2 * confg.audio_ch, confg.dim_f, 256).to(device)
    # create model
    model = DPTDFNet(**confg).to(device)

    pred_detail = model(stft)
    print(pred_detail.shape)


def run_train(src_name):
    yaml = rf"G:\sdx2023fork\configs\model\{src_name}.yaml"
    h = OmegaConf.load(yaml)
    device = torch.device('cuda')
    b = 2
    chunk_size = (h.dim_t - 1) * h.hop_length
    print("before ", chunk_size)
    input1 = torch.rand(b, h.audio_ch, chunk_size).to(device)
    input2 = torch.rand(b, h.audio_ch, chunk_size).to(device)

    tdf = DPTDFNet(**h).to(device)

    output = tdf.training_step((input1, input2))

    print(output["loss"])


def run_val(src_name):
    yaml = rf"G:\sdx2023fork\configs\model\{src_name}.yaml"
    h = OmegaConf.load(yaml)
    device = torch.device('cuda')
    b = 2
    chunk_size = (h.dim_t - 1) * h.hop_length
    print("before ", chunk_size)
    input1 = torch.rand(h.audio_ch, 3*chunk_size)
    input1 = wave_to_batches(input1, chunk_size, h.overlap, b)
    input1 = [i.unsqueeze(0).to(device) for i in input1]
    input2 = torch.rand(1, h.audio_ch, 3*chunk_size).to(device)

    tdf = DPTDFNet(**h).to(device)


    output = tdf.validation_step((input1, input2))

    for k, v in output.items():
        print(k, v)


if __name__ == "__main__":
    src_name = "vocals"
    # src_name = "bass"
    # src_name = "drums"
    # src_name = "other"
    run_once(src_name)
    run_train(src_name)
    with torch.no_grad():
        run_val(src_name)

demos/dtt_demo.py
junyuchen-cjy/DTTNet-Pytorch
import dotenv
import hydra
from omegaconf import DictConfig

import os
from os.path import exists


dotenv.load_dotenv(override=True)

@hydra.main(config_path="../configs/datamodule", config_name="musdb_dev14.yaml", version_base='1.1')
def main(config: DictConfig):
    trainset_path = os.path.join(config.data_dir, 'train')
    validset_path = os.path.join(config.data_dir, 'valid')
    validation_set = config.validation_set
    if not exists(validset_path):
        from shutil import move
        os.mkdir(validset_path)
        for track in validation_set:
            # if exist
            old_path = os.path.join(trainset_path,track)
            if os.path.exists(old_path):
                new_path = os.path.join(validset_path,track)
                move(old_path, new_path)
                print(f"Moved {old_path} to {new_path}")
    return

if __name__ == "__main__":
    main()
demos/split_dataset.py
junyuchen-cjy/DTTNet-Pytorch
import re

import dotenv
import hydra
from omegaconf import DictConfig, OmegaConf
import os
import pandas as pd
from src.utils.omega_resolvers import get_eval_log_dir

dotenv.load_dotenv(override=True)


@hydra.main(config_path="configs/", config_name="evaluation.yaml", version_base='1.1')
def main(config: DictConfig):
    # Imports should be nested inside @hydra.main to optimize tab completion
    # Read more here: https://github.com/facebookresearch/hydra/issues/934

    from src.utils import utils

    # Pretty print config using Rich library
    if config.get("print_config"):
        utils.print_config(config, resolve=True)

    if not config.single:
        from src.evaluation.eval import evaluation
        cSDR, uSDR = evaluation(config)
        df = pd.DataFrame(columns=['cSDR', 'uSDR'])
        df = df._append({'cSDR': cSDR, 'uSDR': uSDR}, ignore_index=True)
        log_dir = os.getcwd()
        df.to_csv(os.path.join(log_dir, "eval.csv"), index=False)
    else:
        from src.evaluation.eval_demo import evaluation
        evaluation(config, 0)

if __name__ == "__main__":
    OmegaConf.register_new_resolver("get_eval_log_dir", get_eval_log_dir)
    main()

run_eval.py
junyuchen-cjy/DTTNet-Pytorch

import os
import hydra
from omegaconf import DictConfig, OmegaConf
import soundfile as sf
from src.utils.utils import load_wav, get_unique_save_path
from src.utils.omega_resolvers import get_eval_log_dir
from pathlib import Path

import dotenv
from src.evaluation.separate import separate_with_ckpt_TDF, no_overlap_inference, overlap_inference
dotenv.load_dotenv(override=True)


@hydra.main(config_path="configs/", config_name="infer.yaml", version_base='1.1')
def main(config: DictConfig):
    # Imports should be nested inside @hydra.main to optimize tab completion
    # Read more here: https://github.com/facebookresearch/hydra/issues/934

    from src.utils import utils

    # Pretty print config using Rich library
    if config.get("print_config"):
        utils.print_config(config, resolve=True)


    model = hydra.utils.instantiate(config.model)
    ckpt_path = Path(config.ckpt_path)
    print(ckpt_path)
    mixture = load_wav(config.mixture_path)
    target_hat = separate_with_ckpt_TDF(config.batch_size, model, ckpt_path, mixture, config.device,
                                        config.double_chunk, config.overlap_add)


    base_name, file_name = os.path.split(config.mixture_path)
    #remove extension
    dir = os.path.splitext(file_name)[0]
    # print(config)

    save_path = os.path.join(config.paths.root_dir, "infer")
    if not os.path.exists(save_path):
        os.mkdir(save_path)
    dir = os.path.join(save_path, dir)
    cur_suffix = get_unique_save_path(dir)
    save_path = f"{dir}_{str(cur_suffix)}"

    if not os.path.exists(save_path):
        os.mkdir(save_path)

    save_path = os.path.join(save_path, model.target_name + ".wav")
    sf.write(save_path, target_hat.T, 44100)

    print(f"Saved to {save_path}")

if __name__ == "__main__":
    main()

run_infer.py
junyuchen-cjy/DTTNet-Pytorch
import os.path
from typing import Dict, Any

import torch
from pytorch_lightning import Callback
import pytorch_lightning as pl
import inspect
from src.models.mdxnet import AbstractMDXNet


class MakeONNXCallback(Callback):
    """Upload all *.py files to wandb as an artifact, at the beginning of the run."""

    def __init__(self, dirpath: str):
        self.dirpath = dirpath
        if not os.path.exists(self.dirpath):
            os.mkdir(self.dirpath)

    def on_save_checkpoint(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule',
                           checkpoint: Dict[str, Any]) -> dict:
        res = super().on_save_checkpoint(trainer, pl_module, checkpoint)

        var = inspect.signature(pl_module.__init__).parameters
        model = pl_module.__class__(**dict((name, pl_module.__dict__[name]) for name in var))
        model.load_state_dict(pl_module.state_dict())

        target_dir = '{}epoch_{}'.format(self.dirpath, pl_module.current_epoch)

        try:
            if not os.path.exists(target_dir):
                os.mkdir(target_dir)

            with torch.no_grad():
                torch.onnx.export(model,
                                  torch.zeros(model.inference_chunk_shape),
                                  '{}/{}.onnx'.format(target_dir, model.target_name),
                                  export_params=True,  # store the trained parameter weights inside the model file
                                  opset_version=13,  # the ONNX version to export the model to
                                  do_constant_folding=True,  # whether to execute constant folding for optimization
                                  input_names=['input'],  # the model's input names
                                  output_names=['output'],  # the model's output names
                                  dynamic_axes={'input': {0: 'batch_size'},  # variable length axes
                                                'output': {0: 'batch_size'}})
        except:
            print('onnx error')
        finally:
            del model

        return res

src/callbacks/onnx_callback.py
junyuchen-cjy/DTTNet-Pytorch
import glob
import os
from typing import List, Optional, Any

import matplotlib.pyplot as plt
import seaborn as sn
import torch
import wandb
from pytorch_lightning import Callback, Trainer
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.utilities.types import STEP_OUTPUT
from sklearn import metrics
from sklearn.metrics import f1_score, precision_score, recall_score


def get_wandb_logger(trainer: Trainer) -> WandbLogger:
    """Safely get Weights&Biases logger from Trainer."""

    if isinstance(trainer.logger, WandbLogger):
        return trainer.logger

    if isinstance(trainer.loggers, list):
        for logger in trainer.loggers:
            if isinstance(logger, WandbLogger):
                return logger

    raise Exception(
        "You are using wandb related callback, but WandbLogger was not found for some reason..."
    )


class UploadValidTrack(Callback):
    def __init__(self, crop: int, upload_after_n_epoch: int):
        self.sample_length = crop * 44100
        self.upload_after_n_epoch = upload_after_n_epoch
        self.len_left_window = self.len_right_window = self.sample_length // 2

    def on_validation_batch_end(
        self,
        trainer: 'pl.Trainer',
        pl_module: 'pl.LightningModule',
        outputs: Optional[STEP_OUTPUT],
        batch: Any,
        batch_idx: int,
        dataloader_idx: int,
    ) -> None:
        if outputs is None:
            return
        track_id = outputs['track_id']
        track = outputs['track']

        logger = get_wandb_logger(trainer=trainer)
        experiment = logger.experiment
        if pl_module.current_epoch < self.upload_after_n_epoch:
            return None

        mid = track.shape[-1]//2
        track = track[:, mid-self.len_left_window:mid+self.len_right_window]

        experiment.log({'track={}_epoch={}'.format(track_id, pl_module.current_epoch):
            [wandb.Audio(track.T, sample_rate=44100)]})


class WatchModel(Callback):
    """Make wandb watch model at the beginning of the run."""

    def __init__(self, log: str = "gradients", log_freq: int = 100):
        self.log = log
        self.log_freq = log_freq

    def on_train_start(self, trainer, pl_module):
        logger = get_wandb_logger(trainer=trainer)
        logger.watch(model=trainer.model, log=self.log, log_freq=self.log_freq)


class UploadCodeAsArtifact(Callback):
    """Upload all *.py files to wandb as an artifact, at the beginning of the run."""

    def __init__(self, code_dir: str):
        self.code_dir = code_dir

    def on_train_start(self, trainer, pl_module):
        logger = get_wandb_logger(trainer=trainer)
        experiment = logger.experiment

        code = wandb.Artifact("project-source", type="code")
        for path in glob.glob(os.path.join(self.code_dir, "**/*.py"), recursive=True):
            code.add_file(path)

        experiment.use_artifact(code)


class UploadCheckpointsAsArtifact(Callback):
    """Upload checkpoints to wandb as an artifact, at the end of run."""

    def __init__(self, ckpt_dir: str = "checkpoints/", upload_best_only: bool = False):
        self.ckpt_dir = ckpt_dir
        self.upload_best_only = upload_best_only

    def on_train_end(self, trainer, pl_module):
        logger = get_wandb_logger(trainer=trainer)
        experiment = logger.experiment

        ckpts = wandb.Artifact("experiment-ckpts", type="checkpoints")

        if self.upload_best_only:
            ckpts.add_file(trainer.checkpoint_callback.best_model_path)
        else:
            for path in glob.glob(os.path.join(self.ckpt_dir, "**/*.ckpt"), recursive=True):
                ckpts.add_file(path)

        experiment.use_artifact(ckpts)


class LogConfusionMatrix(Callback):
    """Generate confusion matrix every epoch and send it to wandb.
    Expects validation step to return predictions and targets.
    """

    def __init__(self):
        self.preds = []
        self.targets = []
        self.ready = True

    def on_sanity_check_start(self, trainer, pl_module) -> None:
        self.ready = False

    def on_sanity_check_end(self, trainer, pl_module):
        """Start executing this callback only after all validation sanity checks end."""
        self.ready = True

    def on_validation_batch_end(
        self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx
    ):
        """Gather data from single batch."""
        if self.ready:
            self.preds.append(outputs["preds"])
            self.targets.append(outputs["targets"])

    def on_validation_epoch_end(self, trainer, pl_module):
        """Generate confusion matrix."""
        if self.ready:
            logger = get_wandb_logger(trainer)
            experiment = logger.experiment

            preds = torch.cat(self.preds).cpu().numpy()
            targets = torch.cat(self.targets).cpu().numpy()

            confusion_matrix = metrics.confusion_matrix(y_true=targets, y_pred=preds)

            # set figure size
            plt.figure(figsize=(14, 8))

            # set labels size
            sn.set(font_scale=1.4)

            # set font size
            sn.heatmap(confusion_matrix, annot=True, annot_kws={"size": 8}, fmt="g")

            # names should be uniqe or else charts from different experiments in wandb will overlap
            experiment.log({f"confusion_matrix/{experiment.name}": wandb.Image(plt)}, commit=False)

            # according to wandb docs this should also work but it crashes
            # experiment.log(f{"confusion_matrix/{experiment.name}": plt})

            # reset plot
            plt.clf()

            self.preds.clear()
            self.targets.clear()


class LogF1PrecRecHeatmap(Callback):
    """Generate f1, precision, recall heatmap every epoch and send it to wandb.
    Expects validation step to return predictions and targets.
    """

    def __init__(self, class_names: List[str] = None):
        self.preds = []
        self.targets = []
        self.ready = True

    def on_sanity_check_start(self, trainer, pl_module):
        self.ready = False

    def on_sanity_check_end(self, trainer, pl_module):
        """Start executing this callback only after all validation sanity checks end."""
        self.ready = True

    def on_validation_batch_end(
        self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx
    ):
        """Gather data from single batch."""
        if self.ready:
            self.preds.append(outputs["preds"])
            self.targets.append(outputs["targets"])

    def on_validation_epoch_end(self, trainer, pl_module):
        """Generate f1, precision and recall heatmap."""
        if self.ready:
            logger = get_wandb_logger(trainer=trainer)
            experiment = logger.experiment

            preds = torch.cat(self.preds).cpu().numpy()
            targets = torch.cat(self.targets).cpu().numpy()
            f1 = f1_score(preds, targets, average=None)
            r = recall_score(preds, targets, average=None)
            p = precision_score(preds, targets, average=None)
            data = [f1, p, r]

            # set figure size
            plt.figure(figsize=(14, 3))

            # set labels size
            sn.set(font_scale=1.2)

            # set font size
            sn.heatmap(
                data,
                annot=True,
                annot_kws={"size": 10},
                fmt=".3f",
                yticklabels=["F1", "Precision", "Recall"],
            )

            # names should be uniqe or else charts from different experiments in wandb will overlap
            experiment.log({f"f1_p_r_heatmap/{experiment.name}": wandb.Image(plt)}, commit=False)

            # reset plot
            plt.clf()

            self.preds.clear()
            self.targets.clear()


class LogImagePredictions(Callback):
    """Logs a validation batch and their predictions to wandb.
    Example adapted from:
        https://wandb.ai/wandb/wandb-lightning/reports/Image-Classification-using-PyTorch-Lightning--VmlldzoyODk1NzY
    """

    def __init__(self, num_samples: int = 8):
        super().__init__()
        self.num_samples = num_samples
        self.ready = True

    def on_sanity_check_start(self, trainer, pl_module):
        self.ready = False

    def on_sanity_check_end(self, trainer, pl_module):
        """Start executing this callback only after all validation sanity checks end."""
        self.ready = True

    def on_validation_epoch_end(self, trainer, pl_module):
        if self.ready:
            logger = get_wandb_logger(trainer=trainer)
            experiment = logger.experiment

            # get a validation batch from the validation dat loader
            val_samples = next(iter(trainer.datamodule.val_dataloader()))
            val_imgs, val_labels = val_samples

            # run the batch through the network
            val_imgs = val_imgs.to(device=pl_module.device)
            logits = pl_module(val_imgs)
            preds = torch.argmax(logits, axis=-1)

            # log the images as wandb Image
            experiment.log(
                {
                    f"Images/{experiment.name}": [
                        wandb.Image(x, caption=f"Pred:{pred}, Label:{y}")
                        for x, pred, y in zip(
                            val_imgs[: self.num_samples],
                            preds[: self.num_samples],
                            val_labels[: self.num_samples],
                        )
                    ]
                }
            )

src/callbacks/wandb_callbacks.py
junyuchen-cjy/DTTNet-Pytorch
import os
from abc import ABCMeta, ABC
from pathlib import Path

import soundfile
from torch.utils.data import Dataset
import torch
import numpy as np
import random
from tqdm import tqdm

from src.utils.utils import load_wav
from src import utils
import numpy as np

log = utils.get_pylogger(__name__)

def check_target_name(target_name, source_names):
    try:
        assert target_name is not None
    except AssertionError:
        print('[ERROR] please identify target name. ex) +datamodule.target_name="vocals"')
        exit(-1)
    try:
        assert target_name in source_names or target_name == 'all'
    except AssertionError:
        print('[ERROR] target name should one of "bass", "drums", "other", "vocals", "all"')
        exit(-1)


def check_sample_rate(sr, sample_track):
    try:
        sample_rate = soundfile.read(sample_track)[1]
        assert sample_rate == sr
    except AssertionError:
        sample_rate = soundfile.read(sample_track)[1]
        print('[ERROR] sampling rate mismatched')
        print('\t=> sr in Config file: {}, but sr of data: {}'.format(sr, sample_rate))
        exit(-1)


class MusdbDataset(Dataset):
    __metaclass__ = ABCMeta

    def __init__(self, data_dir, chunk_size):
        self.source_names = ['bass', 'drums', 'other', 'vocals']
        self.chunk_size = chunk_size
        self.musdb_path = Path(data_dir)


class MusdbTrainDataset(MusdbDataset):
    def __init__(self, data_dir, chunk_size, target_name, aug_params, external_datasets, single_channel, epoch_size):
        super(MusdbTrainDataset, self).__init__(data_dir, chunk_size)

        self.single_channel = single_channel
        self.neg_lst = [x for x in self.source_names if x != target_name]

        self.target_name = target_name
        check_target_name(self.target_name, self.source_names)

        if not self.musdb_path.joinpath('metadata').exists():
            os.mkdir(self.musdb_path.joinpath('metadata'))

        splits = ['train']
        if external_datasets is not None:
            splits += external_datasets

        # collect paths for datasets and metadata (track names and duration)
        datasets, metadata_caches = [], []
        raw_datasets = []    # un-augmented datasets
        for split in splits:
            raw_datasets.append(self.musdb_path.joinpath(split))
            max_pitch, max_tempo = aug_params
            for p in range(-max_pitch, max_pitch+1):
                for t in range(-max_tempo, max_tempo+1, 10):
                    aug_split = split if p==t==0 else split + f'_p={p}_t={t}'
                    datasets.append(self.musdb_path.joinpath(aug_split))
                    metadata_caches.append(self.musdb_path.joinpath('metadata').joinpath(aug_split + '.pkl'))

        # collect all track names and their duration
        self.metadata = []
        raw_track_lengths = []   # for calculating epoch size
        for i, (dataset, metadata_cache) in enumerate(tqdm(zip(datasets, metadata_caches))):
            try:
                metadata = torch.load(metadata_cache)
            except FileNotFoundError:
                print('creating metadata for', dataset)
                metadata = []
                for track_name in sorted(os.listdir(dataset)):
                    track_path = dataset.joinpath(track_name)
                    track_length = load_wav(track_path.joinpath('vocals.wav')).shape[-1]
                    metadata.append((track_path, track_length))
                torch.save(metadata, metadata_cache)

            self.metadata += metadata
            if dataset in raw_datasets:
                raw_track_lengths += [length for path, length in metadata]

        self.epoch_size = sum(raw_track_lengths) // self.chunk_size if epoch_size is None else epoch_size
        log.info(f'epoch size: {self.epoch_size}')

    def __getitem__(self, _):
        sources = []
        for source_name in self.source_names:
            track_path, track_length = random.choice(self.metadata)   # random mixing between tracks
            source = load_wav(track_path.joinpath(source_name + '.wav'),
                              track_length=track_length, chunk_size=self.chunk_size) # (2, times)
            sources.append(source)

        mix = sum(sources)

        if self.target_name == 'all':
            # Targets for models that separate all four sources (ex. Demucs).
            # This adds additional 'source' dimension => batch_shape=[batch, source, channel, time]
            target = sources
        else:
            target = sources[self.source_names.index(self.target_name)]

        mix, target = torch.tensor(mix), torch.tensor(target)
        if self.single_channel:
            mix = torch.mean(mix, dim=0, keepdim=True)
            target = torch.mean(target, dim=0, keepdim=True)
        return mix, target

    def __len__(self):
        return self.epoch_size


class MusdbValidDataset(MusdbDataset):

    def __init__(self, data_dir, chunk_size, target_name, overlap, batch_size, single_channel):
        super(MusdbValidDataset, self).__init__(data_dir, chunk_size)

        self.target_name = target_name
        check_target_name(self.target_name, self.source_names)

        self.overlap = overlap
        self.batch_size = batch_size
        self.single_channel = single_channel

        musdb_valid_path = self.musdb_path.joinpath('valid')
        self.track_paths = [musdb_valid_path.joinpath(track_name)
                            for track_name in os.listdir(musdb_valid_path)]

    def __getitem__(self, index):
        mix = load_wav(self.track_paths[index].joinpath('mixture.wav')) # (2, time)

        if self.target_name == 'all':
            # Targets for models that separate all four sources (ex. Demucs).
            # This adds additional 'source' dimension => batch_shape=[batch, source, channel, time]
            target = [load_wav(self.track_paths[index].joinpath(source_name + '.wav'))
                      for source_name in self.source_names]
        else:
            target = load_wav(self.track_paths[index].joinpath(self.target_name + '.wav'))

        chunk_output_size = self.chunk_size - 2 * self.overlap
        left_pad = np.zeros([2, self.overlap])
        right_pad = np.zeros([2, self.overlap + chunk_output_size - (mix.shape[-1] % chunk_output_size)])
        mix_padded = np.concatenate([left_pad, mix, right_pad], 1)

        num_chunks = mix_padded.shape[-1] // chunk_output_size
        mix_chunks = np.array([mix_padded[:, i * chunk_output_size: i * chunk_output_size + self.chunk_size]
                      for i in range(num_chunks)])
        mix_chunk_batches = torch.tensor(mix_chunks, dtype=torch.float32).split(self.batch_size)
        target = torch.tensor(target)

        if self.single_channel:
            mix_chunk_batches = [torch.mean(t, dim=1, keepdim=True) for t in mix_chunk_batches]
            target = torch.mean(target, dim=0, keepdim=True)

        return mix_chunk_batches, target

    def __len__(self):
        return len(self.track_paths)
src/datamodules/datasets/musdb.py
junyuchen-cjy/DTTNet-Pytorch
import os
from os.path import exists, join
from pathlib import Path
from typing import Optional, Tuple

from pytorch_lightning import LightningDataModule
from torch.utils.data import ConcatDataset, DataLoader, Dataset, random_split

from src.datamodules.datasets.musdb import MusdbTrainDataset, MusdbValidDataset


class MusdbDataModule(LightningDataModule):
    """
    LightningDataModule for Musdb18-HQ dataset.
    A DataModule implements 5 key methods:
        - prepare_data (things to do on 1 GPU/TPU, not on every GPU/TPU in distributed mode)
        - setup (things to do on every accelerator in distributed mode)
        - train_dataloader (the training dataloader)
        - val_dataloader (the validation dataloader(s))
        - test_dataloader (the test dataloader(s))
    This allows you to share a full dataset without explaining how to download,
    split, transform and process the data
    Read the docs:
        https://pytorch-lightning.readthedocs.io/en/latest/extensions/datamodules.html
    """

    def __init__(
            self,
            data_dir: str,
            aug_params,
            target_name: str,
            overlap: int,
            hop_length: int,
            dim_t: int,
            sample_rate: int,
            batch_size: int,
            num_workers: int,
            pin_memory: bool,
            external_datasets,
            audio_ch: int,
            epoch_size,
            **kwargs,
    ):
        super().__init__()

        self.data_dir = Path(data_dir)
        self.target_name = target_name
        self.aug_params = aug_params
        self.external_datasets = external_datasets

        self.batch_size = batch_size
        self.num_workers = num_workers
        self.pin_memory = pin_memory

        # audio-related
        self.hop_length = hop_length
        self.sample_rate = sample_rate
        self.single_channel = audio_ch == 1

        # derived
        self.chunk_size = hop_length * (dim_t - 1)
        self.overlap = overlap

        self.epoch_size = epoch_size

        self.data_train: Optional[Dataset] = None
        self.data_val: Optional[Dataset] = None
        self.data_test: Optional[Dataset] = None

        trainset_path = self.data_dir.joinpath('train')
        validset_path = self.data_dir.joinpath('valid')

        # create validation split
        if not exists(validset_path):
            from shutil import move
            os.mkdir(validset_path)
            for track in kwargs['validation_set']:
                if trainset_path.joinpath(track).exists():
                    move(trainset_path.joinpath(track), validset_path.joinpath(track))
        else:
            valid_files = os.listdir(validset_path)
            assert set(valid_files) == set(kwargs['validation_set'])

    def setup(self, stage: Optional[str] = None):
        """Load data. Set variables: self.data_train, self.data_val, self.data_test."""
        self.data_train = MusdbTrainDataset(self.data_dir,
                                            self.chunk_size,
                                            self.target_name,
                                            self.aug_params,
                                            self.external_datasets,
                                            self.single_channel,
                                            self.epoch_size)

        self.data_val = MusdbValidDataset(self.data_dir,
                                          self.chunk_size,
                                          self.target_name,
                                          self.overlap,
                                          self.batch_size,
                                          self.single_channel)

    def train_dataloader(self):
        return DataLoader(
            dataset=self.data_train,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=self.pin_memory,
            shuffle=True,
        )

    def val_dataloader(self):
        return DataLoader(
            dataset=self.data_val,
            batch_size=1,
            num_workers=self.num_workers,
            pin_memory=self.pin_memory,
            shuffle=False,
        )
src/datamodules/musdb_datamodule.py
junyuchen-cjy/DTTNet-Pytorch
from abc import ABCMeta
from typing import Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from pytorch_lightning import LightningModule
from pytorch_lightning.utilities.types import STEP_OUTPUT

from src.utils.utils import sdr, simplified_msseval


class AbstractModel(LightningModule):
    __metaclass__ = ABCMeta

    def __init__(self, target_name,
                 lr, optimizer,
                  dim_f, dim_t, n_fft, hop_length, overlap,
                 audio_ch,
                  **kwargs):
        super().__init__()
        self.target_name = target_name
        self.lr = lr
        self.optimizer = optimizer
        self.dim_c_in = audio_ch * 2
        self.dim_c_out = audio_ch * 2
        self.dim_f = dim_f
        self.dim_t = dim_t
        self.n_fft = n_fft
        self.n_bins = n_fft // 2 + 1
        self.hop_length = hop_length
        self.audio_ch = audio_ch

        self.chunk_size = hop_length * (self.dim_t - 1)
        self.inference_chunk_size = hop_length * (self.dim_t*2 - 1)
        self.overlap = overlap
        self.window = nn.Parameter(torch.hann_window(window_length=self.n_fft, periodic=True), requires_grad=False)
        self.freq_pad = nn.Parameter(torch.zeros([1, self.dim_c_out, self.n_bins - self.dim_f, 1]), requires_grad=False)
        self.inference_chunk_shape = (self.stft(torch.zeros([1, audio_ch, self.inference_chunk_size]))).shape


    def configure_optimizers(self):
        if self.optimizer == 'rmsprop':
            print("Using RMSprop optimizer")
            return torch.optim.RMSprop(self.parameters(), self.lr)
        elif self.optimizer == 'adamW':
            print("Using AdamW optimizer")
            return torch.optim.AdamW(self.parameters(), self.lr)

    def comp_loss(self, pred_detail, target_wave):
        pred_detail = self.istft(pred_detail)

        comp_loss = F.l1_loss(pred_detail, target_wave)

        self.log("train/comp_loss", comp_loss, sync_dist=True, on_step=False, on_epoch=True, prog_bar=False)

        return comp_loss


    def training_step(self, *args, **kwargs) -> STEP_OUTPUT:
        mix_wave, target_wave = args[0] # (batch, c, 261120)

        # input 1
        stft_44k = self.stft(mix_wave) # (batch, c*2, 1044, 256)
        # forward
        t_est_stft = self(stft_44k) # (batch, c, 1044, 256)

        loss = self.comp_loss(t_est_stft, target_wave)

        self.log("train/loss", loss, sync_dist=True, on_step=True, on_epoch=True, prog_bar=True)

        return {"loss": loss}


    # Validation SDR is calculated on whole tracks and not chunks since
    # short inputs have high possibility of being silent (all-zero signal)
    # which leads to very low sdr values regardless of the model.
    # A natural procedure would be to split a track into chunk batches and
    # load them on multiple gpus, but aggregation was too difficult.
    # So instead we load one whole track on a single device (data_loader batch_size should always be 1)
    # and do all the batch splitting and aggregation on a single device.
    def validation_step(self, *args, **kwargs) -> Optional[STEP_OUTPUT]:
        mix_chunk_batches, target = args[0]

        # remove data_loader batch dimension
        # [(b, c, time)], (c, all_times)
        mix_chunk_batches, target = [batch[0] for batch in mix_chunk_batches], target[0]

        # process whole track in batches of chunks
        target_hat_chunks = []
        for batch in mix_chunk_batches:
            # input
            stft_44k = self.stft(batch)  # (batch, c*2, 1044, 256)
            pred_detail = self(stft_44k) # (batch, c, 1044, 256), irm
            pred_detail = self.istft(pred_detail)

            target_hat_chunks.append(pred_detail[..., self.overlap:-self.overlap])
        target_hat_chunks = torch.cat(target_hat_chunks) # (b*len(ls),c,t)

        # concat all output chunks (c, all_times)
        target_hat = target_hat_chunks.transpose(0, 1).reshape(self.audio_ch, -1)[..., :target.shape[-1]]

        ests = target_hat.detach().cpu().numpy()  # (c, all_times)
        references = target.cpu().numpy()
        score = sdr(ests, references)

        # (src, t, c)
        SDR = simplified_msseval(np.expand_dims(references.T, axis=0), np.expand_dims(ests.T, axis=0), chunk_size=44100)
        # self.log("val/sdr", score, sync_dist=True, on_step=False, on_epoch=True, logger=True)

        return {'song': score, 'chunk': SDR}

    def validation_epoch_end(self, outputs) -> None:
        avg_uSDR = torch.Tensor([x['song'] for x in outputs]).mean()
        self.log("val/usdr", avg_uSDR, sync_dist=True, on_step=False, on_epoch=True, logger=True)

        chunks = [x['chunk'][0, :] for x in outputs]
        # concat np array
        chunks = np.concatenate(chunks, axis=0)
        median_cSDR = np.nanmedian(chunks.flatten(), axis=0)
        median_cSDR = float(median_cSDR)
        self.log("val/csdr", median_cSDR, sync_dist=True, on_step=False, on_epoch=True, logger=True)

    def stft(self, x):
        '''
        Args:
            x: (batch, c, 261120)
        '''
        dim_b = x.shape[0]
        x = x.reshape([dim_b * self.audio_ch, -1]) # (batch*c, 261120)
        x = torch.stft(x, n_fft=self.n_fft, hop_length=self.hop_length, window=self.window, center=True) # (batch*c, 3073, 256, 2)
        x = x.permute([0, 3, 1, 2]) # (batch*c, 2, 3073, 256)
        x = x.reshape([dim_b, self.audio_ch, 2, self.n_bins, -1]).reshape([dim_b, self.audio_ch * 2, self.n_bins, -1]) # (batch, c*2, 3073, 256)
        return x[:, :, :self.dim_f] # (batch, c*2, 2048, 256)

    def istft(self, x):
        '''
        Args:
            x: (batch, c*2, 2048, 256)
        '''
        dim_b = x.shape[0]
        x = torch.cat([x, self.freq_pad.repeat([x.shape[0], 1, 1, x.shape[-1]])], -2) # (batch, c*2, 3073, 256)
        x = x.reshape([dim_b, self.audio_ch, 2, self.n_bins, -1]).reshape([dim_b * self.audio_ch, 2, self.n_bins, -1]) # (batch*c, 2, 3073, 256)
        x = x.permute([0, 2, 3, 1]) # (batch*c, 3073, 256, 2)
        x = torch.istft(x, n_fft=self.n_fft, hop_length=self.hop_length, window=self.window, center=True) # (batch*c, 261120)
        return x.reshape([dim_b, self.audio_ch, -1]) # (batch,c,261120)

    def demix(self, mix, inf_chunk_size, batch_size=5, inf_overf=4):
        '''
        Args:
            mix: (C, L)
        Returns:
            est: (src, C, L)
        '''

        # batch_size = self.config.inference.batch_size
        #  = self.chunk_size
        # self.instruments = ['bass', 'drums', 'other', 'vocals']
        num_instruments = 1

        inf_hop = inf_chunk_size // inf_overf  # hop size
        L = mix.shape[1]
        pad_size = inf_hop - (L - inf_chunk_size) % inf_hop
        mix = torch.cat([torch.zeros(2, inf_chunk_size - inf_hop), torch.Tensor(mix), torch.zeros(2, pad_size + inf_chunk_size - inf_hop)], 1)
        mix = mix.cuda()

        chunks = []
        i = 0
        while i + inf_chunk_size <= mix.shape[1]:
            chunks.append(mix[:, i:i + inf_chunk_size])
            i += inf_hop
        chunks = torch.stack(chunks)

        batches = []
        i = 0
        while i < len(chunks):
            batches.append(chunks[i:i + batch_size])
            i = i + batch_size

        X = torch.zeros(num_instruments, 2, inf_chunk_size - inf_hop) # (src, c, t)
        X = X.cuda()
        with torch.cuda.amp.autocast():
            with torch.no_grad():
                for batch in batches:
                    x = self.stft(batch)
                    x = self(x)
                    x = self.istft(x) # (batch, c, 261120)
                    # insert new axis, the model only predict 1 src so we need to add axis
                    x = x[:,None, ...] # (batch, 1, c, 261120)
                    x = x.repeat([ 1, num_instruments, 1, 1]) # (batch, src, c, 261120)
                    for w in x: # iterate over batch
                        a = X[..., :-(inf_chunk_size - inf_hop)]
                        b = X[..., -(inf_chunk_size - inf_hop):] + w[..., :(inf_chunk_size - inf_hop)]
                        c = w[..., (inf_chunk_size - inf_hop):]
                        X = torch.cat([a, b, c], -1)

        estimated_sources = X[..., inf_chunk_size - inf_hop:-(pad_size + inf_chunk_size - inf_hop)] / inf_overf

        assert L == estimated_sources.shape[-1]

        return estimated_sources


src/dp_tdf/abstract.py
junyuchen-cjy/DTTNet-Pytorch
import torch
import torch.nn as nn

# Original code from https://github.com/amanteur/BandSplitRNN-Pytorch
class RNNModule(nn.Module):
    """
    RNN submodule of BandSequence module
    """

    def __init__(
            self,
            group_num: int,
            input_dim_size: int,
            hidden_dim_size: int,
            rnn_type: str = 'lstm',
            bidirectional: bool = True
    ):
        super(RNNModule, self).__init__()
        self.groupnorm = nn.GroupNorm(group_num, input_dim_size)
        self.rnn = getattr(nn, rnn_type)(
            input_dim_size, hidden_dim_size, batch_first=True, bidirectional=bidirectional # 输出是2*hidden_dim_size，因为是bi
        )
        self.fc = nn.Linear(
            hidden_dim_size * 2 if bidirectional else hidden_dim_size,
            input_dim_size
        )

    def forward(
            self,
            x: torch.Tensor
    ):
        """
        Input shape:
            across T - [batch_size, k_subbands, time, n_features]
            OR
            across K - [batch_size, time, k_subbands, n_features]
        """
        B, K, T, N = x.shape  # across T      across K (keep in mind T->K, K->T)
        # print(x.shape)

        out = x.view(B * K, T, N)  # [BK, T, N]    [BT, K, N]

        # print(out.shape)
        # print(self.groupnorm)
        out = self.groupnorm(
            out.transpose(-1, -2)
        ).transpose(-1, -2)  # [BK, T, N]    [BT, K, N]
        out = self.rnn(out)[0]  # [BK, T, H]    [BT, K, H]， 最后一维是特征
        out = self.fc(out)  # [BK, T, N]    [BT, K, N]

        x = out.view(B, K, T, N) + x  # [B, K, T, N]  [B, T, K, N]

        x = x.permute(0, 2, 1, 3).contiguous()  # [B, T, K, N]  [B, K, T, N]
        return x


class BandSequenceModelModule(nn.Module):
    """
    BandSequence (2nd) Module of BandSplitRNN.
    Runs input through n BiLSTMs in two dimensions - time and subbands.
    """

    def __init__(
            self,
            # group_num,
            input_dim_size: int,
            hidden_dim_size: int,
            rnn_type: str = 'lstm',
            bidirectional: bool = True,
            num_layers: int = 12,
            n_heads: int = 4,
    ):
        super(BandSequenceModelModule, self).__init__()

        self.bsrnn = nn.ModuleList([])
        self.n_heads = n_heads

        input_dim_size = input_dim_size // n_heads
        hidden_dim_size = hidden_dim_size // n_heads
        group_num = input_dim_size // 16
        # print(f"input_dim_size: {input_dim_size}, hidden_dim_size: {hidden_dim_size}, group_num: {group_num}")

        # print(group_num, input_dim_size)

        for _ in range(num_layers):
            rnn_across_t = RNNModule(
                group_num, input_dim_size, hidden_dim_size, rnn_type, bidirectional
            )
            rnn_across_k = RNNModule(
                group_num, input_dim_size, hidden_dim_size, rnn_type, bidirectional
            )
            self.bsrnn.append(
                nn.Sequential(rnn_across_t, rnn_across_k)
            )

    def forward(self, x: torch.Tensor):
        """
        Input shape: [batch_size, k_subbands, time, n_features]
        Output shape: [batch_size, k_subbands, time, n_features]
        """
        # x (b,c,t,f)
        b,c,t,f = x.shape
        x = x.view(b * self.n_heads, c // self.n_heads, t, f) # [b*n_heads, c//n_heads, t, f]

        x = x.permute(0, 3, 2, 1).contiguous()  # [b*n_heads, f, t, c//n_heads]
        for i in range(len(self.bsrnn)):
            x = self.bsrnn[i](x)

        x = x.permute(0, 3, 2, 1).contiguous()  # [b*n_heads, c//n_heads, t, f]
        x = x.view(b, c, t, f)  # [b, c, t, f]
        return x


if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    batch_size, k_subbands, t_timesteps, input_dim = 4, 41, 512, 128
    in_features = torch.rand(batch_size, k_subbands, t_timesteps, input_dim).to(device)

    cfg = {
        # "t_timesteps": t_timesteps,
        "group_num": 32,
        "input_dim_size": 128,
        "hidden_dim_size": 256,
        "rnn_type": "LSTM",
        "bidirectional": True,
        "num_layers": 1
    }
    model = BandSequenceModelModule(**cfg).to(device)
    _ = model.eval()

    with torch.no_grad():
        out_features = model(in_features)

    print(f"In: {in_features.shape}\nOut: {out_features.shape}")
    print(f"Total number of parameters: {sum([p.numel() for p in model.parameters()])}")

src/dp_tdf/bandsequence.py
junyuchen-cjy/DTTNet-Pytorch
import torch.nn as nn
import torch

from src.dp_tdf.modules import TFC_TDF, TFC_TDF_Res1, TFC_TDF_Res2
from src.dp_tdf.bandsequence import BandSequenceModelModule

from src.layers import (get_norm)
from src.dp_tdf.abstract import AbstractModel

class DPTDFNet(AbstractModel):
    def __init__(self, num_blocks, l, g, k, bn, bias, bn_norm, bandsequence, block_type,  **kwargs):

        super(DPTDFNet, self).__init__(**kwargs)
        # self.save_hyperparameters()

        self.num_blocks = num_blocks
        self.l = l
        self.g = g
        self.k = k
        self.bn = bn
        self.bias = bias

        self.n = num_blocks // 2
        scale = (2, 2)

        if block_type == "TFC_TDF":
            T_BLOCK = TFC_TDF
        elif block_type == "TFC_TDF_Res1":
            T_BLOCK = TFC_TDF_Res1
        elif block_type == "TFC_TDF_Res2":
            T_BLOCK = TFC_TDF_Res2
        else:
            raise ValueError(f"Unknown block type {block_type}")

        self.first_conv = nn.Sequential(
            nn.Conv2d(in_channels=self.dim_c_in, out_channels=g, kernel_size=(1, 1)),
            get_norm(bn_norm, g),
            nn.ReLU(),
        )

        f = self.dim_f
        c = g
        self.encoding_blocks = nn.ModuleList()
        self.ds = nn.ModuleList()

        for i in range(self.n):
            c_in = c

            self.encoding_blocks.append(T_BLOCK(c_in, c, l, f, k, bn, bn_norm, bias=bias))
            self.ds.append(
                nn.Sequential(
                    nn.Conv2d(in_channels=c, out_channels=c + g, kernel_size=scale, stride=scale),
                    get_norm(bn_norm, c + g),
                    nn.ReLU()
                )
            )
            f = f // 2
            c += g

        self.bottleneck_block1 = T_BLOCK(c, c, l, f, k, bn, bn_norm, bias=bias)
        self.bottleneck_block2 = BandSequenceModelModule(
            **bandsequence,
            input_dim_size=c,
            hidden_dim_size=2*c
        )

        self.decoding_blocks = nn.ModuleList()
        self.us = nn.ModuleList()
        for i in range(self.n):
            # print(f"i: {i}, in channels: {c}")
            self.us.append(
                nn.Sequential(
                    nn.ConvTranspose2d(in_channels=c, out_channels=c - g, kernel_size=scale, stride=scale),
                    get_norm(bn_norm, c - g),
                    nn.ReLU()
                )
            )

            f = f * 2
            c -= g

            self.decoding_blocks.append(T_BLOCK(c, c, l, f, k, bn, bn_norm, bias=bias))

        self.final_conv = nn.Sequential(
            nn.Conv2d(in_channels=c, out_channels=self.dim_c_out, kernel_size=(1, 1)),
        )

    def forward(self, x):
        '''
        Args:
            x: (batch, c*2, 2048, 256)
        '''
        x = self.first_conv(x)

        x = x.transpose(-1, -2)

        ds_outputs = []
        for i in range(self.n):
            x = self.encoding_blocks[i](x)
            ds_outputs.append(x)
            x = self.ds[i](x)

        # print(f"bottleneck in: {x.shape}")
        x = self.bottleneck_block1(x)
        x = self.bottleneck_block2(x)

        for i in range(self.n):
            x = self.us[i](x)
            # print(f"us{i} in: {x.shape}")
            # print(f"ds{i} out: {ds_outputs[-i - 1].shape}")
            x = x * ds_outputs[-i - 1]
            x = self.decoding_blocks[i](x)

        x = x.transpose(-1, -2)

        x = self.final_conv(x)

        return x
src/dp_tdf/dp_tdf_net.py
junyuchen-cjy/DTTNet-Pytorch
import torch
import torch.nn as nn

from src.layers import (get_norm)

class TFC(nn.Module):
    def __init__(self, c_in, c_out, l, k, bn_norm):
        super(TFC, self).__init__()

        self.H = nn.ModuleList()
        for i in range(l):
            if i == 0:
                c_in = c_in
            else:
                c_in = c_out
            self.H.append(
                nn.Sequential(
                    nn.Conv2d(in_channels=c_in, out_channels=c_out, kernel_size=k, stride=1, padding=k // 2),
                    get_norm(bn_norm, c_out),
                    nn.ReLU(),
                )
            )

    def forward(self, x):
        for h in self.H:
            x = h(x)
        return x


class DenseTFC(nn.Module):
    def __init__(self, c_in, c_out, l, k, bn_norm):
        super(DenseTFC, self).__init__()

        self.conv = nn.ModuleList()
        for i in range(l):
            self.conv.append(
                nn.Sequential(
                    nn.Conv2d(in_channels=c_in, out_channels=c_out, kernel_size=k, stride=1, padding=k // 2),
                    get_norm(bn_norm, c_out),
                    nn.ReLU(),
                )
            )

    def forward(self, x):
        for layer in self.conv[:-1]:
            x = torch.cat([layer(x), x], 1)
        return self.conv[-1](x)


class TFC_TDF(nn.Module):
    def __init__(self, c_in, c_out, l, f, k, bn, bn_norm, dense=False, bias=True):

        super(TFC_TDF, self).__init__()

        self.use_tdf = bn is not None

        self.tfc = DenseTFC(c_in, c_out, l, k, bn_norm) if dense else TFC(c_in, c_out, l, k, bn_norm)

        if self.use_tdf:
            if bn == 0:
                # print(f"TDF={f},{f}")
                self.tdf = nn.Sequential(
                    nn.Linear(f, f, bias=bias),
                    get_norm(bn_norm, c_out),
                    nn.ReLU()
                )
            else:
                # print(f"TDF={f},{f // bn},{f}")
                self.tdf = nn.Sequential(
                    nn.Linear(f, f // bn, bias=bias),
                    get_norm(bn_norm, c_out),
                    nn.ReLU(),
                    nn.Linear(f // bn, f, bias=bias),
                    get_norm(bn_norm, c_out),
                    nn.ReLU()
                )

    def forward(self, x):
        x = self.tfc(x)
        return x + self.tdf(x) if self.use_tdf else x


class TFC_TDF_Res1(nn.Module):
    def __init__(self, c_in, c_out, l, f, k, bn, bn_norm, dense=False, bias=True):

        super(TFC_TDF_Res1, self).__init__()

        self.use_tdf = bn is not None

        self.tfc = DenseTFC(c_in, c_out, l, k, bn_norm) if dense else TFC(c_in, c_out, l, k, bn_norm)

        self.res = TFC(c_in, c_out, 1, k, bn_norm)

        if self.use_tdf:
            if bn == 0:
                # print(f"TDF={f},{f}")
                self.tdf = nn.Sequential(
                    nn.Linear(f, f, bias=bias),
                    get_norm(bn_norm, c_out),
                    nn.ReLU()
                )
            else:
                # print(f"TDF={f},{f // bn},{f}")
                self.tdf = nn.Sequential(
                    nn.Linear(f, f // bn, bias=bias),
                    get_norm(bn_norm, c_out),
                    nn.ReLU(),
                    nn.Linear(f // bn, f, bias=bias),
                    get_norm(bn_norm, c_out),
                    nn.ReLU()
                )

    def forward(self, x):
        res = self.res(x)
        x = self.tfc(x)
        x = x + res
        return x + self.tdf(x) if self.use_tdf else x


class TFC_TDF_Res2(nn.Module):
    def __init__(self, c_in, c_out, l, f, k, bn, bn_norm, dense=False, bias=True):

        super(TFC_TDF_Res2, self).__init__()

        self.use_tdf = bn is not None

        self.tfc1 = TFC(c_in, c_out, l, k, bn_norm)
        self.tfc2 = TFC(c_in, c_out, l, k, bn_norm)

        self.res = TFC(c_in, c_out, 1, k, bn_norm)

        if self.use_tdf:
            if bn == 0:
                # print(f"TDF={f},{f}")
                self.tdf = nn.Sequential(
                    nn.Linear(f, f, bias=bias),
                    get_norm(bn_norm, c_out),
                    nn.ReLU()
                )
            else:
                # print(f"TDF={f},{f // bn},{f}")
                self.tdf = nn.Sequential(
                    nn.Linear(f, f // bn, bias=bias),
                    get_norm(bn_norm, c_out),
                    nn.ReLU(),
                    nn.Linear(f // bn, f, bias=bias),
                    get_norm(bn_norm, c_out),
                    nn.ReLU()
                )

    def forward(self, x):
        res = self.res(x)
        x = self.tfc1(x)
        if self.use_tdf:
            x = x + self.tdf(x)
        x = self.tfc2(x)
        x = x + res
        return x

src/dp_tdf/modules.py
junyuchen-cjy/DTTNet-Pytorch
from os import listdir
from pathlib import Path
from typing import Optional, List

from concurrent import futures
import hydra
import wandb
import os
import shutil
from omegaconf import DictConfig
from pytorch_lightning import LightningDataModule, LightningModule
from pytorch_lightning.loggers import Logger, WandbLogger
import soundfile as sf

from tqdm import tqdm
import numpy as np
from src.callbacks.wandb_callbacks import get_wandb_logger
from src.evaluation.separate import separate_with_onnx_TDF, separate_with_ckpt_TDF
from src.utils import utils
from src.utils.utils import load_wav, sdr, get_median_csdr, save_results, get_metrics

from src.utils import pylogger

log = pylogger.get_pylogger(__name__)


def evaluation(config: DictConfig):

    assert config.split in ['train', 'valid', 'test']

    data_dir = Path(config.get('eval_dir')).joinpath(config['split'])
    assert data_dir.exists()

    # Init Lightning loggers
    loggers: List[Logger] = []
    if "logger" in config:
        for _, lg_conf in config.logger.items():
            if "_target_" in lg_conf:
                log.info(f"Instantiating logger <{lg_conf._target_}>")
                loggers.append(hydra.utils.instantiate(lg_conf))

        if any([isinstance(l, WandbLogger) for l in loggers]):
            utils.wandb_login(key=config.wandb_api_key)

    model = hydra.utils.instantiate(config.model)
    target_name = model.target_name
    ckpt_path = Path(config.ckpt_path)
    is_onnx = os.path.split(ckpt_path)[-1].split('.')[-1] == 'onnx'
    shutil.copy(ckpt_path,os.getcwd()) # copy model

    ssdrs = []
    bss_lst = []
    bss_perms = []
    num_tracks = len(listdir(data_dir))
    target_list = [config.model.target_name,"complement"]


    pool = futures.ProcessPoolExecutor
    with pool(config.pool_workers) as pool:
        datas = sorted(listdir(data_dir))
        if len(datas) > 27: # if not debugging
            # move idx 27 to head
            datas = [datas[27]] + datas[:27] + datas[28:]
        # iterate datas with batchsize 8
        for k in range(0, len(datas), config.pool_workers):
            batch = datas[k:k + config.pool_workers]
            pendings = []
            for i, track in tqdm(enumerate(batch)):
                folder_name = track
                track = data_dir.joinpath(track)
                mixture = load_wav(track.joinpath('mixture.wav')) # (c, t)
                target = load_wav(track.joinpath(target_name + '.wav'))

                if model.audio_ch == 1:
                    mixture = np.mean(mixture, axis=0, keepdims=True)
                    target = np.mean(target, axis=0, keepdims=True)
                #target_hat = {source: separate(config['batch_size'], models[source], onnxs[source], mixture) for source in sources}
                if is_onnx:
                    target_hat = separate_with_onnx_TDF(config.batch_size, model, ckpt_path, mixture)
                else:
                    target_hat = separate_with_ckpt_TDF(config.batch_size, model, ckpt_path, mixture, config.device, config.double_chunk, config.overlap_add)


                pendings.append((folder_name, pool.submit(
                    get_metrics, target_hat, target, mixture, sr=44100,version=config.bss)))

                for wandb_logger in [logger for logger in loggers if isinstance(logger, WandbLogger)]:
                    mid = mixture.shape[-1] // 2
                    track = target_hat[:, mid - 44100 * 3:mid + 44100 * 3]
                    wandb_logger.experiment.log(
                        {f'track={k+i}_target={target_name}': [wandb.Audio(track.T, sample_rate=44100)]})


            for i, (track_name, pending) in tqdm(enumerate(pendings)):
                pending = pending.result()
                bssmetrics, perms, ssdr = pending
                bss_lst.append(bssmetrics)
                bss_perms.append(perms)
                ssdrs.append(ssdr)

                for logger in loggers:
                    logger.log_metrics({'song/ssdr': ssdr}, k+i)
                    logger.log_metrics({'song/csdr': get_median_csdr([bssmetrics])}, k+i)

    log_dir = os.getcwd()
    save_results(log_dir, bss_lst, target_list, bss_perms, ssdrs)

    cSDR = get_median_csdr(bss_lst)
    uSDR = sum(ssdrs)/num_tracks
    for logger in loggers:
        logger.log_metrics({'metrics/mean_sdr_' + target_name: sum(ssdrs)/num_tracks})
        logger.log_metrics({'metrics/median_csdr_' + target_name: get_median_csdr(bss_lst)})
        # get the path of the log dir
        if not isinstance(logger, WandbLogger):
            logger.experiment.close()

    if any([isinstance(logger, WandbLogger) for logger in loggers]):
        wandb.finish()

    return cSDR, uSDR

src/evaluation/eval.py
junyuchen-cjy/DTTNet-Pytorch
from os import listdir
from pathlib import Path
from typing import Optional, List

from concurrent import futures
import hydra
import wandb
import os
import shutil
from omegaconf import DictConfig
from pytorch_lightning import LightningDataModule, LightningModule
from pytorch_lightning.loggers import Logger, WandbLogger

from tqdm import tqdm
import numpy as np
from src.callbacks.wandb_callbacks import get_wandb_logger
from src.evaluation.separate import separate_with_onnx_TDF, separate_with_ckpt_TDF
from src.utils import utils
from src.utils.utils import load_wav, sdr, get_median_csdr, save_results, get_metrics

from src.utils import pylogger
import soundfile as sf
log = pylogger.get_pylogger(__name__)


def evaluation(config: DictConfig, idx):

    assert config.split in ['train', 'valid', 'test']

    data_dir = Path(config.get('eval_dir')).joinpath(config['split'])
    assert data_dir.exists()

    model = hydra.utils.instantiate(config.model)
    target_name = model.target_name
    ckpt_path = Path(config.ckpt_path)
    is_onnx = os.path.split(ckpt_path)[-1].split('.')[-1] == 'onnx'
    shutil.copy(ckpt_path,os.getcwd()) # copy model

    datas = sorted(listdir(data_dir))
    if len(datas) > 27: # if not debugging
        # move idx 27 to head
        datas = [datas[27]] + datas[:27] + datas[28:]


    track = datas[idx]
    track = data_dir.joinpath(track)
    print(track)
    mixture = load_wav(track.joinpath('mixture.wav')) # (c, t)
    target = load_wav(track.joinpath(target_name + '.wav'))
    if model.audio_ch == 1:
        mixture = np.mean(mixture, axis=0, keepdims=True)
        target = np.mean(target, axis=0, keepdims=True)
    #target_hat = {source: separate(config['batch_size'], models[source], onnxs[source], mixture) for source in sources}
    if is_onnx:
        target_hat = separate_with_onnx_TDF(config.batch_size, model, ckpt_path, mixture)
    else:
        target_hat = separate_with_ckpt_TDF(config.batch_size, model, ckpt_path, mixture, config.device, config.double_chunk, overlap_factor=config.overlap_factor)

    bssmetrics, perms, ssdr = get_metrics(target_hat, target, mixture, sr=44100,version=config.bss)
    # dump bssmetrics into pkl
    import pickle
    with open(os.path.join(os.getcwd(),'bssmetrics.pkl'),'wb') as f:
        pickle.dump(bssmetrics,f)

    return bssmetrics







src/evaluation/eval_demo.py
junyuchen-cjy/DTTNet-Pytorch
from os import listdir
from pathlib import Path

import torch
import numpy as np
import onnxruntime as ort
import math
import os
from src.utils.utils import split_nparray_with_overlap, join_chunks

def separate_with_onnx(batch_size, model, onnx_path: Path, mix):
    n_sample = mix.shape[1]

    trim = model.n_fft // 2
    gen_size = model.sampling_size - 2 * trim
    pad = gen_size - n_sample % gen_size
    mix_p = np.concatenate((np.zeros((2, trim)), mix, np.zeros((2, pad)), np.zeros((2, trim))), 1)

    mix_waves = []
    i = 0
    while i < n_sample + pad:
        waves = np.array(mix_p[:, i:i + model.sampling_size], dtype=np.float32)
        mix_waves.append(waves)
        i += gen_size
    mix_waves_batched = torch.tensor(mix_waves, dtype=torch.float32).split(batch_size)

    tar_signals = []

    with torch.no_grad():
        _ort = ort.InferenceSession(str(onnx_path))
        for mix_waves in mix_waves_batched:
            tar_waves = model.istft(torch.tensor(
                _ort.run(None, {'input': model.stft(mix_waves).numpy()})[0]
            ))
            tar_signals.append(tar_waves[:, :, trim:-trim].transpose(0, 1).reshape(2, -1).numpy())
        tar_signal = np.concatenate(tar_signals, axis=-1)[:, :-pad]

    return tar_signal


def separate_with_ckpt(batch_size, model, ckpt_path: Path, mix, device, double_chunk):
    model = model.load_from_checkpoint(ckpt_path).to(device)
    if double_chunk:
        inf_ck = model.inference_chunk_size
    else:
        inf_ck = model.sampling_size
    true_samples = inf_ck - 2 * model.trim

    right_pad = true_samples + model.trim - ((mix.shape[-1]) % true_samples)
    mixture = np.concatenate((np.zeros((2, model.trim), dtype='float32'),
                              mix,
                              np.zeros((2, right_pad), dtype='float32')),
                             1)
    num_chunks = mixture.shape[-1] // true_samples
    mix_waves_batched = [mixture[:, i * true_samples: i * true_samples + inf_ck] for i in
                         range(num_chunks)]
    mix_waves_batched = torch.tensor(mix_waves_batched, dtype=torch.float32).split(batch_size)

    target_wav_hats = []

    with torch.no_grad():
        model.eval()
        for mixture_wav in mix_waves_batched:
            mix_spec = model.stft(mixture_wav.to(device))
            spec_hat = model(mix_spec)
            target_wav_hat = model.istft(spec_hat)
            target_wav_hat = target_wav_hat.cpu().detach().numpy()
            target_wav_hats.append(target_wav_hat)

        target_wav_hat = np.vstack(target_wav_hats)[:, :, model.trim:-model.trim]
        target_wav_hat = np.concatenate(target_wav_hat, axis=-1)[:, :mix.shape[-1]]
    return target_wav_hat




def separate_with_onnx_TDF(batch_size, model, onnx_path: Path, mix):
    n_sample = mix.shape[1]

    overlap = model.n_fft // 2
    gen_size = model.inference_chunk_size - 2 * overlap
    pad = gen_size - n_sample % gen_size
    mix_p = np.concatenate((np.zeros((2, overlap)), mix, np.zeros((2, pad)), np.zeros((2, overlap))), 1)

    mix_waves = []
    i = 0
    while i < n_sample + pad:
        waves = np.array(mix_p[:, i:i + model.inference_chunk_size], dtype=np.float32)
        mix_waves.append(waves)
        i += gen_size
    mix_waves_batched = torch.tensor(mix_waves, dtype=torch.float32).split(batch_size)

    tar_signals = []

    with torch.no_grad():
        _ort = ort.InferenceSession(str(onnx_path), providers=['CUDAExecutionProvider'])
        for mix_waves in mix_waves_batched:
            tar_waves = model.istft(torch.tensor(
                _ort.run(None, {'input': model.stft(mix_waves).numpy()})[0]
            ))
            tar_signals.append(tar_waves[:, :, overlap:-overlap].transpose(0, 1).reshape(2, -1).numpy())
        tar_signal = np.concatenate(tar_signals, axis=-1)[:, :-pad]

    return tar_signal



def separate_with_ckpt_TDF(batch_size, model, ckpt_path: Path, mix, device, double_chunk, overlap_add):
    '''
    Args:
        batch_size: the inference batch size
        model: the model to be used
        ckpt_path: the path to the checkpoint
        mix: (c, t)
        device: the device to be used
        double_chunk: whether to use double chunk size
    Returns:
        target_wav_hat: (c, t)
    '''
    checkpoint = torch.load(ckpt_path)
    model.load_state_dict(checkpoint["state_dict"])
    model = model.to(device)
    # model = model.load_from_checkpoint(ckpt_path).to(device)
    if double_chunk:
        inf_ck = model.inference_chunk_size
    else:
        inf_ck = model.chunk_size

    if overlap_add is None:
        target_wav_hat = no_overlap_inference(model, mix, device, batch_size, inf_ck)
    else:
        if not os.path.exists(overlap_add.tmp_root):
            os.makedirs(overlap_add.tmp_root)
        target_wav_hat = overlap_inference(model, mix, device, batch_size, inf_ck, overlap_add.overlap_rate, overlap_add.tmp_root, overlap_add.samplerate)

    return target_wav_hat

def no_overlap_inference(model, mix, device, batch_size, inf_ck):
    true_samples = inf_ck - 2 * model.overlap

    right_pad = true_samples + model.overlap - ((mix.shape[-1]) % true_samples)
    mixture = np.concatenate((np.zeros((model.audio_ch, model.overlap), dtype='float32'),
                              mix,
                              np.zeros((model.audio_ch, right_pad), dtype='float32')),
                             1)
    num_chunks = mixture.shape[-1] // true_samples
    mix_waves_batched = [mixture[:, i * true_samples: i * true_samples + inf_ck] for i in
                         range(num_chunks)]
    mix_waves_batched = torch.tensor(mix_waves_batched, dtype=torch.float32).split(batch_size)

    target_wav_hats = []

    with torch.no_grad():
        model.eval()
        for mixture_wav in mix_waves_batched:
            mix_spec = model.stft(mixture_wav.to(device))
            spec_hat = model(mix_spec)
            target_wav_hat = model.istft(spec_hat)
            target_wav_hat = target_wav_hat.cpu().detach().numpy()
            target_wav_hats.append(target_wav_hat) # (b, c, t)

        target_wav_hat = np.vstack(target_wav_hats)[:, :, model.overlap:-model.overlap] # (sum(b), c, t)
        target_wav_hat = np.concatenate(target_wav_hat, axis=-1)[:, :mix.shape[-1]]
    return target_wav_hat


def overlap_inference(model, mix, device, batch_size, inf_ck, overlap_rate, tmp_root, samplerate):
    '''
    Args:
        mix: (c, t)
    '''
    hop_length = math.ceil((1 - overlap_rate) * inf_ck)
    overlap_size = inf_ck - hop_length
    step_t = mix.shape[1]
    mix_waves_batched = split_nparray_with_overlap(mix.T, hop_length, overlap_size)

    mix_waves_batched = torch.tensor(mix_waves_batched, dtype=torch.float32).split(batch_size) # [(b, c, t)]

    target_wav_hats = []

    with torch.no_grad():
        model.eval()
        for mixture_wav in mix_waves_batched:
            mix_spec = model.stft(mixture_wav.to(device))
            spec_hat = model(mix_spec)
            target_wav_hat = model.istft(spec_hat)
            target_wav_hat = target_wav_hat.cpu().detach().numpy()
            target_wav_hats.append(target_wav_hat) # (b, c, t)

        target_wav_hat = np.vstack(target_wav_hats) # (sum(b), c, t)
        target_wav_hat = np.transpose(target_wav_hat, (0, 2, 1)) # (sum(b), t, c)
        target_wav_hat = join_chunks(tmp_root, target_wav_hat, samplerate, overlap_size) # (t, c)
    return target_wav_hat[:step_t].T # (c, t)

src/evaluation/separate.py
junyuchen-cjy/DTTNet-Pytorch

from .batch_norm import *
src/layers/__init__.py
junyuchen-cjy/DTTNet-Pytorch
import logging

import torch
import torch.nn.functional as F
from torch import nn

__all__ = ["IBN", "get_norm"]


class BatchNorm(nn.BatchNorm2d):
    def __init__(self, num_features, eps=1e-05, momentum=0.1, weight_freeze=False, bias_freeze=False, weight_init=1.0,
                 bias_init=0.0, **kwargs):
        super().__init__(num_features, eps=eps, momentum=momentum)
        if weight_init is not None: nn.init.constant_(self.weight, weight_init)
        if bias_init is not None: nn.init.constant_(self.bias, bias_init)
        self.weight.requires_grad_(not weight_freeze)
        self.bias.requires_grad_(not bias_freeze)


class SyncBatchNorm(nn.SyncBatchNorm):
    def __init__(self, num_features, eps=1e-05, momentum=0.1, weight_freeze=False, bias_freeze=False, weight_init=1.0,
                 bias_init=0.0):
        super().__init__(num_features, eps=eps, momentum=momentum)
        if weight_init is not None: nn.init.constant_(self.weight, weight_init)
        if bias_init is not None: nn.init.constant_(self.bias, bias_init)
        self.weight.requires_grad_(not weight_freeze)
        self.bias.requires_grad_(not bias_freeze)


class IBN(nn.Module):
    def __init__(self, planes, bn_norm, **kwargs):
        super(IBN, self).__init__()
        half1 = int(planes / 2)
        self.half = half1
        half2 = planes - half1
        self.IN = nn.InstanceNorm2d(half1, affine=True)
        self.BN = get_norm(bn_norm, half2, **kwargs)

    def forward(self, x):
        split = torch.split(x, self.half, 1)
        out1 = self.IN(split[0].contiguous())
        out2 = self.BN(split[1].contiguous())
        out = torch.cat((out1, out2), 1)
        return out


class GhostBatchNorm(BatchNorm):
    def __init__(self, num_features, num_splits=1, **kwargs):
        super().__init__(num_features, **kwargs)
        self.num_splits = num_splits
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))

    def forward(self, input):
        N, C, H, W = input.shape
        if self.training or not self.track_running_stats:
            self.running_mean = self.running_mean.repeat(self.num_splits)
            self.running_var = self.running_var.repeat(self.num_splits)
            outputs = F.batch_norm(
                input.view(-1, C * self.num_splits, H, W), self.running_mean, self.running_var,
                self.weight.repeat(self.num_splits), self.bias.repeat(self.num_splits),
                True, self.momentum, self.eps).view(N, C, H, W)
            self.running_mean = torch.mean(self.running_mean.view(self.num_splits, self.num_features), dim=0)
            self.running_var = torch.mean(self.running_var.view(self.num_splits, self.num_features), dim=0)
            return outputs
        else:
            return F.batch_norm(
                input, self.running_mean, self.running_var,
                self.weight, self.bias, False, self.momentum, self.eps)


class FrozenBatchNorm(nn.Module):
    """
    BatchNorm2d where the batch statistics and the affine parameters are fixed.
    It contains non-trainable buffers called
    "weight" and "bias", "running_mean", "running_var",
    initialized to perform identity transformation.
    The pre-trained backbone models from Caffe2 only contain "weight" and "bias",
    which are computed from the original four parameters of BN.
    The affine transform `x * weight + bias` will perform the equivalent
    computation of `(x - running_mean) / sqrt(running_var) * weight + bias`.
    When loading a backbone model from Caffe2, "running_mean" and "running_var"
    will be left unchanged as identity transformation.
    Other pre-trained backbone models may contain all 4 parameters.
    The forward is implemented by `F.batch_norm(..., training=False)`.
    """

    _version = 3

    def __init__(self, num_features, eps=1e-5, **kwargs):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.register_buffer("weight", torch.ones(num_features))
        self.register_buffer("bias", torch.zeros(num_features))
        self.register_buffer("running_mean", torch.zeros(num_features))
        self.register_buffer("running_var", torch.ones(num_features) - eps)

    def forward(self, x):
        if x.requires_grad:
            # When gradients are needed, F.batch_norm will use extra memory
            # because its backward op computes gradients for weight/bias as well.
            scale = self.weight * (self.running_var + self.eps).rsqrt()
            bias = self.bias - self.running_mean * scale
            scale = scale.reshape(1, -1, 1, 1)
            bias = bias.reshape(1, -1, 1, 1)
            return x * scale + bias
        else:
            # When gradients are not needed, F.batch_norm is a single fused op
            # and provide more optimization opportunities.
            return F.batch_norm(
                x,
                self.running_mean,
                self.running_var,
                self.weight,
                self.bias,
                training=False,
                eps=self.eps,
            )

    def _load_from_state_dict(
            self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
        version = local_metadata.get("version", None)

        if version is None or version < 2:
            # No running_mean/var in early versions
            # This will silent the warnings
            if prefix + "running_mean" not in state_dict:
                state_dict[prefix + "running_mean"] = torch.zeros_like(self.running_mean)
            if prefix + "running_var" not in state_dict:
                state_dict[prefix + "running_var"] = torch.ones_like(self.running_var)

        if version is not None and version < 3:
            logger = logging.getLogger(__name__)
            logger.info("FrozenBatchNorm {} is upgraded to version 3.".format(prefix.rstrip(".")))
            # In version < 3, running_var are used without +eps.
            state_dict[prefix + "running_var"] -= self.eps

        super()._load_from_state_dict(
            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
        )

    def __repr__(self):
        return "FrozenBatchNorm2d(num_features={}, eps={})".format(self.num_features, self.eps)

    @classmethod
    def convert_frozen_batchnorm(cls, module):
        """
        Convert BatchNorm/SyncBatchNorm in module into FrozenBatchNorm.
        Args:
            module (torch.nn.Module):
        Returns:
            If module is BatchNorm/SyncBatchNorm, returns a new module.
            Otherwise, in-place convert module and return it.
        Similar to convert_sync_batchnorm in
        https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/batchnorm.py
        """
        bn_module = nn.modules.batchnorm
        bn_module = (bn_module.BatchNorm2d, bn_module.SyncBatchNorm)
        res = module
        if isinstance(module, bn_module):
            res = cls(module.num_features)
            if module.affine:
                res.weight.data = module.weight.data.clone().detach()
                res.bias.data = module.bias.data.clone().detach()
            res.running_mean.data = module.running_mean.data
            res.running_var.data = module.running_var.data
            res.eps = module.eps
        else:
            for name, child in module.named_children():
                new_child = cls.convert_frozen_batchnorm(child)
                if new_child is not child:
                    res.add_module(name, new_child)
        return res


def get_norm(norm, out_channels, **kwargs):
    """
    Args:
        norm (str or callable): either one of BN, GhostBN, FrozenBN, GN or SyncBN;
            or a callable that takes a channel number and returns
            the normalization layer as a nn.Module
        out_channels: number of channels for normalization layer

    Returns:
        nn.Module or None: the normalization layer
    """
    # return nn.BatchNorm2d(out_channels)

    if isinstance(norm, str):
        if len(norm) == 0:
            return None
        norm = {
            "BN": BatchNorm,
            "syncBN": SyncBatchNorm,
            "GhostBN": GhostBatchNorm,
            "FrozenBN": FrozenBatchNorm,
            "GN": lambda channels, **args: nn.GroupNorm(32, channels),
        }[norm]
    return norm(out_channels, **kwargs)

src/layers/batch_norm.py
junyuchen-cjy/DTTNet-Pytorch
#%%
import numpy as np
import torch



def wave_to_batches(mix, inf_ck, overlap, batch_size):
    '''
    Args:
        mix: (2, N) numpy array
        inf_ck: int, the chunk size as the model input (contains 2*overlap)
                inf_ck = overlap + true_samples + overlap
        overlap: int, the discarded samples at each side
    Returns:
        a tuples of batches, each batch is a (batch, 2, inf_ck) torch tensor
    '''
    true_samples = inf_ck - 2 * overlap
    channels = mix.shape[0]

    right_pad = true_samples + overlap - ((mix.shape[-1]) % true_samples)
    mixture = np.concatenate((np.zeros((channels, overlap), dtype='float32'),
                              mix,
                              np.zeros((channels, right_pad), dtype='float32')),
                             1)

    num_chunks = mixture.shape[-1] // true_samples
    mix_waves_batched = np.array([mixture[:, i * true_samples: i * true_samples + inf_ck] for i in
                         range(num_chunks)]) # (x,2,inf_ck)
    return torch.tensor(mix_waves_batched, dtype=torch.float32).split(batch_size)

def batches_to_wave(target_hat_chunks, overlap, org_len):
    '''
    Args:
        target_hat_chunks: a list of (batch, 2, inf_ck) torch tensors
        overlap: int, the discarded samples at each side
        org_len: int, the original length of the mixture
    Returns:
        (2, N) numpy array
    '''
    target_hat_chunks = [c[..., overlap:-overlap] for c in target_hat_chunks]
    target_hat_chunks = torch.cat(target_hat_chunks)

    # concat all output chunks
    return target_hat_chunks.transpose(0, 1).reshape(2, -1)[..., :org_len].detach().cpu().numpy()

if __name__ == '__main__':
    mix = np.random.rand(2, 14318640)
    inf_ck = 261120
    overlap = 3072
    batch_size = 8
    out = wave_to_batches(mix, inf_ck, overlap, batch_size)
    in_wav = batches_to_wave(out, overlap, mix.shape[-1])
    print(in_wav.shape)
src/layers/chunk_size.py
junyuchen-cjy/DTTNet-Pytorch
from typing import List, Optional

import hydra
import pytorch_lightning as pl
import pyrootutils
import torch
import os
import shutil
from omegaconf import DictConfig
from pytorch_lightning import (
    Callback,
    LightningDataModule,
    LightningModule,
    Trainer,
    seed_everything,
)
from pytorch_lightning.loggers import WandbLogger
from hydra.core.hydra_config import HydraConfig

from src import utils

pyrootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)

log = utils.get_pylogger(__name__)


@utils.task_wrapper
def train(cfg: DictConfig) -> Optional[float]:
    """Contains training pipeline.
    Instantiates all PyTorch Lightning objects from config.

    Args:
        cfg (DictConfig): Configuration composed by Hydra.

    Returns:
        Optional[float]: Metric score for hyperparameter optimization.
    """

    # Set seed for random number generators in pytorch, numpy and python.random
    try:
        if "seed" in cfg:
            # set seed for random number generators in pytorch, numpy and python.random
            if cfg.get("seed"):
                pl.seed_everything(cfg.seed, workers=True)

        else:
            raise ModuleNotFoundError

    except ModuleNotFoundError:
        print('[Error] seed should be fixed for reproducibility \n=> e.g. python run.py +seed=$SEED')
        exit(-1)

    # Init Lightning datamodule
    log.info(f"Instantiating datamodule <{cfg.datamodule._target_}>")
    datamodule: LightningDataModule = hydra.utils.instantiate(cfg.datamodule)

    # Init Lightning model
    log.info(f"Instantiating model <{cfg.model._target_}>")
    model: LightningModule = hydra.utils.instantiate(cfg.model)

    # Init Lightning callbacks
    callbacks: List[Callback] = []
    if "callbacks" in cfg:
        for _, cb_conf in cfg["callbacks"].items():
            if "_target_" in cb_conf:
                log.info(f"Instantiating callback <{cb_conf._target_}>")
                callbacks.append(hydra.utils.instantiate(cb_conf))

    # Init Lightning loggers
    if "resume_from_checkpoint" in cfg.trainer:
        ckpt_path = cfg.trainer.resume_from_checkpoint
        # get the parent directory of the checkpoint path
        log_dir = os.path.dirname(os.path.dirname(ckpt_path))
        tensorboard_dir = os.path.join(log_dir, "tensorboard")
        if os.path.exists(tensorboard_dir):
            # copy tensorboard dir to the parent directory of the checkpoint path
            # HydraConfig.get().run.dir returns new dir so do not use it! (now fixed)
            shutil.copytree(tensorboard_dir,os.path.join(os.getcwd(),"tensorboard"))

        wandb_dir = os.path.join(log_dir, "wandb")
        if os.path.exists(wandb_dir):
            shutil.copytree(wandb_dir,os.path.join(os.getcwd(),"wandb"))


    logger: List = []
    if "logger" in cfg:
        for _, lg_conf in cfg["logger"].items():
            if "_target_" in lg_conf:
                log.info(f"Instantiating logger <{lg_conf._target_}>")
                logger.append(hydra.utils.instantiate(lg_conf))

        for wandb_logger in [l for l in logger if isinstance(l, WandbLogger)]:
            utils.wandb_login(key=cfg.wandb_api_key)
            # utils.wandb_watch_all(wandb_logger, model) # TODO buggy
            break

    # Init Lightning trainer
    log.info(f"Instantiating trainer <{cfg.trainer._target_}>")
    # get env variable use_gloo
    use_gloo = os.environ.get("USE_GLOO", False)
    if use_gloo:
        from pytorch_lightning.strategies import DDPStrategy
        ddp = DDPStrategy(process_group_backend='gloo')
        trainer: Trainer = hydra.utils.instantiate(
            cfg.trainer, strategy=ddp, callbacks=callbacks, logger=logger, _convert_="partial"
        )
    else:
        trainer: Trainer = hydra.utils.instantiate(
            cfg.trainer, callbacks=callbacks, logger=logger, _convert_="partial"
        )

    # Send some parameters from config to all lightning loggers
    log.info("Logging hyperparameters!")
    utils.log_hyperparameters(
        dict(
            cfg=cfg,
            model=model,
            datamodule=datamodule,
            trainer=trainer,
            callbacks=callbacks,
            logger=logger,
        )
    )

    # Train the model
    log.info("Starting training!")
    trainer.fit(model=model, datamodule=datamodule)

    # Evaluate model on test set after training
    # if not cfg.trainer.get("fast_dev_run"):
    #     log.info("Starting testing!")
    #     trainer.test()

    # Make sure everything closed properly
    log.info("Finalizing!")
    # utils.finish(
    #     config=cfg,
    #     model=model,
    #     datamodule=datamodule,
    #     trainer=trainer,
    #     callbacks=callbacks,
    #     logger=logger,
    # )

    # Print path to best checkpoint
    # log.info(f"Best checkpoint path:\n{trainer.checkpoint_callback.best_model_path}")

    # Return metric score for hyperparameter optimization
    # optimized_metric = cfg.get("optimized_metric")
    # if optimized_metric:
    #     return trainer.callback_metrics[optimized_metric]
    return None, None

src/train.py
junyuchen-cjy/DTTNet-Pytorch
from src.utils.pylogger import get_pylogger
from src.utils.rich_utils import enforce_tags, print_config_tree
from src.utils.utils import *

src/utils/__init__.py
junyuchen-cjy/DTTNet-Pytorch
import os
import subprocess as sp
import tempfile
import warnings
from argparse import ArgumentParser
from concurrent import futures

import numpy as np
import soundfile as sf
import torch
from tqdm import tqdm

warnings.simplefilter(action='ignore', category=Warning)
source_names = ['vocals', 'drums', 'bass', 'other']
sample_rate = 44100

def main (args):
    data_root = args.data_dir
    train = args.train
    test = args.test
    valid = args.valid

    musdb_train_path = data_root + 'train/'
    musdb_test_path = data_root + 'test/'
    musdb_valid_path = data_root + 'valid/'
    print(f"train={train}, test={test}, valid={valid}")

    mix_name = 'mixture'

    P = [-3, -2, -1, 0, 1, 2, 3]   # pitch shift amounts (in semitones)
    T = [-30, -20, -10, 0, 10, 20, 30]   # time stretch amounts (10 means 10% slower)

    pool = futures.ProcessPoolExecutor
    pool_workers = 13
    pendings = []
    with pool(pool_workers) as pool:
        for p in P:
            for t in T:
                if not (p==0 and t==0):
                    if train:
                        pendings.append(pool.submit(save_shifted_dataset, p, t, musdb_train_path))
                        # save_shifted_dataset(p, t, musdb_train_path)
                    if valid:
                        save_shifted_dataset(p, t, musdb_valid_path)
                    if test:
                        save_shifted_dataset(p, t, musdb_test_path)
    for pending in pendings:
        pending.result()


def shift(wav, pitch, tempo, voice=False, quick=False, samplerate=44100):
    def i16_pcm(wav):
        if wav.dtype == np.int16:
            return wav
        return (wav * 2 ** 15).clamp_(-2 ** 15, 2 ** 15 - 1).short()

    def f32_pcm(wav):
        if wav.dtype == np.float:
            return wav
        return wav.float() / 2 ** 15

    """
    tempo is a relative delta in percentage, so tempo=10 means tempo at 110%!
    pitch is in semi tones.
    Requires `soundstretch` to be installed, see
    https://www.surina.net/soundtouch/soundstretch.html
    """

    inputfile = tempfile.NamedTemporaryFile(dir="/root/autodl-tmp/tmp", suffix=".wav")
    outfile = tempfile.NamedTemporaryFile(dir="/root/autodl-tmp/tmp", suffix=".wav")

    sf.write(inputfile.name, data=i16_pcm(wav).t().numpy(), samplerate=samplerate, format='WAV')
    command = [
        "soundstretch",
        inputfile.name,
        outfile.name,
        f"-pitch={pitch}",
        f"-tempo={tempo:.6f}",
    ]
    if quick:
        command += ["-quick"]
    if voice:
        command += ["-speech"]
    try:
        sp.run(command, capture_output=True, check=True)
    except sp.CalledProcessError as error:
        raise RuntimeError(f"Could not change bpm because {error.stderr.decode('utf-8')}")
    wav, sr = sf.read(outfile.name, dtype='float32')
    # wav = np.float32(wav)
    # wav = f32_pcm(torch.from_numpy(wav).t())
    assert sr == samplerate
    return wav


def save_shifted_dataset(delta_pitch, delta_tempo, data_path):
    out_path = data_path[:-1] + f'_p={delta_pitch}_t={delta_tempo}/'
    try:
        os.mkdir(out_path)
    except FileExistsError:
        pass
    track_names = list(filter(lambda x: os.path.isdir(f'{data_path}/{x}'), sorted(os.listdir(data_path))))
    for track_name in tqdm(track_names):
        try:
            os.mkdir(f'{out_path}/{track_name}')
        except FileExistsError:
            pass
        for s_name in source_names:
            source = load_wav(f'{data_path}/{track_name}/{s_name}.wav')
            shifted = shift(
                torch.tensor(source),
                delta_pitch,
                delta_tempo,
                voice=s_name == 'vocals')
            sf.write(f'{out_path}/{track_name}/{s_name}.wav', shifted, samplerate=sample_rate, format='WAV')


def load_wav(path, sr=None):
    return sf.read(path, samplerate=sr, dtype='float32')[0].T


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--data_dir', type=str)
    parser.add_argument('--train', type=bool, default=True)
    parser.add_argument('--valid', type=bool, default=False)
    parser.add_argument('--test', type=bool, default=False)

    main(parser.parse_args())
src/utils/data_augmentation.py
junyuchen-cjy/DTTNet-Pytorch
import os
from src.utils import get_unique_save_path

cur_suffix = -1

def get_train_log_dir(a, b, *, _parent_):
    global cur_suffix
    logbase = os.environ.get("LOG_DIR", "dtt_logs")
    if not os.path.exists(logbase):
        os.mkdir(logbase)
    dir_path = os.path.join(logbase, f"{a}_{b}")
    if cur_suffix == -1:
        cur_suffix = get_unique_save_path(dir_path)
        return f"{dir_path}_{str(cur_suffix)}"
    else:
        return f"{dir_path}_{str(cur_suffix)}"

cur_suffix1 = -1

def get_eval_log_dir(ckpt_path, *, _parent_):
    # get environment variable logbase
    global cur_suffix1
    logbase = os.environ.get("LOG_DIR", "dtt_logs")
    if not os.path.exists(logbase):
        os.mkdir(logbase)
    ckpt_name = os.path.basename(ckpt_path)
    # remove the suffix
    ckpt_name = ckpt_name.split(".")[0]
    dir_path = os.path.join(logbase, f"eval_{ckpt_name}")
    if cur_suffix1 == -1:
        cur_suffix1 = get_unique_save_path(dir_path)
        return f"{dir_path}_{str(cur_suffix1)}"
    else:
        return f"{dir_path}_{str(cur_suffix1)}"

cur_suffix2 = -1

def get_sweep_log_dir(a, b, *, _parent_):
    global cur_suffix2
    logbase = os.environ.get("LOG_DIR", "dtt_logs")
    if not os.path.exists(logbase):
        os.mkdir(logbase)
    dir_path = os.path.join(logbase, f"m_{a}_{b}")
    if cur_suffix2 == -1:
        cur_suffix2 = get_unique_save_path(dir_path)
        return f"{dir_path}_{str(cur_suffix2)}"
    else:
        return f"{dir_path}_{str(cur_suffix2)}"
src/utils/omega_resolvers.py
junyuchen-cjy/DTTNet-Pytorch
import os
import re
import traceback
import pandas as pd
import wandb
import numpy as np
from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

class AbstractLoggerQuery:
    def get_value_steps(self, tag):
        raise NotImplementedError

    # def steps_to_epoch(self, steps):
    #     raise NotImplementedError

    def get_best(self, tag, method="min"):
        v, s = self.get_value_steps(tag)
        indexes = [i for i in range(len(s))]
        # sort the indexes by v
        if method == "min":
            indexes.sort(key=lambda x: v[x])
        else:
            indexes.sort(key=lambda x: -v[x])
        # get the top 5 indexes
        top5_idx = indexes[:3]
        # the the values
        top5_values = [v[i] for i in top5_idx]
        top5_steps = [s[i] for i in top5_idx]
        return top5_idx, top5_values, top5_steps

    def get_best_epochs(self, tag, method="min"):
        v, e = self.get_value_epochs(tag)
        indexes = [i for i in range(len(e))]
        # sort the indexes by v
        if method == "min":
            indexes.sort(key=lambda x: v[x])
        else:
            indexes.sort(key=lambda x: -v[x])
        # get the top 5 indexes
        top5_idx = indexes[:3]
        # the the values
        top5_values = [v[i] for i in top5_idx]
        top5_epochs = [e[i] for i in top5_idx]
        return top5_idx, top5_values, top5_epochs

    def get_last_ep(self, tag):
        v, s = self.get_value_steps(tag)
        return len(s) - 1

    def _bsearch(self, arr, target_step):
        l = 0
        r = len(arr) - 1
        while l < r:
            mid = (l + r + 1) // 2
            if arr[mid] <= target_step:
                l = mid
            else:
                r = mid - 1
        return l

    def steps_to_epoch(self, steps):
        ep_ep, ep_step = self.get_value_steps('epoch')
        mapped = []
        for vs in steps:
            idx = self._bsearch(ep_step, vs)
            mapped.append(ep_ep[idx])
        return mapped

    def get_value_epochs(self, tag):
        raise NotImplementedError


class TFLoggerQuery(AbstractLoggerQuery):
    def __init__(self, log_dir):
        self.event_acc = EventAccumulator(log_dir)
        self.event_acc.Reload()

    def get_value_steps(self, tag):
        event_list = self.event_acc.Scalars(tag)
        values = list(map(lambda x: x.value, event_list))
        step = list(map(lambda x: x.step, event_list))
        return values, step

    # def steps_to_epoch(self, steps):
    #     v, s = self.get_value_steps('epoch')
    #     # zip the two list
    #     return [int(v[s.index(step)]) for step in steps]

    def get_tags(self):
        return self.event_acc.Tags()["scalars"]


class WandbLoggerQuery(AbstractLoggerQuery):
    def __init__(self, api_key, path):
        self.api_key = api_key
        self.path = path

    def _get_run(self):
        return wandb.Api(api_key=self.api_key).run(self.path)

    def get_exp_name(self):
        run = self._get_run()
        return run.name

    def get_value_steps(self, tag):
        run = self._get_run()
        # run.history returns sampled data, do not use it
        # history = run.history(keys=[tag])
        history = run.scan_history(keys=[tag, "_step"])
        losses = [row for row in history]
        df = pd.DataFrame(losses)
        df = df.sort_values(by="_step")
        values = df[tag].values
        step = df["_step"].values
        return values, step

    def get_value_epochs(self, tag, sort_by="epoch"):
        run = self._get_run()
        history = run.scan_history(keys=[tag, "epoch"])
        losses = [row for row in history]
        df = pd.DataFrame(losses)
        df = df.sort_values(by=sort_by)
        values = df[tag].values
        epoch = df["epoch"].values
        return values, epoch



    def steps_to_runtime(self, steps):
        rt_ep, rt_step = self.get_value_steps('_runtime')
        mapped = []
        for vs in steps:
            idx = self._bsearch(rt_step, vs)
            secs = rt_ep[idx]
            # seconds to days hours minutes seconds
            m, s = divmod(secs, 60)
            h, m = divmod(m, 60)
            d, h = divmod(h, 24)
            mapped.append(f"{d:.0f}d {h:.0f}h {m:.0f}m {s:.0f}s")
        return mapped

    def epochs_to_runtime(self, epochs):
        rt_ep, ep_lst = self.get_value_epochs('_runtime')
        mapped = []
        for vs in epochs:
            idx = self._bsearch(ep_lst, vs)
            secs = rt_ep[idx]
            # seconds to days hours minutes seconds
            m, s = divmod(secs, 60)
            h, m = divmod(m, 60)
            d, h = divmod(h, 24)
            mapped.append(f"{d:.0f}d {h:.0f}h {m:.0f}m {s:.0f}s")
        return mapped

    def _secs_to_string(self, secs):
        # seconds to days hours minutes seconds
        m, s = divmod(secs, 60)
        h, m = divmod(m, 60)
        d, h = divmod(h, 24)
        return f"{d:.0f}d {h:.0f}h {m:.0f}m {s:.0f}s"

    def avg_runtime_per_ep(self):
        epochs = [i for i in range(1, 100)]
        rt_ep, ep_lst = self.get_value_epochs('_runtime', sort_by="_runtime") # 坑：这里必须要按照runtime排序，因为epochs有很多相同的值
        mapped = []
        for vs in epochs:
            idx = self._bsearch(ep_lst, vs)
            secs = rt_ep[idx]
            mapped.append(secs)
        new = []
        for i in range(1, len(mapped)):
            new.append(mapped[i] - mapped[i-1])
        meansecs = np.mean(new)
        print(new)
        return self._secs_to_string(meansecs)


if __name__ == "__main__":
    tag = "val/usdr"
    api_key = "xxxxxxxxxxxx"
    path = "username/dtt_vocals/xxxxxxx"

    w_query = WandbLoggerQuery(api_key=api_key, path=path)
    best_sdr_idx, best_sdr_values, best_sdr_epochs = w_query.get_best_epochs(tag, method="max")
    print("idx", best_sdr_idx)
    print("values", best_sdr_values)
    print("epochs", best_sdr_epochs)
    print("runtime", w_query.epochs_to_runtime(best_sdr_epochs))
src/utils/pick_best.py
junyuchen-cjy/DTTNet-Pytorch
import warnings
from copy import deepcopy
from importlib.util import find_spec
from typing import Callable, List

import re
import hydra
import museval
import numpy as np
import wandb
from omegaconf import DictConfig, OmegaConf
from pytorch_lightning import Callback
from pytorch_lightning.loggers import Logger
import pytorch_lightning as pl
import tempfile
import soundfile as sf
import shutil
import subprocess


from pytorch_lightning.utilities import rank_zero_only

from src.utils import pylogger, rich_utils

log = pylogger.get_pylogger(__name__)


def task_wrapper(task_func: Callable) -> Callable:
    """Optional decorator that wraps the task function in extra utilities.

    Makes multirun more resistant to failure.

    Utilities:
    - Calling the `utils.extras()` before the task is started
    - Calling the `utils.close_loggers()` after the task is finished or failed
    - Logging the exception if occurs
    - Logging the output dir
    """

    def wrap(cfg: DictConfig):

        # execute the task
        try:

            # apply extra utilities
            extras(cfg)

            metric_dict, object_dict = task_func(cfg=cfg)

        # things to do if exception occurs
        except Exception as ex:

            # save exception to `.log` file
            log.exception("")

            # when using hydra plugins like Optuna, you might want to disable raising exception
            # to avoid multirun failure
            raise ex

        # things to always do after either success or exception
        finally:

            # display output dir path in terminal
            log.info(f"Output dir: {cfg.paths.output_dir}")

            # close loggers (even if exception occurs so multirun won't fail)
            close_loggers()

        return metric_dict, object_dict

    return wrap


def extras(cfg: DictConfig) -> None:
    """Applies optional utilities before the task is started.

    Utilities:
    - Ignoring python warnings
    - Setting tags from command line
    - Rich config printing
    """

    # return if no `extras` config
    if not cfg.get("extras"):
        log.warning("Extras config not found! <cfg.extras=null>")
        return

    # disable python warnings
    if cfg.extras.get("ignore_warnings"):
        log.info("Disabling python warnings! <cfg.extras.ignore_warnings=True>")
        warnings.filterwarnings("ignore")

    # prompt user to input tags from command line if none are provided in the config
    if cfg.extras.get("enforce_tags"):
        log.info("Enforcing tags! <cfg.extras.enforce_tags=True>")
        rich_utils.enforce_tags(cfg, save_to_file=True)

    # pretty print config tree using Rich library
    if cfg.extras.get("print_config"):
        log.info("Printing config tree with Rich! <cfg.extras.print_config=True>")
        rich_utils.print_config_tree(cfg, resolve=True, save_to_file=True)


def instantiate_callbacks(callbacks_cfg: DictConfig) -> List[Callback]:
    """Instantiates callbacks from config."""
    callbacks: List[Callback] = []

    if not callbacks_cfg:
        log.warning("No callback configs found! Skipping..")
        return callbacks

    if not isinstance(callbacks_cfg, DictConfig):
        raise TypeError("Callbacks config must be a DictConfig!")

    for _, cb_conf in callbacks_cfg.items():
        if isinstance(cb_conf, DictConfig) and "_target_" in cb_conf:
            log.info(f"Instantiating callback <{cb_conf._target_}>")
            callbacks.append(hydra.utils.instantiate(cb_conf))

    return callbacks


def instantiate_loggers(logger_cfg: DictConfig) -> List[Logger]:
    """Instantiates loggers from config."""
    Logger: List[pl.loggers.Logger] = []

    if not logger_cfg:
        log.warning("No Logger configs found! Skipping...")
        return Logger

    if not isinstance(logger_cfg, DictConfig):
        raise TypeError("Logger config must be a DictConfig!")

    for _, lg_conf in logger_cfg.items():
        if isinstance(lg_conf, DictConfig) and "_target_" in lg_conf:
            log.info(f"Instantiating Logger <{lg_conf._target_}>")
            Logger.append(hydra.utils.instantiate(lg_conf))

    return Logger


@rank_zero_only
def log_hyperparameters(object_dict: dict) -> None:
    """Controls which config parts are saved by lightning loggers.

    Additionally saves:
    - Number of model parameters
    """

    hparams = {}

    cfg = object_dict["cfg"]
    model = object_dict["model"]
    trainer = object_dict["trainer"]

    if not trainer.logger:
        log.warning("trainer.logger not found! Skipping hyperparameter logging...")
        return

    hparams["model"] = cfg["model"]

    # save number of model parameters
    hparams["model/params/total"] = sum(p.numel() for p in model.parameters())
    hparams["model/params/trainable"] = sum(
        p.numel() for p in model.parameters() if p.requires_grad
    )
    hparams["model/params/non_trainable"] = sum(
        p.numel() for p in model.parameters() if not p.requires_grad
    )
    #
    _cfg = deepcopy(cfg)
    OmegaConf.resolve(_cfg)
    hparams["datamodule"] = _cfg["datamodule"]
    hparams["trainer"] = _cfg["trainer"]

    hparams["callbacks"] = _cfg.get("callbacks")
    # hparams["extras"] = cfg.get("extras")

    hparams["task_name"] = _cfg.get("task_name")
    hparams["tags"] = _cfg.get("tags")
    hparams["ckpt_path"] = _cfg.get("ckpt_path")
    hparams["seed"] = _cfg.get("seed")

    # send hparams to all loggers
    for Logger in trainer.loggers:
        Logger.log_hyperparams(hparams)


def get_metric_value(metric_dict: dict, metric_name: str) -> float:
    """Safely retrieves value of the metric logged in LightningModule."""

    if not metric_name:
        log.info("Metric name is None! Skipping metric value retrieval...")
        return None

    if metric_name not in metric_dict:
        raise Exception(
            f"Metric value not found! <metric_name={metric_name}>\n"
            "Make sure metric name logged in LightningModule is correct!\n"
            "Make sure `optimized_metric` name in `hparams_search` config is correct!"
        )

    metric_value = metric_dict[metric_name].item()
    log.info(f"Retrieved metric value! <{metric_name}={metric_value}>")

    return metric_value


def close_loggers() -> None:
    """Makes sure all loggers closed properly (prevents logging failure during multirun)."""

    log.info("Closing loggers...")

    if find_spec("wandb"):  # if wandb is installed
        import wandb

        if wandb.run:
            log.info("Closing wandb!")
            wandb.finish()


@rank_zero_only
def save_file(path: str, content: str) -> None:
    """Save file in rank zero mode (only on one process in multi-GPU setup)."""
    with open(path, "w+") as file:
        file.write(content)


@rank_zero_only
def wandb_login(key):
    wandb.login(key=key)

from functools import reduce
import pandas as pd

def print_single_result(x, target_list, metrics):
    df = pd.DataFrame(x)
    df.columns = target_list
    df.index = metrics
    print(df)


def concat_results(results, reduce_per_song):
    '''
    return: metrics_num x len(target_list) x no.
    '''
    if reduce_per_song:
        medians = []
        for i in range(len(results)):
            medians.append(np.nanmedian(results[i], axis=2, keepdims=True))

        total = reduce(lambda x, y: np.concatenate((x, y), axis=2), medians)
    else:
        total = reduce(lambda x, y: np.concatenate((x, y), axis=2), results)
    return total

import pickle
import os

def save_results(model_path,results,target_list, perms, ssdrs):

    # if the path is dir
    if os.path.isdir(model_path):
        base_name = os.path.join(model_path, "bss_metrics")
    else: # remove suffix
        # base_name = os.path.splitext(model_path)[0]
        base_name = model_path

    # make path unique by assigning id
    id = 0
    save_path = f"{base_name}_{str(id)}.pkl"
    while os.path.isfile(save_path):
        save_path = f"{base_name}_{str(id)}.pkl"
        id = id + 1

    # save into file
    with open(save_path, "wb") as f:
        d = {'results': results,
             'target_list':target_list,
             'metrics':["SDR", "ISR", "SIR", "SAR"],
             'perms':perms,
             'ssdrs':ssdrs}
        pickle.dump(d, f)
        # np.savez(**d)

def get_metrics(target_hat, target, mixture, sr, version="naive"):
    ssdr = sdr(target_hat, target)  # c, t

    target_hat = target_hat.T # (t,c)
    target = target.T # (t,c)
    mixture = mixture.T # (t,c)
    references = np.stack([target, mixture - target]) # (nsrc, nsampl, nchan)

    ests = np.stack([target_hat, mixture - target_hat]) # (nsrc, nsampl, nchan)

    # evaluate
    if version == "official":
        SDR, ISR, SIR, SAR, perms = museval.metrics.bss_eval(references, ests, window=sr, hop=sr)
        bss_ret = np.stack((SDR, ISR, SIR, SAR), axis=0)
    elif version == "fast":
        SDR = simplified_msseval(references, ests, chunk_size=sr)
        bss_ret = np.expand_dims(SDR, axis=0)
        perms = []
    else:
        raise ValueError("version must be official or fast")

    return bss_ret, perms, ssdr

def get_median_csdr(results):
    total = concat_results(results, False)
    return np.nanmedian(total, axis=2)[0,0]

def sdr(est, ref):
    '''
    Args:
        est: (c, t)
        ref: (c, t)
    '''
    ratio = np.sum(ref ** 2) / np.sum((ref - est) ** 2)
    if ratio == 0:
        return np.nan
    return 10 * np.log10(ratio)

def simplified_msseval(references, ests, chunk_size):
    '''
    Args:
        ests: (n, t, c)
        references: (n, t, c)
    Returns:
        results: (src, chunks)
    '''
    results = []
    for i in range(ests.shape[0]):
        e = ests[i] # (t, c)
        r = references[i] # (t, c)
        # pad to multiple of chunk_size
        # e = np.pad(e, ((0, chunk_size - e.shape[0] % chunk_size), (0, 0)))
        # r = np.pad(r, ((0, chunk_size - r.shape[0] % chunk_size), (0, 0)))
        # remove the last chunk if it's not full (to align with museval)
        if e.shape[0] % chunk_size != 0:
            e = e[:-(e.shape[0] % chunk_size)]
            r = r[:-(r.shape[0] % chunk_size)]
        # split into chunks
        e_chunks = np.split(e, e.shape[0] // chunk_size)
        r_chunks = np.split(r, r.shape[0] // chunk_size)
        sdr_ls = []
        for j in range(len(e_chunks)):
            e_chunk = e_chunks[j]
            r_chunk = r_chunks[j]
            sdr_chunk = sdr(e_chunk, r_chunk)
            sdr_ls.append(sdr_chunk)
        results.append(np.array(sdr_ls))
    results = np.array(results) # (src, chunks)

    # remove silent chunks SDR (i.e. to nan, to algin with museval)
    is_nan = np.isnan(results)
    is_nan = np.sum(is_nan, axis=0, keepdims=True)  # to count
    is_nan = is_nan > 0  # to boolean
    is_nan = np.repeat(is_nan,results.shape[0], axis=0)  # (src,chunks)
    results[is_nan == 1] = np.nan
    return results

def load_wav(path, track_length=None, chunk_size=None):
    if track_length is None:
        return sf.read(path, dtype='float32')[0].T # (c, t)
    else:
        s = np.random.randint(track_length - chunk_size)
        return sf.read(path, dtype='float32', start=s, frames=chunk_size)[0].T

def get_unique_save_path(dir_path):
    '''
    Given dir_path, add number suffix to make if unique
    :param dir_path: the directory path or file path without suffix
    :return:
    '''
    # get the parent directory
    parent_dir = os.path.dirname(dir_path)
    # ls the parent directory
    ls = os.listdir(parent_dir)
    # get the folder name
    folder_prefix = os.path.basename(dir_path)
    # filter out the files with the same prefix

    ls = [f for f in ls if re.match(folder_prefix + "_\d+", f)]
    if len(ls) == 0:
        return 0
    else:
        # get the number suffix
        ls = [int(f.split("_")[-1]) for f in ls]
        return max(ls) + 1

# https://github.com/RetroCirce/Zero_Shot_Audio_Source_Separation/blob/main/utils.py
def split_nparray_with_overlap(array, hop_length, overlap_size):
    '''
    Args:
        array: np.array (t, c)
        array_size: n_segs
        overlap_size: int
    Returns:
        [(c, t_seg + overlap_size)]
    '''
    result = []
    n, k, s = array.shape[0], hop_length + overlap_size, hop_length
    p = s - (n - k) % s
    array = np.pad(array, ((0, p), (0, 0)), 'constant', constant_values=0)
    array_size = (n + p - k) // s + 1
    # print(f"padded array size = {array.shape}")
    # no_ovelap_size = int(len(array) / array_size)
    for i in range(array_size):
        offset = int(i * hop_length)
        # last_loop = i == array_size
        chunk = array[offset: offset + k]
        # chunk = chunk.copy()
        # if chunk.shape[0] < no_ovelap_size + overlap_size:
        #     print(f"zero padding = {no_ovelap_size + overlap_size - chunk.shape[0]}")
        #     chunk.resize((no_ovelap_size + overlap_size, 2), refcheck=False)
        # chunk.resize((element_size + overlap_size, 2), refcheck=False)  # zero padding
        result.append(chunk.T)
    return result

# https://github.com/RetroCirce/Zero_Shot_Audio_Source_Separation/blob/main/models/asp_model.py
def join_chunks(tmp_root, chunk_ls, samplerate, overlap_size):
    '''
    Args:
        chunk_ls: list of np.array (chunks, chunk_size, c)
        overlap_size: int
    '''
    # tmp_root = "G:\\sdx2023fork\\tmp"

    # song_dir = tmp_root
    # src_filename = "test.wav"
    # out_path = os.path.join(song_dir, src_filename)

    array_size = len(chunk_ls)
    tmpfolder = tempfile.TemporaryDirectory(dir=tmp_root)
    tmpfolder_name = tmpfolder.name
    tmpfolder.cleanup()

    out_path = os.path.join(tmpfolder_name, "out.wav")

    # print(tmpfolder_name)
    os.mkdir(tmpfolder_name)


    filters = []
    args = ["ffmpeg", "-y", "-loglevel", "quiet"]
    for i in range(array_size):
        # tmpfile = tempfile.NamedTemporaryFile(dir=tmpfolder.name, suffix=".wav",delete=True)
        # os.remove(tmpfile.name)
        tmpfile_name = os.path.join(tmpfolder_name, f"chunk_{i}.wav")
        args.extend(["-i", tmpfile_name])
        sf.write(tmpfile_name, data=chunk_ls[i], samplerate=samplerate, format='WAV')

        if i < array_size - 1:
            filter_cmd = "[" + ("a" if i != 0 else "") + "{0}][{1}]acrossfade=ns={2}:c1=tri:c2=tri".format(i, i + 1,
                                                                                                           overlap_size)

            if i != array_size - 2:
                filter_cmd += "[a{0}];".format(i + 1)

            filters.append(filter_cmd)

    args.extend([
        "-filter_complex",
        "".join(filters),
        "-y",
        out_path
    ])

    # print(" ".join(args))

    try:
        out = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # shell=False, otherwise cmd length exceed
        # print(out)
        # print(out.stderr.decode("gbk"))
    except:
        raise Exception("ffmpeg does not exist. Install ffmpeg or set config.overlap_rate to zero.")



    # remove zero padding
    test, sr = sf.read(out_path)
    # print(f"test.shape = {test.shape}")
    # os.remove(out_path)
    shutil.rmtree(tmpfolder_name)
    return test


@rank_zero_only
def wandb_watch_all(wandb_logger, model):
    wandb_logger.watch(model, 'all')

if __name__ == "__main__":
    import  os
    cur_dir = r"/root/unmix_plus/scripts/open-unmix"
    # ls all the files in the dir
    files = os.listdir(cur_dir)
    # filter the files with name start with "vocals"
    files = [file for file in files if file.startswith("vocals")]
    # get the number suffix
    files = [int(file.split("_")[-1].split(".")[0]) for file in files]

    # rename the files with suffix decrease by 1
    for file in files:
        os.rename(os.path.join(cur_dir, f"vocals_{file}.wav"), os.path.join(cur_dir, f"vocals_{file-1}.wav"))

src/utils/utils.py
junyuchen-cjy/DTTNet-Pytorch
import dotenv
import hydra
from omegaconf import DictConfig, OmegaConf
from src.utils.omega_resolvers import get_train_log_dir, get_sweep_log_dir
# load environment variables from `.env` file if it exists
# recursively searches for `.env` in all folders starting from work dir
from pytorch_lightning.utilities import rank_zero_info

from src.utils import print_config_tree

dotenv.load_dotenv(override=True)

@hydra.main(config_path="configs/", config_name="config.yaml", version_base='1.1')
def main(config: DictConfig):

    # Imports should be nested inside @hydra.main to optimize tab completion
    # Read more here: https://github.com/facebookresearch/hydra/issues/934
    from src.train import train
    from src.utils import utils

    rank_zero_info(OmegaConf.to_yaml(config))

    # A couple of optional utilities:
    # - disabling python warnings
    # - easier access to debug mode
    # - forcing debug friendly configuration
    # - forcing multi-gpu friendly configuration
    # You can safely get rid of this line if you don't want those
    utils.extras(config)

    # Pretty print config using Rich library
    if config.get("print_config"):
        print_config_tree(config, resolve=True)

    # Train model
    train(config)



if __name__ == "__main__":
    # register resolvers with hydra key run.dir
    OmegaConf.register_new_resolver("get_train_log_dir", get_train_log_dir)
    OmegaConf.register_new_resolver("get_sweep_log_dir", get_sweep_log_dir)
    main()

train.py
junyuchen-cjy/DTTNet-Pytorch