refactor dataset

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
 
-class EnhancerDataset(IterableDataset):
+    def __iter__(self):
+        return self.dataset.train__iter__()
+
+    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):
+        super().__init__(
-            raise ValueError(f"{clean_dir} is not a valid directory")
+            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.files = files
-        self.noisy_dir = noisy_dir
         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)
+    def setup(self, stage:Optional[str]=None):
-        self.valid_files = fp.prepare_matching_dict()
         
-    def __iter__(self):
+        super().setup(stage=stage)
+
+    def train__iter__(self):
 
         rng = create_unique_rng(12) ##pass epoch number here
         
         while True:
 
-            file_dict,*_ = rng.choices(self.valid_files,k=1,
+            file_dict,*_ = rng.choices(self.train_data,k=1,
-                        weights=[self.valid_files[file]['duration'] for file in self.valid_files])
+                        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)
 
+    def val__len__(self):
+        return len(self._validation)
 
 
-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 valid_dataloader(self):
-        return DataLoader(self.valid_dataset, batch_size = self.batch_size)