added typecheck workflow and fixed the errors

.github/workflows/lint.yml

@@ -31,3 +31,40 @@ jobs:
           python-version: "3.12"
 
       - uses: pre-commit/action@v3.0.1
+
+  # mypy is a CI job rather than a pre-commit hook because it needs the project's
+  # dependencies importable to be worth anything: without kornia and torch installed
+  # every annotation resolves to Any and the check passes while checking nothing.
+  # That install is too heavy for a hook contributors run on every commit.
+  #
+  # Config lives in [tool.mypy] in pyproject.toml. `smauglab/` only -- unit_tests/
+  # and scripts/ are excluded there.
+  typecheck:
+    name: mypy
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          # poetry-dynamic-versioning derives the version from the git tag;
+          # a shallow clone has no tags and would build as 0.0.0.
+          fetch-depth: 0
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: pip
+          cache-dependency-path: pyproject.toml
+
+      - name: Install CPU-only PyTorch
+        # Must come first, same as in tests.yml: the default index serves the CUDA
+        # build (several GB of nvidia-* wheels), which is slow and can fill the
+        # runner disk. Nothing here needs a GPU.
+        run: |
+          python -m pip install --upgrade pip
+          pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+
+      - name: Install SmaugLab
+        run: pip install -e ".[dev]"
+
+      - name: Run mypy
+        run: mypy smauglab/

pyproject.toml

@@ -76,6 +76,7 @@ wandb = { version = "*", optional = true }
 nnunetv2 = { version = "*", optional = true }
 build = { version = "*", optional = true }
 coverage = { version = ">=7", optional = true }
+mypy = { version = ">=1.10", optional = true }
 pre-commit = { version = "*", optional = true }
 pytest = { version = ">=8", optional = true }
 pytest-cov = { version = "*", optional = true }
@@ -88,7 +89,7 @@ twine = { version = "*", optional = true }
 [tool.poetry.extras]
 nnunetv2 = ["nnunetv2"]
 all = ["monai", "tqdm", "wandb"]
-dev = ["build", "coverage", "pre-commit", "pytest", "pytest-cov", "ruff", "twine"]
+dev = ["build", "coverage", "mypy", "pre-commit", "pytest", "pytest-cov", "ruff", "twine"]
 
 [tool.poetry.scripts]
 smauglab_add_nnunettrainer = "smauglab.add_trainer:main"
@@ -238,6 +239,28 @@ docstring-code-format = true
 # Like Black, automatically detect the appropriate line ending.
 line-ending = "auto"
 
