Dataloader updations

Author: saisha17

data/direct_dataset.py

@@ -1,4 +1,3 @@
-#python3 train.py  --dataroot ./datasets/datasets/data-six/ --name endtoend_no_shadow_ --checkpoints_dir all_checkpoints/endtoend_checkpoints/ --model pix2pix_endtoend_3 --dataset_mode endtoend_no_shadow --input_nc 1 --output_nc 1  --crop_size 256     --gpu_id 1 --netG unet_256 
 import os.path
 from data.base_dataset import BaseDataset, get_params, get_transform
 from data.image_folder import make_dataset
@@ -7,11 +6,6 @@
 import json
 
 class DirectDataset(BaseDataset):
-    """A dataset class for paired image dataset.
-
-    It assumes that the directory '/path/to/data/train' contains image pairs in the form of {A,B}.
-    During test time, you need to prepare a directory '/path/to/data/test'.
-    """
 
     def __init__(self, opt):
         """Initialize this dataset class.
@@ -23,18 +17,19 @@ def __init__(self, opt):
         self.C_paths=[]
         self.G_paths=[]
         self.K_paths=[]
+        self.texture_paths={}
 
         self.dir_AB = opt.dataroot  # get the image directory
         for i in os.listdir(self.dir_AB):
           T_path = self.dir_AB +'/'+ i+'/params/meta.ndjson';
           with open(T_path) as f:
             for line in f:
               j_content = json.loads(line)
-              texture=j_content["texture"]["name"]
-          if texture!='sha' and texture!='som':
-            self.C_paths += make_dataset(self.dir_AB+'/'+i+'/r_contour/', opt.max_dataset_size)  # get image paths
-            self.G_paths += make_dataset(self.dir_AB+'/'+i+'/r_gnomon/', opt.max_dataset_size)  # get image paths
-            self.K_paths += make_dataset(self.dir_AB+'/'+i+'/r_sketch/', opt.max_dataset_size)  # get image paths
+              tex=j_content["texture"]["name"]
+          self.C_paths += make_dataset(self.dir_AB+'/'+i+'/r_contour/', opt.max_dataset_size)  # get image paths
+          self.G_paths += make_dataset(self.dir_AB+'/'+i+'/r_gnomon/', opt.max_dataset_size)  # get image paths
+          self.K_paths += make_dataset(self.dir_AB+'/'+i+'/r_sketch/', opt.max_dataset_size)  # get image paths
+          self.texture_paths[self.dir_AB+'/'+i+'/r_contour'] = opt.texture+'/'+tex+'/'  # get image paths
         self.C_paths=sorted(self.C_paths)
         self.G_paths=sorted(self.G_paths)
         self.K_paths=sorted(self.K_paths)
@@ -43,17 +38,6 @@ def __init__(self, opt):
         self.output_nc = self.opt.output_nc 
 
     def __getitem__(self, index):
-        """Return a data point and its metadata information.
-
-        Parameters:
-            index - - a random integer for data indexing
-
-        Returns a dictionary that contains A, B, A_paths and B_paths
-            A (tensor) - - an image in the input domain
-            B (tensor) - - its corresponding image in the target domain
-            A_paths (str) - - image paths
-            B_paths (str) - - image paths (same as A_paths)
-        """
         # read a image given a random integer index
 
         C_path = self.C_paths[index]
@@ -73,14 +57,14 @@ def __getitem__(self, index):
             for line in f:
               j_content = json.loads(line)
               texture=j_content["texture"]["name"]
-       
-        C=Image.open("/nfs/151/gpu/raghav/data/shadegan/brushes_v2/"+texture+"/high.png").convert('RGB')
+
+        C=Image.open(self.texture_paths[C_path.rsplit('/',1)[0]]+"high.png").convert('RGB')
 
-        D=Image.open("/nfs/151/gpu/raghav/data/shadegan/brushes_v2/"+texture+"/mid.png").convert('RGB')
+        D=Image.open(self.texture_paths[C_path.rsplit('/',1)[0]]+"mid.png").convert('RGB')
 
-        E=Image.open("/nfs/151/gpu/raghav/data/shadegan/brushes_v2/"+texture+"/shade.png").convert('RGB')
+        E=Image.open(self.texture_paths[C_path.rsplit('/',1)[0]]+"shade.png").convert('RGB')
 
-        F=Image.open("/nfs/151/gpu/raghav/data/shadegan/brushes_v2/"+texture+"/shadow.png").convert('RGB')
+        F=Image.open(self.texture_paths[C_path.rsplit('/',1)[0]]+"shadow.png").convert('RGB')
         # apply the same transform to both A and B
         transform_params = get_params(self.opt, A.size)
         A_transform = get_transform(self.opt, transform_params, grayscale=(self.input_nc == 1))

data/sp_dataset.py

@@ -1,4 +1,3 @@
-#python3 train.py  --dataroot ./datasets/datasets/data-six/ --name endtoend_no_shadow_ --checkpoints_dir all_checkpoints/endtoend_checkpoints/ --model pix2pix_endtoend_3 --dataset_mode endtoend_no_shadow --input_nc 1 --output_nc 1  --crop_size 256     --gpu_id 1 --netG unet_256 
 import os.path
 from data.base_dataset import BaseDataset, get_params, get_transform
 from data.image_folder import make_dataset
@@ -7,18 +6,8 @@
 import json
 
 class SpDataset(BaseDataset):
-    """A dataset class for paired image dataset.
-
-    It assumes that the directory '/path/to/data/train' contains image pairs in the form of {A,B}.
-    During test time, you need to prepare a directory '/path/to/data/test'.
-    """
 
     def __init__(self, opt):
-        """Initialize this dataset class.
-
-        Parameters:
-            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
-        """
         BaseDataset.__init__(self, opt)
         self.C_paths=[]
         self.G_paths=[]
@@ -27,22 +16,23 @@ def __init__(self, opt):
         self.S_paths=[]
         self.K_paths=[]
         self.Sh_paths=[]
+        self.texture_paths={}
 
         self.dir_AB = opt.dataroot  # get the image directory
         for i in os.listdir(self.dir_AB):
           T_path = self.dir_AB +'/'+ i+'/params/meta.ndjson';
           with open(T_path) as f:
             for line in f:
               j_content = json.loads(line)
-              texture=j_content["texture"]["name"]
-          if texture!='sha' and texture!='som':
-            self.C_paths += make_dataset(self.dir_AB+'/'+i+'/r_contour/', opt.max_dataset_size)  # get image paths
-            self.G_paths += make_dataset(self.dir_AB+'/'+i+'/r_gnomon/', opt.max_dataset_size)  # get image paths
-            self.H_paths += make_dataset(self.dir_AB+'/'+i+'/r_highlights/', opt.max_dataset_size)  # get image paths
-            self.M_paths += make_dataset(self.dir_AB+'/'+i+'/r_midtones/', opt.max_dataset_size)  # get image paths
-            self.S_paths += make_dataset(self.dir_AB+'/'+i+'/r_shades/', opt.max_dataset_size)  # get image paths
-            self.Sh_paths += make_dataset(self.dir_AB+'/'+i+'/r_shadow/', opt.max_dataset_size)  # get image paths
-            self.K_paths += make_dataset(self.dir_AB+'/'+i+'/r_sketch/', opt.max_dataset_size)  # get image paths
+              tex=j_content["texture"]["name"]
+          self.C_paths += make_dataset(self.dir_AB+'/'+i+'/r_contour/', opt.max_dataset_size)  # get image paths
+          self.G_paths += make_dataset(self.dir_AB+'/'+i+'/r_gnomon/', opt.max_dataset_size)  # get image paths
+          self.H_paths += make_dataset(self.dir_AB+'/'+i+'/r_highlights/', opt.max_dataset_size)  # get image paths
+          self.M_paths += make_dataset(self.dir_AB+'/'+i+'/r_midtones/', opt.max_dataset_size)  # get image paths
+          self.S_paths += make_dataset(self.dir_AB+'/'+i+'/r_shades/', opt.max_dataset_size)  # get image paths
+          self.Sh_paths += make_dataset(self.dir_AB+'/'+i+'/r_shadow/', opt.max_dataset_size)  # get image paths
+          self.K_paths += make_dataset(self.dir_AB+'/'+i+'/r_sketch/', opt.max_dataset_size)  # get image paths
+          self.texture_paths[self.dir_AB+'/'+i+'/r_contour'] = opt.texture+'/'+tex+'/'  # get image paths
         self.C_paths=sorted(self.C_paths)
         self.G_paths=sorted(self.G_paths)
         self.H_paths=sorted(self.H_paths)
@@ -55,17 +45,6 @@ def __init__(self, opt):
         self.output_nc = self.opt.output_nc 
 
     def __getitem__(self, index):
-        """Return a data point and its metadata information.
-
-        Parameters:
-            index - - a random integer for data indexing
-
-        Returns a dictionary that contains A, B, A_paths and B_paths
-            A (tensor) - - an image in the input domain
-            B (tensor) - - its corresponding image in the target domain
-            A_paths (str) - - image paths
-            B_paths (str) - - image paths (same as A_paths)
-        """
         # read a image given a random integer index
 
         C_path = self.C_paths[index]
@@ -93,14 +72,14 @@ def __getitem__(self, index):
             for line in f:
               j_content = json.loads(line)
               texture=j_content["texture"]["name"]
-       
-        C=Image.open("/nfs/151/gpu/raghav/data/shadegan/brushes_v2/"+texture+"/high.png").convert('RGB')
+
+        C=Image.open(self.texture_paths[C_path.rsplit('/',1)[0]]+"high.png").convert('RGB')
 
-        D=Image.open("/nfs/151/gpu/raghav/data/shadegan/brushes_v2/"+texture+"/mid.png").convert('RGB')
+        D=Image.open(self.texture_paths[C_path.rsplit('/',1)[0]]+"mid.png").convert('RGB')
 
-        E=Image.open("/nfs/151/gpu/raghav/data/shadegan/brushes_v2/"+texture+"/shade.png").convert('RGB')
+        E=Image.open(self.texture_paths[C_path.rsplit('/',1)[0]]+"shade.png").convert('RGB')
 
-        F=Image.open("/nfs/151/gpu/raghav/data/shadegan/brushes_v2/"+texture+"/shadow.png").convert('RGB')
+        F=Image.open(self.texture_paths[C_path.rsplit('/',1)[0]]+"shadow.png").convert('RGB')
 
             # apply the same transform to both A and B
         transform_params = get_params(self.opt, A.size)

data/sp_ws_dataset.py

@@ -6,12 +6,6 @@
 import json
 
 class SpWsDataset(BaseDataset):
-    """A dataset class for paired image dataset.
-
-    It assumes that the directory '/path/to/data/train' contains image pairs in the form of {A,B}.
-    During test time, you need to prepare a directory '/path/to/data/test'.
-    """
-
     def __init__(self, opt):
         """Initialize this dataset class.
 
