Add mixed-device LabelStats handling and tests (#1)

* Add mixed-device LabelStats coverage

* Autofix to mirror CI: Restrict versions of isort, ruff, and black.

Signed-off-by: R. Garcia-Dias <rafaelagd@gmail.com>

---------

Signed-off-by: R. Garcia-Dias <rafaelagd@gmail.com>

Author: garciadias

monai/apps/detection/transforms/dictionary.py

@@ -125,8 +125,10 @@ def __init__(self, box_keys: KeysCollection, box_ref_image_keys: str, allow_miss
         super().__init__(box_keys, allow_missing_keys)
         box_ref_image_keys_tuple = ensure_tuple(box_ref_image_keys)
         if len(box_ref_image_keys_tuple) > 1:
-            raise ValueError("Please provide a single key for box_ref_image_keys.\
-                All boxes of box_keys are attached to box_ref_image_keys.")
+            raise ValueError(
+                "Please provide a single key for box_ref_image_keys.\
+                All boxes of box_keys are attached to box_ref_image_keys."
+            )
         self.box_ref_image_keys = box_ref_image_keys
 
     def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
@@ -287,8 +289,10 @@ def __init__(
         super().__init__(box_keys, allow_missing_keys)
         box_ref_image_keys_tuple = ensure_tuple(box_ref_image_keys)
         if len(box_ref_image_keys_tuple) > 1:
-            raise ValueError("Please provide a single key for box_ref_image_keys.\
-                All boxes of box_keys are attached to box_ref_image_keys.")
+            raise ValueError(
+                "Please provide a single key for box_ref_image_keys.\
+                All boxes of box_keys are attached to box_ref_image_keys."
+            )
         self.box_ref_image_keys = box_ref_image_keys
         self.image_meta_key = image_meta_key or f"{box_ref_image_keys}_{image_meta_key_postfix}"
         self.converter_to_image_coordinate = AffineBox()
@@ -306,8 +310,10 @@ def extract_affine(self, data: Mapping[Hashable, torch.Tensor]) -> tuple[Ndarray
         else:
             raise ValueError(f"{meta_key} is not found. Please check whether it is the correct the image meta key.")
         if "affine" not in meta_dict:
-            raise ValueError(f"'affine' is not found in {meta_key}. \
-                Please check whether it is the correct the image meta key.")
+            raise ValueError(
+                f"'affine' is not found in {meta_key}. \
+                Please check whether it is the correct the image meta key."
+            )
         affine: NdarrayOrTensor = meta_dict["affine"]
 
         if self.affine_lps_to_ras:  # RAS affine
@@ -809,12 +815,16 @@ def __init__(
     ) -> None:
         box_keys_tuple = ensure_tuple(box_keys)
         if len(box_keys_tuple) != 1:
-            raise ValueError("Please provide a single key for box_keys.\
-                All label_keys are attached to this box_keys.")
+            raise ValueError(
+                "Please provide a single key for box_keys.\
+                All label_keys are attached to this box_keys."
+            )
         box_ref_image_keys_tuple = ensure_tuple(box_ref_image_keys)
         if len(box_ref_image_keys_tuple) != 1:
-            raise ValueError("Please provide a single key for box_ref_image_keys.\
-                All box_keys and label_keys are attached to this box_ref_image_keys.")
+            raise ValueError(
+                "Please provide a single key for box_ref_image_keys.\
+                All box_keys and label_keys are attached to this box_ref_image_keys."
+            )
         self.label_keys = ensure_tuple(label_keys)
         super().__init__(box_keys_tuple, allow_missing_keys)
 
@@ -1081,8 +1091,10 @@ def __init__(
 
         box_keys_tuple = ensure_tuple(box_keys)
         if len(box_keys_tuple) != 1:
-            raise ValueError("Please provide a single key for box_keys.\
-                All label_keys are attached to this box_keys.")
+            raise ValueError(
+                "Please provide a single key for box_keys.\
+                All label_keys are attached to this box_keys."
+            )
         self.box_keys = box_keys_tuple[0]
         self.label_keys = ensure_tuple(label_keys)
 

monai/apps/detection/utils/anchor_utils.py

@@ -124,8 +124,10 @@ def __init__(
             aspect_ratios = (aspect_ratios,) * len(self.sizes)
 
         if len(self.sizes) != len(aspect_ratios):
-            raise ValueError("len(sizes) and len(aspect_ratios) should be equal. \
-                It represents the number of feature maps.")
+            raise ValueError(
+                "len(sizes) and len(aspect_ratios) should be equal. \
+                It represents the number of feature maps."
+            )
 
         spatial_dims = len(ensure_tuple(aspect_ratios[0][0])) + 1
         spatial_dims = look_up_option(spatial_dims, [2, 3])
@@ -170,12 +172,16 @@ def generate_anchors(
         scales_t = torch.as_tensor(scales, dtype=dtype, device=device)  # sized (N,)
         aspect_ratios_t = torch.as_tensor(aspect_ratios, dtype=dtype, device=device)  # sized (M,) or (M,2)
         if (self.spatial_dims >= 3) and (len(aspect_ratios_t.shape) != 2):
-            raise ValueError(f"In {self.spatial_dims}-D image, aspect_ratios for each level should be \
-                {len(aspect_ratios_t.shape) - 1}-D. But got aspect_ratios with shape {aspect_ratios_t.shape}.")
+            raise ValueError(
+                f"In {self.spatial_dims}-D image, aspect_ratios for each level should be \
+                {len(aspect_ratios_t.shape) - 1}-D. But got aspect_ratios with shape {aspect_ratios_t.shape}."
+            )
 
         if (self.spatial_dims >= 3) and (aspect_ratios_t.shape[1] != self.spatial_dims - 1):
-            raise ValueError(f"In {self.spatial_dims}-D image, aspect_ratios for each level should has \
-                shape (_,{self.spatial_dims - 1}). But got aspect_ratios with shape {aspect_ratios_t.shape}.")
+            raise ValueError(
+                f"In {self.spatial_dims}-D image, aspect_ratios for each level should has \
+                shape (_,{self.spatial_dims - 1}). But got aspect_ratios with shape {aspect_ratios_t.shape}."
+            )
 
         # if 2d, w:h = 1:aspect_ratios
         if self.spatial_dims == 2:

monai/apps/reconstruction/transforms/array.py

@@ -61,8 +61,10 @@ def __init__(
                 real/imaginary parts.
         """
         if len(center_fractions) != len(accelerations):
-            raise ValueError("Number of center fractions \
-                should match number of accelerations")
+            raise ValueError(
+                "Number of center fractions \
+                should match number of accelerations"
+            )
 
         self.center_fractions = center_fractions
         self.accelerations = accelerations

monai/auto3dseg/analyzer.py

@@ -476,6 +476,12 @@ def __call__(self, data: Mapping[Hashable, MetaTensor]) -> dict[Hashable, MetaTe
         restore_grad_state = torch.is_grad_enabled()
         torch.set_grad_enabled(False)
 
+        if isinstance(image_tensor, (MetaTensor, torch.Tensor)) and isinstance(
+            label_tensor, (MetaTensor, torch.Tensor)
+        ):
+            if label_tensor.device != image_tensor.device:
+                label_tensor = label_tensor.to(image_tensor.device)
+
         ndas: list[MetaTensor] = [image_tensor[i] for i in range(image_tensor.shape[0])]  # type: ignore
         ndas_label: MetaTensor = label_tensor.astype(torch.int16)  # (H,W,D)
 

monai/bundle/utils.py

@@ -124,8 +124,10 @@
         "run_name": None,
         # may fill it at runtime
         "save_execute_config": True,
-        "is_not_rank0": ("$torch.distributed.is_available() \
-                and torch.distributed.is_initialized() and torch.distributed.get_rank() > 0"),
+        "is_not_rank0": (
+            "$torch.distributed.is_available() \
+                and torch.distributed.is_initialized() and torch.distributed.get_rank() > 0"
+        ),
         # MLFlowHandler config for the trainer
         "trainer": {
             "_target_": "MLFlowHandler",

monai/losses/dice.py

@@ -203,9 +203,11 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
                 self.class_weight = torch.as_tensor([self.class_weight] * num_of_classes)
             else:
                 if self.class_weight.shape[0] != num_of_classes:
-                    raise ValueError("""the length of the `weight` sequence should be the same as the number of classes.
+                    raise ValueError(
+                        """the length of the `weight` sequence should be the same as the number of classes.
                         If `include_background=False`, the weight should not include
-                        the background category class 0.""")
+                        the background category class 0."""
+                    )
             if self.class_weight.min() < 0:
                 raise ValueError("the value/values of the `weight` should be no less than 0.")
             # apply class_weight to loss

monai/losses/focal_loss.py

@@ -183,9 +183,11 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
                 self.class_weight = torch.as_tensor([self.class_weight] * num_of_classes)
             else:
                 if self.class_weight.shape[0] != num_of_classes:
-                    raise ValueError("""the length of the `weight` sequence should be the same as the number of classes.
+                    raise ValueError(
+                        """the length of the `weight` sequence should be the same as the number of classes.
                         If `include_background=False`, the weight should not include
-                        the background category class 0.""")
+                        the background category class 0."""
+                    )
             if self.class_weight.min() < 0:
                 raise ValueError("the value/values of the `weight` should be no less than 0.")
             # apply class_weight to loss

requirements-dev.txt

@@ -14,9 +14,9 @@ mccabe
 pep8-naming
 pycodestyle
 pyflakes
-black>=25.1.0
-isort>=5.1, !=6.0.0
-ruff
+black==25.1.0
+isort>=5.1, <6, !=6.0.0
+ruff>=0.14.11,<0.15
 pytype>=2020.6.1, <=2024.4.11; platform_system != "Windows"
 types-setuptools
 mypy>=1.5.0, <1.12.0

tests/apps/test_auto3dseg.py

@@ -53,7 +53,7 @@
     SqueezeDimd,
     ToDeviced,
 )
-from monai.utils.enums import DataStatsKeys
+from monai.utils.enums import DataStatsKeys, LabelStatsKeys
 from tests.test_utils import skip_if_no_cuda
 
 device = "cpu"
@@ -78,6 +78,13 @@
 
 SIM_GPU_TEST_CASES = [[{"sim_dim": (32, 32, 32), "label_key": "label"}], [{"sim_dim": (32, 32, 32), "label_key": None}]]
 
+LABEL_STATS_DEVICE_TEST_CASES = [
+    [{"image_device": "cpu", "label_device": "cpu", "image_meta": False}],
+    [{"image_device": "cuda", "label_device": "cuda", "image_meta": True}],
+    [{"image_device": "cpu", "label_device": "cuda", "image_meta": True}],
+    [{"image_device": "cuda", "label_device": "cpu", "image_meta": False}],
+]
+
 
 def create_sim_data(dataroot: str, sim_datalist: dict, sim_dim: tuple, image_only: bool = False, **kwargs) -> None:
     """
@@ -360,6 +367,50 @@ def test_label_stats_case_analyzer(self):
             report_format = analyzer.get_report_format()
             assert verify_report_format(d["label_stats"], report_format)
 
+    @parameterized.expand(LABEL_STATS_DEVICE_TEST_CASES)
+    def test_label_stats_mixed_device_analyzer(self, input_params):
+        image_device = torch.device(input_params["image_device"])
+        label_device = torch.device(input_params["label_device"])
+
+        if (image_device.type == "cuda" or label_device.type == "cuda") and not torch.cuda.is_available():
+            self.skipTest("CUDA is not available for mixed-device LabelStats tests.")
+
+        analyzer = LabelStats(image_key="image", label_key="label")
+
+        image_tensor = torch.tensor(
+            [
+                [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]],
+                [[[11.0, 12.0], [13.0, 14.0]], [[15.0, 16.0], [17.0, 18.0]]],
+            ],
+            dtype=torch.float32,
+        ).to(image_device)
+        label_tensor = torch.tensor([[[0, 1], [1, 0]], [[0, 1], [0, 1]]], dtype=torch.int64).to(label_device)
+
+        if input_params["image_meta"]:
+            image_tensor = MetaTensor(image_tensor)
+        label_tensor = MetaTensor(label_tensor)
+
+        result = analyzer({"image": image_tensor, "label": label_tensor})
+        report = result["label_stats"]
+
+        assert verify_report_format(report, analyzer.get_report_format())
+        assert report[LabelStatsKeys.LABEL_UID] == [0, 1]
+
+        label_stats = report[LabelStatsKeys.LABEL]
+        self.assertAlmostEqual(label_stats[0][LabelStatsKeys.PIXEL_PCT], 0.5)
+        self.assertAlmostEqual(label_stats[1][LabelStatsKeys.PIXEL_PCT], 0.5)
+
+        label0_intensity = label_stats[0][LabelStatsKeys.IMAGE_INTST]
+        label1_intensity = label_stats[1][LabelStatsKeys.IMAGE_INTST]
+        self.assertAlmostEqual(label0_intensity[0]["mean"], 4.25)
+        self.assertAlmostEqual(label1_intensity[0]["mean"], 4.75)
+        self.assertAlmostEqual(label0_intensity[1]["mean"], 14.25)
+        self.assertAlmostEqual(label1_intensity[1]["mean"], 14.75)
+
+        foreground_stats = report[LabelStatsKeys.IMAGE_INTST]
+        self.assertAlmostEqual(foreground_stats[0]["mean"], 4.75)
+        self.assertAlmostEqual(foreground_stats[1]["mean"], 14.75)
+
     def test_filename_case_analyzer(self):
         analyzer_image = FilenameStats("image", DataStatsKeys.BY_CASE_IMAGE_PATH)
         analyzer_label = FilenameStats("label", DataStatsKeys.BY_CASE_IMAGE_PATH)

versioneer.py

@@ -429,7 +429,9 @@ def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False, env=
     return stdout, process.returncode
 
 
-LONG_VERSION_PY["git"] = r'''
+LONG_VERSION_PY[
+    "git"
+] = r'''
 # This file helps to compute a version number in source trees obtained from
 # git-archive tarball (such as those provided by githubs download-from-tag
 # feature). Distribution tarballs (built by setup.py sdist) and build