Add options to skip operations for RestoreLabeld Transform

monai/apps/deepgrow/transforms.py

@@ -819,6 +819,10 @@ def __init__(
  original_shape_key: str = "foreground_original_shape",
  cropped_shape_key: str = "foreground_cropped_shape",
  allow_missing_keys: bool = False,
+ restore_resize: bool = True,
+ restore_crop: bool = True,
+ restore_spacing: bool = True,
+ restore_slicing: bool = True,
  ) -> None:
  super().__init__(keys, allow_missing_keys)
  self.ref_image = ref_image
@@ -833,6 +837,10 @@ def __init__(
  self.end_coord_key = end_coord_key
  self.original_shape_key = original_shape_key
  self.cropped_shape_key = cropped_shape_key
+ self.restore_resize = restore_resize
+ self.restore_crop = restore_crop
+ self.restore_spacing = restore_spacing
+ self.restore_slicing = restore_slicing
 
  def __call__(self, data: Any) -> dict:
  d = dict(data)
@@ -842,38 +850,45 @@ def __call__(self, data: Any) -> dict:
  image = d[key]
 
  # Undo Resize
- current_shape = image.shape
- cropped_shape = meta_dict[self.cropped_shape_key]
- if np.any(np.not_equal(current_shape, cropped_shape)):
- resizer = Resize(spatial_size=cropped_shape[1:], mode=mode)
- image = resizer(image, mode=mode, align_corners=align_corners)
+ if self.restore_resize:
+ current_shape = image.shape
+ cropped_shape = meta_dict[self.cropped_shape_key]
+ if np.any(np.not_equal(current_shape, cropped_shape)):
+ resizer = Resize(spatial_size=cropped_shape[1:], mode=mode)
+ image = resizer(image, mode=mode, align_corners=align_corners)
 
  # Undo Crop
- original_shape = meta_dict[self.original_shape_key]
- result = np.zeros(original_shape, dtype=np.float32)
- box_start = meta_dict[self.start_coord_key]
- box_end = meta_dict[self.end_coord_key]
-
- spatial_dims = min(len(box_start), len(image.shape[1:]))
- slices = tuple(
- [slice(None)] + [slice(s, e) for s, e in zip(box_start[:spatial_dims], box_end[:spatial_dims])]
- )
- result[slices] = image
+ if self.restore_crop:
+ original_shape = meta_dict[self.original_shape_key]
+ result = np.zeros(original_shape, dtype=np.float32)
+ box_start = meta_dict[self.start_coord_key]
+ box_end = meta_dict[self.end_coord_key]
+
+ spatial_dims = min(len(box_start), len(image.shape[1:]))
+ slices = tuple(
+ [slice(None)] + [slice(s, e) for s, e in zip(box_start[:spatial_dims], box_end[:spatial_dims])]
+ )
+ result[slices] = image
+ else:
+ result = image
 
  # Undo Spacing
- current_size = result.shape[1:]
- # change spatial_shape from HWD to DHW
- spatial_shape = list(np.roll(meta_dict["spatial_shape"], 1))
- spatial_size = spatial_shape[-len(current_size) :]
+ if self.restore_spacing:
+ current_size = result.shape[1:]
+ # change spatial_shape from HWD to DHW
+ spatial_shape = list(np.roll(meta_dict["spatial_shape"], 1))
+ spatial_size = spatial_shape[-len(current_size) :]
 
- if np.any(np.not_equal(current_size, spatial_size)):
- resizer = Resize(spatial_size=spatial_size, mode=mode)
- result = resizer(result, mode=mode, align_corners=align_corners) # type: ignore
+  if np.any(np.not_equal(current_size, spatial_size)):
+  resizer = Resize(spatial_size=spatial_size, mode=mode)
+  result = resizer(result, mode=mode, align_corners=align_corners) # type: ignore
 
  # Undo Slicing
  slice_idx = meta_dict.get("slice_idx")
  final_result: NdarrayOrTensor
- if slice_idx is None or self.slice_only:
+ if self.restore_slicing == False:
+ final_result = result
+ elif slice_idx is None or self.slice_only:
  final_result = result if len(result.shape) <= 3 else result[0]
  else:
  slice_idx = meta_dict["slice_idx"][0]

tests/test_deepgrow_transforms.py

@@ -312,7 +312,7 @@
 ]
 
 RESTORE_LABEL_TEST_CASE_1 = [
- {"keys": ["pred"], "ref_image": "image", "mode": "nearest"},
+ {"keys": ["pred"], "ref_image": "image", "mode": "nearest", "restore_resize": True, "restore_crop": True, "restore_spacing": True, "restore_slicing": True},
  DATA_10,
  np.array([[[1, 2, 3], [4, 5, 6], [7, 8, 9]]]),
 ]
@@ -327,7 +327,36 @@
  ]
 )
 
-RESTORE_LABEL_TEST_CASE_2 = [{"keys": ["pred"], "ref_image": "image", "mode": "nearest"}, DATA_11, RESULT]
+RESTORE_LABEL_TEST_CASE_2 = [
+ {"keys": ["pred"], "ref_image": "image", "mode": "nearest"},
+ DATA_11,
+ RESULT
+]
+
+RESTORE_LABEL_TEST_CASE_3_RESULT = np.zeros((10, 20, 20))
+for layer in range(5):
+ RESTORE_LABEL_TEST_CASE_3_RESULT[layer, 0:10, 0:10] = 1
+ RESTORE_LABEL_TEST_CASE_3_RESULT[layer, 0:10, 10:20] = 2
+ RESTORE_LABEL_TEST_CASE_3_RESULT[layer, 10:20, 0:10] = 3
+ RESTORE_LABEL_TEST_CASE_3_RESULT[layer, 10:20, 10:20] = 4
+
+for layer in range(5, 10):
+ RESTORE_LABEL_TEST_CASE_3_RESULT[layer, 0:10, 0:10] = 5
+ RESTORE_LABEL_TEST_CASE_3_RESULT[layer, 0:10, 10:20] = 6
+ RESTORE_LABEL_TEST_CASE_3_RESULT[layer, 10:20, 0:10] = 7
+ RESTORE_LABEL_TEST_CASE_3_RESULT[layer, 10:20, 10:20] = 8
+
+RESTORE_LABEL_TEST_CASE_3 = [
+ {"keys": ["pred"], "ref_image": "image", "mode": "nearest", "restore_crop": False},
+ DATA_11,
+ RESTORE_LABEL_TEST_CASE_3_RESULT,
+]
+
+RESTORE_LABEL_TEST_CASE_4 = [
+ {"keys": ["pred"], "ref_image": "image", "mode": "nearest", "restore_resize": False, "restore_spacing": False, "restore_slicing": False, "restore_crop": False},
+ DATA_11,
+ np.array([[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]]),
+]
 
 FETCH_2D_SLICE_TEST_CASE_1 = [
  {"keys": ["image"], "guidance": "guidance"},
@@ -445,7 +474,7 @@ def test_correct_results(self, arguments, input_data, expected_result):
 
 class TestRestoreLabeld(unittest.TestCase):
 
- @parameterized.expand([RESTORE_LABEL_TEST_CASE_1, RESTORE_LABEL_TEST_CASE_2])
+ @parameterized.expand([RESTORE_LABEL_TEST_CASE_1, RESTORE_LABEL_TEST_CASE_2, RESTORE_LABEL_TEST_CASE_3, RESTORE_LABEL_TEST_CASE_4])
  def test_correct_results(self, arguments, input_data, expected_result):
  result = RestoreLabeld(**arguments)(input_data)
  np.testing.assert_allclose(result["pred"], expected_result)