@@ -25,21 +19,23 @@ def __init__(self, opt):
         self.M_paths=[]
         self.S_paths=[]
         self.K_paths=[]
+        self.texture_paths={}
 
         self.dir_AB = opt.dataroot  # get the image directory
         for i in os.listdir(self.dir_AB):
           T_path = self.dir_AB +'/'+ i+'/params/meta.ndjson';
           with open(T_path) as f:
             for line in f:
               j_content = json.loads(line)
-              texture=j_content["texture"]["name"]
-          if texture!='sha' and texture!='som':
+              tex=j_content["texture"]["name"]
+            
             self.C_paths += make_dataset(self.dir_AB+'/'+i+'/r_contour/', opt.max_dataset_size)  # get image paths
             self.G_paths += make_dataset(self.dir_AB+'/'+i+'/r_gnomon/', opt.max_dataset_size)  # get image paths
             self.H_paths += make_dataset(self.dir_AB+'/'+i+'/r_highlights/', opt.max_dataset_size)  # get image paths
             self.M_paths += make_dataset(self.dir_AB+'/'+i+'/r_midtones/', opt.max_dataset_size)  # get image paths
             self.S_paths += make_dataset(self.dir_AB+'/'+i+'/r_shades/', opt.max_dataset_size)  # get image paths
             self.K_paths += make_dataset(self.dir_AB+'/'+i+'/r_sketch_geom/', opt.max_dataset_size)  # get image paths
