diff --git a/monai/transforms/intensity/array.py b/monai/transforms/intensity/array.py index 5c8c77408b..25110ac0dd 100644 --- a/monai/transforms/intensity/array.py +++ b/monai/transforms/intensity/array.py @@ -2293,6 +2293,7 @@ def __call__(self, image: NdarrayOrTensor): class ComputeHoVerMaps(Transform): """Compute horizontal and vertical maps from an instance mask It generates normalized horizontal and vertical distances to the center of mass of each region. + Input data with the size of [1xHxW[xD]], which channel dim will temporarily removed for calculating coordinates. Args: dtype: the data type of output Tensor. Defaults to `"float32"`. @@ -2311,6 +2312,7 @@ def __call__(self, mask: NdarrayOrTensor): h_map = instance_mask.astype(self.dtype, copy=True) v_map = instance_mask.astype(self.dtype, copy=True) + instance_mask = instance_mask.squeeze(0) # remove channel dim for region in skimage.measure.regionprops(instance_mask): v_dist = region.coords[:, 0] - region.centroid[0] @@ -2325,5 +2327,5 @@ def __call__(self, mask: NdarrayOrTensor): h_map[h_map == region.label] = h_dist v_map[v_map == region.label] = v_dist - hv_maps = convert_to_tensor(np.stack([h_map, v_map]), track_meta=get_track_meta()) + hv_maps = convert_to_tensor(np.concatenate([h_map, v_map]), track_meta=get_track_meta()) return hv_maps diff --git a/tests/test_compute_ho_ver_maps.py b/tests/test_compute_ho_ver_maps.py index ca8576488a..f5091a57af 100644 --- a/tests/test_compute_ho_ver_maps.py +++ b/tests/test_compute_ho_ver_maps.py @@ -22,11 +22,11 @@ _, has_skimage = optional_import("skimage", "0.19.0", min_version) -INSTANCE_MASK = np.zeros((16, 16), dtype="int16") -INSTANCE_MASK[5:8, 4:11] = 1 -INSTANCE_MASK[3:5, 6:9] = 1 -INSTANCE_MASK[8:10, 6:9] = 1 -INSTANCE_MASK[13:, 13:] = 2 +INSTANCE_MASK = np.zeros((1, 16, 16), dtype="int16") +INSTANCE_MASK[:, 5:8, 4:11] = 1 +INSTANCE_MASK[:, 3:5, 6:9] = 1 +INSTANCE_MASK[:, 8:10, 6:9] = 1 +INSTANCE_MASK[:, 13:, 13:] = 2 H_MAP = torch.zeros((16, 16), dtype=torch.float32) H_MAP[5:8, 4] = -1.0 H_MAP[5:8, 5] = -2.0 / 3.0 diff --git a/tests/test_compute_ho_ver_maps_d.py b/tests/test_compute_ho_ver_maps_d.py index cf7b2ee1ec..3c20c7f200 100644 --- a/tests/test_compute_ho_ver_maps_d.py +++ b/tests/test_compute_ho_ver_maps_d.py @@ -22,11 +22,11 @@ _, has_skimage = optional_import("skimage", "0.19.0", min_version) -INSTANCE_MASK = np.zeros((16, 16), dtype="int16") -INSTANCE_MASK[5:8, 4:11] = 1 -INSTANCE_MASK[3:5, 6:9] = 1 -INSTANCE_MASK[8:10, 6:9] = 1 -INSTANCE_MASK[13:, 13:] = 2 +INSTANCE_MASK = np.zeros((1, 16, 16), dtype="int16") +INSTANCE_MASK[:, 5:8, 4:11] = 1 +INSTANCE_MASK[:, 3:5, 6:9] = 1 +INSTANCE_MASK[:, 8:10, 6:9] = 1 +INSTANCE_MASK[:, 13:, 13:] = 2 H_MAP = torch.zeros((16, 16), dtype=torch.float32) H_MAP[5:8, 4] = -1.0 H_MAP[5:8, 5] = -2.0 / 3.0