Set layout in Python Op pipelines

Signed-off-by: Joaquin Anton <janton@nvidia.com>

dali/test/python/test_operator_crop.py

@@ -119,11 +119,13 @@ def __init__(self, function, batch_size, layout, iterator, num_threads=1, device
         self.layout = layout
         self.iterator = iterator
         self.inputs = ops.ExternalSource()
-        self.crop = ops.PythonFunction(function=function)
+        self.crop = ops.PythonFunction(function = function)
+        self.set_layout = ops.Reshape(layout = layout)
 
     def define_graph(self):
         self.data = self.inputs()
         out = self.crop(self.data)
+        out = self.set_layout(out)
         return out
 
     def iter_setup(self):
@@ -290,11 +292,13 @@ def __init__(self, function, batch_size, iterator, data_shape, data_layout, num_
         def crop_func(image):
             return function(image, layout=self.data_layout, shape=self.data_shape)
 
-        self.crop = ops.PythonFunction(function=crop_func)
+        self.crop = ops.PythonFunction(function = crop_func)
+        self.set_layout = ops.Reshape(layout = data_layout)
 
     def define_graph(self):
         self.data = self.inputs()
         out = self.crop(self.data)
+        out = self.set_layout(out)
         return out
 
     def iter_setup(self):

dali/test/python/test_operator_crop_mirror_normalize.py

@@ -120,17 +120,19 @@ def test_cmn_no_crop_args_vs_decoder_only():
             yield check_cmn_no_crop_args_vs_decoder_only, device, batch_size
 
 class PythonOpPipeline(Pipeline):
-    def __init__(self, batch_size, function, num_threads=1, device_id=0, num_gpus=1):
+    def __init__(self, batch_size, function, output_layout, num_threads=1, device_id=0, num_gpus=1):
 
         super(PythonOpPipeline, self).__init__(batch_size, num_threads, device_id, seed=7865, exec_async=False, exec_pipelined=False)
         self.input = ops.CaffeReader(path = caffe_db_folder, shard_id = device_id, num_shards = num_gpus)
         self.decode = ops.ImageDecoder(device = "cpu", output_type = types.RGB)
         self.cmn = ops.PythonFunction(function=function)
+        self.set_layout = ops.Reshape(layout=output_layout)
 
     def define_graph(self):
         inputs, labels = self.input(name="Reader")
         images = self.decode(inputs)
         images = self.cmn(images)
+        images = self.set_layout(images)
         return images
 
 # Those are hardcoded coin flip results when using `seed=7865`
@@ -261,7 +263,7 @@ def check_cmn_vs_numpy(device, batch_size, dtype, output_layout,
     compare_pipelines(CropMirrorNormalizePipeline(device, batch_size, dtype=dtype,
                                                   output_layout=output_layout, mirror_probability=mirror_probability,
                                                   mean=mean, std=std, pad_output=should_pad),
-                      PythonOpPipeline(batch_size, function),
+                      PythonOpPipeline(batch_size, function, output_layout),
                       batch_size=batch_size, N_iterations=iterations)
 
 def test_cmn_vs_numpy():
@@ -317,7 +319,7 @@ def iter_setup(self):
         self.feed_input(self.data, data, layout=self.layout)
 
 class CMNRandomDataPythonOpPipeline(Pipeline):
-    def __init__(self, function, batch_size, layout, iterator, num_threads=1, device_id=0):
+    def __init__(self, function, batch_size, layout, output_layout, iterator, num_threads=1, device_id=0):
         super(CMNRandomDataPythonOpPipeline, self).__init__(batch_size,
                                                             num_threads,
                                                             device_id,
@@ -327,10 +329,12 @@ def __init__(self, function, batch_size, layout, iterator, num_threads=1, device
         self.iterator = iterator
         self.inputs = ops.ExternalSource()
         self.cmn = ops.PythonFunction(function=function)
+        self.set_layout = ops.Reshape(layout=output_layout)
 
     def define_graph(self):
         self.data = self.inputs()
         out = self.cmn(self.data)
+        out = self.set_layout(out)
         return out
 
     def iter_setup(self):
@@ -359,7 +363,7 @@ def check_cmn_random_data_vs_numpy(device, batch_size, dtype, input_layout, inpu
                               dtype = dtype, output_layout = output_layout,
                               mirror_probability = mirror_probability, mean = mean, std= std,
                               pad_output = should_pad),
-        CMNRandomDataPythonOpPipeline(function, batch_size, input_layout, iter(eii2)),
+        CMNRandomDataPythonOpPipeline(function, batch_size, input_layout, output_layout, iter(eii2)),
                                       batch_size=batch_size, N_iterations=1)
 
 def test_cmn_random_data_vs_numpy():

dali/test/python/test_operator_flip.py

@@ -110,10 +110,12 @@ def __init__(self, batch_size, layout, data_iterator):
         h_dim, v_dim, d_dim = find_dims(layout)
         fun = lambda d, hor, ver, depth: numpy_flip(d, h_dim, v_dim, d_dim, hor, ver, depth)
         self.python_flip = ops.PythonFunction(function=fun)
+        self.set_layout = ops.Reshape(layout=layout)
 
     def define_graph(self):
         self.data = self.input()
         flipped = self.python_flip(self.data, self.coin(), self.coin(), self.coin())
+        flipped = self.set_layout(flipped)
         return flipped
 
     def iter_setup(self):

dali/test/python/test_operator_lookup_table.py

@@ -80,10 +80,12 @@ def lookup_table_func(input_data):
                             dictionary=dictionary,
                             default_value=default_value)
         self.lookup = ops.PythonFunction(function=lookup_table_func)
+        self.set_layout = ops.Reshape(layout = data_layout)
 
     def define_graph(self):
         self.data = self.inputs()
         out = self.lookup(self.data)
+        out = self.set_layout(out)
         return out
 
     def iter_setup(self):

dali/test/python/test_operator_rotate.py

@@ -125,14 +125,16 @@ def __init__(self, batch_size, output_type, input_type, fixed_size, num_threads=
         self.input = ops.CaffeReader(path = caffe_db_folder, shard_id = device_id, num_shards = num_gpus)
         self.decode = ops.ImageDecoder(device = "cpu", output_type = types.RGB)
         self.rotate = ops.PythonFunction(function=CVRotate(output_type, input_type, fixed_size))
-        self.uniform = ops.Uniform(range = (-180.0, 180.0), seed = 42);
+        self.set_layout = ops.Reshape(layout="HWC")
+        self.uniform = ops.Uniform(range = (-180.0, 180.0), seed = 42)
         self.iter = 0
 
     def define_graph(self):
         self.jpegs, self.labels = self.input(name = "Reader")
         images = self.decode(self.jpegs)
         angles = self.uniform()
-        outputs = self.rotate(images, angles);
+        outputs = self.rotate(images, angles)
+        outputs = self.set_layout(outputs)
         return outputs
 
 

dali/test/python/test_operator_slice.py

@@ -269,12 +269,14 @@ def __init__(self, batch_size, layout, iterator, pos_size_iter,
             slice_func_helper, axes, axis_names, self.layout,
             normalized_anchor, normalized_shape)
         self.slice = ops.PythonFunction(function=function)
+        self.set_layout = ops.Reshape(layout=layout)
 
     def define_graph(self):
         self.data = self.inputs()
         self.crop_pos = self.input_crop_pos()
         self.crop_size = self.input_crop_size()
         out = self.slice(self.data, self.crop_pos, self.crop_size)
+        out = self.set_layout(out)
         return out
 
     def iter_setup(self):
@@ -308,13 +310,15 @@ def __init__(self, batch_size, pos_size_iter,
             slice_func_helper, axes, axis_names, self.layout,
             normalized_anchor, normalized_shape)
         self.slice = ops.PythonFunction(function=function)
+        self.set_layout = ops.Reshape(layout="HWC")
 
     def define_graph(self):
         imgs, _ = self.input()
         imgs = self.decode(imgs)
         self.crop_pos = self.input_crop_pos()
         self.crop_size = self.input_crop_size()
         out = self.slice(imgs, self.crop_pos, self.crop_size)
+        out = self.set_layout(out)
         return out
 
     def iter_setup(self):

dali/test/python/test_operator_warp.py

@@ -131,6 +131,7 @@ def __init__(self, batch_size, output_type, input_type, use_input, num_threads=3
           self.warp = ops.PythonFunction(function=CVWarp(output_type, input_type))
         else:
           self.warp = ops.PythonFunction(function=CVWarp(output_type, input_type, [[0.1, 0.9, 10], [0.8, -0.2, -20]]))
+        self.set_layout = ops.Reshape(layout="HWC")
         self.iter = 0
 
     def define_graph(self):
@@ -141,6 +142,7 @@ def define_graph(self):
           outputs = self.warp(images, self.transform)
         else:
           outputs = self.warp(images)
+        outputs = self.set_layout(outputs)
         return outputs
 
 def compare(pipe1, pipe2, eps):

dali/test/python/test_operator_water.py

@@ -73,13 +73,15 @@ def __init__(self,  batch_size,function,  num_threads=1, device_id=0, num_gpus=1
         self.input = ops.CaffeReader(path = caffe_db_folder, shard_id = device_id, num_shards = num_gpus)
         self.decode = ops.ImageDecoder(device = "cpu", output_type = types.RGB)
         self.water = ops.PythonFunction(function=function)
+        self.set_layout = ops.Reshape(layout="HWC")
 
 
     def define_graph(self):
         inputs, labels = self.input(name="Reader")
 
         images = self.decode(inputs)
         images = self.water(images)
+        images = self.set_layout(images)
         return images
 
 def check_water_cpu_vs_gpu(batch_size):

dali/test/python/test_pipeline_multichannel.py

@@ -135,10 +135,12 @@ def __init__(self, function, batch_size, layout, iterator, num_threads=1, device
         self.iterator = iterator
         self.inputs = ops.ExternalSource()
         self.oper = ops.PythonFunction(function=function)
+        self.set_layout = ops.Reshape(layout=layout)
 
     def define_graph(self):
         self.data = self.inputs()
         out = self.oper(self.data)
+        out = self.set_layout(out)
         return out
 
     def iter_setup(self):
@@ -224,11 +226,13 @@ def __init__(self, function, batch_size, num_threads=1, device_id=0):
         self.reader = ops.FileReader(file_root=multichannel_tiff_root)
         self.decoder = ops.ImageDecoder(device = 'cpu', output_type = types.ANY_DATA)
         self.oper = ops.PythonFunction(function=function)
+        self.set_layout = ops.Reshape(layout="HWC")
 
     def define_graph(self):
         encoded_data, _ = self.reader()
         decoded_data = self.decoder(encoded_data)
         out = self.oper(decoded_data)
+        out = self.set_layout(out)
         return out