+            self.texture_paths[self.dir_AB+'/'+i+'/r_contour'] = opt.texture+'/'+tex+'/'  # get image paths
         self.C_paths=sorted(self.C_paths)
         self.G_paths=sorted(self.G_paths)
         self.H_paths=sorted(self.H_paths)
@@ -52,17 +48,6 @@ def __init__(self, opt):
         self.output_nc = self.opt.output_nc 
 
     def __getitem__(self, index):
-        """Return a data point and its metadata information.
-
-        Parameters:
-            index - - a random integer for data indexing
-
-        Returns a dictionary that contains A, B, A_paths and B_paths
-            A (tensor) - - an image in the input domain
-            B (tensor) - - its corresponding image in the target domain
-            A_paths (str) - - image paths
-            B_paths (str) - - image paths (same as A_paths)
-        """
         # read a image given a random integer index
 
         C_path = self.C_paths[index]
@@ -83,19 +68,18 @@ def __getitem__(self, index):
 
         except:
           print("EXCEPTION:"+str(AB_path))
-        texture="zero"
         with open(T_path) as f:
             for line in f:
               j_content = json.loads(line)
               texture=j_content["texture"]["name"]
 
-        C=Image.open("/nfs/151/gpu/raghav/data/shadegan/brushes_v2/"+texture+"/high.png").convert('RGB')
+        C=Image.open(self.texture_paths[C_path.rsplit('/',1)[0]]+"high.png").convert('RGB')
 
