Fixing problem with torch.stack non-empty elemnt

Author: FabianHoerst

cellvit/inference/inference.py

@@ -475,7 +475,7 @@ def _import_postprocessing(self) -> Tuple[Callable, Callable]:
             Tuple[Callable, Callable]: Postprocessing module
         """
         if self.system_configuration["cupy"]:
-            from cellvit.inference.postprocessing_numpy import (
+            from cellvit.inference.postprocessing_cupy import (
                 DetectionCellPostProcessor,
                 create_batch_pooling_actor,
             )

cellvit/inference/postprocessing_cupy.py

@@ -653,22 +653,23 @@ def convert_batch_to_graph_nodes(
                 batch_cell_positions = batch_cell_positions + patch_cell_positions
 
             if self.detection_cell_postprocessor.classifier is not None:
-                batch_cell_tokens_pt = torch.stack(batch_cell_tokens)
-                updated_preds = self.detection_cell_postprocessor.classifier(
-                    batch_cell_tokens_pt
-                )
-                updated_preds = F.softmax(updated_preds, dim=1)
-                updated_classes = torch.argmax(updated_preds, dim=1)
-                updated_class_preds = updated_preds[
-                    torch.arange(updated_classes.shape[0]), updated_classes
-                ]
-
-                for f, z in zip(batch_complete, updated_classes):
-                    f["type"] = int(z)
-                for f, z in zip(batch_complete, updated_class_preds):
-                    f["type_prob"] = int(z)
-                for f, z in zip(batch_detection, updated_classes):
-                    f["type"] = int(z)
+                if len(batch_cell_tokens) > 0:
+                    batch_cell_tokens_pt = torch.stack(batch_cell_tokens)
+                    updated_preds = self.detection_cell_postprocessor.classifier(
+                        batch_cell_tokens_pt
+                    )
+                    updated_preds = F.softmax(updated_preds, dim=1)
+                    updated_classes = torch.argmax(updated_preds, dim=1)
+                    updated_class_preds = updated_preds[
+                        torch.arange(updated_classes.shape[0]), updated_classes
+                    ]
+
+                    for f, z in zip(batch_complete, updated_classes):
+                        f["type"] = int(z)
+                    for f, z in zip(batch_complete, updated_class_preds):
+                        f["type_prob"] = int(z)
+                    for f, z in zip(batch_detection, updated_classes):
+                        f["type"] = int(z)
             if self.detection_cell_postprocessor.binary:
                 for f in batch_complete:
                     f["type"] = 1

cellvit/inference/postprocessing_numpy.py

@@ -653,22 +653,23 @@ def convert_batch_to_graph_nodes(
                 batch_cell_positions = batch_cell_positions + patch_cell_positions
 
             if self.detection_cell_postprocessor.classifier is not None:
-                batch_cell_tokens_pt = torch.stack(batch_cell_tokens)
-                updated_preds = self.detection_cell_postprocessor.classifier(
-                    batch_cell_tokens_pt
-                )
-                updated_preds = F.softmax(updated_preds, dim=1)
-                updated_classes = torch.argmax(updated_preds, dim=1)
-                updated_class_preds = updated_preds[
-                    torch.arange(updated_classes.shape[0]), updated_classes
-                ]
-
-                for f, z in zip(batch_complete, updated_classes):
-                    f["type"] = int(z)
-                for f, z in zip(batch_complete, updated_class_preds):
-                    f["type_prob"] = int(z)
-                for f, z in zip(batch_detection, updated_classes):
-                    f["type"] = int(z)
+                if len(batch_cell_tokens) > 0:
+                    batch_cell_tokens_pt = torch.stack(batch_cell_tokens)
+                    updated_preds = self.detection_cell_postprocessor.classifier(
+                        batch_cell_tokens_pt
+                    )
+                    updated_preds = F.softmax(updated_preds, dim=1)
+                    updated_classes = torch.argmax(updated_preds, dim=1)
+                    updated_class_preds = updated_preds[
+                        torch.arange(updated_classes.shape[0]), updated_classes
+                    ]
+
+                    for f, z in zip(batch_complete, updated_classes):
+                        f["type"] = int(z)
+                    for f, z in zip(batch_complete, updated_class_preds):
+                        f["type_prob"] = int(z)
+                    for f, z in zip(batch_detection, updated_classes):
+                        f["type"] = int(z)
             if self.detection_cell_postprocessor.binary:
                 for f in batch_complete:
                     f["type"] = 1

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "cellvit"
-version = "1.0.1b"
+version = "1.0.2b"
 description = "CellViT Inference Pipeline for Whole Slide Images (WSI) in Memory"
 authors = [
     { name = "Fabian Hörst", email = "[email protected]" }

tests/test_cli/test_cli_yaml.py

@@ -125,8 +125,8 @@ def test_invalid_batch_size(self, mock_device_count):
             InferenceConfiguration(invalid_config)
         self.assertEqual(
             str(context.exception),
-            "Batch size must be between 2 and 32",
-            "Batch size must be between 2 and 32",
+            "Batch size must be between 2 and 48",
+            "Batch size must be between 2 and 48",
         )
 
     @patch("torch.cuda.device_count")

tests/test_dataclass/test_wsi_meta.py

@@ -27,7 +27,7 @@ def test_load_wsi_meta_with_svs_file_1(self):
         self.assertEqual(
             slide_properties["magnification"], 20.0
         )  # Fill in expected value
-        self.assertEqual(target_mpp, 0.25)  # Fill in expected value
+        self.assertEqual(target_mpp, 0.2495)  # Fill in expected value
 
     def test_load_wsi_meta_with_svs_file_2(self):
         """Test loading WSI metadata from SVS file 2."""