Add model output and labels to training view

Author: C-Achard

napari_cellseg3d/code_models/worker_training.py

@@ -44,7 +44,6 @@
     PRETRAINED_WEIGHTS_DIR,
     LogSignal,
     QuantileNormalizationd,
-    RemapTensor,
     TrainingReport,
     WeightsDownloader,
 )
@@ -640,14 +639,16 @@ def get_loader_func(num_samples):
                             # TODO : more parameters/flexibility
                             post_pred = Compose(
                                 [
-                                    RemapTensor(new_max=1, new_min=0),
+                                    # RemapTensor(new_max=1, new_min=0),
                                     AsDiscrete(threshold=0.5),  # needed ?
                                     EnsureType(),
                                 ]
                             )  #
                             post_label = EnsureType()
 
-                            output_raw = [RemapTensor(0, 1)(t) for t in pred]
+                            # output_raw = [RemapTensor(0, 1)(t) for t in pred]
+                            output_raw = pred
+
                             val_outputs = [
                                 post_pred(res_tensor) for res_tensor in pred
                             ]
@@ -658,7 +659,7 @@ def get_loader_func(num_samples):
 
                             # logger.debug(len(val_outputs))
                             # logger.debug(len(val_labels))
-                            dice_test = np.array(
+                            dice_test = np.array(  # TODO(cyril): remove
                                 [
                                     utils.dice_coeff(i, j)
                                     for i, j in zip(val_outputs, val_labels)
@@ -673,6 +674,7 @@ def get_loader_func(num_samples):
                         checkpoint_output.append(
                             [
                                 output_raw[0].detach().cpu(),
+                                val_outputs[0].detach().cpu(),
                                 val_inputs[0].detach().cpu(),
                                 val_labels[0].detach().cpu(),
                             ]
@@ -682,7 +684,7 @@ def get_loader_func(num_samples):
                             for channel in checkpoint_output
                             for item in channel
                         ]
-                        checkpoint_output[2] = checkpoint_output[2].astype(
+                        checkpoint_output[3] = checkpoint_output[2].astype(
                             np.uint16
                         )
 

napari_cellseg3d/code_plugins/plugin_model_training.py

@@ -996,13 +996,21 @@ def _display_results(self, images, names, complete_missing=False):
             layer_output = self._viewer.add_image(
                 data=images[0], name=names[0], colormap="turbo"
             )
+            layer_output_discrete = self._viewer.add_image(
+                data=images[1], name=names[1], colormap="bop orange"
+            )
             layer_image = self._viewer.add_image(
-                data=images[1], name=names[1], colormap="inferno"
+                data=images[2], name=names[2], colormap="inferno"
             )
             layer_labels = self._viewer.add_labels(
-                data=images[2], name=names[2]
+                data=images[3], name=names[3]
             )
-            self.result_layers += [layer_output, layer_image, layer_labels]
+            self.result_layers += [
+                layer_output,
+                layer_output_discrete,
+                layer_image,
+                layer_labels,
+            ]
             self._viewer.grid.enabled = True
             self._viewer.dims.ndisplay = 3
             self._viewer.reset_view()
@@ -1018,15 +1026,22 @@ def _display_results(self, images, names, complete_missing=False):
                         )
                         self.result_layers[0] = layer_output
                     elif i == 1:
+                        layer_output_discrete = self._viewer.add_image(
+                            data=images[i],
+                            name=names[i],
+                            colormap="bop orange",
+                        )
+                        self.result_layers[1] = layer_output_discrete
+                    elif i == 2:
                         layer_image = self._viewer.add_image(
                             data=images[i], name=names[i], colormap="inferno"
                         )
-                        self.result_layers[1] = layer_image
+                        self.result_layers[2] = layer_image
                     else:
                         layer_labels = self._viewer.add_labels(
                             data=images[i], name=names[i]
                         )
-                        self.result_layers[2] = layer_labels
+                        self.result_layers[3] = layer_labels
                 self.result_layers[i].data = images[i]
                 self.result_layers[i].refresh()
 
@@ -1042,6 +1057,7 @@ def on_yield(self, report: TrainingReport):
             try:
                 layer_names = [
                     "Validation output",
+                    "Validation output (discrete)",
                     "Validation image",
                     "Validation labels",
                 ]