fix(datasets): ret_binary & ret_inst opts

Author: okunator

cellseg_models_pytorch/datasets/_base_dataset.py

@@ -27,7 +27,8 @@ def __init__(
         img_transforms: List[str],
         inst_transforms: List[str],
         normalization: str = None,
-        return_inst: bool = True,
+        return_inst: bool = False,
+        return_binary: bool = True,
         return_type: bool = True,
         return_sem: bool = False,
         return_weight: bool = False,
@@ -48,7 +49,9 @@ def __init__(
             normalization : str, optional
                 Apply img normalization after all the transformations. One of "minmax",
                 "norm", "percentile", None.
-            return_inst : bool, default=True
+            return_inst : bool, default=False
+                If True, returns the instance labelled mask. (If the db contains these.)
+            return_binary : bool, default=True
                 If True, returns a binarized instance mask. (If the db contains these.)
             return_type : bool, default=True
                 If True, returns a type mask. (If the db contains these.)
@@ -79,6 +82,7 @@ def __init__(
             )
 
         # Return masks
+        self.return_binary = return_binary
         self.return_inst = return_inst
         self.return_type = return_type
         self.return_sem = return_sem
@@ -89,7 +93,7 @@ def __init__(
             img_transforms.append(NORM_TRANSFORMS[normalization]())
 
         inst_transforms = [INST_TRANSFORMS[tr](**kwargs) for tr in inst_transforms]
-        if return_inst:
+        if return_binary:
             inst_transforms.append(INST_TRANSFORMS["binarize"]())
 
         if return_weight:
@@ -142,10 +146,14 @@ def _getitem(
             out[n] = aux_map
 
         # remove redundant target (not needed in downstream).
+
         if self.return_inst:
-            out["inst"] = out["binary"]
-            del out["binary"]
+            out["inst_map"] = out["inst"]
         else:
             del out["inst"]
 
+        if self.return_binary:
+            out["inst"] = out["binary"]
+            del out["binary"]
+
         return out

cellseg_models_pytorch/datasets/folder_dataset_train.py

@@ -21,7 +21,8 @@ def __init__(
         img_transforms: List[str],
         inst_transforms: List[str],
         normalization: str = None,
-        return_inst: bool = True,
+        return_inst: bool = False,
+        return_binary: bool = True,
         return_type: bool = True,
         return_sem: bool = False,
         return_weight: bool = False,
@@ -46,7 +47,9 @@ def __init__(
             normalization : str, optional
                 Apply img normalization after all the transformations. One of "minmax",
                 "norm", "percentile", None.
-            return_inst : bool, default=True
+            return_inst : bool, default=False
+                If True, returns the instance labelled mask. (If the db contains these.)
+            return_binary : bool, default=True
                 If True, returns a binarized instance mask. (If the db contains these.)
             return_type : bool, default=True
                 If True, returns a type mask. (If the db contains these.)
@@ -67,6 +70,7 @@ def __init__(
             return_type=return_type,
             return_sem=return_sem,
             return_weight=return_weight,
+            return_binary=return_binary,
             **kwargs,
         )
 

cellseg_models_pytorch/datasets/hdf5_dataset.py

@@ -24,7 +24,8 @@ def __init__(
         img_transforms: List[str],
         inst_transforms: List[str],
         normalization: str = None,
-        return_inst: bool = True,
+        return_inst: bool = False,
+        return_binary: bool = True,
         return_type: bool = True,
         return_sem: bool = False,
         return_weight: bool = False,
@@ -49,7 +50,9 @@ def __init__(
             normalization : str, optional
                 Apply img normalization after all the transformations. One of "minmax",
                 "norm", "percentile", None.
-            return_inst : bool, default=True
+            return_inst : bool, default=False
+                If True, returns the instance labelled mask. (If the db contains these.)
+            return_binary : bool, default=True
                 If True, returns a binarized instance mask. (If the db contains these.)
             return_type : bool, default=True
                 If True, returns a type mask. (If the db contains these.)
@@ -66,6 +69,7 @@ def __init__(
             return_type=return_type,
             return_sem=return_sem,
             return_weight=return_weight,
+            return_binary=return_binary,
         )
 
         self.path = Path(path)

cellseg_models_pytorch/training/lit/lightning_experiment.py

@@ -140,52 +140,66 @@ def step(
         loss = self.criterion(yhats=soft_masks, targets=targets)
         metrics = self.compute_metrics(soft_masks, targets, phase)
 
+        if batch_idx % self.log_freq == 0:
+            ret_masks = soft_masks
+        elif phase == "test":
+            ret_masks = soft_masks
+        else:
+            ret_masks = None
+
         ret = {
-            "soft_masks": soft_masks if batch_idx % self.log_freq == 0 else None,
+            "soft_masks": ret_masks,
             "loss": loss,
         }
 
         return {**ret, **metrics}
 
-    def training_step(
-        self, batch: Dict[str, torch.Tensor], batch_idx: int
+    def log_step(
+        self, batch: Dict[str, torch.Tensor], batch_idx: int, phase: str
     ) -> Dict[str, torch.Tensor]:
-        """Training step + train metric logs."""
-        res = self.step(batch, batch_idx, "train")
+        """Do the logging."""
+        on_epoch = phase in ("val", "test")
+        on_step = phase == "train"
+        prog_bar = phase == "train"
+
+        res = self.step(batch, batch_idx, phase)
 
         # log all the metrics
-        self.log("train_loss", res["loss"], prog_bar=True, on_epoch=False, on_step=True)
+        self.log(
+            f"{phase}_loss",
+            res["loss"],
+            prog_bar=prog_bar,
+            on_epoch=on_epoch,
+            on_step=on_step,
+        )
+
         for k, val in res.items():
             if k not in ("loss", "soft_masks"):
-                self.log(f"train_{k}", val, prog_bar=True, on_epoch=False, on_step=True)
+                self.log(
+                    f"{phase}_{k}",
+                    val,
+                    prog_bar=prog_bar,
+                    on_epoch=on_epoch,
+                    on_step=on_step,
+                )
 
         return res
 
+    def training_step(
+        self, batch: Dict[str, torch.Tensor], batch_idx: int
+    ) -> Dict[str, torch.Tensor]:
+        """Training step + train metric logs."""
+        return self.log_step(batch, batch_idx, "train")
+
     def validation_step(
         self, batch: Dict[str, torch.Tensor], batch_idx: int
     ) -> Dict[str, torch.Tensor]:
         """Validate step + validation metric logs + example outputs for logging."""
-        res = self.step(batch, batch_idx, "val")
-
-        # log all the metrics
-        self.log("val_loss", res["loss"], prog_bar=False, on_epoch=True, on_step=False)
-        for k, val in res.items():
-            if k not in ("loss", "soft_masks"):
-                self.log(f"val_{k}", val, prog_bar=False, on_epoch=True, on_step=False)
-
-        return res
+        return self.log_step(batch, batch_idx, "val")
 
     def test_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> None:
         """Test step + test metric logs."""
-        res = self.step(batch, batch_idx, "test")
-
-        del res["soft_masks"]  # soft masks not needed for logging
-        loss = res.pop("loss")
-
-        # log all the metrics
-        self.log("test_loss", loss, prog_bar=False, on_epoch=True, on_step=False)
-        for k, val in res.items():
-            self.log(f"test_{k}", val, prog_bar=False, on_epoch=True, on_step=False)
+        return self.log_step(batch, batch_idx, "test")
 
     def compute_metrics(
         self,