fix: fix minor bugs and more verbose errors

Author: okunator

cellseg_models_pytorch/datasets/folder_dataset_train.py

@@ -142,7 +142,8 @@ def read_img_mask(
             out["inst"] = masks["inst_map"]
         except KeyError:
             raise KeyError(
-                f"The file {self.fnames_masks[ix]} does not contain key `inst_map`."
+                f"The file {self.fnames_masks[ix]} does not contain key `inst_map`. "
+                "Try setting `return_inst=False`."
             )
 
         if return_type:
@@ -151,6 +152,7 @@ def read_img_mask(
             except KeyError:
                 raise KeyError(
                     f"The file {self.fnames_masks[ix]} does not contain key `type_map`."
+                    " Try setting `return_type=False`."
                 )
 
         if return_sem:
@@ -159,6 +161,7 @@ def read_img_mask(
             except KeyError:
                 raise KeyError(
                     f"The file {self.fnames_masks[ix]} does not contain key `sem_map`."
+                    "Try setting `return_sem=False`."
                 )
 
         return out

cellseg_models_pytorch/datasets/hdf5_dataset.py

@@ -119,20 +119,27 @@ def read_h5_patch(
                 out["inst"] = h5.root.insts[ix, ...]
             except Exception:
                 raise IOError(
-                    "The HDF5 database does not contain instance labelled masks."
+                    "The HDF5 database does not contain instance labelled masks. Try "
+                    "setting `return_inst=False`"
                 )
 
             if return_type:
                 try:
                     out["type"] = h5.root.types[ix, ...]
                 except Exception:
-                    raise IOError("The HDF5 database does not contain type masks.")
+                    raise IOError(
+                        "The HDF5 database does not contain type masks. Try setting "
+                        "`return_type = False` "
+                    )
 
             if return_sem:
                 try:
                     out["sem"] = h5.root.areas[ix, ...]
                 except Exception:
-                    raise IOError("The HDF5 database does not contain semantic masks.")
+                    raise IOError(
+                        "The HDF5 database does not contain semantic masks. Try "
+                        "setting `return_sem = False`"
+                    )
 
             return out
 

cellseg_models_pytorch/losses/weighted_base_loss.py

@@ -71,6 +71,8 @@ def apply_ls_to_target(
     ) -> torch.Tensor:
         """Apply regular label smoothing to the target map.
 
+        https://arxiv.org/abs/1512.00567
+
         Parameters
         ----------
             target : torch.Tensor

cellseg_models_pytorch/training/callbacks/wandb_callbacks.py

@@ -147,7 +147,7 @@ def train_batch_end(
         outputs: Dict[str, torch.Tensor],
         batch: Dict[str, torch.Tensor],
         batch_idx: int,
-        dataloader_idx: int,
+        dataloader_idx: int = 0,
     ) -> None:
         """Log the inputs and outputs of the model to wandb."""
         self.batch_end(trainer, outputs["soft_masks"], batch, batch_idx, phase="train")
@@ -159,7 +159,7 @@ def validation_batch_end(
         outputs: Dict[str, torch.Tensor],
         batch: Dict[str, torch.Tensor],
         batch_idx: int,
-        dataloader_idx: int,
+        dataloader_idx: int = 0,
     ) -> None:
         """Log the inputs and outputs of the model to wandb."""
         self.batch_end(trainer, outputs["soft_masks"], batch, batch_idx, phase="val")
@@ -171,7 +171,7 @@ def test_batch_end(
         outputs: Dict[str, torch.Tensor],
         batch: Dict[str, torch.Tensor],
         batch_idx: int,
-        dataloader_idx: int,
+        dataloader_idx: int = 0,
     ) -> None:
         """Log the inputs and outputs of the model to wandb."""
         self.batch_end(trainer, outputs["soft_masks"], batch, batch_idx, phase="test")
@@ -198,6 +198,7 @@ def batch_end(
         batch: Dict[str, torch.Tensor],
         batch_idx: int,
         phase: str,
+        dataloader_idx: int = None,
     ) -> None:
         """Log metrics at every 100th step to wandb."""
         if batch_idx % self.freq == 0:
@@ -223,7 +224,7 @@ def on_validation_batch_end(
         outputs: Dict[str, torch.Tensor],
         batch: Dict[str, torch.Tensor],
         batch_idx: int,
-        dataloader_idx: int,
+        dataloader_idx: int = 0,
     ) -> None:
         """Call the callback at val time."""
         self.validation_batch_end(
@@ -237,7 +238,7 @@ def on_train_batch_end(
         outputs: Dict[str, torch.Tensor],
         batch: Dict[str, torch.Tensor],
         batch_idx: int,
-        dataloader_idx: int,
+        dataloader_idx: int = 0,
     ) -> None:
         """Call the callback at val time."""
         self.train_batch_end(

cellseg_models_pytorch/training/tests/test_training.py

@@ -9,8 +9,7 @@
 from cellseg_models_pytorch.training.lit import SegmentationExperiment
 
 
-# @pytest.mark.parametrize
-def test_training(img_patch_dir, mask_patch_dir):
+def test_training(img_patch_dir, mask_patch_dir, auto_lr):
     train_ds = SegmentationFolderDataset(
         path=img_patch_dir.as_posix(),
         mask_path=mask_patch_dir.as_posix(),