DRY

- Refactored inference on folder/layer to not have horrible duplicates

Author: C-Achard

napari_cellseg3d/model_workers.py

@@ -311,9 +311,8 @@ def load_folder(self):
 
         # TODO : better solution than loading first image always ?
         data_check = LoadImaged(keys=["image"])(images_dict[0])
-        # print(data)
+
         check = data_check["image"].shape
-        # print(check)
         # TODO remove
         # z_aniso = 5 / 1.5
         # if zoom is not None :
@@ -384,7 +383,14 @@ def load_layer(self):
         self.log("Done")
         return input_image
 
-    def model_output(self, inputs, model, post_process_transforms):
+    def model_output(
+        self,
+        inputs,
+        model,
+        post_process_transforms,
+        post_process=True,
+        aniso_transform=None,
+    ):
 
         inputs = inputs.to("cpu")
 
@@ -412,7 +418,103 @@ def model_output(self, inputs, model, post_process_transforms):
         )
 
         out = outputs.detach().cpu()
-        return out
+
+        if aniso_transform is not None:
+            out = aniso_transform(out)
+
+        if post_process:
+            out = np.array(out).astype(np.float32)
+            out = np.squeeze(out)
+            return out
+        else:
+            return out
+
+    def create_result_dict(
+        self,
+        semantic_labels,
+        instance_labels,
+        from_layer: bool,
+        original=None,
+        data_dict=None,
+        i=0,
+    ):
+
+        if not from_layer and original is None:
+            raise ValueError(
+                "If the image is not from a layer, an original should always be available"
+            )
+
+        if from_layer:
+            if i != 0:
+                raise ValueError(
+                    "A layer's ID should always be 0 (default value)"
+                )
+            semantic_labels = np.swapaxes(semantic_labels, 0, 2)
+
+        return {
+            "image_id": i + 1,
+            "original": original,
+            "instance_labels": instance_labels,
+            "object stats": data_dict,
+            "result": semantic_labels,
+            "model_name": self.model_dict["name"],
+        }
+
+    def get_original_filename(self, i):
+        return os.path.basename(self.images_filepaths[i]).split(".")[0]
+
+    def get_instance_result(self, semantic_labels, from_layer=False, i=-1):
+
+        if not from_layer and i == -1:
+            raise ValueError(
+                "An ID should be provided when running from a file"
+            )
+
+        if self.instance_params["do_instance"]:
+            instance_labels = self.instance_seg(
+                semantic_labels,
+                i + 1,
+            )
+            if from_layer:
+                instance_labels = np.swapaxes(instance_labels, 0, 2)
+            data_dict = self.stats_csv(instance_labels)
+        else:
+            instance_labels = None
+            data_dict = None
+        return instance_labels, data_dict
+
+    def save_image(
+        self,
+        image,
+        from_layer=False,
+        i=0,
+    ):
+
+        if not from_layer:
+            original_filename = "_" + self.get_original_filename(i) + "_"
+        else:
+            original_filename = "_"
+
+        time = utils.get_date_time()
+
+        file_path = (
+            self.results_path
+            + "/"
+            + f"Prediction_{i+1}"
+            + original_filename
+            + self.model_dict["name"]
+            + f"_{time}_"
+            + self.filetype
+        )
+
+        imwrite(file_path, image)
+
+        if from_layer:
+            self.log(f"\nLayer prediction saved as :")
+        else:
+            self.log(f"\nFile n°{i+1} saved as :")
+            filename = os.path.split(file_path)[1]
+            self.log(filename)
 
     def aniso_transform(self, image):
         zoom = self.transforms["zoom"][1]
@@ -468,71 +570,35 @@ def method(image):
         return instance_labels
         # print(self.stats_to_csv)
 
-    def inference_on_list(self, inf_data, i, model, post_process_transforms):
+    def inference_on_folder(self, inf_data, i, model, post_process_transforms):
 
         self.log("-" * 10)
         self.log(f"Inference started on image n°{i + 1}...")
 
         inputs = inf_data["image"]
-        # print(inputs.shape)
-
-        out = self.model_output(inputs, model, post_process_transforms)
-        out = self.aniso_transform(out)
 
