aet.py

import argparse
import pathlib
import yaml
import torch
import torchaudio
from torch.utils.data import DataLoader
import numpy as np
import random
import librosa
from dataset import Dataset
import pickle
from lightning_module import (
    SSLStepLightningModule,
    SSLDualLightningModule,
)
from utils import plot_and_save_mels
import os
import tqdm


def get_arg():
    parser = argparse.ArgumentParser()
    parser.add_argument("--stage", required=True, type=str)
    parser.add_argument("--config_path", required=True, type=pathlib.Path)
    parser.add_argument("--exist_src_aux", action="store_true")
    parser.add_argument("--run_name", required=True, type=str)
    return parser.parse_args()


class AETDataset(Dataset):
    def __init__(self, filetxt, src_config, tar_config):
        self.config = src_config

        self.preprocessed_dir_src = pathlib.Path(
            src_config["general"]["preprocessed_path"]
        )
        self.preprocessed_dir_tar = pathlib.Path(
            tar_config["general"]["preprocessed_path"]
        )
        for item in [
            "sampling_rate",
            "fft_length",
            "frame_length",
            "frame_shift",
            "fmin",
            "fmax",
            "n_mels",
        ]:
            assert src_config["preprocess"][item] == tar_config["preprocess"][item]

        self.spec_module = torchaudio.transforms.MelSpectrogram(
            sample_rate=src_config["preprocess"]["sampling_rate"],
            n_fft=src_config["preprocess"]["fft_length"],
            win_length=src_config["preprocess"]["frame_length"],
            hop_length=src_config["preprocess"]["frame_shift"],
            f_min=src_config["preprocess"]["fmin"],
            f_max=src_config["preprocess"]["fmax"],
            n_mels=src_config["preprocess"]["n_mels"],
            power=1,
            center=True,
            norm="slaney",
            mel_scale="slaney",
        )

        with open(self.preprocessed_dir_src / filetxt, "r") as fr:
            self.filelist_src = [pathlib.Path(path.strip("\n")) for path in fr]
        with open(self.preprocessed_dir_tar / filetxt, "r") as fr:
            self.filelist_tar = [pathlib.Path(path.strip("\n")) for path in fr]

        self.d_out = {"src": {}, "tar": {}}
        for item in ["wavs", "wavsaux"]:
            self.d_out["src"][item] = []
            self.d_out["tar"][item] = []

        for swp in self.filelist_src:
            if src_config["general"]["corpus_type"] == "single":
                basename = str(swp.stem)
            else:
                basename = str(swp.parent.name) + "-" + str(swp.stem)
            with open(
                self.preprocessed_dir_src / "{}.pickle".format(basename), "rb"
            ) as fw:
                d_preprocessed = pickle.load(fw)
            for item in ["wavs", "wavsaux"]:
                try:
                    self.d_out["src"][item].extend(d_preprocessed[item])
                except:
                    pass

        for twp in self.filelist_tar:
            if tar_config["general"]["corpus_type"] == "single":
                basename = str(twp.stem)
            else:
                basename = str(twp.parent.name) + "-" + str(twp.stem)
            with open(
                self.preprocessed_dir_tar / "{}.pickle".format(basename), "rb"
            ) as fw:
                d_preprocessed = pickle.load(fw)
            for item in ["wavs", "wavsaux"]:
                try:
                    self.d_out["tar"][item].extend(d_preprocessed[item])
                except:
                    pass

        min_len = min(len(self.d_out["src"]["wavs"]), len(self.d_out["tar"]["wavs"]))
        for spk in ["src", "tar"]:
            for item in ["wavs", "wavsaux"]:
                if self.d_out[spk][item] != None:
                    self.d_out[spk][item] = np.asarray(self.d_out[spk][item][:min_len])

    def __len__(self):
        return len(self.d_out["src"]["wavs"])

    def __getitem__(self, idx):
        d_batch = {}

        for spk in ["src", "tar"]:
            for item in ["wavs", "wavsaux"]:
                if self.d_out[spk][item].size > 0:
                    d_batch["{}_{}".format(item, spk)] = torch.from_numpy(
                        self.d_out[spk][item][idx]
                    )
                    d_batch["{}_{}".format(item, spk)] = self.normalize_waveform(
                        d_batch["{}_{}".format(item, spk)], db=-3
                    )

        d_batch["melspecs_src"] = self.calc_spectrogram(d_batch["wavs_src"])
        return d_batch


