Remove deprecated call to AddChannel

Author: C-Achard

napari_cellseg3d/code_models/worker_inference.py

@@ -8,7 +8,7 @@
 from monai.data import DataLoader, Dataset
 from monai.inferers import sliding_window_inference
 from monai.transforms import (
-    AddChannel,
+    # AddChannel,
     # AsDiscrete,
     Compose,
     EnsureChannelFirstd,
@@ -231,9 +231,8 @@ def load_layer(self):
                 f"Data array is not 3-dimensional but {volume_dims}-dimensional,"
                 f" please check for extra channel/batch dimensions"
             )
-        volume = np.swapaxes(
-            volume, 0, 2
-        )  # for dims to be monai-like, i.e. xyz, from napari zyx
+        volume = utils.correct_rotation(volume)
+        # volume = np.reshape(volume, newshape=(1, 1, *volume.shape))
 
         dims_check = volume.shape
 
@@ -252,9 +251,9 @@ def load_layer(self):
                     normalization,
                     ToTensor(),
                     # anisotropic_transform,
-                    AddChannel(),
+                    # AddChannel(),
                     # SpatialPad(spatial_size=pad),
-                    AddChannel(),
+                    # AddChannel(),
                     EnsureType(),
                 ],
                 map_items=False,
@@ -269,9 +268,9 @@ def load_layer(self):
                     normalization,
                     ToTensor(),
                     # anisotropic_transform,
-                    AddChannel(),
+                    # AddChannel(),
                     SpatialPad(spatial_size=pad),
-                    AddChannel(),
+                    # AddChannel(),
                     EnsureType(),
                 ],
                 map_items=False,
@@ -541,6 +540,14 @@ def run_crf(self, image, labels, aniso_transform, image_id=0):
         try:
             if aniso_transform is not None:
                 image = aniso_transform(image)
+
+            if image.shape[-3:] != labels.shape[-3:]:
+                image = utils.correct_rotation(image)
+                if image.shape[-3:] != labels.shape[-3:]:
+                    logger.warning(
+                        f"Labels shape mismatch: target {image.shape}, got {labels.shape}. CRF will likely fail."
+                    )
+
             crf_results = crf_with_config(
                 image, labels, config=self.config.crf_config, log=self.log
             )
@@ -572,13 +579,24 @@ def stats_csv(self, instance_labels):
     def inference_on_layer(self, image, model, post_process_transforms):
         self.log("-" * 10)
         self.log("Inference started on layer...")
-
+        image = image.view((1, 1, *image.shape))
+        logger.debug(f"Layer shape @ inference input: {image.shape}")
         out = self.model_output(
             image,
             model,
             post_process_transforms,
             aniso_transform=self.aniso_transform,
         )
+        logger.debug(f"Inference on layer result shape : {out.shape}")
+        out = utils.correct_rotation(out)
+        extra_dims = len(image.shape) - 3
+        layer_shape_corrected = np.swapaxes(
+            image, extra_dims, 2 + extra_dims
+        ).shape
+        if out.shape[-3:] != layer_shape_corrected[-3:]:
+            logger.debug(
+                f"Output shape {out.shape[-3:]} does not match input shape {layer_shape_corrected[-3:]} on HWD dims even after rotation"
+            )
         self.save_image(out, from_layer=True)
 
         instance_labels, stats = self.get_instance_result(