add ci

.devcontainer/devcontainer.json

@@ -1,33 +0,0 @@
-{
-  "name": "Python 3",
-  // Or use a Dockerfile or Docker Compose file. More info: https://containers.dev/guide/dockerfile
-  "image": "mcr.microsoft.com/devcontainers/python:1-3.11-bookworm",
-  "customizations": {
-    "codespaces": {
-      "openFiles": [
-        "README.md",
-        "app.py"
-      ]
-    },
-    "vscode": {
-      "settings": {},
-      "extensions": [
-        "ms-python.python",
-        "ms-python.vscode-pylance"
-      ]
-    }
-  },
-  "updateContentCommand": "[ -f packages.txt ] && sudo apt update && sudo apt upgrade -y && sudo xargs apt install -y <packages.txt; [ -f requirements.txt ] && pip3 install --user -r requirements.txt; pip3 install --user streamlit; echo '✅ Packages installed and Requirements met'",
-  "postAttachCommand": {
-    "server": "streamlit run app.py --server.enableCORS false --server.enableXsrfProtection false"
-  },
-  "portsAttributes": {
-    "8501": {
-      "label": "Application",
-      "onAutoForward": "openPreview"
-    }
-  },
-  "forwardPorts": [
-    8501
-  ]
-}

.github/workflows/ci.yaml

@@ -0,0 +1,25 @@
+name: CI
+
+on:
+  push
+
+jobs:
+  quality-assurance:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.9'
+          cache: 'pip' # Speeds up subsequent runs
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements.txt
+
+      - name: Run CI Suite
+        run: make check

.gitignore

@@ -6,4 +6,6 @@ temp.ipynb
 __pycache__/
 **/__pycache__/
 *.pyc
-.ruff_cache
+.ruff_cache/
+.mypy_cache/
+.pytest_cache/

Makefile

@@ -0,0 +1,30 @@
+.PHONY: format lint typecheck test check
+
+format:
+	ruff format .
+
+lint:
+	ruff check . --fix
+
+typecheck:
+	mypy .
+
+test:
+	pytest
+
+clean:
+	@echo "Cleaning project..."
+	# Delete compiled bytecode
+	@python3 -Bc "import pathlib; [p.unlink() for p in pathlib.Path('.').rglob('*.py[co]')]"
+	# Delete directory-based caches
+	@python3 -Bc "import shutil, pathlib; \
+		[shutil.rmtree(p) for p in pathlib.Path('.').rglob('__pycache__')]; \
+		[shutil.rmtree(p) for p in pathlib.Path('.').rglob('.ipynb_checkpoints')]; \
+		[shutil.rmtree(p) for p in pathlib.Path('.').rglob('.monai-cache')]; \
+		[shutil.rmtree(p) for p in pathlib.Path('.').rglob('.mypy_cache')]; \
+		[shutil.rmtree(p) for p in pathlib.Path('.').rglob('.ruff_cache')]; \
+		[shutil.rmtree(p) for p in pathlib.Path('.').rglob('.pytest_cache')]"
+
+# Updated 'check' to clean before running (optional)
+# This ensures you are testing from a "blank slate"
+check: format lint typecheck test clean

preprocess_main.py

@@ -1,14 +1,15 @@
-import os
-from src.preprocessing.register_and_crop import register_files
-from src.preprocessing.prostate_mask import get_segmask
-from src.preprocessing.histogram_match import histmatch
-from src.preprocessing.generate_heatmap import get_heatmap
-import logging
-from src.utils import setup_logging
-from src.utils import validate_steps
 import argparse
+import logging
+import os
+
 import yaml
 
+from src.preprocessing.generate_heatmap import get_heatmap
+from src.preprocessing.histogram_match import histmatch
+from src.preprocessing.prostate_mask import get_segmask
+from src.preprocessing.register_and_crop import register_files
+from src.utils import setup_logging, validate_steps
+
 
 def parse_args():
     parser = argparse.ArgumentParser(description="File preprocessing")
@@ -37,7 +38,7 @@ def parse_args():
 
     args = parser.parse_args()
     if args.config:
-        with open(args.config, "r") as config_file:
+        with open(args.config) as config_file:
             config = yaml.safe_load(config_file)
             args.__dict__.update(config)
     return args

pyproject.toml

@@ -1,9 +1,51 @@
 [tool.ruff]
 line-length = 100
+target-version = "py39" # Ensures ruff uses 3.9 compatible syntax
 
 [tool.ruff.lint]
-select = ["E", "W"]
+select = ["E", "W", "F", "I", "B", "N", "UP"]
 ignore = ["E501"]
 
+[tool.ruff.lint.isort]
+# This makes your imports look like a professional library
+combine-as-imports = true
+force-wrap-aliases = true
+known-first-party = ["src"] # Treats your 'src' folder as a local package
+
 [tool.ruff.format]
-quote-style = "double"
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+
+[tool.ruff.lint.pep8-naming]
+# Add your custom class name here
+ignore-names = ["sitk", "NormalizeIntensity_custom", "NormalizeIntensity_customd"]
+
+[tool.mypy]
+ignore_missing_imports = true
+disable_error_code = ["override", "import-untyped"]
+mypy_path = "."
+pretty = true
+show_error_codes = true
+
+[[tool.mypy.overrides]]
+# These settings apply specifically to these external libraries
+module = [
+    "yaml.*",
+    "nrrd.*",
+    "nibabel.*",
+    "scipy.*",
+    "sklearn.*"
+]
+ignore_missing_imports = true
+
+[tool.pytest.ini_options]
+# Automatically adds these flags every time you run 'pytest'
+addopts = "-v --showlocals --durations=5"
+
+# Where pytest should look for tests
+testpaths = ["tests"]
+
+# Patterns to identify test files
+python_files = "test_*.py"
+python_functions = "test_*"

run_cspca.py

@@ -1,22 +1,23 @@
 import argparse
+import logging
 import os
 import shutil
-import yaml
 import sys
-import torch
 from pathlib import Path
+
+import torch
+import yaml
 from monai.utils import set_determinism
-import logging
-from src.model.MIL import MILModel_3D
-from src.model.csPCa_model import csPCa_Model
+
 from src.data.data_loader import get_dataloader
-from src.utils import save_cspca_checkpoint, get_metrics, setup_logging
+from src.model.cspca_model import CSPCAModel
+from src.model.mil import MILModel3D
 from src.train.train_cspca import train_epoch, val_epoch
-import random
+from src.utils import get_metrics, save_cspca_checkpoint, setup_logging
 
 
 def main_worker(args):
-    mil_model = MILModel_3D(num_classes=args.num_classes, mil_mode=args.mil_mode)
+    mil_model = MILModel3D(num_classes=args.num_classes, mil_mode=args.mil_mode)
     cache_dir_path = Path(os.path.join(args.logdir, "cache"))
 
     if args.mode == "train":
@@ -31,7 +32,7 @@ def main_worker(args):
 
         train_loader = get_dataloader(args, split="train")
         valid_loader = get_dataloader(args, split="test")
-        cspca_model = csPCa_Model(backbone=mil_model).to(args.device)
+        cspca_model = CSPCAModel(backbone=mil_model).to(args.device)
         for submodule in [
             cspca_model.backbone.net,
             cspca_model.backbone.myfc,
@@ -45,23 +46,17 @@ def main_worker(args):
         )
 
         old_loss = float("inf")
-        old_auc = 0.0
         for epoch in range(args.epochs):
             train_loss, train_auc = train_epoch(
                 cspca_model, train_loader, optimizer, epoch=epoch, args=args
             )
-            logging.info(
-                f"EPOCH {epoch} TRAIN loss: {train_loss:.4f} AUC: {train_auc:.4f}"
-            )
+            logging.info(f"EPOCH {epoch} TRAIN loss: {train_loss:.4f} AUC: {train_auc:.4f}")
             val_metric = val_epoch(cspca_model, valid_loader, epoch=epoch, args=args)
             logging.info(
                 f"EPOCH {epoch} VAL loss: {val_metric['loss']:.4f} AUC: {val_metric['auc']:.4f}"
             )
             if val_metric["loss"] < old_loss:
                 old_loss = val_metric["loss"]