class AETModule(torch.nn.Module):
    """
    src: Dataset from which we extract the channel features
    tar: Dataset to which the src channel features are added
    """

    def __init__(self, args, chmatch_config, src_config, tar_config):
        super().__init__()
        if args.stage == "ssl-step":
            LModule = SSLStepLightningModule
        elif args.stage == "ssl-dual":
            LModule = SSLDualLightningModule
        else:
            raise NotImplementedError()

        src_model = LModule(src_config).load_from_checkpoint(
            checkpoint_path=chmatch_config["general"]["source"]["ckpt_path"],
            config=src_config,
        )
        self.src_config = src_config

        self.encoder_src = src_model.encoder
        if src_config["general"]["use_gst"]:
            self.gst_src = src_model.gst
        else:
            self.channelfeats_src = src_model.channelfeats
        self.channel_src = src_model.channel

    def forward(self, melspecs_src, wavsaux_tar):
        if self.src_config["general"]["use_gst"]:
            chfeats_src = self.gst_src(melspecs_src.transpose(1, 2))
        else:
            _, enc_hidden_src = self.encoder_src(
                melspecs_src.unsqueeze(1).transpose(2, 3)
            )
            chfeats_src = self.channelfeats_src(enc_hidden_src)
        wavschmatch_tar = self.channel_src(wavsaux_tar, chfeats_src)
        return wavschmatch_tar


def calc_deg_baseline(wav, char_vector, tar_config):
    wav = wav[0, ...].cpu().detach().numpy()
    spec = librosa.stft(
        wav,
        n_fft=tar_config["preprocess"]["fft_length"],
        hop_length=tar_config["preprocess"]["frame_shift"],
        win_length=tar_config["preprocess"]["frame_length"],
    )
    spec_converted = spec * char_vector.reshape(-1, 1)
    wav_converted = librosa.istft(
        spec_converted,
        hop_length=tar_config["preprocess"]["frame_shift"],
        win_length=tar_config["preprocess"]["frame_length"],
    )
    wav_converted = torch.from_numpy(wav_converted).to(torch.float32).unsqueeze(0)
    return wav_converted


