From c34746ac098559bf37719beff360d76fd66c6263 Mon Sep 17 00:00:00 2001 From: shahules786 Date: Sat, 3 Sep 2022 10:48:59 +0530 Subject: [PATCH] refactor dataset --- enhancer/data/dataset.py | 152 ++++++++++++++++++++++++++------------- 1 file changed, 103 insertions(+), 49 deletions(-) diff --git a/enhancer/data/dataset.py b/enhancer/data/dataset.py index 2807aed..adcdea1 100644 --- a/enhancer/data/dataset.py +++ b/enhancer/data/dataset.py @@ -1,95 +1,149 @@ +from dataclasses import dataclass import glob import math import os +from typing_extensions import dataclass_transform import pytorch_lightning as pl -from torch.utils.data import IterableDataset, DataLoader +from torch.utils.data import IterableDataset, DataLoader, Dataset import torch.nn.functional as F from typing import Optional from enhancer.utils.random import create_unique_rng from enhancer.utils.io import Audio -from enhancer.utils import Fileprocessor +from enhancer.utils import Fileprocessor, check_files from enhancer.utils.config import Files +class TrainDataset(IterableDataset): + + def __init__(self,dataset): + self.dataset = dataset + def __iter__(self): + return self.dataset.train__iter__() -class EnhancerDataset(IterableDataset): + def __len__(self): + return self.dataset.train__len__() + +class ValidDataset(Dataset): + + def __init__(self,dataset): + self.dataset = dataset + + def __getitem__(self,idx): + return self.dataset.val__getitem__(idx) + + def __len__(self): + return self.dataset.val__len__() + +class TaskDataset(pl.LightningDataModule): + + def __init__( + self, + name:str, + root_dir:str, + files:Files, + duration:float=1.0, + sampling_rate:int=48000, + matching_function = None, + batch_size=32): + super().__init__() + + self.name = name + self.files,self.root_dir = check_files(root_dir,files) + self.duration = duration + self.sampling_rate = sampling_rate + self.batch_size = batch_size + self.matching_function = matching_function + + def setup(self, stage: Optional[str] = None): + + if stage in ("fit",None): + + fp = Fileprocessor.from_name(self.name,self.files.train_clean,self.files.train_noisy,self.matching_function) + self.train_data = fp.prepare_matching_dict() + + fp = Fileprocessor.from_name(self.name,self.files.test_clean,self.files.test_noisy,self.matching_function) + val_data = fp.prepare_matching_dict() + + for item in val_data: + clean,noisy,total_dur = item.values() + if total_dur < self.duration: + continue + num_segments = round(total_dur/self.duration) + for index in range(num_segments): + start_time = index * self.duration + self._validation.append(({"clean_file":clean,"noisy_file":noisy}, + start_time)) + def train_dataloader(self): + return DataLoader(TrainDataset(self), batch_size = self.batch_size) + + def val_dataloader(self): + return DataLoader(ValidDataset(self), batch_size = self.batch_size) + +class EnhancerDataset(TaskDataset): """Dataset object for creating clean-noisy speech enhancement datasets""" - def __init__(self,name:str,clean_dir,noisy_dir,duration=1.0,sampling_rate=48000, matching_function=None): + def __init__( + self, + name:str, + root_dir:str, + files:Files, + duration=1.0, + sampling_rate=48000, + matching_function=None, + batch_size=32): - if not os.path.isdir(clean_dir): - raise ValueError(f"{clean_dir} is not a valid directory") + super().__init__( + name=name, + root_dir=root_dir, + files=files, + sampling_rate=sampling_rate, + duration=duration, + matching_function = matching_function, + batch_size=batch_size - if not os.path.isdir(noisy_dir): - raise ValueError(f"{clean_dir} is not a valid directory") + ) self.sampling_rate = sampling_rate - self.clean_dir = clean_dir - self.noisy_dir = noisy_dir + self.files = files self.duration = max(1.0,duration) self.audio = Audio(self.sampling_rate,mono=True,return_tensor=True) - fp = Fileprocessor.from_name(name,clean_dir,noisy_dir,matching_function) - self.valid_files = fp.prepare_matching_dict() + def setup(self, stage:Optional[str]=None): + + super().setup(stage=stage) - def __iter__(self): + def train__iter__(self): rng = create_unique_rng(12) ##pass epoch number here while True: - file_dict,*_ = rng.choices(self.valid_files,k=1, - weights=[self.valid_files[file]['duration'] for file in self.valid_files]) + file_dict,*_ = rng.choices(self.train_data,k=1, + weights=[file["duration"] for file in self.train_data]) file_duration = file_dict['duration'] start_time = round(rng.uniform(0,file_duration- self.duration),2) data = self.prepare_segment(file_dict,start_time) yield data + def val__getitem__(self,idx): + return self.prepare_segment(*self._validation[idx]) + def prepare_segment(self,file_dict:dict, start_time:float): - clean_segment = self.audio(file_dict.keys()[0], + clean_segment = self.audio(file_dict["clean"], offset=start_time,duration=self.duration) - noisy_segment = self.audio(file_dict['noisy'], + noisy_segment = self.audio(file_dict["noisy"], offset=start_time,duration=self.duration) clean_segment = F.pad(clean_segment,(0,int(self.duration*self.sampling_rate-clean_segment.shape[-1]))) noisy_segment = F.pad(noisy_segment,(0,int(self.duration*self.sampling_rate-noisy_segment.shape[-1]))) return {"clean": clean_segment,"noisy":noisy_segment} - def __len__(self): - + def train__len__(self): return math.ceil(sum([file["duration"] for file in self.valid_files])/self.duration) - - -class Dataset(pl.LightningDataModule): - - def __init__(self,name:str,root_dir:str, files:Files, - duration:float=1.0, sampling_rate:int=48000, batch_size=32): - super().__init__() - - self.train_clean = os.path.join(root_dir, files.train_clean) - self.train_noisy = os.path.join(root_dir,files.train_noisy) - self.valid_clean = os.path.join(root_dir,files.test_clean) - self.valid_noisy = os.path.join(root_dir,files.test_noisy) - self.name = name - self.duration = duration - self.sampling_rate = sampling_rate - self.batch_size = batch_size - - def setup(self, stage: Optional[str] = None): - - if stage in (None,"fit"): - self.train_dataset = EnhancerDataset(self.name, self.train_clean, - self.train_noisy, self.duration, self.sampling_rate) - - self.valid_dataset = EnhancerDataset(self.name, self.valid_clean, - self.valid_noisy, self.duration, self.sampling_rate) - - def train_dataloader(self): - return DataLoader(self.train_dataset, batch_size = self.batch_size) + def val__len__(self): + return len(self._validation) - def valid_dataloader(self): - return DataLoader(self.valid_dataset, batch_size = self.batch_size) \ No newline at end of file