inference.py: update 2d patching assignment and warnings

ivadomed · Jul 11, 2022 · 246282c · 246282c
1 parent f04d778
commit 246282c
Showing 1 changed file with 35 additions and 34 deletions.
diff --git a/ivadomed/inference.py b/ivadomed/inference.py
@@ -349,8 +349,8 @@ def segment_volume(folder_model: str, fname_images: list, gpu_id: int = 0, optio
                             Length equals 2 [PixelSizeX, PixelSizeY] for 2D or 3 [PixelSizeX, PixelSizeY, PixelSizeZ] for 3D, \
                             where X is the width, Y the height and Z the depth of the image.
             * 'pixel_size_units': (str) Units of pixel size (Must be either "mm", "um" or "nm")
-            * 'patch': (bool) Indicate if 2d patching is used while segmenting (command '--segment'), otherwise inference is \
-                       performed on the entire image at once.
+            * 'no_patch': (bool) If True, 2D patches are not used while segmenting with models that includes the \
+                          "length_2D" parameter.
             * 'overlap_2D': (list of int) List of overlaps in pixels for 2D patching. Length equals 2 [OverlapX, OverlapY], \
                             where X is the width and Y the height of the image.
             * 'metadata': (str) Film metadata.
@@ -408,39 +408,40 @@ def segment_volume(folder_model: str, fname_images: list, gpu_id: int = 0, optio
     kernel_3D = bool(ConfigKW.MODIFIED_3D_UNET in context and context[ConfigKW.MODIFIED_3D_UNET][ModelParamsKW.APPLIED]) or \
                 not context[ConfigKW.DEFAULT_MODEL][ModelParamsKW.IS_2D]
 
-    if (options is not None) and ('patch' in options):
-        is_2d_patch = options['patch']
-    else:
-        is_2d_patch = False
-
-    if is_2d_patch:
-        # Assign length_2D and stride_2D for 2D patching
-        length_2D = context[ConfigKW.DEFAULT_MODEL][ModelParamsKW.LENGTH_2D] if \
-            ModelParamsKW.LENGTH_2D in context[ConfigKW.DEFAULT_MODEL] else []
-        stride_2D = context[ConfigKW.DEFAULT_MODEL][ModelParamsKW.STRIDE_2D] if \
-            ModelParamsKW.STRIDE_2D in context[ConfigKW.DEFAULT_MODEL] else []
-        if not (length_2D and stride_2D):
-            is_2d_patch = False
-            logger.warning(f"'patch' option is provided but the default model has no 'length_2D' and 'stride_2D' "
-                           f"parameters in the configuration file '{fname_model_metadata.split('/')[-1]}'. "
-                           f"2D patching is ignored and segmentation is done without patches.")
-    else:
-        length_2D = []
-        stride_2D = []
-
-    if (options is not None) and (OptionKW.OVERLAP_2D in options):
-        if (length_2D and stride_2D):
-            overlap_2D = options.get(OptionKW.OVERLAP_2D)
-            # Swap OverlapX and OverlapY resulting in an array in order [OverlapY, OverlapX]
-            # to match length_2D and stride_2D in [Height, Width] orientation.
-            overlap_2D[1], overlap_2D[0] = overlap_2D[0], overlap_2D[1]
-            # Adjust stride_2D with overlap_2D
-            stride_2D = [x1 - x2 for (x1, x2) in zip(length_2D, overlap_2D)]
+    # Assign length_2D and stride_2D for 2D patching
+    length_2D = context[ConfigKW.DEFAULT_MODEL][ModelParamsKW.LENGTH_2D] if \
+        ModelParamsKW.LENGTH_2D in context[ConfigKW.DEFAULT_MODEL] else []
+    stride_2D = context[ConfigKW.DEFAULT_MODEL][ModelParamsKW.STRIDE_2D] if \
+        ModelParamsKW.STRIDE_2D in context[ConfigKW.DEFAULT_MODEL] else []
+
+    is_2d_patch = bool(length_2D)
+    if (options is not None) and (OptionKW.NO_PATCH in options):
+        if is_2d_patch:
+            is_2d_patch = not options.get(OptionKW.NO_PATCH)
+            length_2D = []
+            stride_2D = []
         else:
-            logger.warning(f"'overlap_2d' option is provided but 'patch' option is not, or the default model has no "
-                           f"'length_2D' and 'stride_2D' parameters in the configuration file "
-                           f"'{fname_model_metadata.split('/')[-1]}'. 2D patching is ignored and the segmentation is "
-                           f"done without patches.")
+            logger.warning(f"The 'no_patch' option is provided but the model has no 'length_2D' and "
+                               f"'stride_2D' parameters in its configuration file "
+                               f"'{fname_model_metadata.split('/')[-1]}'. 2D patching is ignored, the segmentation "
+                               f"'is done on the entire image without patches.")
+        if OptionKW.OVERLAP_2D in options:
+            logger.warning(f"The 'no_patch' option is provided along with the 'overlap_2D' option. "
+                           f"2D patching is ignored, the segmentation is done on the entire image without patches.")
+    else:
+        if (options is not None) and (OptionKW.OVERLAP_2D in options):
+            if (length_2D and stride_2D):
+                overlap_2D = options.get(OptionKW.OVERLAP_2D)
+                # Swap OverlapX and OverlapY resulting in an array in order [OverlapY, OverlapX]
+                # to match length_2D and stride_2D in [Height, Width] orientation.
+                overlap_2D[1], overlap_2D[0] = overlap_2D[0], overlap_2D[1]
+                # Adjust stride_2D with overlap_2D
+                stride_2D = [x1 - x2 for (x1, x2) in zip(length_2D, overlap_2D)]
+            else:
+                logger.warning(f"The 'overlap_2D' option is provided but the model has no 'length_2D' and "
+                               f"'stride_2D' parameters in its configuration file "
+                               f"'{fname_model_metadata.split('/')[-1]}'. 2D patching is ignored, the segmentation "
+                               f"is done on the entire image without patches.")
 
     # Add microscopy pixel size and pixel size units from options to metadata for filenames_pairs
     if (options is not None) and (OptionKW.PIXEL_SIZE in options):