-        D=Image.open("/nfs/151/gpu/raghav/data/shadegan/brushes_v2/"+texture+"/mid.png").convert('RGB')
+        D=Image.open(self.texture_paths[C_path.rsplit('/',1)[0]]+"mid.png").convert('RGB')
 
-        E=Image.open("/nfs/151/gpu/raghav/data/shadegan/brushes_v2/"+texture+"/shade.png").convert('RGB')
+        E=Image.open(self.texture_paths[C_path.rsplit('/',1)[0]]+"shade.png").convert('RGB')
 
-        F=Image.open("/nfs/151/gpu/raghav/data/shadegan/brushes_v2/"+texture+"/shadow.png").convert('RGB')
+        F=Image.open(self.texture_paths[C_path.rsplit('/',1)[0]]+"shadow.png").convert('RGB')
 
         transform_params = get_params(self.opt, A.size)
         A_transform = get_transform(self.opt, transform_params, grayscale=(self.input_nc == 1))

options/base_options.py

@@ -54,6 +54,7 @@ def initialize(self, parser):
         parser.add_argument('--load_iter', type=int, default='0', help='which iteration to load? if load_iter > 0, the code will load models by iter_[load_iter]; otherwise, the code will load models by [epoch]')
         parser.add_argument('--verbose', action='store_true', help='if specified, print more debugging information')
         parser.add_argument('--suffix', default='', type=str, help='customized suffix: opt.name = opt.name + suffix: e.g., {model}_{netG}_size{load_size}')
+        parser.add_argument('--texture', default='', type=str, help='path to texture images')
         self.initialized = True
         return parser
 

test_commands

@@ -1,11 +1,11 @@
 direct:
