diff --git a/monai/data/wsi_reader.py b/monai/data/wsi_reader.py index 2930676be3..7f1bbfefca 100644 --- a/monai/data/wsi_reader.py +++ b/monai/data/wsi_reader.py @@ -824,7 +824,7 @@ def _get_patch( f"The image is expected to have three or four color channels in '{mode}' mode but has " f"{patch.shape[self.channel_dim]}. " ) - patch = patch[:3] + patch = np.take(patch, [0, 1, 2], self.channel_dim) return patch @@ -1183,6 +1183,6 @@ def _get_patch( f"The image is expected to have three or four color channels in '{mode}' mode but has " f"{patch.shape[self.channel_dim]}. " ) - patch = patch[:3] + patch = np.take(patch, [0, 1, 2], self.channel_dim) return patch diff --git a/tests/test_wsireader.py b/tests/test_wsireader.py index 68de71d021..aae2b0dbaf 100644 --- a/tests/test_wsireader.py +++ b/tests/test_wsireader.py @@ -81,15 +81,27 @@ TEST_CASE_3 = [ WSI_GENERIC_TIFF_PATH, {"channel_dim": -1}, - {"location": (WSI_GENERIC_TIFF_HEIGHT // 2, WSI_GENERIC_TIFF_WIDTH // 2), "size": (2, 1), "level": 0}, - np.moveaxis(np.array([[[246], [246]], [[246], [246]], [[246], [246]]], dtype=np.uint8), 0, -1), + {"location": (WSI_GENERIC_TIFF_HEIGHT // 2, WSI_GENERIC_TIFF_WIDTH // 2), "size": (4, 1), "level": 0}, + np.moveaxis( + np.array( + [[[246], [246], [246], [246]], [[246], [246], [246], [246]], [[246], [246], [246], [246]]], dtype=np.uint8 + ), + 0, + -1, + ), ] TEST_CASE_4 = [ WSI_GENERIC_TIFF_PATH, {"channel_dim": 2}, - {"location": (0, 0), "size": (2, 1), "level": 8}, - np.moveaxis(np.array([[[242], [242]], [[242], [242]], [[242], [242]]], dtype=np.uint8), 0, -1), + {"location": (0, 0), "size": (4, 1), "level": 8}, + np.moveaxis( + np.array( + [[[242], [242], [242], [242]], [[242], [242], [242], [242]], [[242], [242], [242], [242]]], dtype=np.uint8 + ), + 0, + -1, + ), ] TEST_CASE_5 = [