Spaces:
Running on Zero
Running on Zero
File size: 2,129 Bytes
2e1dc7f | 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 | import torch
import torch.utils.data
from pathlib import Path
from data.stemmed_dataset import StemmedDataset
import config as cfg
class MixDataset(torch.utils.data.Dataset):
def __init__(self, root_dir: Path, *args, **kwargs):
for dirname in [
"moisesdb_v0.1", "lag-moisesdb", "lag_moisesdb", "moisesdb",
"moises"
]:
if (root_dir / dirname).exists():
self.root_dir_moises: Path = root_dir / dirname
break
else:
raise FileNotFoundError(
f"Couldn't find subdirectory for moisesdb under {root_dir}")
self.root_dir_mus: Path = root_dir / "musdb"
self.moisesdb_dataset: StemmedDataset = StemmedDataset(
self.root_dir_moises,
*args,
**kwargs,
)
self.musdb_dataset: StemmedDataset = StemmedDataset(
self.root_dir_mus,
*args,
**kwargs,
)
# self.n_samples = len(self.moisesdb_dataset) + len(self.musdb_dataset)
self.n_samples_moises = len(self.moisesdb_dataset)
self.n_samples_mus = len(self.musdb_dataset)
def __len__(self):
return self.n_samples_moises + self.n_samples_mus
def __getitem__(self, x):
if x < self.n_samples_moises:
return self.moisesdb_dataset[x]
else:
return self.musdb_dataset[x - self.n_samples_moises]
if __name__ == "__main__":
from data.stem import Stem
moises_root = cfg.moises_path()
mus_root = cfg.mus_path()
stems = {
Stem.DRUMS, Stem.GUITAR, Stem.BASS, Stem.PIANO, Stem.KEYBOARD,
Stem.STRINGS
}
d = MixDataset(
cfg.mixdata_path(),
stems,
target_stem=Stem.DRUMS,
single_stem=True,
add_click=False,
bpm_in_caption=False,
sync_chunks=False,
train=True,
sample_rate=32_000,
chunk_size_samples=32_000 * 10,
speed_transform_p=1,
pitch_transform_p=1,
stereo=False,
n_samples_per_epoch=2000,
)
sample = next(iter(d))
|