File size: 3,360 Bytes
1002053 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | from dataclasses import dataclass
from pathlib import Path
import numpy as np
import pandas as pd
import soundfile as sf
from simple_parsing import Serializable
from torch.utils.data import Dataset as TorchDataset
from audiozen.acoustics.io import load_audio
@dataclass
class DatasetArgs(Serializable):
librispeech_dir: str
librispeech_metadata_fpath: str
duration: float = 6.0
sr: int = 24000
num_samples: int = 40000
class Dataset(TorchDataset):
def __init__(self, args: DatasetArgs):
super().__init__()
librispeech_dir = Path(args.librispeech_dir).expanduser().resolve()
librispeech_metadata_fpath = Path(args.librispeech_metadata_fpath).expanduser().resolve()
self.librispeech_metadata = pd.read_csv(librispeech_metadata_fpath, engine="python")
self.librispeech_dir = librispeech_dir
self.num_samples = args.num_samples
self.duration = args.duration
self.sr = args.sr
def __len__(self) -> int:
return self.num_samples
def __getitem__(self, index):
spk_1_info = self.librispeech_metadata.sample(n=1).iloc[0]
spk_1_id = spk_1_info["speaker_ID"]
spk_2_info = self.librispeech_metadata.sample(n=1).iloc[0]
spk_2_id = spk_2_info["speaker_ID"]
while spk_1_id == spk_2_id:
spk_2_info = self.librispeech_metadata.sample(n=1).iloc[0]
spk_2_id = spk_2_info["speaker_ID"]
spk_1_path = self.librispeech_dir / spk_1_info["origin_path"]
spk_2_path = self.librispeech_dir / spk_2_info["origin_path"]
spk_1_y, _ = load_audio(spk_1_path, duration=self.duration, sr=self.sr)
spk_2_y, _ = load_audio(spk_2_path, duration=self.duration, sr=self.sr)
# Randomly select a delay size from 1s to 5s
delay_size = np.random.randint(1, 5)
delay_frames = int(delay_size * self.sr)
# Offset the second one and add it to the first one.
spk_2_y = spk_2_y[:-delay_frames]
spk_2_y = np.pad(spk_2_y, (delay_frames, 0), mode="constant")
assert spk_1_y.shape == spk_2_y.shape
# Mix the two files with the delay size
mix_y = spk_1_y + spk_2_y
# Normalize the mixture
max_val = np.max(np.abs(mix_y))
mix_y /= max_val
spk_1_y /= max_val
spk_2_y /= max_val
file_id = f"{spk_1_id}_{spk_2_id}_{delay_size}.wav"
spk_y = np.stack([spk_1_y, spk_2_y], axis=0)
return mix_y, spk_y, file_id
if __name__ == "__main__":
dist_dir = Path("./tmp")
dataset = Dataset(
librispeech_dir="/nfs/xhao/data/LibriSpeech",
librispeech_metadata_fpath="/home/xhao/proj/audiozen/recipes/librimix_sot/local/metadata/LibriSpeech/train-clean-100-24K.csv",
duration=6.0,
sr=24000,
num_samples=500,
)
for i in range(len(dataset)):
mix_y, spk_y, file_id = dataset[i]
(dist_dir / "mix").mkdir(parents=True, exist_ok=True)
(dist_dir / "spk_1").mkdir(parents=True, exist_ok=True)
(dist_dir / "spk_2").mkdir(parents=True, exist_ok=True)
sf.write(dist_dir / "mix" / f"{file_id}.wav", mix_y, samplerate=24000)
sf.write(dist_dir / "spk_1" / f"{file_id}.wav", spk_y[0], samplerate=24000)
sf.write(dist_dir / "spk_2" / f"{file_id}.wav", spk_y[1], samplerate=24000)
|