import pickle
import pathlib
import torch
from torch.utils.data.dataloader import DataLoader
import pytorch_lightning as pl
import numpy as np
import yaml
import torchaudio
import pyworld
import pysptk
import random
class DataModule(pl.LightningDataModule):
def __init__(self, config):
super().__init__()
self.config = config
self.batchsize = config["train"]["batchsize"]
self.preprocessed_dir = pathlib.Path(config["general"]["preprocessed_path"])
def setup(self, stage):
if not self.preprocessed_dir.exists():
raise RuntimeError("Preprocessed directory was not be found")
if "dual" in self.config:
if self.config["dual"]["enable"]:
task_config = yaml.load(
open(self.config["dual"]["config_path"], "r"),
Loader=yaml.FullLoader,
)
task_preprocessed_dir = (
self.preprocessed_dir.parent
/ pathlib.Path(task_config["general"]["preprocessed_path"]).name
)
if not task_preprocessed_dir.exists():
raise RuntimeError(
"Preprocessed directory for multi-task learning was not be found"
)
self.flnames = {
"train": "train.txt",
"val": "val.txt",
"test": "test.txt",
}
def get_ds(self, phase):
ds = Dataset(self.flnames[phase], self.config)
return ds
def get_loader(self, phase):
ds = self.get_ds(phase)
dl = DataLoader(
ds,
self.batchsize,
shuffle=True if phase == "train" else False,
num_workers=self.config["train"]["num_workers"],
drop_last=True,
)
return dl
def train_dataloader(self):
return self.get_loader(phase="train")
def val_dataloader(self):
return self.get_loader(phase="val")
def test_dataloader(self):
return self.get_loader(phase="test")
class Dataset(torch.utils.data.Dataset):
def __init__(self, filetxt, config):
self.preprocessed_dir = pathlib.Path(config["general"]["preprocessed_path"])
self.config = config
self.spec_module = torchaudio.transforms.MelSpectrogram(
sample_rate=config["preprocess"]["sampling_rate"],
n_fft=config["preprocess"]["fft_length"],
win_length=config["preprocess"]["frame_length"],
hop_length=config["preprocess"]["frame_shift"],
f_min=config["preprocess"]["fmin"],
f_max=config["preprocess"]["fmax"],
n_mels=config["preprocess"]["n_mels"],
power=1,
center=True,
norm="slaney",
mel_scale="slaney",
)
self.resample_candidate = [8000, 11025, 12000, 16000]
self.quantization_candidate = range(2 ** 6, 2 ** 10 + 2, 2)
self.segment_length = config["preprocess"]["segment_length"]
with open(self.preprocessed_dir / filetxt, "r") as fr:
self.filelist = [pathlib.Path(path.strip("\n")) for path in fr]
self.d_out = dict()
for item in ["wavs", "wavsaux"]:
self.d_out[item] = []
for wp in self.filelist:
if config["general"]["corpus_type"] == "single":
basename = str(wp.stem)
else:
basename = str(wp.parent.name) + "-" + str(wp.stem)
with open(self.preprocessed_dir / "{}.pickle".format(basename), "rb") as fw:
d_preprocessed = pickle.load(fw)
for item in ["wavs", "wavsaux"]:
try:
self.d_out[item].extend(d_preprocessed[item])
except:
pass
for item in ["wavs", "wavsaux"]:
if self.d_out[item] != None:
self.d_out[item] = np.asarray(self.d_out[item])
if "dual" in self.config:
if self.config["dual"]["enable"]:
task_config = yaml.load(
open(config["dual"]["config_path"], "r"),
Loader=yaml.FullLoader,
)
task_preprocessed_dir = (
self.preprocessed_dir.parent
/ pathlib.Path(task_config["general"]["preprocessed_path"]).name
)
with open(task_preprocessed_dir / filetxt, "r") as fr:
task_filelist = [pathlib.Path(path.strip("\n")) for path in fr]
self.d_out["wavstask"] = []
for wp in task_filelist:
if task_config["general"]["corpus_type"] == "single":
basename = str(wp.stem)
else:
basename = str(wp.parent.name) + "-" + str(wp.stem)
with open(
task_preprocessed_dir / "{}.pickle".format(basename), "rb"
) as fw:
d_preprocessed = pickle.load(fw)
self.d_out["wavstask"].extend(d_preprocessed["wavs"])
self.d_out["wavstask"] = np.asarray(self.d_out["wavstask"])
def __len__(self):
return len(self.d_out["wavs"])
def __getitem__(self, idx):
d_batch = {}
if self.d_out["wavs"].size > 0:
d_batch["wavs"] = torch.from_numpy(self.d_out["wavs"][idx])
if self.segment_length > 0:
d_batch["wavs"] = self.get_segment(d_batch["wavs"], self.segment_length)
if self.d_out["wavsaux"].size > 0:
d_batch["wavsaux"] = torch.from_numpy(self.d_out["wavsaux"][idx])
if self.segment_length > 0:
d_batch["wavsaux"] = self.get_segment(
d_batch["wavsaux"], self.segment_length
)
if self.config["general"]["stage"] == "pretrain":
if self.config["train"]["augment"]:
d_batch["wavs"] = self.augmentation(d_batch["wavsaux"])
d_batch["wavs"] = self.normalize_waveform(d_batch["wavs"], db=-3)
d_batch["wavsaux"] = self.normalize_waveform(d_batch["wavsaux"], db=-3)
if len(d_batch["wavs"]) != len(d_batch["wavsaux"]):
min_seq_len = min(len(d_batch["wavs"]), len(d_batch["wavsaux"]))
d_batch["wavs"] = d_batch["wavs"][:min_seq_len]
d_batch["wavsaux"] = d_batch["wavsaux"][:min_seq_len]
d_batch["melspecs"] = self.calc_spectrogram(d_batch["wavs"])
if self.config["general"]["feature_type"] == "melspec":
d_batch["melspecsaux"] = self.calc_spectrogram(d_batch["wavsaux"])
elif self.config["general"]["feature_type"] == "vocfeats":
d_batch["melceps"] = self.calc_melcep(d_batch["wavsaux"])
d_batch["f0s"] = self.calc_f0(d_batch["wavs"])
d_batch["melcepssrc"] = self.calc_melcep(d_batch["wavs"])
else:
raise NotImplementedError()
elif self.config["general"]["stage"].startswith("ssl"):
d_batch["wavs"] = self.normalize_waveform(d_batch["wavs"], db=-3)
d_batch["melspecs"] = self.calc_spectrogram(d_batch["wavs"])
if self.config["general"]["feature_type"] == "vocfeats":
d_batch["f0s"] = self.calc_f0(d_batch["wavs"])
d_batch["melcepssrc"] = self.calc_melcep(d_batch["wavs"])
if self.d_out["wavsaux"].size > 0:
d_batch["wavsaux"] = self.normalize_waveform(d_batch["wavsaux"], db=-3)
if self.config["general"]["feature_type"] == "melspec":
d_batch["melspecsaux"] = self.calc_spectrogram(d_batch["wavsaux"])
elif self.config["general"]["feature_type"] == "vocfeats":
d_batch["melceps"] = self.calc_melcep(d_batch["wavsaux"])
if "dual" in self.config:
if self.config["dual"]["enable"]:
d_batch["wavstask"] = torch.from_numpy(self.d_out["wavstask"][idx])
d_batch["wavstask"] = self.get_segment(
d_batch["wavstask"], self.segment_length
)
d_batch["wavstask"] = self.normalize_waveform(
d_batch["wavstask"], db=-3
)
if self.config["general"]["feature_type"] == "melspec":
d_batch["melspecstask"] = self.calc_spectrogram(
d_batch["wavstask"]
)
elif self.config["general"]["feature_type"] == "vocfeats":
d_batch["melcepstask"] = self.calc_melcep(d_batch["wavstask"])
else:
raise NotImplementedError()
else:
raise NotImplementedError()
return d_batch
def calc_spectrogram(self, wav):
specs = self.spec_module(wav)
log_spec = torch.log(
torch.clamp_min(specs, self.config["preprocess"]["min_magnitude"])
* self.config["preprocess"]["comp_factor"]
).to(torch.float32)
return log_spec
def calc_melcep(self, wav):
wav = wav.numpy()
_, sp, _ = pyworld.wav2world(
wav.astype(np.float64),
self.config["preprocess"]["sampling_rate"],
fft_size=self.config["preprocess"]["fft_length"],
frame_period=(
self.config["preprocess"]["frame_shift"]
/ self.config["preprocess"]["sampling_rate"]
* 1000
),
)
melcep = pysptk.sp2mc(
sp,
order=self.config["preprocess"]["cep_order"],
alpha=pysptk.util.mcepalpha(self.config["preprocess"]["sampling_rate"]),
).transpose(1, 0)
melcep = torch.from_numpy(melcep).to(torch.float32)
return melcep
def calc_f0(self, wav):
if self.config["preprocess"]["f0_extractor"] == "dio":
return self.calc_f0_dio(wav)
elif self.config["preprocess"]["f0_extractor"] == "harvest":
return self.calc_f0_harvest(wav)
elif self.config["preprocess"]["f0_extractor"] == "swipe":
return self.calc_f0_swipe(wav)
else:
raise NotImplementedError()
def calc_f0_dio(self, wav):
wav = wav.numpy()
_f0, _t = pyworld.dio(
wav.astype(np.float64),
self.config["preprocess"]["sampling_rate"],
frame_period=(
self.config["preprocess"]["frame_shift"]
/ self.config["preprocess"]["sampling_rate"]
* 1000
),
)
f0 = pyworld.stonemask(
wav.astype(np.float64), _f0, _t, self.config["preprocess"]["sampling_rate"]
)
f0 = torch.from_numpy(f0).to(torch.float32)
return f0
def calc_f0_harvest(self, wav):
wav = wav.numpy()
_f0, _t = pyworld.harvest(
wav.astype(np.float64),
self.config["preprocess"]["sampling_rate"],
frame_period=(
self.config["preprocess"]["frame_shift"]
/ self.config["preprocess"]["sampling_rate"]
* 1000
),
)
f0 = pyworld.stonemask(
wav.astype(np.float64), _f0, _t, self.config["preprocess"]["sampling_rate"]
)
f0 = torch.from_numpy(f0).to(torch.float32)
return f0
def calc_f0_swipe(self, wav):
wav = wav.numpy()
f0 = pysptk.sptk.swipe(
wav.astype(np.float64),
fs=self.config["preprocess"]["sampling_rate"],
min=71,
max=800,
hopsize=self.config["preprocess"]["frame_shift"],
otype="f0",
)
f0 = torch.from_numpy(f0).to(torch.float32)
return f0
def augmentation(self, wav):
wav /= torch.max(torch.abs(wav))
new_freq = random.choice(self.resample_candidate)
new_quantization = random.choice(self.quantization_candidate)
mulaw_encoder = torchaudio.transforms.MuLawEncoding(
quantization_channels=new_quantization
)
wav_quantized = mulaw_encoder(wav) / new_quantization * 2.0 - 1.0
downsampler = torchaudio.transforms.Resample(
orig_freq=self.config["preprocess"]["sampling_rate"],
new_freq=new_freq,
resampling_method="sinc_interpolation",
lowpass_filter_width=6,
dtype=torch.float32,
)
upsampler = torchaudio.transforms.Resample(
orig_freq=new_freq,
new_freq=self.config["preprocess"]["sampling_rate"],
resampling_method="sinc_interpolation",
lowpass_filter_width=6,
dtype=torch.float32,
)
wav_processed = upsampler(downsampler(wav_quantized))
return wav_processed
def normalize_waveform(self, wav, db=-3):
wav, _ = torchaudio.sox_effects.apply_effects_tensor(
wav.unsqueeze(0),
self.config["preprocess"]["sampling_rate"],
[["norm", "{}".format(db)]],
)
return wav.squeeze(0)
def get_segment(self, wav, segment_length):
seg_size = self.config["preprocess"]["sampling_rate"] * segment_length
if len(wav) >= seg_size:
max_wav_start = len(wav) - seg_size
wav_start = random.randint(0, max_wav_start)
wav = wav[wav_start : wav_start + seg_size]
else:
wav = torch.nn.functional.pad(wav, (0, seg_size - len(wav)), "constant")
return wav