Better dimension handling

NOTE :
Layer orientation can be hard to determine, anisotropy ends up being wrong due to this depending on whether layer is loaded image or "original" from previous inference

Author: C-Achard

napari_cellseg3d/model_workers.py

@@ -344,7 +344,17 @@ def load_folder(self):
 
     def load_layer(self):
 
-        volume = np.array(self.layer.data, dtype=np.int16)
+        data = np.squeeze(self.layer.data)
+
+        volume = np.array(data, dtype=np.int16)
+
+        volume_dims = len(volume.shape)
+        if volume_dims != 3:
+            raise ValueError(
+                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 anisotropy to be monai-like, i.e. zyx
@@ -537,7 +547,7 @@ def inference_on_layer(self, image, model, post_process_transforms):
 
         # print(inputs.shape)
         image = image.type(torch.FloatTensor)
-        out = self.model_output(image,model, post_process_transforms)
+        out = self.model_output(image, model, post_process_transforms)
         out = self.aniso_transform(out)
         # if self.transforms["zoom"][0]:
         #     zoom = self.transforms["zoom"][1]
@@ -596,9 +606,9 @@ def inference_on_layer(self, image, model, post_process_transforms):
         result = {
             "image_id": 0,
             "original": None,
-            "instance_labels": np.swapaxes(instance_labels, 0,2),
+            "instance_labels": np.swapaxes(instance_labels, 0, 2),
             "object stats": data_dict,
-            "result": np.swapaxes(out, 0,2),
+            "result": np.swapaxes(out, 0, 2),
             "model_name": self.model_dict["name"],
         }
         return result