-                old_auc = val_metric["auc"]
-                sensitivity = val_metric["sensitivity"]
-                specificity = val_metric["specificity"]
                 save_cspca_checkpoint(cspca_model, val_metric, model_dir)
 
         args.checkpoint_cspca = os.path.join(model_dir, "cspca_model.pth")
@@ -69,7 +64,7 @@ def main_worker(args):
             shutil.rmtree(cache_dir_path)
 
 
-    cspca_model = csPCa_Model(backbone=mil_model).to(args.device)
+    cspca_model = CSPCAModel(backbone=mil_model).to(args.device)
     checkpt = torch.load(args.checkpoint_cspca, map_location="cpu")
     cspca_model.load_state_dict(checkpt["state_dict"])
     cspca_model = cspca_model.to(args.device)
@@ -96,6 +91,7 @@ def main_worker(args):
     get_metrics(metrics_dict)
 
 
+
 def parse_args():
     parser = argparse.ArgumentParser(
         description="Multiple Instance Learning (MIL) for csPCa risk prediction."
@@ -164,7 +160,7 @@ def parse_args():
     )
     args = parser.parse_args()
     if args.config:
-        with open(args.config, "r") as config_file:
+        with open(args.config) as config_file:
             config = yaml.safe_load(config_file)
             args.__dict__.update(config)
 
@@ -172,12 +168,13 @@ def parse_args():
 
 
 if __name__ == "__main__":
-
     args = parse_args()
     if args.project_dir is None:
-        args.project_dir = Path(__file__).resolve().parent # Set project directory
+        args.project_dir = Path(__file__).resolve().parent  # Set project directory
 
-    slurm_job_name = os.getenv('SLURM_JOB_NAME') # If the script is submitted via slurm, job name is the run name
+    slurm_job_name = os.getenv(
+        "SLURM_JOB_NAME"
+    )  # If the script is submitted via slurm, job name is the run name
     if slurm_job_name:
         args.run_name = slurm_job_name
 
@@ -207,10 +204,15 @@ if __name__ == "__main__":
 
     if args.dry_run:
         logging.info("Dry run mode enabled.")
-        args.epochs = 2
+        args.epochs = 1
         args.batch_size = 2
         args.workers = 0
-        args.num_seeds = 2
+        args.num_seeds = 1
         args.wandb = False
+        args.tile_size = 10
+        args.tile_count = 5
 
     main_worker(args)
+
+    if args.dry_run:
+        shutil.rmtree(args.logdir)

run_inference.py

@@ -1,18 +1,20 @@
 import argparse
+import json
+import logging
 import os
-import yaml
+
 import torch
-import logging
-from src.model.MIL import MILModel_3D
-from src.model.csPCa_model import csPCa_Model
-from src.utils import setup_logging, get_parent_image, get_patch_coordinate
-from src.preprocessing.register_and_crop import register_files
-from src.preprocessing.prostate_mask import get_segmask
-from src.preprocessing.histogram_match import histmatch
-from src.preprocessing.generate_heatmap import get_heatmap
-from src.data.data_loader import data_transform, list_data_collate
+import yaml
 from monai.data import Dataset
-import json
+
+from src.data.data_loader import data_transform, list_data_collate
+from src.model.cspca_model import CSPCAModel
+from src.model.mil import MILModel3D
+from src.preprocessing.generate_heatmap import get_heatmap
+from src.preprocessing.histogram_match import histmatch
+from src.preprocessing.prostate_mask import get_segmask
+from src.preprocessing.register_and_crop import register_files
+from src.utils import get_parent_image, get_patch_coordinate, setup_logging
 
 
 def parse_args():
@@ -36,7 +38,7 @@ def parse_args():
 
     args = parser.parse_args()
     if args.config:
-        with open(args.config, "r") as config_file:
+        with open(args.config) as config_file:
             config = yaml.safe_load(config_file)
             args.__dict__.update(config)
     return args
@@ -64,14 +66,14 @@ if __name__ == "__main__":
     args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     logging.info("Loading PIRADS model")
-    pirads_model = MILModel_3D(num_classes=args.num_classes, mil_mode=args.mil_mode)
+    pirads_model = MILModel3D(num_classes=args.num_classes, mil_mode=args.mil_mode)
     pirads_checkpoint = torch.load(
         os.path.join(args.project_dir, "models", "pirads.pt"), map_location="cpu"
     )
     pirads_model.load_state_dict(pirads_checkpoint["state_dict"])
     pirads_model.to(args.device)
     logging.info("Loading csPCa model")
-    cspca_model = csPCa_Model(backbone=pirads_model).to(args.device)
+    cspca_model = CSPCAModel(backbone=pirads_model).to(args.device)
     checkpt = torch.load(
         os.path.join(args.project_dir, "models", "cspca_model.pth"), map_location="cpu"
     )
@@ -109,7 +111,7 @@ if __name__ == "__main__":
     cspca_model.eval()
     patches_top_5_list = []
     with torch.no_grad():
-        for idx, batch_data in enumerate(loader):
+        for _, batch_data in enumerate(loader):
             data = batch_data["image"].as_subclass(torch.Tensor).to(args.device)
             logits = pirads_model(data)
             pirads_score = torch.argmax(logits, dim=1)
@@ -137,10 +139,10 @@ if __name__ == "__main__":
                 patches_top_5.append(patch_temp)
             patches_top_5_list.append(patches_top_5)
     coords_list = []
-    for i in args.data_list:
+    for j, i in enumerate(args.data_list):
         parent_image = get_parent_image([i], args)
 
-        coords = get_patch_coordinate(patches_top_5, parent_image)
+        coords = get_patch_coordinate(patches_top_5_list[j], parent_image)
         coords_list.append(coords)
     output_dict = {}
 

run_pirads.py

@@ -1,19 +1,21 @@
 import argparse
+import logging
 import os
 import shutil
-import time
-import yaml
 import sys
+import time
+from pathlib import Path
+
 import numpy as np
 import torch
-from torch.utils.tensorboard import SummaryWriter
-from monai.utils import set_determinism
 import wandb
-import logging
-from pathlib import Path
+import yaml
+from monai.utils import set_determinism
+from torch.utils.tensorboard import SummaryWriter
+
 from src.data.data_loader import get_dataloader
+from src.model.mil import MILModel3D
 from src.train.train_pirads import train_epoch, val_epoch
-from src.model.MIL import MILModel_3D
 from src.utils import save_pirads_checkpoint, setup_logging
 
 
@@ -21,7 +23,7 @@ def main_worker(args):
     if args.device == torch.device("cuda"):
         torch.backends.cudnn.benchmark = True
 
-    model = MILModel_3D(num_classes=args.num_classes, mil_mode=args.mil_mode)
+    model = MILModel3D(num_classes=args.num_classes, mil_mode=args.mil_mode)
     start_epoch = 0
     best_acc = 0.0
     if args.checkpoint is not None:
@@ -250,19 +252,20 @@ def parse_args():
     )
     args = parser.parse_args()
     if args.config:
-        with open(args.config, "r") as config_file:
+        with open(args.config) as config_file:
             config = yaml.safe_load(config_file)
             args.__dict__.update(config)
     return args
 
 
 if __name__ == "__main__":
-
     args = parse_args()
     if args.project_dir is None:
-        args.project_dir = Path(__file__).resolve().parent # Set project directory
+        args.project_dir = Path(__file__).resolve().parent  # Set project directory
 
-    slurm_job_name = os.getenv('SLURM_JOB_NAME') # If the script is submitted via slurm, job name is the run name
+    slurm_job_name = os.getenv(
+        "SLURM_JOB_NAME"
+    )  # If the script is submitted via slurm, job name is the run name
     if slurm_job_name:
         args.run_name = slurm_job_name
 
@@ -288,11 +291,13 @@ if __name__ == "__main__":
 
     if args.dry_run:
         logging.info("Dry run mode enabled.")
-        args.epochs = 2
+        args.epochs = 1
         args.batch_size = 2
         args.workers = 0