def calc_deg_charactaristics(chmatch_config):
    src_config = yaml.load(
        open(chmatch_config["general"]["source"]["config_path"], "r"),
        Loader=yaml.FullLoader,
    )
    tar_config = yaml.load(
        open(chmatch_config["general"]["target"]["config_path"], "r"),
        Loader=yaml.FullLoader,
    )
    # configs
    preprocessed_dir = pathlib.Path(src_config["general"]["preprocessed_path"])
    n_train = src_config["preprocess"]["n_train"]
    SR = src_config["preprocess"]["sampling_rate"]

    os.makedirs(preprocessed_dir, exist_ok=True)

    sourcepath = pathlib.Path(src_config["general"]["source_path"])

    if src_config["general"]["corpus_type"] == "single":
        fulllist = list(sourcepath.glob("*.wav"))
        random.seed(0)
        random.shuffle(fulllist)
        train_filelist = fulllist[:n_train]
    elif src_config["general"]["corpus_type"] == "multi-seen":
        fulllist = list(sourcepath.glob("*/*.wav"))
        random.seed(0)
        random.shuffle(fulllist)
        train_filelist = fulllist[:n_train]
    elif src_config["general"]["corpus_type"] == "multi-unseen":
        spk_list = list(set([x.parent for x in sourcepath.glob("*/*.wav")]))
        train_filelist = []
        random.seed(0)
        random.shuffle(spk_list)
        for i, spk in enumerate(spk_list):
            sourcespkpath = sourcepath / spk
            if i < n_train:
                train_filelist.extend(list(sourcespkpath.glob("*.wav")))
    else:
        raise NotImplementedError(
            "corpus_type specified in config.yaml should be {single, multi-seen, multi-unseen}"
        )

    specs_all = np.zeros((tar_config["preprocess"]["fft_length"] // 2 + 1, 1))

    for wp in tqdm.tqdm(train_filelist):
        wav, _ = librosa.load(wp, sr=SR)
        spec = np.abs(
            librosa.stft(
                wav,
                n_fft=src_config["preprocess"]["fft_length"],
                hop_length=src_config["preprocess"]["frame_shift"],
                win_length=src_config["preprocess"]["frame_length"],
            )
        )

        auxpath = pathlib.Path(src_config["general"]["aux_path"])
        if src_config["general"]["corpus_type"] == "single":
            wav_aux, _ = librosa.load(auxpath / wp.name, sr=SR)
        else:
            wav_aux, _ = librosa.load(auxpath / wp.parent.name / wp.name, sr=SR)
        spec_aux = np.abs(
            librosa.stft(
                wav_aux,
                n_fft=src_config["preprocess"]["fft_length"],
                hop_length=src_config["preprocess"]["frame_shift"],
                win_length=src_config["preprocess"]["frame_length"],
            )
        )
        min_len = min(spec.shape[1], spec_aux.shape[1])
        spec_diff = spec[:, :min_len] / (spec_aux[:, :min_len] + 1e-10)
        specs_all = np.hstack([specs_all, np.mean(spec_diff, axis=1).reshape(-1, 1)])

    char_vector = np.mean(specs_all, axis=1)
    char_vector = char_vector / (np.sum(char_vector) + 1e-10)
    return char_vector


def normalize_waveform(wav, tar_config, db=-3):
    wav, _ = torchaudio.sox_effects.apply_effects_tensor(
        wav,
        tar_config["preprocess"]["sampling_rate"],
        [["norm", "{}".format(db)]],
    )
    return wav


def main(args, chmatch_config, device):
    src_config = yaml.load(
        open(chmatch_config["general"]["source"]["config_path"], "r"),
        Loader=yaml.FullLoader,
    )
    tar_config = yaml.load(
        open(chmatch_config["general"]["target"]["config_path"], "r"),
        Loader=yaml.FullLoader,
    )
    output_path = pathlib.Path(chmatch_config["general"]["output_path"]) / args.run_name
    dataset = AETDataset("test.txt", src_config, tar_config)
    loader = DataLoader(dataset, batch_size=1, shuffle=False)
    chmatch_module = AETModule(args, chmatch_config, src_config, tar_config).to(device)

    if args.exist_src_aux:
        char_vector = calc_deg_charactaristics(chmatch_config)

    for idx, batch in enumerate(tqdm.tqdm(loader)):
        melspecs_src = batch["melspecs_src"].to(device)
        wavsdeg_src = batch["wavs_src"].to(device)
        wavsaux_tar = batch["wavsaux_tar"].to(device)
        if args.exist_src_aux:
            wavsdegbaseline_tar = calc_deg_baseline(
                batch["wavsaux_tar"], char_vector, tar_config
            )
            wavsdegbaseline_tar = normalize_waveform(wavsdegbaseline_tar, tar_config)
            wavsdeg_tar = batch["wavs_tar"].to(device)
        wavsmatch_tar = normalize_waveform(
            chmatch_module(melspecs_src, wavsaux_tar).cpu().detach(), tar_config
        )
        torchaudio.save(
            output_path / "test_wavs" / "{}-src_wavsdeg.wav".format(idx),
            wavsdeg_src.cpu(),
            src_config["preprocess"]["sampling_rate"],
        )
        torchaudio.save(
            output_path / "test_wavs" / "{}-tar_wavsaux.wav".format(idx),
            wavsaux_tar.cpu(),
            tar_config["preprocess"]["sampling_rate"],
        )
        if args.exist_src_aux:
            torchaudio.save(
                output_path / "test_wavs" / "{}-tar_wavsdegbaseline.wav".format(idx),
                wavsdegbaseline_tar.cpu(),
                tar_config["preprocess"]["sampling_rate"],
            )
            torchaudio.save(
                output_path / "test_wavs" / "{}-tar_wavsdeg.wav".format(idx),
                wavsdeg_tar.cpu(),
                tar_config["preprocess"]["sampling_rate"],
            )
        torchaudio.save(
            output_path / "test_wavs" / "{}-tar_wavsmatch.wav".format(idx),
            wavsmatch_tar.cpu(),
            tar_config["preprocess"]["sampling_rate"],
        )
        plot_and_save_mels(
            wavsdeg_src[0, ...].cpu().detach(),
            output_path / "test_mels" / "{}-src_melsdeg.png".format(idx),
            src_config,
        )
        plot_and_save_mels(
            wavsaux_tar[0, ...].cpu().detach(),
            output_path / "test_mels" / "{}-tar_melsaux.png".format(idx),
            tar_config,
        )
        if args.exist_src_aux:
            plot_and_save_mels(
                wavsdegbaseline_tar[0, ...].cpu().detach(),
                output_path / "test_mels" / "{}-tar_melsdegbaseline.png".format(idx),
                tar_config,
            )
            plot_and_save_mels(
                wavsdeg_tar[0, ...].cpu().detach(),
                output_path / "test_mels" / "{}-tar_melsdeg.png".format(idx),
                tar_config,
            )
        plot_and_save_mels(
            wavsmatch_tar[0, ...].cpu().detach(),
            output_path / "test_mels" / "{}-tar_melsmatch.png".format(idx),
            tar_config,
        )


if __name__ == "__main__":
    args = get_arg()
    chmatch_config = yaml.load(open(args.config_path, "r"), Loader=yaml.FullLoader)
    output_path = pathlib.Path(chmatch_config["general"]["output_path"]) / args.run_name
    os.makedirs(output_path, exist_ok=True)
    os.makedirs(output_path / "test_wavs", exist_ok=True)
    os.makedirs(output_path / "test_mels", exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    main(args, chmatch_config, device)