-        out = np.array(out).astype(np.float32)
-        out = np.squeeze(out)
-        to_instance = out  # avoid post processing since thresholding is done there anyway
-        image_id = i + 1
-        time = utils.get_date_time()
-        # print(time)
-
-        original_filename = os.path.basename(self.images_filepaths[i]).split(
-            "."
-        )[0]
-
-        # File output save name : original-name_model_date+time_number.filetype
-        file_path = (
-            self.results_path
-            + "/"
-            + f"Prediction_{image_id}_"
-            + original_filename
-            + "_"
-            + self.model_dict["name"]
-            + f"_{time}_"
-            + self.filetype
+        out = self.model_output(
+            inputs,
+            model,
+            post_process_transforms,
+            aniso_transform=self.aniso_transform,
         )
 
-        # print(filename)
-        imwrite(file_path, out)
-
-        self.log(f"\nFile n°{image_id} saved as :")
-        filename = os.path.split(file_path)[1]
-        self.log(filename)
-
-        #################
-        #################
-        #################
-        if self.instance_params["do_instance"]:
-            instance_labels = self.instance_seg(
-                to_instance, image_id, original_filename
-            )
-            data_dict = self.stats_csv(instance_labels)
-        else:
-            instance_labels = None
-            data_dict = None
+        self.save_image(out, i=i)
+        instance_labels, data_dict = self.get_instance_result(out, i=i)
 
         original = np.array(inf_data["image"]).astype(np.float32)
 
         self.log(f"Inference completed on layer")
-        result = {
-            "image_id": i + 1,
-            "original": original,
-            "instance_labels": instance_labels,
-            "object stats": data_dict,
-            "result": out,
-            "model_name": self.model_dict["name"],
-        }
-        return result
+
+        return self.create_result_dict(
+            out,
+            instance_labels,
+            from_layer=False,
+            original=original,
+            data_dict=data_dict,
+            i=i,
+        )
 
     def stats_csv(self, instance_labels):
         if self.stats_to_csv:
@@ -554,74 +620,20 @@ def inference_on_layer(self, image, model, post_process_transforms):
         self.log("-" * 10)
         self.log(f"Inference started on layer...")
 
-        # print(inputs.shape)
         image = image.type(torch.FloatTensor)
-        out = self.model_output(image, model, post_process_transforms)
-        out = self.aniso_transform(out)
-        # if self.transforms["zoom"][0]:
-        #     zoom = self.transforms["zoom"][1]
-        #     anisotropic_transform = Zoom(
-        #         zoom=zoom,
-        #         keep_size=False,
-        #         padding_mode="empty",
-        #     )
-        #     out = anisotropic_transform(out[0])
-        ##################
-        ##################
-        ##################
-        out = post_process_transforms(out)
-        out = np.array(out).astype(np.float32)
-        out = np.squeeze(out)
-        to_instance = out  # avoid post processing since thresholding is done there anyway
-
-        # batch_len = out.shape[1]
-        # print("trying to check len")
-        # print(batch_len)
-        # if batch_len != 1 :
-        #     sum  = np.sum(out, axis=1)
-        #     print(sum.shape)
-        #     out = sum
-        #     print(out.shape)
 
-        time = utils.get_date_time()
-        # print(time)
-        # File output save name : original-name_model_date+time_number.filetype
-        file_path = os.path.join(
-            self.results_path,
-            f"Prediction_layer"
-            + "_"
-            + self.model_dict["name"]
-            + f"_{time}_"
-            + self.filetype,
+        out = self.model_output(
+            image,
+            model,
+            post_process_transforms,
+            aniso_transform=self.aniso_transform,
         )
 
-        # print(filename)
-        imwrite(file_path, out)
-
-        self.log(f"\nLayer prediction saved as :")
-        filename = os.path.split(file_path)[1]
-        self.log(filename)
+        self.save_image(out, from_layer=True)
 
-        #################
-        #################
-        #################
-        if self.instance_params["do_instance"]:
-            instance_labels = self.instance_seg(to_instance)
-            instance_labels = np.swapaxes(instance_labels, 0, 2)
-            data_dict = self.stats_csv(instance_labels)
-        else:
-            instance_labels = None
-            data_dict = None
+        instance_labels, data_dict = self.get_instance_result(out,from_layer=True)
 
-        result = {
-            "image_id": 0,
-            "original": None,
-            "instance_labels": instance_labels,
-            "object stats": data_dict,
-            "result": np.swapaxes(out, 0, 2),
-            "model_name": self.model_dict["name"],
-        }
-        return result
+        return self.create_result_dict(out, instance_labels, from_layer=True, data_dict=data_dict)
 
     def inference(self):
         """
@@ -789,7 +801,7 @@ def inference(self):
                 ################################
                 if is_folder:
                     for i, inf_data in enumerate(inference_loader):
-                        yield self.inference_on_list(
+                        yield self.inference_on_folder(
                             inf_data, i, model, post_process_transforms
                         )
                 elif is_layer:
@@ -1328,7 +1340,7 @@ def train(self):
                 f"at epoch: {best_metric_epoch}"
             )
             model.to("cpu")
-            
+
         except Exception as e:
             self.log(f"Error : {e}")
             self.quit()