-python3 test.py  --dataroot /nfs/154/dataset2/shadegan/data-test/renders/  --name direct --checkpoints_dir test_checkpoints --model direct  --dataset_mode direct --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1 --results_dir test_results
+python3 test.py  --dataroot /nfs/154/dataset2/shadegan/data-test/renders/  --name direct --checkpoints_dir test_checkpoints --model direct  --dataset_mode direct --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1 --results_dir test_results --texture /nfs/151/gpu/raghav/data/shadegan/brushes_v2
 
 sp:
-python3 test.py  --dataroot /nfs/154/dataset2/shadegan/data-test/renders/  --name sp --checkpoints_dir test_checkpoints --model sp  --dataset_mode sp --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1 --results_dir test_results
+python3 test.py  --dataroot /nfs/154/dataset2/shadegan/data-test/renders/  --name sp --checkpoints_dir test_checkpoints --model sp  --dataset_mode sp --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1 --results_dir test_results --texture /nfs/151/gpu/raghav/data/shadegan/brushes_v2
 
 sp-ws:
-python3 test.py  --dataroot /nfs/154/dataset2/shadegan/data-test/renders/  --name sp_ws --checkpoints_dir test_checkpoints --model sp_ws  --dataset_mode sp_ws --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1 --results_dir test_results
+python3 test.py  --dataroot /nfs/154/dataset2/shadegan/data-test/renders/  --name sp_ws --checkpoints_dir test_checkpoints --model sp_ws  --dataset_mode sp_ws --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1 --results_dir test_results --texture /nfs/151/gpu/raghav/data/shadegan/brushes_v2
 
 se:
-python3 test.py  --dataroot /nfs/154/dataset2/shadegan/data-test/renders/  --name se --checkpoints_dir test_checkpoints --model se  --dataset_mode sp --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1 --results_dir test_results
+python3 test.py  --dataroot /nfs/154/dataset2/shadegan/data-test/renders/  --name se --checkpoints_dir test_checkpoints --model se  --dataset_mode sp --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1 --results_dir test_results --texture /nfs/151/gpu/raghav/data/shadegan/brushes_v2

train_commands

@@ -1,11 +1,11 @@
 direct:
-python3 train.py  --dataroot /nfs/154/dataset2/shadegan/data-single/renders/  --name direct --checkpoints_dir test_checkpoints --model direct  --dataset_mode direct --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1
+python3 train.py  --dataroot /nfs/154/dataset2/shadegan/data-single/renders/  --name direct --checkpoints_dir test_checkpoints --model direct  --dataset_mode direct --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1 --texture /nfs/151/gpu/raghav/data/shadegan/brushes_v2
 
 sp:
-python3 train.py  --dataroot /nfs/154/dataset2/shadegan/data-single/renders/  --name sp --checkpoints_dir test_checkpoints --model sp  --dataset_mode sp --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1
+python3 train.py  --dataroot /nfs/154/dataset2/shadegan/data-single/renders/  --name sp --checkpoints_dir test_checkpoints --model sp  --dataset_mode sp --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1 --texture /nfs/151/gpu/raghav/data/shadegan/brushes_v2
 
 sp-ws:
-python3 train.py  --dataroot /nfs/154/dataset2/shadegan/data-six/renders/  --name sp_ws --checkpoints_dir test_checkpoints --model sp_ws  --dataset_mode sp_ws --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1
+python3 train.py  --dataroot /nfs/154/dataset2/shadegan/data-six/renders/  --name sp_ws --checkpoints_dir test_checkpoints --model sp_ws  --dataset_mode sp_ws --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1 --texture /nfs/151/gpu/raghav/data/shadegan/brushes_v2
 
 se:
-python3 train.py  --dataroot /nfs/154/dataset2/shadegan/data-six/renders/  --name se --checkpoints_dir test_checkpoints --model se  --dataset_mode sp --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1 --gpu_id 1
+python3 train.py  --dataroot /nfs/154/dataset2/shadegan/data-six/renders/  --name se --checkpoints_dir test_checkpoints --model se  --dataset_mode sp --input_nc 1 --output_nc 1 --load_size 256 --crop_size 256  --batch_size 1 --texture /nfs/151/gpu/raghav/data/shadegan/brushes_v2