-        args.num_seeds = 2
+        args.num_seeds = 1
         args.wandb = False
+        args.tile_size = 10
+        args.tile_count = 5
 
     mode_wandb = "online" if args.wandb and args.mode != "test" else "disabled"
 
@@ -314,3 +319,6 @@ if __name__ == "__main__":
     main_worker(args)
 
     wandb.finish()
+
+    if args.dry_run:
+        shutil.rmtree(args.logdir)

src/data/custom_transforms.py

@@ -1,18 +1,19 @@
+from collections.abc import Hashable, Mapping, Sequence
+from typing import Union
+
+import cv2
 import numpy as np
 import torch
-from typing import Union
-from monai.transforms import MapTransform
 from monai.config import DtypeLike, KeysCollection
 from monai.config.type_definitions import NdarrayOrTensor
 from monai.data.meta_obj import get_track_meta
+from monai.transforms import MapTransform
 from monai.transforms.transform import Transform
 from monai.transforms.utils import soft_clip
 from monai.transforms.utils_pytorch_numpy_unification import clip, percentile
 from monai.utils.enums import TransformBackends
 from monai.utils.type_conversion import convert_data_type, convert_to_dst_type, convert_to_tensor
 from scipy.ndimage import binary_dilation
-import cv2
-from collections.abc import Hashable, Mapping, Sequence
 
 
 class DilateAndSaveMaskd(MapTransform):
@@ -100,7 +101,7 @@ class ClipMaskIntensityPercentiles(Transform):
         self.channel_wise = channel_wise
         self.dtype = dtype
 
-    def _clip(self, img: NdarrayOrTensor, mask_data: NdarrayOrTensor) -> NdarrayOrTensor:
+    def _clip(self, img: NdarrayOrTensor, mask_data: NdarrayOrTensor) -> torch.Tensor:
         masked_img = img * (mask_data > 0)
         if self.sharpness_factor is not None:
             lower_percentile = (
@@ -125,8 +126,8 @@ class ClipMaskIntensityPercentiles(Transform):
             )
             img = clip(img, lower_percentile, upper_percentile)
 
-        img = convert_to_tensor(img, track_meta=False)
-        return img
+        img_tensor = convert_to_tensor(img, track_meta=False)
+        return img_tensor
 
     def __call__(self, img: NdarrayOrTensor, mask_data: NdarrayOrTensor) -> NdarrayOrTensor:
         """

src/data/data_loader.py

@@ -1,30 +1,33 @@
+import argparse
 import os
+from typing import Literal
+
 import numpy as np
-from monai.data import load_decathlon_datalist, ITKReader, PersistentDataset
+import torch
+from monai.data import PersistentDataset, load_decathlon_datalist
 from monai.transforms import (
     Compose,
-    LoadImaged,
-    ToTensord,
     ConcatItemsd,
     DeleteItemsd,
     EnsureTyped,
-    RandCropByPosNegLabeld,
+    LoadImaged,
     NormalizeIntensityd,
-    Transposed,
+    RandCropByPosNegLabeld,
     RandWeightedCropd,
+    ToTensord,
+    Transform,
+    Transposed,
 )
+from torch.utils.data.dataloader import default_collate
+
 from .custom_transforms import (
-    NormalizeIntensity_customd,
     ClipMaskIntensityPercentilesd,
     ElementwiseProductd,
+    NormalizeIntensity_customd,
 )
-import torch
-from torch.utils.data.dataloader import default_collate
-from typing import Literal
-import collections.abc
 
 
-def list_data_collate(batch: collections.abc.Sequence):
+def list_data_collate(batch: list):
     """
     Combine instances from a list of dicts into a single dict, by stacking them along first dim
     [{'image' : 3xHxW}, {'image' : 3xHxW}, {'image' : 3xHxW}...] - > {'image' : Nx3xHxW}
@@ -42,13 +45,13 @@ def list_data_collate(batch: collections.abc.Sequence):
     return default_collate(batch)
 
 
