Repository: Project-MONAI/MONAI

Issue:
`num_workers` for cucim in `WSIReader`
Currently CuCIM relies on the explicit set of `num_workers` when extracting patches, which is essentially different than `num_workers` provided to data loader. Although in most cases where dealing with loading small patches instead of the whole slide the default setting is fine, to gain the maximum performance of cucim in loading whole slides, we should include this argument explicitly in WSIReader.


Discussion:

Patch:
diff --git a/monai/data/wsi_reader.py b/monai/data/wsi_reader.py
--- a/monai/data/wsi_reader.py
+++ b/monai/data/wsi_reader.py
@@ -262,6 +262,7 @@ class WSIReader(BaseWSIReader):
         backend: the name of backend whole slide image reader library, the default is cuCIM.
         level: the level at which patches are extracted.
         channel_dim: the desired dimension for color channel. Default to 0 (channel first).
+        num_workers: number of workers for multi-thread image loading (cucim backend only).
         kwargs: additional arguments to be passed to the backend library
 
     """
@@ -356,6 +357,7 @@ class CuCIMWSIReader(BaseWSIReader):
         level: the whole slide image level at which the image is extracted. (default=0)
             This is overridden if the level argument is provided in `get_data`.
         channel_dim: the desired dimension for color channel. Default to 0 (channel first).
+        num_workers: number of workers for multi-thread image loading
         kwargs: additional args for `cucim.CuImage` module:
             https://github.com/rapidsai/cucim/blob/main/cpp/include/cucim/cuimage.h
 
@@ -364,8 +366,9 @@ class CuCIMWSIReader(BaseWSIReader):
     supported_suffixes = ["tif", "tiff", "svs"]
     backend = "cucim"
 
-    def __init__(self, level: int = 0, channel_dim: int = 0, **kwargs):
+    def __init__(self, level: int = 0, channel_dim: int = 0, num_workers: int = 0, **kwargs):
         super().__init__(level, channel_dim, **kwargs)
+        self.num_workers = num_workers
 
     @staticmethod
     def get_level_count(wsi) -> int:
@@ -448,7 +451,9 @@ def get_patch(
         """
         # Extract a patch or the entire image
         # (reverse the order of location and size to become WxH for cuCIM)
-        patch: np.ndarray = wsi.read_region(location=location[::-1], size=size[::-1], level=level)
+        patch: np.ndarray = wsi.read_region(
+            location=location[::-1], size=size[::-1], level=level, num_workers=self.num_workers
+        )
 
         # Convert to numpy
         patch = np.asarray(patch, dtype=dtype)


Test patch:
diff --git a/tests/test_wsireader.py b/tests/test_wsireader.py
--- a/tests/test_wsireader.py
+++ b/tests/test_wsireader.py
@@ -385,6 +385,37 @@ def test_with_dataloader_batch(self, file_path, level, expected_spatial_shape, e
                 assert_allclose(s, expected_spatial_shape, type_test=False)
             self.assertTupleEqual(data["image"].shape, (batch_size, *expected_shape[1:]))
 
+        @parameterized.expand([TEST_CASE_0])
+        def test_read_whole_image_multi_thread(self, file_path, level, expected_shape):
+            if self.backend == "cucim":
+                reader = WSIReader(self.backend, level=level, num_workers=4)
+                with reader.read(file_path) as img_obj:
+                    img, meta = reader.get_data(img_obj)
+                self.assertTupleEqual(img.shape, expected_shape)
+                self.assertEqual(meta["backend"], self.backend)
+                self.assertEqual(meta["path"], str(os.path.abspath(file_path)))
+                self.assertEqual(meta["patch_level"], level)
+                assert_array_equal(meta["patch_size"], expected_shape[1:])
+                assert_array_equal(meta["patch_location"], (0, 0))
+
+        @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4])
+        def test_read_region_multi_thread(self, file_path, kwargs, patch_info, expected_img):
+            if self.backend == "cucim":
+                reader = WSIReader(backend=self.backend, num_workers=2, **kwargs)
+                with reader.read(file_path) as img_obj:
+                    # Read twice to check multiple calls
+                    img, meta = reader.get_data(img_obj, **patch_info)
+                    img2 = reader.get_data(img_obj, **patch_info)[0]
+                    self.assertTupleEqual(img.shape, img2.shape)
+                    self.assertIsNone(assert_array_equal(img, img2))
+                    self.assertTupleEqual(img.shape, expected_img.shape)
+                    self.assertIsNone(assert_array_equal(img, expected_img))
+                    self.assertEqual(meta["backend"], self.backend)
+                    self.assertEqual(meta["path"], str(os.path.abspath(file_path)))
+                    self.assertEqual(meta["patch_level"], patch_info["level"])
+                    assert_array_equal(meta["patch_size"], patch_info["size"])
+                    assert_array_equal(meta["patch_location"], patch_info["location"])
+
 
 @skipUnless(has_cucim, "Requires cucim")
 class TestCuCIMDeprecated(WSIReaderDeprecatedTests.Tests):