reformat and fix

neptune-ai · May 9, 2018 · c21f4a5 · c21f4a5
1 parent a910c64
commit c21f4a5
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/pipelines.py b/pipelines.py
@@ -19,15 +19,18 @@ def unet(config, train_mode):
 
     loader = preprocessing(config, model_type='single', is_train=train_mode)
     unet = Step(name='unet',
-                transformer= PyTorchUNetStream(**config.unet) if config.execution.stream_mode else PyTorchUNet(**config.unet),
+                transformer=PyTorchUNetStream(**config.unet) if config.execution.stream_mode else PyTorchUNet(
+                    **config.unet),
                 input_steps=[loader],
                 cache_dirpath=config.env.cache_dirpath,
                 save_output=save_output, load_saved_output=load_saved_output)
 
     mask_postprocessed = mask_postprocessing(unet, config, save_output=save_output)
     if config.postprocessor["dilate_selem_size"] > 0:
         mask_postprocessed = Step(name='mask_dilation',
-                                  transformer=MaskDilatorStream(**config.postprocessor) if config.execution.stream_mode else MaskDilator(**config.unet),
+                                  transformer=MaskDilatorStream(
+                                      **config.postprocessor) if config.execution.stream_mode else MaskDilator(
+                                      **config.postprocessor),
                                   input_steps=[mask_postprocessed],
                                   adapter={'images': ([(mask_postprocessed.name, 'categorized_images')]),
                                            },

diff --git a/postprocessing.py b/postprocessing.py
@@ -77,12 +77,15 @@ def _transform(self, images):
 
 
 class MaskDilatorStream(BaseTransformer):
+    def __init__(self, dilate_selem_size):
+        self.selem_size = dilate_selem_size
+
     def transform(self, images):
         return {'categorized_images': self._transform(images)}
 
     def _transform(self, images):
         for image in tqdm(images):
-            yield dilate_image(image)
+            yield dilate_image(image, self.selem_size)
 
 
 def label(mask):