Reintroduced best Dice channel seeking + refacto

Author: C-Achard

napari_cellseg3d/code_models/worker_training.py

@@ -779,33 +779,17 @@ def eval(self, model, epoch) -> TrainingReport:
                     f"Val decoder outputs shape: {val_decoder_outputs.shape}"
                 )
 
-                # dices = []
-                # Find in which channel the labels are (avoid background)
-                # for channel in range(val_outputs.shape[1]):
-                #     dices.append(
-                #         utils.dice_coeff(
-                #             y_pred=val_outputs[
-                #                 0, channel : (channel + 1), :, :, :
-                #             ],
-                #             y_true=val_labels[0],
-                #         )
-                #     )
-                # logger.debug(f"DICE COEFF: {dices}")
-                # max_dice_channel = torch.argmax(
-                #     torch.Tensor(dices)
-                # )
-                # logger.debug(
-                #     f"MAX DICE CHANNEL: {max_dice_channel}"
-                # )
+                max_dice_channel = utils.seek_best_dice_coeff_channel(
+                    y_pred=val_outputs, y_true=val_labels
+                )
                 self.dice_metric(
-                    y_pred=val_outputs,
-                    # [
-                    #     :,
-                    #     max_dice_channel : (max_dice_channel + 1),
-                    #     :,
-                    #     :,
-                    #     :,
-                    # ],
+                    y_pred=val_outputs[
+                        :,
+                        max_dice_channel : (max_dice_channel + 1),
+                        :,
+                        :,
+                        :,
+                    ],
                     y=val_labels,
                 )
 
@@ -1282,7 +1266,7 @@ def get_patch_loader_func(num_samples):
                 batch_size=self.config.batch_size,
                 num_workers=self.config.num_workers,
             )
-            logger.info("\nDone")
+            logger.debug("\nDone")
 
             logger.debug("Optimizer")
             optimizer = (

napari_cellseg3d/utils.py

@@ -229,6 +229,29 @@ def dice_coeff(
     )
 
 
+def seek_best_dice_coeff_channel(y_pred, y_true) -> torch.Tensor:
+    """Compute Dice-Sorensen coefficient between unsupervised model output and ground truth labels;
+    returns the channel with the highest dice coefficient.
+    Args:
+        y_true: Ground truth label
+        y_pred: Prediction label
+    Returns: best Dice coefficient channel
+    """
+    dices = []
+    # Find in which channel the labels are (to avoid background)
+    for channel in range(y_pred.shape[1]):
+        dices.append(
+            dice_coeff(
+                y_pred=y_pred[0, channel : (channel + 1), :, :, :],
+                y_true=y_true[0],
+            )
+        )
+    LOGGER.debug(f"DICE COEFF: {dices}")
+    max_dice_channel = torch.argmax(torch.Tensor(dices))
+    LOGGER.debug(f"MAX DICE CHANNEL: {max_dice_channel}")
+    return max_dice_channel
+
+
 def correct_rotation(image):
     """Rotates the exes 0 and 2 in [DHW] section of image array"""
     extra_dims = len(image.shape) - 3