added option for user to define data transforms

torch_enhance/datasets/bsds300.py

@@ -15,12 +15,17 @@ def __init__(
         image_size=256,
         color_space='RGB',
         train=True,
-        data_dir=None
+        data_dir=None,
+        lr_transforms=None,
+        hr_transforms=None
     ):
         super(BSDS300, self).__init__()
+
         self.scale_factor = scale_factor
         self.image_size = image_size
         self.color_space = color_space
+        self.lr_transforms = lr_transforms
+        self.hr_transforms = hr_transforms
 
         if data_dir is None:
             data_dir = os.path.join(os.getcwd(), self.base_dir)
@@ -30,8 +35,10 @@ def __init__(
         self.set_dir = os.path.join(self.root_dir, 'train' if train else 'test')
         self.file_names = self.get_files(self.set_dir)
 
-        self.lr_transform = self.get_lr_transforms()
-        self.hr_transform = self.get_hr_transforms()
+        if self.lr_transforms is None:
+            self.lr_transform = self.get_lr_transforms()
+        if self.hr_transforms is None:
+            self.hr_transform = self.get_hr_transforms()
 
     def download(self, data_dir):
 

torch_enhance/datasets/bsds500.py

@@ -16,12 +16,17 @@ def __init__(
         image_size=256,
         color_space='RGB',
         set_type='train',
-        data_dir=None
+        data_dir=None,
+        lr_transforms=None,
+        hr_transforms=None
     ):
         super(BSDS500, self).__init__()
+
         self.scale_factor = scale_factor
         self.image_size = image_size
         self.color_space = color_space
+        self.lr_transforms = lr_transforms
+        self.hr_transforms = hr_transforms
 
         if data_dir is None:
             data_dir = os.path.join(os.getcwd(), self.base_dir)
@@ -31,8 +36,10 @@ def __init__(
         self.set_dir = os.path.join(self.root_dir, set_type)
         self.file_names = self.get_files(self.set_dir)
 
-        self.lr_transform = self.get_lr_transforms()
-        self.hr_transform = self.get_hr_transforms()
+        if self.lr_transforms is None:
+            self.lr_transform = self.get_lr_transforms()
+        if self.hr_transforms is None:
+            self.hr_transform = self.get_hr_transforms()
 
     def download(self, data_dir):
 

torch_enhance/datasets/historical.py

@@ -14,13 +14,17 @@ def __init__(
         scale_factor=2,
         image_size=256,
         color_space='L',
-        data_dir=None
+        data_dir=None,
+        lr_transforms=None,
+        hr_transforms=None
     ):
         super(Historical, self).__init__()
 
         self.scale_factor = scale_factor
         self.image_size = image_size
         self.color_space = color_space
+        self.lr_transforms = lr_transforms
+        self.hr_transforms = hr_transforms
 
         if data_dir is None:
             data_dir = os.path.join(os.getcwd(), self.base_dir)
@@ -29,8 +33,10 @@ def __init__(
         self.download(data_dir)
         self.file_names = self.get_files(self.root_dir)
 
-        self.lr_transform = self.get_lr_transforms()
-        self.hr_transform = self.get_hr_transforms()
+        if self.lr_transforms is None:
+            self.lr_transform = self.get_lr_transforms()
+        if self.hr_transforms is None:
+            self.hr_transform = self.get_hr_transforms()
 
     def download(self, data_dir):
 

torch_enhance/datasets/set14.py

@@ -14,23 +14,29 @@ def __init__(
         scale_factor=2,
         image_size=256,
         color_space='RGB',
-        data_dir=None
+        data_dir=None,
+        lr_transforms=None,
+        hr_transforms=None
     ):
         super(Set14, self).__init__()
 
         self.scale_factor = scale_factor
         self.image_size = image_size
         self.color_space = color_space
-
+        self.lr_transforms = lr_transforms
+        self.hr_transforms = hr_transforms
+
         if data_dir is None:
             data_dir = os.path.join(os.getcwd(), self.base_dir)
 
         self.root_dir = os.path.join(data_dir, 'Set14')
         self.download(data_dir)
         self.file_names = self.get_files(self.root_dir)
 
-        self.lr_transform = self.get_lr_transforms()
-        self.hr_transform = self.get_hr_transforms()
+        if self.lr_transforms is None:
+            self.lr_transform = self.get_lr_transforms()
+        if self.hr_transforms is None:
+            self.hr_transform = self.get_hr_transforms()
 
     def download(self, data_dir):
 

torch_enhance/datasets/set5.py

@@ -14,13 +14,17 @@ def __init__(
         scale_factor=2,
         image_size=256,
         color_space='RGB',
-        data_dir=None
+        data_dir=None,
+        lr_transforms=None,
+        hr_transforms=None
     ):
         super(Set5, self).__init__()
 
         self.scale_factor = scale_factor
         self.image_size = image_size
         self.color_space = color_space
+        self.lr_transforms = lr_transforms
+        self.hr_transforms = hr_transforms
 
         if data_dir is None:
             data_dir = os.path.join(os.getcwd(), self.base_dir)
@@ -29,8 +33,10 @@ def __init__(
         self.download(data_dir)
         self.file_names = self.get_files(self.root_dir)
 
-        self.lr_transform = self.get_lr_transforms()
-        self.hr_transform = self.get_hr_transforms()
+        if self.lr_transforms is None:
+            self.lr_transform = self.get_lr_transforms()
+        if self.hr_transforms is None:
+            self.hr_transform = self.get_hr_transforms()
 
     def download(self, data_dir):
 

torch_enhance/datasets/t91.py

@@ -14,22 +14,29 @@ def __init__(
         scale_factor=2,
         image_size=256,
         color_space='RGB',
-        data_dir=None
+        data_dir=None,
+        lr_transforms=None,
+        hr_transforms=None
     ):
         super(T91, self).__init__()
+
         self.scale_factor = scale_factor
         self.image_size = image_size
         self.color_space = color_space
-
+        self.lr_transforms = lr_transforms
+        self.hr_transforms = hr_transforms
+
         if data_dir is None:
             data_dir = os.path.join(os.getcwd(), self.base_dir)
 
         self.root_dir = os.path.join(data_dir, 'T91')
         self.download(data_dir)
         self.file_names = self.get_files(self.root_dir)
 
-        self.lr_transform = self.get_lr_transforms()
-        self.hr_transform = self.get_hr_transforms()
+        if self.lr_transforms is None:
+            self.lr_transform = self.get_lr_transforms()
+        if self.hr_transforms is None:
+            self.hr_transform = self.get_hr_transforms()
 
     def download(self, data_dir):