Initial commit

Functional prototype for single layer inference

Features :
- functional basic inference on selected layer
- improvements to inference code layout

MISSING :
- Proper button behaviour
- Error if no layer selected
- Full test of all functionalities

Author: C-Achard

napari_cellseg3d/model_workers.py

@@ -17,7 +17,6 @@
 from monai.transforms import AddChannel
 from monai.transforms import AsDiscrete
 from monai.transforms import Compose
-from monai.transforms import EnsureChannelFirst
 from monai.transforms import EnsureChannelFirstd
 from monai.transforms import EnsureType
 from monai.transforms import EnsureTyped
@@ -246,33 +245,32 @@ def load_folder(self):
     def load_layer(self):
 
         volume = np.array(self.layer.data, dtype=np.int16)
+        volume = np.swapaxes(
+            volume, 0, 2
+        )  # for anisotropy to be monai-like, i.e. zyx
         print("Loading layer")
-        print(volume.shape)
-
         dims_check = volume.shape
         self.log("\nChecking dimensions...")
         pad = utils.get_padding_dim(dims_check)
-        print(f"padding: {pad}")
+        print(volume.shape)
+        print(volume.dtype)
         load_transforms = Compose(
             [
-                AddChannel(),
-                AddChannel(),
                 ToTensor(),
-                EnsureType(),
-                # EnsureChannelFirst(),
-                # Orientationd(keys=["image"], axcodes="PLI"),
                 # anisotropic_transform,
+                AddChannel(),
                 SpatialPad(spatial_size=pad),
-            ]
+                AddChannel(),
+                EnsureType(),
+            ],
+            map_items=False,
+            log_stats=True,
         )
 
         self.log("\nLoading dataset...")
-        inference_ds = Dataset(data=volume, transform=load_transforms)
-        inference_loader = DataLoader(
-            inference_ds, batch_size=1, num_workers=2
-        )
+        input_image = load_transforms(volume)
         self.log("Done")
-        return inference_loader
+        return input_image
 
     def model_output(self, inputs, model, post_process_transforms):
 
@@ -357,12 +355,7 @@ def method(image):
         self.log(os.path.split(instance_filepath)[1])
 
         # print(self.stats_to_csv)
-        if self.stats_to_csv:
-            data_dict = volume_stats(
-                instance_labels
-            )  # TODO test with area mesh function
-            return data_dict
-        return None
+        return instance_labels
 
     def inference_on_list(self, inf_data, i, model, post_process_transforms):
 
@@ -538,16 +531,13 @@ def inference_on_list(self, inf_data, i, model, post_process_transforms):
         # print(result)
         return result
 
-    def inference_on_array(self, image, model, post_process_transforms):
+    def inference_on_layer(self, image, model, post_process_transforms):
 
         self.log("-" * 10)
         self.log(f"Inference started on layer...")
 
         # print(inputs.shape)
-
-        image = ToTensor()(image)
-        image = AddChannel()(image)
-        image = AddChannel()(image)
+        image = image.type(torch.FloatTensor)
         inputs = image.to("cpu")
 
         model_output = lambda inputs: post_process_transforms(
@@ -632,7 +622,13 @@ def inference_on_array(self, image, model, post_process_transforms):
         #################
         #################
         if self.instance_params["do_instance"]:
-            data_dict = self.instance_seg(to_instance)
+            instance_labels = self.instance_seg(to_instance)
+            if self.stats_to_csv:
+                data_dict = volume_stats(
+                    instance_labels
+                )  # TODO test with area mesh function
+            else:
+                data_dict = None
         #     self.log(
         #         f"\nRunning instance segmentation for image n°{image_id}"
         #     )
@@ -691,16 +687,14 @@ def inference_on_array(self, image, model, post_process_transforms):
             instance_labels = None
             data_dict = None
 
-        # logging(f"Inference completed on image {i+1}")
         result = {
             "image_id": 0,
             "original": None,
-            "instance_labels": instance_labels,
+            "instance_labels": np.swapaxes(instance_labels, 0,2),
             "object stats": data_dict,
-            "result": out,
+            "result": np.swapaxes(out, 0,2),
             "model_name": self.model_dict["name"],
         }
-        # print(result)
         return result
 
     def inference(self):
@@ -838,8 +832,6 @@ def inference(self):
             )
         elif is_folder:
             inference_loader = self.load_folder()
-        elif is_layer:
-            inference_loader = self.load_layer()
             ##################
             ##################
             # DEBUG
@@ -850,6 +842,10 @@ def inference(self):
             print(image.shape)
             ##################
             ##################
+        elif is_layer:
+            input_image = self.load_layer()
+            print(input_image.shape)
+
         else:
             raise ValueError("No data has been provided. Aborting.")
 
@@ -861,13 +857,11 @@ def inference(self):
             if is_folder:
                 for i, inf_data in enumerate(inference_loader):
                     yield self.inference_on_list(
-                        inf_data,i, model, post_process_transforms
+                        inf_data, i, model, post_process_transforms
                     )
             elif is_layer:
-                image = self.layer.data
-                print(image.shape)
-                yield self.inference_on_array(
-                    image, model, post_process_transforms
+                yield self.inference_on_layer(
+                    input_image, model, post_process_transforms
                 )
             # for i, inf_data in enumerate(inference_loader):
             #

napari_cellseg3d/plugin_model_inference.py

@@ -15,6 +15,9 @@
 from napari_cellseg3d.model_workers import InferenceWorker
 
 
+# TODO for layer inference : button behaviour/visibility, error if no layer selected, test all funcs
+
+
 class Inferer(ModelFramework):
     """A plugin to run already trained models in evaluation mode to preform inference and output a label on all
     given volumes."""
@@ -605,7 +608,7 @@ def start(self, on_layer=False):
                     keep_on_cpu=self.keep_on_cpu,
                     stats_csv=self.stats_to_csv,
                 )
-            elif on_layer:
+            else:
                 layer = self._viewer.layers.selection.active
                 self.worker = InferenceWorker(
                     device=device,
@@ -663,7 +666,7 @@ def on_error(self):
         """Catches errors and tries to clean up. TODO : upgrade"""
         self.log.print_and_log("Worker errored...")
         self.log.print_and_log("Trying to clean up...")
-        self.btn_start.setText("Start")
+        self.btn_start.setText("Start on folder")
         self.btn_close.setVisible(True)
 
         self.worker = None
@@ -673,7 +676,7 @@ def on_finish(self):
         """Catches finished signal from worker, resets workspace for next run."""
         self.log.print_and_log(f"\nWorker finished at {utils.get_time()}")
         self.log.print_and_log("*" * 20)
-        self.btn_start.setText("Start")
+        self.btn_start.setText("Start on folder")
         self.btn_close.setVisible(True)
 
         self.worker = None
@@ -704,13 +707,13 @@ def on_yield(data, widget):
 
             zoom = widget.zoom
 
-            print(data["original"].shape)
+            # print(data["original"].shape)
             print(data["result"].shape)
 
             viewer.dims.ndisplay = 3
             viewer.scale_bar.visible = True
 
-            if widget.show_original:
+            if widget.show_original and data["original"] is not None:
                 original_layer = viewer.add_image(
                     data["original"],
                     colormap="inferno",