-def data_transform(args):
+def data_transform(args: argparse.Namespace) -> Transform:
     if args.use_heatmap:
         transform = Compose(
             [
                 LoadImaged(
                     keys=["image", "mask", "dwi", "adc", "heatmap"],
-                    reader=ITKReader(),
+                    reader="ITKReader",
                     ensure_channel_first=True,
                     dtype=np.float32,
                 ),
@@ -75,7 +78,7 @@ def data_transform(args):
             [
                 LoadImaged(
                     keys=["image", "mask", "dwi", "adc"],
-                    reader=ITKReader(),
+                    reader="ITKReader",
                     ensure_channel_first=True,
                     dtype=np.float32,
                 ),
@@ -101,14 +104,16 @@ def data_transform(args):
     return transform
 
 
-def get_dataloader(args, split: Literal["train", "test"]):
+def get_dataloader(
+    args: argparse.Namespace, split: Literal["train", "test"]
+) -> torch.utils.data.DataLoader:
     data_list = load_decathlon_datalist(
         data_list_file_path=args.dataset_json,
         data_list_key=split,
         base_dir=args.data_root,
     )
     if args.dry_run:
-        data_list = data_list[:8]  # Use only 8 samples for dry run
+        data_list = data_list[:2]  # Use only 8 samples for dry run
     cache_dir_ = os.path.join(args.logdir, "cache")
     os.makedirs(os.path.join(cache_dir_, split), exist_ok=True)
     transform = data_transform(args)

src/model/{csPCa_model.py → cspca_model.py}

@@ -1,4 +1,5 @@
 from __future__ import annotations
+
 import torch
 import torch.nn as nn
 from monai.utils.module import optional_import
@@ -17,8 +18,8 @@ class SimpleNN(nn.Module):
         input_dim (int): The number of input features.
     """
 
-    def __init__(self, input_dim):
-        super(SimpleNN, self).__init__()
+    def __init__(self, input_dim: int) -> None:
+        super().__init__()
         self.net = nn.Sequential(
             nn.Linear(input_dim, 256),
             nn.ReLU(),
@@ -42,7 +43,7 @@ class SimpleNN(nn.Module):
         return self.net(x)
 
 
-class csPCa_Model(nn.Module):
+class CSPCAModel(nn.Module):
     """
     Clinically Significant Prostate Cancer (csPCa) risk prediction model using a MIL backbone.
 
@@ -67,7 +68,7 @@ class csPCa_Model(nn.Module):
         backbone: The MIL based PI-RADS classifier.
     """
 
-    def __init__(self, backbone):
+    def __init__(self, backbone: nn.Module) -> None:
         super().__init__()
         self.backbone = backbone
         self.fc_dim = backbone.myfc.in_features

src/model/{MIL.py → mil.py}

@@ -1,14 +1,16 @@
 from __future__ import annotations
+
 from typing import cast
+
 import torch
 import torch.nn as nn
-from monai.utils.module import optional_import
 from monai.networks.nets import resnet
+from monai.utils.module import optional_import
 
 models, _ = optional_import("torchvision.models")
 
 
-class MILModel_3D(nn.Module):
+class MILModel3D(nn.Module):
     """
     Multiple Instance Learning (MIL) model, with a backbone classification model.
     Adapted from MONAI, modified for 3D images. The expected shape of input data is `[B, N, C, D, H, W]`,
@@ -56,6 +58,7 @@ class MILModel_3D(nn.Module):
         self.mil_mode = mil_mode.lower()
         self.attention = nn.Sequential()
         self.transformer: nn.Module | None = None
+        net: nn.Module
 
         if backbone is None:
             net = resnet.resnet18(
@@ -63,8 +66,9 @@ class MILModel_3D(nn.Module):
                 n_input_channels=3,
                 num_classes=5,
             )
+            assert net.fc is not None
             nfc = net.fc.in_features  # save the number of final features
-            net.fc = torch.nn.Identity()  # remove final linear layer
+            net.fc = torch.nn.Identity()  # type: ignore[assignment]
 
             self.extra_outputs: dict[str, torch.Tensor] = {}
 
@@ -90,7 +94,7 @@ class MILModel_3D(nn.Module):
 
             if getattr(net, "fc", None) is not None:
                 nfc = net.fc.in_features  # save the number of final features
-                net.fc = torch.nn.Identity()  # remove final linear layer
+                net.fc = torch.nn.Identity()  # type: ignore[assignment]
             else:
                 raise ValueError(
                     "Unable to detect FC layer for the torchvision model " + str(backbone),
@@ -100,8 +104,13 @@ class MILModel_3D(nn.Module):
         elif isinstance(backbone, nn.Module):
             # use a custom backbone
             net = backbone
+
+            if backbone_num_features is None:
+                raise ValueError(
+                    "Number of endencoder features must be provided for a custom backbone model"
+                )
             nfc = backbone_num_features
-            net.fc = torch.nn.Identity()  # remove final linear layer
+            net.fc = torch.nn.Identity()  # type: ignore[assignment]
 
             if mil_mode == "att_trans_pyramid":
                 # register hooks to capture outputs of intermediate layers
@@ -109,11 +118,6 @@ class MILModel_3D(nn.Module):
                     "Cannot use att_trans_pyramid with custom backbone. Have to use the default ResNet 18 backbone."
                 )
 
-            if backbone_num_features is None:
-                raise ValueError(
-                    "Number of endencoder features must be provided for a custom backbone model"
-                )
-
         else:
             raise ValueError("Unsupported backbone")
 

src/preprocessing/center_crop.py

@@ -12,7 +12,7 @@
 
 
 # import argparse
-from typing import Union
+from typing import Union, cast
 
 import SimpleITK as sitk  # noqa N813
 
@@ -21,7 +21,11 @@ def _flatten(t):
     return [item for sublist in t for item in sublist]
 
 
-def crop(image: sitk.Image, margin: Union[int, float], interpolator=sitk.sitkLinear):
+def crop(
+    image: sitk.Image,
+    margin_: Union[int, float, list[Union[int, float]]],
+    interpolator=sitk.sitkLinear,
+):
     """
     Crops a sitk.Image while retaining correct spacing. Negative margins will lead to zero padding
 
@@ -30,13 +34,14 @@ def crop(image: sitk.Image, margin: Union[int, float], interpolator=sitk.sitkLin
         margin: margins to crop. Single integer or float (percentage crop),
                 lists of int/float or nestes lists are supported.
     """
-    if isinstance(margin, (list, tuple)):
-        assert len(margin) == 3, "expected margin to be of length 3"
+    if isinstance(margin_, (list, tuple)):
+        margin_temp = margin_
+        assert len(margin_) == 3, "expected margin to be of length 3"
     else:
-        assert isinstance(margin, (int, float)), "expected margin to be a float value"
-        margin = [margin, margin, margin]
+        assert isinstance(margin_, (int, float)), "expected margin to be a float value"
+        margin_temp = [margin_, margin_, margin_]
 
-    margin = [m if isinstance(m, (tuple, list)) else [m, m] for m in margin]
+    margin = [m if isinstance(m, (tuple, list)) else [m, m] for m in margin_temp]
     old_size = image.GetSize()
 
     # calculate new origin and new image size
@@ -46,7 +51,7 @@ def crop(image: sitk.Image, margin: Union[int, float], interpolator=sitk.sitkLin
         )
         to_crop = [[int(sz * _m) for _m in m] for sz, m in zip(old_size, margin)]
     elif all([isinstance(m, int) for m in _flatten(margin)]):
-        to_crop = margin
+        to_crop = cast(list[list[int]], margin)
     else:
         raise ValueError("Wrong format of margins.")
 

src/preprocessing/generate_heatmap.py

@@ -1,11 +1,13 @@
+import argparse
+import logging
 import os
-import numpy as np
+
 import nrrd
+import numpy as np
 from tqdm import tqdm
-import logging
 
 
-def get_heatmap(args):
+def get_heatmap(args: argparse.Namespace) -> argparse.Namespace:
     """
     Generate heatmaps from DWI (Diffusion Weighted Imaging) and ADC (Apparent Diffusion Coefficient) medical imaging data.
     This function processes medical imaging files (DWI and ADC) along with their corresponding

src/preprocessing/histogram_match.py

@@ -1,8 +1,10 @@
+import argparse
+import logging
 import os
+
 import nrrd
-from skimage import exposure
-import logging
 import numpy as np
+from skimage import exposure
 from tqdm import tqdm
 
 
@@ -34,7 +36,7 @@ def get_histmatched(
     return matched_img
 
 
-def histmatch(args):
+def histmatch(args: argparse.Namespace) -> argparse.Namespace:
     files = os.listdir(args.t2_dir)
 
     t2_histmatched_dir = os.path.join(args.output_dir, "t2_histmatched")

src/preprocessing/prostate_mask.py

@@ -1,31 +1,29 @@
+import argparse
+import logging
 import os
-import numpy as np
+
 import nrrd
-from tqdm import tqdm
-from monai.bundle import ConfigParser
+import numpy as np
 import torch
-
-
+from monai.bundle import ConfigParser
+from monai.data import MetaTensor
 from monai.transforms import (
     Compose,
+    EnsureChannelFirstd,
+    EnsureTyped,
     LoadImaged,
-    ScaleIntensityd,
     NormalizeIntensityd,
-)
-from monai.utils import set_determinism
-from monai.transforms import (
-    EnsureChannelFirstd,
     Orientationd,
+    ScaleIntensityd,
     Spacingd,
-    EnsureTyped,
 )
-from monai.data import MetaTensor
-import logging
+from monai.utils import set_determinism
+from tqdm import tqdm
 
 set_determinism(43)
 
 
-def get_segmask(args):
+def get_segmask(args: argparse.Namespace) -> argparse.Namespace:
     """
     Generate prostate segmentation masks using a pre-trained deep learning model.
     This function performs inference on T2-weighted MRI images to segment the prostate gland.

src/preprocessing/register_and_crop.py

@@ -1,12 +1,15 @@
-import SimpleITK as sitk
+import argparse
+import logging
 import os
-from tqdm import tqdm
+
+import SimpleITK as sitk
 from picai_prep.preprocessing import PreprocessingSettings, Sample
+from tqdm import tqdm
+
 from .center_crop import crop
-import logging
 
 
-def register_files(args):
+def register_files(args: argparse.Namespace) -> argparse.Namespace:
     """
     Register and crop medical images (T2, DWI, and ADC) to a standardized spacing and size.
     This function reads medical images from specified directories, resamples them to a
@@ -55,9 +58,9 @@ def register_files(args):
 
         pat_case = Sample(
             scans=[
-                images_to_preprocess.get("t2"),
-                images_to_preprocess.get("hbv"),
-                images_to_preprocess.get("adc"),
+                images_to_preprocess["t2"],
+                images_to_preprocess["hbv"],
+                images_to_preprocess["adc"],
             ],
             settings=PreprocessingSettings(
                 spacing=[3.0, 0.4, 0.4], matrix_size=[new_size[2], new_size[1], new_size[0]]

src/train/train_cspca.py

@@ -1,8 +1,7 @@
 import torch
 import torch.nn as nn
 from monai.metrics import Cumulative, CumulativeAverage
-from sklearn.metrics import confusion_matrix
-from sklearn.metrics import roc_auc_score
+from sklearn.metrics import confusion_matrix, roc_auc_score
 
 
 def train_epoch(cspca_model, loader, optimizer, epoch, args):
@@ -10,10 +9,10 @@ def train_epoch(cspca_model, loader, optimizer, epoch, args):
     criterion = nn.BCELoss()
     loss = 0.0
     run_loss = CumulativeAverage()
-    TARGETS = Cumulative()
-    PREDS = Cumulative()
+    targets_cumulative = Cumulative()
+    preds_cumulative = Cumulative()
 
-    for idx, batch_data in enumerate(loader):
+    for _, batch_data in enumerate(loader):
         data = batch_data["image"].as_subclass(torch.Tensor).to(args.device)
         target = batch_data["label"].as_subclass(torch.Tensor).to(args.device)
 
@@ -24,13 +23,13 @@ def train_epoch(cspca_model, loader, optimizer, epoch, args):
         loss.backward()
         optimizer.step()
 
-        TARGETS.extend(target.detach().cpu())
-        PREDS.extend(output.detach().cpu())
+        targets_cumulative.extend(target.detach().cpu())
+        preds_cumulative.extend(output.detach().cpu())
         run_loss.append(loss.item())
 
     loss_epoch = run_loss.aggregate()
-    target_list = TARGETS.get_buffer().cpu().numpy()
-    pred_list = PREDS.get_buffer().cpu().numpy()
+    target_list = targets_cumulative.get_buffer().cpu().numpy()
+    pred_list = preds_cumulative.get_buffer().cpu().numpy()
     auc_epoch = roc_auc_score(target_list, pred_list)
 
     return loss_epoch, auc_epoch
@@ -41,10 +40,10 @@ def val_epoch(cspca_model, loader, epoch, args):
     criterion = nn.BCELoss()
     loss = 0.0
     run_loss = CumulativeAverage()
-    TARGETS = Cumulative()
-    PREDS = Cumulative()
+    targets_cumulative = Cumulative()
+    preds_cumulative = Cumulative()
     with torch.no_grad():
-        for idx, batch_data in enumerate(loader):
+        for _, batch_data in enumerate(loader):
             data = batch_data["image"].as_subclass(torch.Tensor).to(args.device)
             target = batch_data["label"].as_subclass(torch.Tensor).to(args.device)
 
@@ -52,13 +51,13 @@ def val_epoch(cspca_model, loader, epoch, args):
             output = output.squeeze(1)
             loss = criterion(output, target)
 
-            TARGETS.extend(target.detach().cpu())
-            PREDS.extend(output.detach().cpu())
+            targets_cumulative.extend(target.detach().cpu())
+            preds_cumulative.extend(output.detach().cpu())
             run_loss.append(loss.item())
 
     loss_epoch = run_loss.aggregate()
-    target_list = TARGETS.get_buffer().cpu().numpy()
-    pred_list = PREDS.get_buffer().cpu().numpy()
+    target_list = targets_cumulative.get_buffer().cpu().numpy()
+    pred_list = preds_cumulative.get_buffer().cpu().numpy()
     auc_epoch = roc_auc_score(target_list, pred_list)
     y_pred_categoric = pred_list >= 0.5
     tn, fp, fn, tp = confusion_matrix(target_list, y_pred_categoric).ravel()

src/train/train_pirads.py

@@ -1,20 +1,27 @@
+import argparse
+import logging
 import time
+
 import numpy as np
 import torch
 import torch.nn as nn
 from monai.metrics import Cumulative, CumulativeAverage
 from sklearn.metrics import cohen_kappa_score
-import logging
 
 
-def get_lambda_att(epoch, max_lambda=2.0, warmup_epochs=10):
+def get_lambda_att(epoch: int, max_lambda: float = 2.0, warmup_epochs: int = 10) -> float:
     if epoch < warmup_epochs:
         return (epoch / warmup_epochs) * max_lambda
     else:
         return max_lambda
 
 
-def get_attention_scores(data, target, heatmap, args):
+def get_attention_scores(
+    data: torch.Tensor,
+    target: torch.Tensor,
+    heatmap: torch.Tensor,
+    args: argparse.Namespace,
+) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Compute attention scores from heatmaps and shuffle data accordingly.
     This function generates attention scores based on spatial heatmaps, applies
@@ -136,17 +143,7 @@ def train_epoch(model, loader, optimizer, scaler, epoch, args):
             run_loss.append(loss.detach().cpu())
             run_acc.append(acc.detach().cpu())
             logging.info(
-                "Epoch {}/{} {}/{} loss: {:.4f} attention loss: {:.4f} acc: {:.4f} grad norm: {:.4f} time {:.2f}s".format(
-                    epoch,
-                    args.epochs,
-                    idx,
-                    len(loader),
-                    loss.item(),
-                    attn_loss.item(),
-                    acc,
-                    total_norm,
-                    time.time() - start_time,
-                )
+                f"Epoch {epoch}/{args.epochs} {idx}/{len(loader)} loss: {loss.item():.4f} attention loss: {attn_loss.item():.4f} acc: {acc:.4f} grad norm: {total_norm:.4f} time {time.time() - start_time:.2f}s"
             )
             start_time = time.time()
 
@@ -164,8 +161,8 @@ def val_epoch(model, loader, epoch, args):
 
     run_loss = CumulativeAverage()
     run_acc = CumulativeAverage()
-    PREDS = Cumulative()
-    TARGETS = Cumulative()
+    preds_cumulative = Cumulative()
+    targets_cumulative = Cumulative()
 
     start_time = time.time()
     loss, acc = 0.0, 0.0
@@ -188,12 +185,10 @@ def val_epoch(model, loader, epoch, args):
 
             run_loss.append(loss.detach().cpu())
             run_acc.append(acc.detach().cpu())
-            PREDS.extend(pred.detach().cpu())
-            TARGETS.extend(target.detach().cpu())
+            preds_cumulative.extend(pred.detach().cpu())
+            targets_cumulative.extend(target.detach().cpu())
             logging.info(
-                "Val epoch {}/{} {}/{} loss: {:.4f} acc: {:.4f} time {:.2f}s".format(
-                    epoch, args.epochs, idx, len(loader), loss, acc, time.time() - start_time
-                )
+                f"Val epoch {epoch}/{args.epochs} {idx}/{len(loader)} loss: {loss:.4f} acc: {acc:.4f} time {time.time() - start_time:.2f}s"
             )
             start_time = time.time()
 
@@ -201,9 +196,11 @@ def val_epoch(model, loader, epoch, args):
             torch.cuda.empty_cache()
 
         # Calculate QWK metric (Quadratic Weigted Kappa) https://en.wikipedia.org/wiki/Cohen%27s_kappa
-        PREDS = PREDS.get_buffer().cpu().numpy()
-        TARGETS = TARGETS.get_buffer().cpu().numpy()
+        preds_cumulative = preds_cumulative.get_buffer().cpu().numpy()
+        targets_cumulative = targets_cumulative.get_buffer().cpu().numpy()
         loss_epoch = run_loss.aggregate()
         acc_epoch = run_acc.aggregate()
-        qwk = cohen_kappa_score(TARGETS.astype(np.float64), PREDS.astype(np.float64))
+        qwk = cohen_kappa_score(
+            targets_cumulative.astype(np.float64), preds_cumulative.astype(np.float64)
+        )
     return loss_epoch, acc_epoch, qwk

src/utils.py

@@ -1,21 +1,31 @@
+import argparse
+import logging
 import os
 import sys
+from pathlib import Path
+from typing import Any, Union
+
+import cv2
 import numpy as np
 import torch
+from monai.data import Dataset
 from monai.transforms import (
     Compose,
+    EnsureTyped,
     LoadImaged,
     ToTensord,
-    EnsureTyped,
 )
+
 from .data.custom_transforms import ClipMaskIntensityPercentilesd, NormalizeIntensity_customd
-from monai.data import Dataset, ITKReader
-import logging
-from pathlib import Path
-import cv2
 
 
-def save_pirads_checkpoint(model, epoch, args, filename="model.pth", best_acc=0):
+def save_pirads_checkpoint(
+    model: torch.nn.Module,
+    epoch: int,
+    args: argparse.Namespace,
+    filename: str = "model.pth",
+    best_acc: float = 0,
+) -> None:
     """Save checkpoint for the PI-RADS model"""
 
     state_dict = model.state_dict()
@@ -25,7 +35,11 @@ def save_pirads_checkpoint(model, epoch, args, filename="model.pth", best_acc=0)
     logging.info(f"Saving checkpoint {filename}")
 
 
-def save_cspca_checkpoint(model, val_metric, model_dir):
+def save_cspca_checkpoint(
+    model: torch.nn.Module,
+    val_metric: dict[str, Any],
+    model_dir: str,
+) -> None:
     """Save checkpoint for the csPCa model"""
 
     state_dict = model.state_dict()
@@ -41,7 +55,7 @@ def save_cspca_checkpoint(model, val_metric, model_dir):
     logging.info(f"Saving model with auc: {val_metric['auc']}")
 
 
-def get_metrics(metric_dict: dict):
+def get_metrics(metric_dict: dict) -> None:
     for metric_name, metric_list in metric_dict.items():
         metric_list = np.array(metric_list)
         lower = np.percentile(metric_list, 2.5)
@@ -51,7 +65,7 @@ def get_metrics(metric_dict: dict):
         logging.info(f"95% CI: ({lower:.3f}, {upper:.3f})")
 
 
-def setup_logging(log_file):
+def setup_logging(log_file: Union[str, Path]) -> None:
     log_file = Path(log_file)
     log_file.parent.mkdir(parents=True, exist_ok=True)
     if log_file.exists():
@@ -66,13 +80,13 @@ def setup_logging(log_file):
 
 
 def validate_steps(steps):
-    REQUIRES = {
+    requires = {
         "get_segmentation_mask": ["register_and_crop"],
         "histogram_match": ["get_segmentation_mask", "register_and_crop"],
         "get_heatmap": ["get_segmentation_mask", "histogram_match", "register_and_crop"],
     }
     for i, step in enumerate(steps):
-        required = REQUIRES.get(step, [])
+        required = requires.get(step, [])
         for req in required:
             if req not in steps[:i]:
                 logging.error(
@@ -81,7 +95,10 @@ def validate_steps(steps):
                 sys.exit(1)
 
 
-def get_patch_coordinate(patches_top_5, parent_image):
+def get_patch_coordinate(
+    patches_top_5: list[np.ndarray],
+    parent_image: np.ndarray,
+) -> list[tuple[int, int, int]]:
     """
     Locate the coordinates of top-5 patches within a parent image.
 
@@ -90,7 +107,7 @@ def get_patch_coordinate(patches_top_5, parent_image):
     coordinates (row, column) and the slice index where each patch is found.
 
     Args:
-        patches_top_5 (list): List of top-5 patch tensors, each with shape (C, H, W)
+        patches_top_5 (list): List of top-5 patches as np arrays, each with shape (C, H, W)
                              where C is channels, H is height, W is width.
         parent_image (np.ndarray): 3D image volume with shape (height, width, slices)
                                    to search within.
@@ -137,12 +154,12 @@ def get_patch_coordinate(patches_top_5, parent_image):
     return coords
 
 
-def get_parent_image(temp_data_list, args):
+def get_parent_image(temp_data_list, args: argparse.Namespace) -> np.ndarray:
     transform_image = Compose(
         [
             LoadImaged(
                 keys=["image", "mask"],
-                reader=ITKReader(),
+                reader="ITKReader",
                 ensure_channel_first=True,
                 dtype=np.float32,
             ),
@@ -169,9 +186,9 @@ def visualise_patches():
     for i in range(rows):
         for j in range(slices):
             ax = axes[i, j]
-            
+
             if j == 0:
-            
+
                 for k in range(parent_image.shape[2]):
                     img_temp = parent_image[:, :, k]
                     H, W = img_temp.shape
@@ -187,11 +204,11 @@ def visualise_patches():
                                 break
                         if bool1:
                             break
-                        
+
                     if bool1:
                         break
 
-                
+
 
 
             ax.imshow(parent_image[:, :, k+j], cmap='gray')
@@ -199,7 +216,7 @@ def visualise_patches():
                                     linewidth=2, edgecolor='red', facecolor='none')
             ax.add_patch(rect)
             ax.axis('off')
-        
+
 
     plt.tight_layout()
     plt.show()

temp.ipynb

@@ -160,7 +160,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "id": "c91a5802",
    "metadata": {},
    "outputs": [
@@ -168,9 +168,18 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "  0%|          | 0/1 [00:16<?, ?it/s]\n",
+      "  0%|          | 0/1 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 1/1 [00:02<00:00,  2.45s/it]\n",
       "You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
-      "  0%|          | 0/1 [00:18<?, ?it/s]\n"
+      "100%|██████████| 1/1 [00:38<00:00, 38.82s/it]\n",
+      "100%|██████████| 1/1 [00:00<00:00,  4.11it/s]\n",
+      "100%|██████████| 1/1 [00:00<00:00,  9.82it/s]\n"
      ]
     }
    ],
@@ -188,8 +197,8 @@
     "\n",
     "args = register_files(args)\n",
     "args = get_segmask(args)\n",
-    "#args = histmatch(args)\n",
-    "#args = get_heatmap(args)\n"
+    "args = histmatch(args)\n",
+    "args = get_heatmap(args)\n"
    ]
   },
   {
@@ -370,7 +379,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "id": "8b5d382e",
    "metadata": {},
    "outputs": [],
@@ -385,7 +394,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 6,
    "id": "4cf061ec",
    "metadata": {},
    "outputs": [
@@ -419,7 +428,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 7,
    "id": "fac15515",
    "metadata": {},
    "outputs": [],
@@ -452,7 +461,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 8,
    "id": "eb80047b",
    "metadata": {},
    "outputs": [],
@@ -493,227 +502,23 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 9,
    "id": "dbcfc97f",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "[array([[[-0.43333182, -0.52243334, -1.4245864 , ..., -1.7698549 ,\n",
-       "          -1.6027895 , -1.257521  ],\n",
-       "         [-0.7786003 , -1.2797965 , -2.0817103 , ..., -2.0371597 ,\n",
-       "          -2.0817103 , -2.059435  ],\n",
-       "         [-1.4802749 , -1.658478  , -1.7809926 , ..., -1.4691372 ,\n",
-       "          -2.1596742 , -1.9591957 ],\n",
-       "         ...,\n",
-       "         [ 0.8697783 ,  0.34630668,  0.6358867 , ...,  0.624749  ,\n",
-       "           0.24606745,  0.03445129],\n",
-       "         [ 0.44654593,  0.23492976,  0.7806767 , ...,  1.2373221 ,\n",
-       "           0.2906182 , -1.0124918 ],\n",
-       "         [ 0.3017559 , -0.2551287 ,  0.27948052, ...,  1.1036698 ,\n",
-       "          -0.16602719, -1.2463834 ]],\n",
-       " \n",
-       "        [[-0.84542644, -1.3800358 , -1.4357241 , ..., -2.5383558 ,\n",
-       "          -2.0817103 , -0.99021643],\n",
-       "         [-0.4444695 , -1.3466226 , -1.8255434 , ..., -1.335485  ,\n",
-       "          -1.6362027 , -1.4579996 ],\n",
-       "         [-1.2129703 , -1.7921304 , -2.0371597 , ...,  0.14582822,\n",
-       "          -0.04351256, -0.82315105],\n",
-       "         ...,\n",
-       "         [-0.4556072 ,  1.4378005 ,  1.5603153 , ...,  0.2906182 ,\n",
-       "          -0.85656416, -1.1906949 ],\n",
-       "         [-1.2352457 , -0.35536796,  1.1816336 , ..., -0.39991874,\n",
-       "          -1.3800358 , -1.2129703 ],\n",
-       "         [-1.2909342 , -1.1127311 ,  0.5467852 , ..., -0.9568034 ,\n",
-       "          -1.4023111 , -0.9011149 ]],\n",
-       " \n",
-       "        [[-2.1596742 , -1.914645  , -1.7809926 , ..., -1.0124918 ,\n",
-       "           0.03445129,  0.44654593],\n",
-       "         [-2.0037465 , -1.8923696 , -1.8700942 , ..., -1.2129703 ,\n",
-       "          -1.1572819 , -0.50015795],\n",
-       "         [-1.8478189 , -1.9369203 , -1.9814711 , ..., -0.65608567,\n",
-       "          -1.2352457 , -1.4134488 ],\n",
-       "         ...,\n",
-       "         [-0.6338103 , -1.0124918 , -0.16602719, ...,  0.41313285,\n",
-       "           0.13469052, -0.80087566],\n",
-       "         [-0.7451872 , -1.3577603 , -0.2996795 , ...,  0.34630668,\n",
-       "           0.68043745,  0.45768362],\n",
-       "         [ 0.0121759 , -0.7674626 , -0.33309257, ...,  0.09013975,\n",
-       "           0.46882132,  1.0034306 ]]], dtype=float32),\n",
-       " array([[[-0.64494795, -0.789738  , -0.5447087 , ...,  0.0233136 ,\n",
-       "           0.14582822, -0.35536796],\n",
-       "         [-0.8120134 , -0.7340495 , -0.42219412, ..., -0.31081718,\n",
-       "           0.46882132,  0.15696591],\n",
-       "         [-0.08806333,  0.03445129,  0.12355283, ..., -0.2774041 ,\n",
-       "           0.73612595,  0.7249882 ],\n",
-       "         ...,\n",
-       "         [ 0.09013975,  0.0233136 , -0.24399103, ...,  0.34630668,\n",
-       "           0.914329  ,  0.6358867 ],\n",
-       "         [ 0.2906182 ,  0.335169  ,  0.624749  , ...,  0.24606745,\n",
-       "           0.9254667 ,  0.9922929 ],\n",
-       "         [ 0.44654593,  0.70271283,  1.0479814 , ..., -0.09920102,\n",
-       "           0.37971976,  0.70271283]],\n",
-       " \n",
-       "        [[-0.23285334, -0.01009948, -0.1326141 , ...,  1.1593583 ,\n",
-       "           1.5603153 ,  1.5603153 ],\n",
-       "         [-0.23285334,  0.13469052,  0.1792413 , ...,  0.98115516,\n",
-       "           1.5603153 ,  1.5603153 ],\n",
-       "         [-0.36650565, -0.01009948,  0.190379  , ...,  0.8697783 ,\n",
-       "           1.5603153 ,  1.5603153 ],\n",
-       "         ...,\n",
-       "         [-0.16602719, -0.06578795,  0.190379  , ..., -1.4357241 ,\n",
-       "          -1.368898  , -1.6027895 ],\n",
-       "         [-0.2662664 , -0.35536796,  0.190379  , ..., -1.513688  ,\n",
-       "          -1.3800358 , -1.6362027 ],\n",
-       "         [-0.789738  , -0.7786003 , -0.17716487, ..., -1.7141665 ,\n",
-       "          -1.5693765 , -1.9591957 ]],\n",
-       " \n",
-       "        [[-0.12147641, -0.01009948,  0.0233136 , ...,  0.769539  ,\n",
-       "           0.8140898 ,  0.2906182 ],\n",
-       "         [-0.35536796, -0.2774041 , -0.16602719, ...,  0.55792284,\n",
-       "           0.68043745,  0.335169  ],\n",
-       "         [-0.18830256,  0.1681036 ,  0.7918144 , ...,  0.70271283,\n",
-       "           0.4910967 ,  0.10127745],\n",
-       "         ...,\n",
-       "         [-0.97907877, -0.97907877, -0.9011149 , ..., -0.43333182,\n",
-       "          -0.5447087 , -0.1437518 ],\n",
-       "         [-1.1238688 , -1.1461442 , -1.0347673 , ..., -0.7117741 ,\n",
-       "          -0.12147641,  0.479959  ],\n",
-       "         [-0.8788395 , -0.62267256, -0.9122526 , ..., -0.92339027,\n",
-       "          -0.42219412,  0.22379206]]], dtype=float32),\n",
-       " array([[[ 2.7948052e-01,  1.0368437e+00,  1.5603153e+00, ...,\n",
-       "          -1.2797965e+00, -8.3428878e-01, -2.4399103e-01],\n",
-       "         [ 6.7864366e-02,  3.4630668e-01,  1.1704960e+00, ...,\n",
-       "          -7.7860028e-01, -8.3428878e-01, -3.1081718e-01],\n",
-       "         [ 2.6834285e-01,  9.0139754e-02,  2.3492976e-01, ...,\n",
-       "          -6.3381028e-01, -9.5680338e-01, -4.4446951e-01],\n",
-       "         ...,\n",
-       "         [-1.3688980e+00, -1.2241080e+00, -9.4566566e-01, ...,\n",
-       "          -1.0570426e+00, -2.6626641e-01,  1.2707351e+00],\n",
-       "         [-1.0793180e+00, -8.8997722e-01, -1.1015934e+00, ...,\n",
-       "          -7.6746261e-01, -5.4650255e-02,  1.0257059e+00],\n",
-       "         [-7.3404950e-01, -6.5608567e-01, -1.0124918e+00, ...,\n",
-       "          -5.3357106e-01,  2.6834285e-01,  8.5864055e-01]],\n",
-       " \n",
-       "        [[ 1.1593583e+00,  1.5603153e+00,  1.5603153e+00, ...,\n",
-       "          -4.5560721e-01, -1.1033872e-01,  8.6977828e-01],\n",
-       "         [ 8.4750289e-01,  1.5603153e+00,  1.5603153e+00, ...,\n",
-       "          -1.3020718e+00, -1.0459049e+00,  6.3588673e-01],\n",
-       "         [ 7.9002060e-02,  1.3469052e-01, -3.1081718e-01, ...,\n",
-       "          -1.3688980e+00, -1.0570426e+00,  4.9109671e-01],\n",
-       "         ...,\n",
-       "         [ 1.5603153e+00,  6.4702439e-01, -4.3512560e-02, ...,\n",
-       "           5.2450979e-01,  8.9205366e-01, -2.9967949e-01],\n",
-       "         [ 1.2930106e+00,  6.1361134e-01, -4.6674490e-01, ...,\n",
-       "           2.7948052e-01,  8.4750289e-01,  7.8067672e-01],\n",
-       "         [ 1.0382106e-03,  2.5720516e-01, -4.1105643e-01, ...,\n",
-       "           6.2474900e-01,  1.5603153e+00,  1.5603153e+00]],\n",
-       " \n",
-       "        [[-9.2339027e-01, -1.0236295e+00, -1.0347673e+00, ...,\n",
-       "           2.6834285e-01, -1.9944026e-01, -1.1033872e-01],\n",
-       "         [-8.2315105e-01, -1.1127311e+00, -9.7907877e-01, ...,\n",
-       "           1.2818729e+00,  3.2403129e-01,  6.7864366e-02],\n",
-       "         [-8.0087566e-01, -9.7907877e-01, -8.0087566e-01, ...,\n",
-       "           1.5046268e+00,  6.1361134e-01, -1.7716487e-01],\n",
-       "         ...,\n",
-       "         [ 1.3932499e+00,  6.8043745e-01,  2.3492976e-01, ...,\n",
-       "          -1.8144057e+00, -1.2129703e+00,  2.3313597e-02],\n",
-       "         [ 1.5603153e+00,  1.5603153e+00,  1.5603153e+00, ...,\n",
-       "          -1.7587173e+00, -1.1572819e+00,  5.6906056e-01],\n",
-       "         [ 1.5603153e+00,  1.5603153e+00,  1.5603153e+00, ...,\n",
-       "          -1.8366811e+00, -8.7883949e-01,  1.3598367e+00]]], dtype=float32),\n",
-       " array([[[-1.4245864 , -1.3466226 , -1.079318  , ..., -0.36650565,\n",
-       "          -0.66722333, -1.079318  ],\n",
-       "         [-1.1350064 , -1.2352457 , -1.4357241 , ..., -0.01009948,\n",
-       "          -0.41105643, -0.7563249 ],\n",
-       "         [-1.0459049 , -0.84542644, -1.368898  , ..., -0.22171564,\n",
-       "          -0.4778826 , -0.82315105],\n",
-       "         ...,\n",
-       "         [-2.2599134 , -2.1151235 , -1.6139272 , ...,  0.5245098 ,\n",
-       "           0.41313285,  0.37971976],\n",
-       "         [-2.560631  , -2.204225  , -1.7921304 , ...,  0.10127745,\n",
-       "           0.3128936 ,  0.37971976],\n",
-       "         [-2.2933266 , -2.1596742 , -2.0705726 , ...,  0.0233136 ,\n",
-       "           0.3017559 ,  0.4910967 ]],\n",
-       " \n",
-       "        [[-1.0570426 , -1.3911734 , -1.658478  , ...,  0.769539  ,\n",
-       "           0.7806767 ,  0.95887977],\n",
-       "         [-1.6362027 , -1.4134488 , -1.5693765 , ...,  1.0145682 ,\n",
-       "           1.1816336 ,  1.1370829 ],\n",
-       "         [-1.914645  , -1.1127311 , -1.224108  , ...,  0.68043745,\n",
-       "           0.9922929 ,  0.85864055],\n",
-       "         ...,\n",
-       "         [-2.1262612 , -2.026022  , -1.7921304 , ...,  0.4910967 ,\n",
-       "           0.9031913 ,  1.0702567 ],\n",
-       "         [-2.3712904 , -2.5494936 , -2.304464  , ...,  0.70271283,\n",
-       "           0.914329  ,  1.0479814 ],\n",
-       "         [-2.2710512 , -2.3267395 , -2.4715295 , ...,  1.0145682 ,\n",
-       "           0.7806767 ,  0.7472636 ]],\n",
-       " \n",
-       "        [[-0.94566566, -0.6003972 , -0.85656416, ...,  1.2373221 ,\n",
-       "           1.5603153 ,  1.5603153 ],\n",
-       "         [-1.1795572 , -0.38878104, -0.37764335, ...,  1.5603153 ,\n",
-       "           1.5603153 ,  1.5603153 ],\n",
-       "         [-1.335485  , -1.3577603 , -0.32195488, ...,  1.5603153 ,\n",
-       "           1.5603153 ,  1.5603153 ],\n",
-       "         ...,\n",
-       "         [-1.224108  , -1.1795572 , -1.4914126 , ..., -0.04351256,\n",
-       "          -0.5112957 , -0.64494795],\n",
-       "         [-1.3132095 , -1.8589565 , -1.7587173 , ..., -0.03237487,\n",
-       "          -0.42219412, -0.1437518 ],\n",
-       "         [-1.6473403 , -2.4381166 , -2.1262612 , ...,  0.37971976,\n",
-       "           0.34630668,  0.46882132]]], dtype=float32),\n",
-       " array([[[-0.84542644, -0.8120134 , -0.6894987 , ...,  0.46882132,\n",
-       "          -0.01009948, -0.39991874],\n",
-       "         [-0.65608567, -0.8120134 , -0.6338103 , ...,  0.59133595,\n",
-       "           0.04558898, -0.5112957 ],\n",
-       "         [-0.96794105, -0.8342888 , -0.35536796, ...,  0.5801982 ,\n",
-       "           0.13469052, -0.37764335],\n",
-       "         ...,\n",
-       "         [-0.09920102, -0.11033872, -1.224108  , ..., -0.92339027,\n",
-       "          -1.2686588 , -0.85656416],\n",
-       "         [ 0.0121759 , -0.01009948, -1.0236295 , ..., -1.2129703 ,\n",
-       "          -1.4023111 , -0.7340495 ],\n",
-       "         [ 0.23492976, -0.07692564, -0.70063645, ..., -0.70063645,\n",
-       "          -0.934528  , -1.0124918 ]],\n",
-       " \n",
-       "        [[-1.6139272 , -0.97907877, -0.7340495 , ..., -0.8342888 ,\n",
-       "          -0.6894987 , -1.0347673 ],\n",
-       "         [-1.0459049 , -0.5892595 , -0.789738  , ..., -1.257521  ,\n",
-       "          -1.0124918 , -1.4357241 ],\n",
-       "         [-0.42219412, -0.5892595 , -1.4134488 , ..., -1.5025504 ,\n",
-       "          -1.2463834 , -1.4802749 ],\n",
-       "         ...,\n",
-       "         [ 0.10127745,  0.40199515,  0.13469052, ...,  0.25720516,\n",
-       "           0.55792284,  0.12355283],\n",
-       "         [-0.2885418 , -0.2551287 ,  0.41313285, ...,  0.46882132,\n",
-       "          -0.33309257, -1.2797965 ],\n",
-       "         [-0.2996795 , -0.64494795,  0.9366044 , ...,  0.56906056,\n",
-       "          -0.41105643, -1.3466226 ]],\n",
-       " \n",
-       "        [[-1.1238688 , -1.4134488 , -1.7364419 , ..., -1.1350064 ,\n",
-       "          -1.6918911 , -2.3156018 ],\n",
-       "         [-1.4914126 , -1.3020718 , -0.99021643, ..., -1.658478  ,\n",
-       "          -1.7698549 , -1.8812319 ],\n",
-       "         [-1.5582387 , -0.85656416, -0.43333182, ..., -1.6027895 ,\n",
-       "          -1.914645  , -1.7698549 ],\n",
-       "         ...,\n",
-       "         [ 0.6358867 ,  0.9031913 ,  1.0702567 , ...,  0.11241514,\n",
-       "           0.07900206,  0.34630668],\n",
-       "         [ 0.70271283,  1.1593583 ,  0.9254667 , ...,  0.335169  ,\n",
-       "           0.41313285,  0.23492976],\n",
-       "         [ 0.35744438,  1.1482205 ,  0.8697783 , ..., -0.2662664 ,\n",
-       "           0.56906056,  0.624749  ]]], dtype=float32)]"
+       "list"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 9,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "patches_top_5"
+    "type(patches_top_5)"
    ]
   },
   {

tests/test_run.py

@@ -20,7 +20,15 @@ def test_run_pirads_training():
 
     # Run the script with the config
     result = subprocess.run(
-        [sys.executable, str(script_path), "--mode", "train", "--config", str(config_path), "--dry_run", "True" ],
+        [
+            sys.executable,
+            str(script_path),
+            "--mode",
+            "train",
+            "--config",
+            str(config_path),
+            "--dry_run",
+        ],
         capture_output=True,
         text=True,
     )
@@ -28,6 +36,7 @@ def test_run_pirads_training():
     # Check that it ran without errors
     assert result.returncode == 0, f"Script failed with:\n{result.stderr}"
 
+
 def test_run_pirads_inference():
     """
     Test that run_cspca.py runs without crashing using an existing YAML config.
@@ -45,14 +54,23 @@ def test_run_pirads_inference():
 
     # Run the script with the config
     result = subprocess.run(
-        [sys.executable, str(script_path), "--mode", "test", "--config", str(config_path), "--dry_run", "True" ],
+        [
+            sys.executable,
+            str(script_path),
+            "--mode",
+            "test",
+            "--config",
+            str(config_path),
+            "--dry_run",
+        ],
         capture_output=True,
         text=True,
     )
 
     # Check that it ran without errors
     assert result.returncode == 0, f"Script failed with:\n{result.stderr}"
-    
+
+
 def test_run_cspca_training():
     """
     Test that run_cspca.py runs without crashing using an existing YAML config.
@@ -70,14 +88,23 @@ def test_run_cspca_training():
 
     # Run the script with the config
     result = subprocess.run(
-        [sys.executable, str(script_path), "--mode", "train", "--config", str(config_path), "--dry_run", "True" ],
+        [
+            sys.executable,
+            str(script_path),
+            "--mode",
+            "train",
+            "--config",
+            str(config_path),
+            "--dry_run",
+        ],
         capture_output=True,
         text=True,
     )
 
     # Check that it ran without errors
     assert result.returncode == 0, f"Script failed with:\n{result.stderr}"
-    
+
+
 def test_run_cspca_inference():
     """
     Test that run_cspca.py runs without crashing using an existing YAML config.
@@ -95,12 +122,18 @@ def test_run_cspca_inference():
 
     # Run the script with the config
     result = subprocess.run(
-        [sys.executable, str(script_path), "--mode", "test", "--config", str(config_path), "--dry_run", "True" ],
+        [
+            sys.executable,
+            str(script_path),
+            "--mode",
+            "test",
+            "--config",
+            str(config_path),
+            "--dry_run",
+        ],
         capture_output=True,
         text=True,
     )
 
     # Check that it ran without errors
     assert result.returncode == 0, f"Script failed with:\n{result.stderr}"
-
-