+[tool.mypy]
+# Checked in CI by the `typecheck` job in .github/workflows/lint.yml, which has to
+# install the real dependencies first: with kornia and torch absent every type
+# collapses to Any and the run passes while checking nothing.
+#
+# The floor of the supported range, not the version CI happens to run on -- this is
+# what `requires-python = ">=3.10"` promises.
+python_version = "3.10"
+# Pragmatic rather than strict. Nothing in the dependency stack ships stubs
+# (kornia, batchgeneratorsv2, torchio, nnunetv2, nibabel, progress), and full
+# annotation coverage is not enforced -- tighten with disallow_untyped_defs
+# incrementally rather than in one pass.
+ignore_missing_imports = true
+warn_redundant_casts = true
+# smauglab/, smauglab/transforms/ and smauglab/utils/ have no __init__.py (PEP 420,
+# see the `packages` note above), so mypy cannot derive module names without these.
+namespace_packages = true
+explicit_package_bases = true
+# Only the shipped package is gated. Tests and standalone scripts are not installed
+# and assert against literals, which reads badly to a type checker.
+exclude = ["unit_tests/", "scripts/"]
+
 [tool.pytest.ini_options]
 testpaths = ["unit_tests"]
 filterwarnings = [

smauglab/add_trainer.py

@@ -1,6 +1,7 @@
 import argparse
 import importlib.resources
 import shutil
+from pathlib import Path
 
 import nnunetv2
 
@@ -30,8 +31,13 @@ def main():
 
 def add_trainer(trainer_name: str, overwrite: bool = False):
 
-    # Find trainer path
-    trainers_path = importlib.resources.files(trainers)
+    # Find trainer path.
+    # importlib.resources returns a Traversable, which only promises open()/read_bytes()
+    # -- not .exists(), and not something shutil.copy accepts. Copying *into* the
+    # installed nnunetv2 package needs a real directory on disk regardless (nnU-Net
+    # cannot run from a zipped install), so resolve both ends to concrete paths here.
+    # Same idiom as unit_tests/helpers.py.
+    trainers_path = Path(str(importlib.resources.files(trainers)))
     if trainer_name == "nnUNetTrainerDAExt":
         source_trainer = trainers_path / "nnUNetTrainerDAExt.py"
     elif trainer_name == "nnUNetTrainerTest":
@@ -40,7 +46,7 @@ def add_trainer(trainer_name: str, overwrite: bool = False):
         raise ValueError(f"Trainer {trainer_name} not recognized.")
 
     # Find nnUNet path
-    nnunetv2_path = importlib.resources.files(nnunetv2)
+    nnunetv2_path = Path(str(importlib.resources.files(nnunetv2)))
     nnunet_trainers_path = nnunetv2_path / "training" / "nnUNetTrainer"
 
     # Copy trainer

smauglab/trainers/nnUNetTrainerDAExt.py

@@ -44,7 +44,7 @@ def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self):
 
     @staticmethod
     def get_training_transforms(
-        patch_size: Union[np.ndarray, tuple[int]],
+        patch_size: Union[np.ndarray, tuple[int, ...]],
         rotation_for_DA: RandomScalar,
         deep_supervision_scales: Union[list, tuple, None],
         mirror_axes: tuple[int, ...],
@@ -173,7 +173,7 @@ def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self):
 
     @staticmethod
     def get_training_transforms(
-        patch_size: Union[np.ndarray, tuple[int]],
+        patch_size: Union[np.ndarray, tuple[int, ...]],
         rotation_for_DA: RandomScalar,
         deep_supervision_scales: Union[list, tuple, None],
         mirror_axes: tuple[int, ...],
@@ -371,7 +371,7 @@ def configure_rotation_dummyDA_mirroring_and_inital_patch_size(self):
 
     @staticmethod
     def get_training_transforms(
-        patch_size: Union[np.ndarray, tuple[int]],
+        patch_size: Union[np.ndarray, tuple[int, ...]],
         rotation_for_DA: RandomScalar,
         deep_supervision_scales: Union[list, tuple, None],
         mirror_axes: tuple[int, ...],

smauglab/trainers/nnUNetTrainerTest.py

@@ -32,7 +32,7 @@ def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dic
 
     @staticmethod
     def get_training_transforms(
-        patch_size: Union[np.ndarray, tuple[int]],
+        patch_size: Union[np.ndarray, tuple[int, ...]],
         rotation_for_DA: RandomScalar,
         deep_supervision_scales: Union[list, tuple, None],
         mirror_axes: tuple[int, ...],
@@ -136,7 +136,7 @@ def __init__(self, plans: dict, configuration: str, fold: int, dataset_json: dic
 
     @staticmethod
     def get_training_transforms(
-        patch_size: Union[np.ndarray, tuple[int]],
+        patch_size: Union[np.ndarray, tuple[int, ...]],
         rotation_for_DA: RandomScalar,
         deep_supervision_scales: Union[list, tuple, None],
         mirror_axes: tuple[int, ...],

smauglab/trainers/utils.py

@@ -1,3 +1,4 @@
+from collections.abc import Sequence
 from typing import Union
 
 import torch
@@ -27,6 +28,8 @@ def __call__(self, segmentation: torch.Tensor) -> list[torch.Tensor]:
             List of downsampled tensors, each with shape [batch, channels, spatial_dims...]
         """
         results = []
+        # Per-axis scale factors: either broadcast from a scalar or taken as given.
+        s: Sequence[float]
         for ds_scale in self.ds_scales:
             if not isinstance(ds_scale, (tuple, list)):
                 # If single scale value, apply to all spatial dimensions

smauglab/transforms/cpu/artifact.py

@@ -49,7 +49,7 @@ def apply(self, data_dict: dict, **params) -> dict:
             data_dict["image"], data_dict["segmentation"] = self._apply_to_image(data_dict["image"], data_dict["segmentation"], **params)
         return data_dict
 
-    def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> torch.Tensor:
+    def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> tuple[torch.Tensor, torch.Tensor]:
         if params["motion"]:
             img, seg = aug_motion(img, seg)
         if params["ghosting"]:

smauglab/transforms/cpu/contrast.py

@@ -1,3 +1,5 @@
+from typing import Union
+
 import torch
 import torch.nn.functional as F
 from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
@@ -20,6 +22,9 @@ def __init__(self, kernel_type: str = "Laplace", absolute: bool = False, retain_
         self.retain_stats = retain_stats
 
     def get_parameters(self, **data_dict) -> dict:
+        # Scharr yields one kernel per spatial direction, Laplace a single kernel;
+        # _apply_to_image dispatches on kernel_type to tell the two apart.
+        kernel: Union[torch.Tensor, list[torch.Tensor]]
         spatial_dims = len(data_dict["image"].shape) - 1
         if spatial_dims == 2:
             if self.kernel_type == "Laplace":

smauglab/transforms/cpu/fromSeg.py

@@ -27,7 +27,7 @@ def apply(self, data_dict: dict, **params) -> dict:
             data_dict["image"], data_dict["segmentation"] = self._apply_to_image(data_dict["image"], data_dict["segmentation"], **params)
         return data_dict
 
-    def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> torch.Tensor:
+    def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> tuple[torch.Tensor, torch.Tensor]:
         for c in range(1):  # Works on the first channel only
             img[c], seg[c] = aug_redistribute_seg(
                 img[c], seg[c], classes=params["classes"], in_seg=params["in_seg"], retain_stats=params["retain_stats"]

smauglab/transforms/cpu/spatial.py

@@ -37,7 +37,7 @@ def apply(self, data_dict: dict, **params) -> dict:
             data_dict["image"], data_dict["segmentation"] = self._apply_to_image(data_dict["image"], data_dict["segmentation"], **params)
         return data_dict
 
-    def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> torch.Tensor:
+    def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> tuple[torch.Tensor, torch.Tensor]:
         if params["flip"]:
             img, seg = aug_flip(img, seg)
         if params["affine"]:
@@ -153,7 +153,7 @@ def apply(self, data_dict: dict, **params) -> dict:
             data_dict["image"], data_dict["segmentation"] = self._apply_to_image(data_dict["image"], data_dict["segmentation"], **params)
         return data_dict
 
-    def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> torch.Tensor:
+    def _apply_to_image(self, img: torch.Tensor, seg: torch.Tensor, **params) -> tuple[torch.Tensor, torch.Tensor]:
         # Compute random shape
         img_shape = img.shape[1:]
         new_shape = [random.randint(params["shape_min"], s) if i not in params["ignore_axes"] else s for i, s in enumerate(img_shape)]

smauglab/transforms/cpu/transforms.py

@@ -28,9 +28,14 @@ def __init__(
         self,
         json_path: str,
         do_dummy_2d_data_aug: bool,
-        patch_size: Union[np.ndarray, tuple[int]],
+        # tuple[int, ...], not tuple[int]: these are 3D patch shapes and axis
+        # tuples, so the one-element form was never what callers pass.
+        patch_size: Union[np.ndarray, tuple[int, ...]],
         rotation_for_DA: RandomScalar,
-        mirror_axes: tuple[int],
+        # Accepted for signature compatibility with the nnU-Net trainers but
+        # unused -- the mirror axes are read from the JSON config instead, see
+        # _build_transforms below. Callers legitimately pass None.
+        mirror_axes: tuple[int, ...] | None,
     ):
         # Load transform parameters from JSON
         config_path = os.path.join(json_path)
@@ -48,7 +53,11 @@ def __init__(
         super().__init__(transforms=self.transforms)
 
     def _build_transforms(
-        self, do_dummy_2d_data_aug: bool, patch_size: Union[np.ndarray, tuple[int]], rotation_for_DA: RandomScalar, mirror_axes: tuple[int]
+        self,
+        do_dummy_2d_data_aug: bool,
+        patch_size: Union[np.ndarray, tuple[int, ...]],
+        rotation_for_DA: RandomScalar,
+        mirror_axes: tuple[int, ...] | None,
     ):
         transform_params = self.transform_params
         transforms = []
@@ -341,7 +350,7 @@ def _build_transforms(self):
     from smauglab.utils.utils import normalize
 
     configs_path = importlib.resources.files(configs)
-    json_path = configs_path / "transform_params_hybrid_TAGE.json"
+    json_path = str(configs_path / "transform_params_hybrid_TAGE.json")
 
     # Load images and masks tensors
     img_path = "/home/GRAMES.POLYMTL.CA/p118739/data_nvme_p118739/data/datasets/data-multi-subject/sub-amu02/anat/sub-amu02_T1w.nii.gz"
@@ -384,8 +393,8 @@ def _build_transforms(self):
     nb_col = 6
     for key in ["image", "segmentation"]:
         output = []
-        line = []
-        aug = [[]]
+        line: list[np.ndarray] = []
+        aug: list[list[str]] = [[]]
         for _idx, (augment, _dic) in enumerate(tensor_dict.items()):
             if len(line) < nb_col:
                 img = 255 * normalize(np.sum(tensor_dict[augment][key].detach().numpy(), axis=0, keepdims=True)[0, 64])

smauglab/transforms/gpu/base.py

@@ -29,6 +29,10 @@
 
 DataType = Union[Tensor, list[Tensor], Boxes, Keypoints]
 SequenceDataType = Union[list[Tensor], list[list[Tensor]], list[Boxes], list[Keypoints]]
+# Anything AugmentationSequentialCustom accepts as a stage. Every smauglab GPU
+# transform qualifies, via ImageOnlyTransform or RigidAffineAugmentationBase3D.
+# Narrower than nn.Module, which is what lets the `*transforms` splat type-check.
+TransformType = Union[_AugmentationBase, ImageSequential]
 
 
 class ImageOnlyTransform(RigidAffineAugmentationBase3D):
@@ -305,7 +309,14 @@ def transform(
         keys = [dk.name for dk in _data_keys]
         if "MASK" in keys:
             mask_index = keys.index("MASK")
-            param.data["seg"] = arg[mask_index]
+            # kornia types ParamItem.data as dict | list[ParamItem] | None and the
+            # inputs as the wider DataType, but the MASK entry of a leaf
+            # augmentation is always a params dict holding a plain tensor. Asserted
+            # rather than ignored so a violated assumption still fails loudly.
+            mask = arg[mask_index]
+            assert isinstance(param.data, dict)
+            assert isinstance(mask, Tensor)
+            param.data["seg"] = mask
 
         outputs = []
         for inp, dcate in zip(arg, _data_keys):