Seich
/
mayavoz
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor dataset

This commit is contained in:
shahules786 2022-09-03 10:48:59 +05:30
parent 761683fd41
commit c34746ac09
1 changed files with 103 additions and 49 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

152
enhancer/data/dataset.py
View File

@ -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)