cleanup

.gitignore

@@ -5,4 +5,5 @@ temp_data/
 temp.ipynb
 __pycache__/
 **/__pycache__/
-*.pyc
+*.pyc
+.ruff_cache

config/config_cspca_test.yaml

@@ -1,6 +1,5 @@
-project_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/new_folder/
 data_root: /sc-projects/sc-proj-cc06-ag-ki-radiologie/pirad_model_test_PICAI/registered/t2_hist_matched/
-dataset_json: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/new_folder/dataset/PICAI_cspca.json
+dataset_json: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/PICAI_cspca.json
 num_classes: !!int 4
 mil_mode: att_trans
 tile_count: !!int 24
@@ -8,10 +7,7 @@ tile_size: !!int 64
 depth: !!int 3
 use_heatmap: !!bool True
 workers: !!int 6
-checkpoint_cspca: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/new_folder/models/cspca_model.pth
-num_seeds: !!int 2
-batch_size: !!int 1
-
-
+checkpoint_cspca: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/models/cspca_model.pth
+batch_size: !!int 8
 
 

config/config_cspca_train.yaml

@@ -1,6 +1,5 @@
-project_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/new_folder/
 data_root: /sc-projects/sc-proj-cc06-ag-ki-radiologie/pirad_model_test_PICAI/registered/t2_hist_matched/
-dataset_json: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/new_folder/dataset/PICAI_cspca.json
+dataset_json: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/PICAI_cspca.json
 num_classes: !!int 4
 mil_mode: att_trans
 tile_count: !!int 24
@@ -8,12 +7,11 @@ tile_size: !!int 64
 depth: !!int 3
 use_heatmap: !!bool True
 workers: !!int 6
-checkpoint_pirads: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/new_folder/models/pirads.pt
-epochs: !!int 1
+checkpoint_pirads: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/models/pirads.pt
+epochs: !!int 80
 batch_size: !!int 8
 optim_lr: !!float 2e-4
 
 
 
 
-

config/config_pirads_test.yaml

@@ -1,18 +1,15 @@
 run_name: pirads_test_run
-project_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/new_folder/
 data_root: /sc-projects/sc-proj-cc06-ag-ki-radiologie/prostate-foundation/PICAI_registered/t2_hist_matched
-dataset_json: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/new_folder/dataset/PI-RADS_data.json
+dataset_json: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/PI-RADS_data.json
 num_classes: !!int 4
 mil_mode: att_trans
 tile_count: !!int 24
 tile_size: !!int 64
 depth: !!int 3
 use_heatmap: !!bool True
-workers: !!int 0
-checkpoint: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/new_folder/models/pirads.pt
+workers: !!int 8
+checkpoint: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/models/pirads.pt
 amp: !!bool True
-dry_run: !!bool True
-
 
 
 

config/config_pirads_train.yaml

@@ -1,21 +1,18 @@
-project_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/new_folder/
 data_root: /sc-projects/sc-proj-cc06-ag-ki-radiologie/prostate-foundation/PICAI_registered/t2_hist_matched
-dataset_json: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/new_folder/dataset/PI-RADS_data.json
+dataset_json: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/PI-RADS_data.json
 num_classes: !!int 4
 mil_mode: att_trans
 tile_count: !!int 24
 tile_size: !!int 64
 depth: !!int 3
 use_heatmap: !!bool True
-workers: !!int 0
-checkpoint_pirads: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/new_folder/models/pirads.pt
-epochs: !!int 2
+workers: !!int 4
+epochs: !!int 100
 batch_size: !!int 8
 optim_lr: !!float 2e-4
 weight_decay: !!float 1e-5
 amp: !!bool True
 wandb: !!bool True
-dry_run: !!bool True
 
 
 

job_scripts/train_pirads.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+#SBATCH --job-name=pirads_training       # Specify job name
+#SBATCH --partition=gpu               # Specify partition name
+#SBATCH --mem=128G   
+#SBATCH --gres=gpu:1             
+#SBATCH --time=48:00:00               # Set a limit on the total run time
+#SBATCH --output=/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/logs/%x/log.o%j      # File name for standard output
+#SBATCH --error=/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/logs/%x/log.e%j       # File name for standard error output
+#SBATCH --mail-user=anirudh.balaraman@charite.de
+#SBATCH --mail-type=END,FAIL
+
+
+source /etc/profile.d/conda.sh
+conda activate foundation
+
+RUNDIR="/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate"
+
+
+srun python -u $RUNDIR/run_pirads.py --mode train --config $RUNDIR/config/config_pirads_train.yaml

preprocess_main.py

@@ -1,68 +1,61 @@
-import SimpleITK as sitk 
 import os
-import numpy as np 
-import nrrd
-from AIAH_utility.viewer import BasicViewer
-from tqdm import tqdm
-import pandas as pd
-from picai_prep.preprocessing import PreprocessingSettings, Sample
-import multiprocessing
-import sys
 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 pathlib import Path
 from src.utils import setup_logging
 from src.utils import validate_steps
 import argparse
-import yaml 
+import yaml
+
 
 def parse_args():
-    FUNCTIONS = {
-    "register_and_crop": register_files,
-    "histogram_match": histmatch,
-    "get_segmentation_mask": get_segmask,
-    "get_heatmap": get_heatmap,
-    }
     parser = argparse.ArgumentParser(description="File preprocessing")
     parser.add_argument("--config", type=str, help="Path to YAML config file")
     parser.add_argument(
         "--steps",
-        nargs="+",                      # ← list of strings
-        choices=FUNCTIONS.keys(),       # ← restrict allowed values
+        nargs="+",  # ← list of strings
+        choices=[
+            "register_and_crop",
+            "histogram_match",
+            "get_segmentation_mask",
+            "get_heatmap",
+        ],  # ← restrict allowed values
         required=True,
-        help="Steps to execute (one or more)"
+        help="Steps to execute (one or more)",
     )
     parser.add_argument("--t2_dir", default=None, help="Path to T2W files")
     parser.add_argument("--dwi_dir", default=None, help="Path to DWI files")
     parser.add_argument("--adc_dir", default=None, help="Path to ADC files")
     parser.add_argument("--seg_dir", default=None, help="Path to segmentation masks")
     parser.add_argument("--output_dir", default=None, help="Path to output folder")
-    parser.add_argument("--margin", default=0.2, type=float, help="Margin to center crop the images")
+    parser.add_argument(
+        "--margin", default=0.2, type=float, help="Margin to center crop the images"
+    )
     parser.add_argument("--project_dir", default=None, help="Project directory")
-    
+
     args = parser.parse_args()
     if args.config:
-        with open(args.config, 'r') as config_file:
+        with open(args.config, "r") as config_file:
             config = yaml.safe_load(config_file)
             args.__dict__.update(config)
     return args
 
+
 if __name__ == "__main__":
     args = parse_args()
     FUNCTIONS = {
-    "register_and_crop": register_files,
-    "histogram_match": histmatch,
-    "get_segmentation_mask": get_segmask,
-    "get_heatmap": get_heatmap,
+        "register_and_crop": register_files,
+        "histogram_match": histmatch,
+        "get_segmentation_mask": get_segmask,
+        "get_heatmap": get_heatmap,
     }
 
-    args.logfile = os.path.join(args.output_dir, f"preprocessing.log")
+    args.logfile = os.path.join(args.output_dir, "preprocessing.log")
     setup_logging(args.logfile)
     logging.info("Starting preprocessing")
     validate_steps(args.steps)
     for step in args.steps:
         func = FUNCTIONS[step]
-        args = func(args) 
+        args = func(args)

pyproject.toml

@@ -0,0 +1,9 @@
+[tool.ruff]
+line-length = 100
+
+[tool.ruff.lint]
+select = ["E", "W"]
+ignore = ["E501"]
+
+[tool.ruff.format]
+quote-style = "double"

run_cspca.py

@@ -1,120 +1,102 @@
 import argparse
 import os
 import shutil
-import time
 import yaml
 import sys
-import gdown
-import numpy as np
 import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-import torch.nn as nn
-import torch.nn.functional as F
-from monai.config import KeysCollection
-from monai.metrics import Cumulative, CumulativeAverage
-from monai.networks.nets import milmodel, resnet, MILModel
-
-from sklearn.metrics import cohen_kappa_score
-from torch.cuda.amp import GradScaler, autocast
-from torch.utils.data.dataloader import default_collate
-from torchvision.models.resnet import ResNet50_Weights
-import shutil
 from pathlib import Path
-from torch.utils.data.distributed import DistributedSampler
-from torch.utils.tensorboard import SummaryWriter
 from monai.utils import set_determinism
-
-import matplotlib.pyplot as plt
-
-import wandb
-import math
 import logging
-from pathlib import Path
-
-
 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.train.train_cspca import train_epoch, val_epoch
+import random
 
-def main_worker(args):
 
-    mil_model = MILModel_3D(
-        num_classes=args.num_classes,  
-        mil_mode=args.mil_mode 
-    ).to(args.device)
+def main_worker(args):
+    mil_model = MILModel_3D(num_classes=args.num_classes, mil_mode=args.mil_mode)
     cache_dir_path = Path(os.path.join(args.logdir, "cache"))
 
-    if args.mode == 'train':
-
+    if args.mode == "train":
         checkpoint = torch.load(args.checkpoint_pirads, weights_only=False, map_location="cpu")
         mil_model.load_state_dict(checkpoint["state_dict"])
         mil_model = mil_model.to(args.device)
-        
-        model_dir = os.path.join(args.project_dir,'models')
-        metrics_dict = {'auc':[], 'sensitivity':[], 'specificity':[]}
+
+        model_dir = os.path.join(args.logdir, "models")
+        os.makedirs(model_dir, exist_ok=True)
+
+        metrics_dict = {"auc": [], "sensitivity": [], "specificity": []}
         for st in list(range(args.num_seeds)):
             set_determinism(seed=st)
-            
+
             train_loader = get_dataloader(args, split="train")
             valid_loader = get_dataloader(args, split="test")
             cspca_model = csPCa_Model(backbone=mil_model).to(args.device)
-            for submodule in [cspca_model.backbone.net, 
-                            cspca_model.backbone.myfc, 
-                            cspca_model.backbone.transformer]:
+            for submodule in [
+                cspca_model.backbone.net,
+                cspca_model.backbone.myfc,
+                cspca_model.backbone.transformer,
+            ]:
                 for param in submodule.parameters():
                     param.requires_grad = False
 
-            optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, cspca_model.parameters()), lr=args.optim_lr)
+            optimizer = torch.optim.AdamW(
+                filter(lambda p: p.requires_grad, cspca_model.parameters()), lr=args.optim_lr
+            )
 
-            old_loss = float('inf')
+            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"STATE {st} EPOCH {epoch} TRAIN loss: {train_loss:.4f} AUC: {train_auc:.4f}")
+                train_loss, train_auc = train_epoch(
+                    cspca_model, train_loader, optimizer, epoch=epoch, args=args
+                )
+                logging.info(
+                    f"STATE {st} 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"STATE {st} EPOCH {epoch} VAL loss: {val_metric['loss']:.4f} AUC: {val_metric['auc']:.4f}")
-                val_metric['state'] = st
-                if val_metric['loss'] < old_loss:
-                    old_loss = val_metric['loss']
-                    old_auc = val_metric['auc']
-                    sensitivity = val_metric['sensitivity']
-                    specificity = val_metric['specificity']
-                    if len(metrics_dict['auc']) == 0:
-                        save_cspca_checkpoint(cspca_model, val_metric, model_dir)
-                    elif val_metric['auc'] >= max(metrics_dict['auc']):
+                logging.info(
+                    f"STATE {st} EPOCH {epoch} VAL loss: {val_metric['loss']:.4f} AUC: {val_metric['auc']:.4f}"
+                )
+                val_metric["state"] = st
+                if val_metric["loss"] < old_loss:
+                    old_loss = val_metric["loss"]
+                    old_auc = val_metric["auc"]
+                    sensitivity = val_metric["sensitivity"]
+                    specificity = val_metric["specificity"]
+                    if not metrics_dict["auc"] or val_metric["auc"] >= max(metrics_dict["auc"]):
                         save_cspca_checkpoint(cspca_model, val_metric, model_dir)
 
-            metrics_dict['auc'].append(old_auc)
-            metrics_dict['sensitivity'].append(sensitivity)
-            metrics_dict['specificity'].append(specificity)
+            metrics_dict["auc"].append(old_auc)
+            metrics_dict["sensitivity"].append(sensitivity)
+            metrics_dict["specificity"].append(specificity)
             if cache_dir_path.exists() and cache_dir_path.is_dir():
                 shutil.rmtree(cache_dir_path)
 
         get_metrics(metrics_dict)
 
-    elif args.mode == 'test':
-
+    elif args.mode == "test":
         cspca_model = csPCa_Model(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.load_state_dict(checkpt["state_dict"])
         cspca_model = cspca_model.to(args.device)
-        if 'auc' in checkpt and 'sensitivity' in checkpt and 'specificity' in checkpt:
-            auc, sens, spec = checkpt['auc'], checkpt['sensitivity'], checkpt['specificity']
-            logging.info(f"csPCa Model loaded from {args.checkpoint_cspca} with AUC: {auc}, Sensitivity: {sens}, Specificity: {spec} on the test set.")
+        if "auc" in checkpt and "sensitivity" in checkpt and "specificity" in checkpt:
+            auc, sens, spec = checkpt["auc"], checkpt["sensitivity"], checkpt["specificity"]
+            logging.info(
+                f"csPCa Model loaded from {args.checkpoint_cspca} with AUC: {auc}, Sensitivity: {sens}, Specificity: {spec} on the test set."
+            )
         else:
             logging.info(f"csPCa Model loaded from {args.checkpoint_cspca}.")
-        
-        metrics_dict = {'auc':[], 'sensitivity':[], 'specificity':[]}
+
+        metrics_dict = {"auc": [], "sensitivity": [], "specificity": []}
         for st in list(range(args.num_seeds)):
             set_determinism(seed=st)
             test_loader = get_dataloader(args, split="test")
             test_metric = val_epoch(cspca_model, test_loader, epoch=0, args=args)
-            metrics_dict['auc'].append(test_metric['auc'])
-            metrics_dict['sensitivity'].append(test_metric['sensitivity'])
-            metrics_dict['specificity'].append(test_metric['specificity'])
+            metrics_dict["auc"].append(test_metric["auc"])
+            metrics_dict["sensitivity"].append(test_metric["sensitivity"])
+            metrics_dict["specificity"].append(test_metric["specificity"])
 
             if cache_dir_path.exists() and cache_dir_path.is_dir():
                 shutil.rmtree(cache_dir_path)
@@ -122,43 +104,58 @@ def main_worker(args):
         get_metrics(metrics_dict)
 
 
-
-
 def parse_args():
-    parser = argparse.ArgumentParser(description="Multiple Instance Learning (MIL) for csPCa risk prediction.")
-    parser.add_argument('--mode', type=str, choices=['train', 'test'], required=True, help='Operation mode: train or infer')
-    parser.add_argument('--run_name', type=str, default='train_cspca', help='run name for log file')
-    parser.add_argument('--config', type=str, help='Path to YAML config file')
-    parser.add_argument(
-        "--project_dir", default=None, help="path to project firectory"
+    parser = argparse.ArgumentParser(
+        description="Multiple Instance Learning (MIL) for csPCa risk prediction."
     )
     parser.add_argument(
-        "--data_root", default=None, help="path to root folder of images"
+        "--mode",
+        type=str,
+        choices=["train", "test"],
+        required=True,
+        help="Operation mode: train or infer",
     )
+    parser.add_argument("--run_name", type=str, default="train_cspca", help="run name for log file")
+    parser.add_argument("--config", type=str, help="Path to YAML config file")
+    parser.add_argument("--project_dir", default=None, help="path to project firectory")
+    parser.add_argument("--data_root", default=None, help="path to root folder of images")
     parser.add_argument("--dataset_json", default=None, type=str, help="path to dataset json file")
     parser.add_argument("--num_classes", default=4, type=int, help="number of output classes")
-    parser.add_argument("--mil_mode", default="att_trans", help="MIL algorithm: choose either att_trans or att_pyramid")
     parser.add_argument(
-        "--tile_count", default=24, type=int, help="number of patches (instances) to extract from MRI input"
+        "--mil_mode",
+        default="att_trans",
+        help="MIL algorithm: choose either att_trans or att_pyramid",
     )
-    parser.add_argument("--tile_size", default=64, type=int, help="size of square patch (instance) in pixels")
-    parser.add_argument("--depth", default=3, type=int, help="number of slices in each 3D patch (instance)")
     parser.add_argument(
-        "--use_heatmap", action="store_true",
-        help="enable weak attention heatmap guided patch generation"
+        "--tile_count",
+        default=24,
+        type=int,
+        help="number of patches (instances) to extract from MRI input",
     )
     parser.add_argument(
-        "--no_heatmap", dest="use_heatmap", action="store_false",
-        help="disable heatmap"
+        "--tile_size", default=64, type=int, help="size of square patch (instance) in pixels"
+    )
+    parser.add_argument(
+        "--depth", default=3, type=int, help="number of slices in each 3D patch (instance)"
+    )
+    parser.add_argument(
+        "--use_heatmap",
+        action="store_true",
+        help="enable weak attention heatmap guided patch generation",
+    )
+    parser.add_argument(
+        "--no_heatmap", dest="use_heatmap", action="store_false", help="disable heatmap"
     )
     parser.set_defaults(use_heatmap=True)
     parser.add_argument("--workers", default=2, type=int, help="number of workers for data loading")
-    #parser.add_argument("--dry-run", action="store_true")
+    # parser.add_argument("--dry-run", action="store_true")
     parser.add_argument("--checkpoint_pirads", default=None, help="Load PI-RADS model")
-    parser.add_argument("--epochs", "--max_epochs", default=30, type=int, help="number of training epochs")
+    parser.add_argument(
+        "--epochs", "--max_epochs", default=30, type=int, help="number of training epochs"
+    )
     parser.add_argument("--batch_size", default=32, type=int, help="number of MRI scans per batch")
     parser.add_argument("--optim_lr", default=2e-4, type=float, help="initial learning rate")
-    #parser.add_argument("--amp", action="store_true", help="use AMP, recommended")
+    # parser.add_argument("--amp", action="store_true", help="use AMP, recommended")
     parser.add_argument(
         "--val_every",
         "--val_interval",
@@ -166,42 +163,50 @@ def parse_args():
         type=int,
         help="run validation after this number of epochs, default 1 to run every epoch",
     )
-    parser.add_argument("--dry_run", action="store_true", help="Run the script in dry-run mode (default: False)")
+    parser.add_argument(
+        "--dry_run", action="store_true", help="Run the script in dry-run mode (default: False)"
+    )
     parser.add_argument("--checkpoint_cspca", default=None, help="load existing checkpoint")
-    parser.add_argument("--num_seeds", default=20, type=int, help="number of seeds to be run to build CI")
+    parser.add_argument(
+        "--num_seeds", default=20, type=int, help="number of seeds to be run to build CI"
+    )
     args = parser.parse_args()
     if args.config:
-        with open(args.config, 'r') as config_file:
+        with open(args.config, "r") 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
+
+    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
+
     args.logdir = os.path.join(args.project_dir, "logs", args.run_name)
     os.makedirs(args.logdir, exist_ok=True)
     args.logfile = os.path.join(args.logdir, f"{args.run_name}.log")
     setup_logging(args.logfile)
 
-
     logging.info("Argument values:")
     for k, v in vars(args).items():
         logging.info(f"{k} => {v}")
     logging.info("-----------------")
 
     if args.dataset_json is None:
-        logging.error('Dataset path not provided. Quitting.')
+        logging.error("Dataset path not provided. Quitting.")
         sys.exit(1)
-    if args.checkpoint_pirads is None and args.mode == 'train':
-        logging.error('PI-RADS checkpoint path not provided. Quitting.')
+    if args.checkpoint_pirads is None and args.mode == "train":
+        logging.error("PI-RADS checkpoint path not provided. Quitting.")
         sys.exit(1)
-    elif args.checkpoint_cspca is None and args.mode == 'test':
-        logging.error('csPCa checkpoint path not provided. Quitting.')
+    elif args.checkpoint_cspca is None and args.mode == "test":
+        logging.error("csPCa checkpoint path not provided. Quitting.")
         sys.exit(1)
     args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -216,5 +221,4 @@ if __name__ == "__main__":
         args.num_seeds = 2
         args.wandb = False
 
-
     main_worker(args)

run_inference.py

@@ -1,66 +1,21 @@
 import argparse
 import os
-import shutil
-import time
 import yaml
-import sys
-import gdown
-import numpy as np
 import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-import torch.nn as nn
-import torch.nn.functional as F
-from monai.config import KeysCollection
-from monai.metrics import Cumulative, CumulativeAverage
-from monai.networks.nets import milmodel, resnet, MILModel
-
-from sklearn.metrics import cohen_kappa_score
-from torch.cuda.amp import GradScaler, autocast
-from torch.utils.data.dataloader import default_collate
-from torchvision.models.resnet import ResNet50_Weights
-import shutil
-from pathlib import Path
-from torch.utils.data.distributed import DistributedSampler
-from torch.utils.tensorboard import SummaryWriter
-from monai.utils import set_determinism
-import matplotlib.pyplot as plt
-import wandb
-import math
 import logging
-from pathlib import Path
-
-
 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, save_pirads_checkpoint, get_parent_image, get_patch_coordinate
-from src.train import train_cspca, train_pirads
-import SimpleITK as sitk 
-
-import nrrd
-
-from tqdm import tqdm
-import pandas as pd
-from picai_prep.preprocessing import PreprocessingSettings, Sample
-import multiprocessing
-import sys
+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
-import logging
-from pathlib import Path
-from src.utils import setup_logging
-from src.utils import validate_steps
-import argparse
-import yaml 
 from src.data.data_loader import data_transform, list_data_collate
-from monai.data import Dataset, load_decathlon_datalist, ITKReader, NumpyReader, PersistentDataset
+from monai.data import Dataset
 import json
 
-def parse_args():
 
+def parse_args():
     parser = argparse.ArgumentParser(description="File preprocessing")
     parser.add_argument("--config", type=str, help="Path to YAML config file")
     parser.add_argument("--t2_dir", default=None, help="Path to T2W files")
@@ -68,7 +23,9 @@ def parse_args():
     parser.add_argument("--adc_dir", default=None, help="Path to ADC files")
     parser.add_argument("--seg_dir", default=None, help="Path to segmentation masks")
     parser.add_argument("--output_dir", default=None, help="Path to output folder")
-    parser.add_argument("--margin", default=0.2, type=float, help="Margin to center crop the images")
+    parser.add_argument(
+        "--margin", default=0.2, type=float, help="Margin to center crop the images"
+    )
     parser.add_argument("--num_classes", default=4, type=int)
     parser.add_argument("--mil_mode", default="att_trans", type=str)
     parser.add_argument("--use_heatmap", default=True, type=bool)
@@ -76,47 +33,49 @@ def parse_args():
     parser.add_argument("--tile_count", default=24, type=int)
     parser.add_argument("--depth", default=3, type=int)
     parser.add_argument("--project_dir", default=None, help="Project directory")
-    
+
     args = parser.parse_args()
     if args.config:
-        with open(args.config, 'r') as config_file:
+        with open(args.config, "r") as config_file:
             config = yaml.safe_load(config_file)
             args.__dict__.update(config)
     return args
 
+
 if __name__ == "__main__":
     args = parse_args()
     FUNCTIONS = {
-    "register_and_crop": register_files,
-    "histogram_match": histmatch,
-    "get_segmentation_mask": get_segmask,
-    "get_heatmap": get_heatmap,
+        "register_and_crop": register_files,
+        "histogram_match": histmatch,
+        "get_segmentation_mask": get_segmask,
+        "get_heatmap": get_heatmap,
     }
 
-    args.logfile = os.path.join(args.output_dir, f"inference.log")
+    args.logfile = os.path.join(args.output_dir, "inference.log")
     setup_logging(args.logfile)
     logging.info("Starting preprocessing")
     steps = ["register_and_crop", "get_segmentation_mask", "histogram_match", "get_heatmap"]
     for step in steps:
         func = FUNCTIONS[step]
-        args = func(args) 
+        args = func(args)
 
     logging.info("Preprocessing completed.")
 
     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 = MILModel_3D(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_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)
-    checkpt = torch.load(os.path.join(args.project_dir, 'models', 'cspca_model.pth'), map_location="cpu")
-    cspca_model.load_state_dict(checkpt['state_dict'])
+    checkpt = torch.load(
+        os.path.join(args.project_dir, "models", "cspca_model.pth"), map_location="cpu"
+    )
+    cspca_model.load_state_dict(checkpt["state_dict"])
     cspca_model = cspca_model.to(args.device)
 
     transform = data_transform(args)
@@ -124,12 +83,12 @@ if __name__ == "__main__":
     args.data_list = []
     for file in files:
         temp = {}
-        temp['image'] = os.path.join(args.t2_dir, file)
-        temp['dwi'] = os.path.join(args.dwi_dir, file)
-        temp['adc'] = os.path.join(args.adc_dir, file)
-        temp['heatmap'] = os.path.join(args.heatmapdir, file)
-        temp['mask'] = os.path.join(args.seg_dir, file)
-        temp['label'] = 0  # dummy label
+        temp["image"] = os.path.join(args.t2_dir, file)
+        temp["dwi"] = os.path.join(args.dwi_dir, file)
+        temp["adc"] = os.path.join(args.adc_dir, file)
+        temp["heatmap"] = os.path.join(args.heatmapdir, file)
+        temp["mask"] = os.path.join(args.seg_dir, file)
+        temp["label"] = 0  # dummy label
         args.data_list.append(temp)
 
     dataset = Dataset(data=args.data_list, transform=transform)
@@ -139,7 +98,7 @@ if __name__ == "__main__":
         shuffle=False,
         num_workers=0,
         pin_memory=True,
-        multiprocessing_context= None,
+        multiprocessing_context=None,
         sampler=None,
         collate_fn=list_data_collate,
     )
@@ -153,7 +112,7 @@ if __name__ == "__main__":
         for idx, 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)
+            pirads_score = torch.argmax(logits, dim=1)
             pirads_list.append(pirads_score.item())
 
             output = cspca_model(data)
@@ -181,17 +140,19 @@ if __name__ == "__main__":
     for i in args.data_list:
         parent_image = get_parent_image([i], args)
 
-        coords = get_patch_coordinate(patches_top_5, parent_image, args)
+        coords = get_patch_coordinate(patches_top_5, parent_image)
         coords_list.append(coords)
     output_dict = {}
 
-    for i,j in enumerate(files):
-        logging.info(f"File: {j}, PIRADS score: {pirads_list[i]}, csPCa risk score: {cspca_risk_list[i]:.4f}")
-
+    for i, j in enumerate(files):
+        logging.info(
+            f"File: {j}, PIRADS score: {pirads_list[i] + 2.0}, csPCa risk score: {cspca_risk_list[i]:.4f}"
+        )
         output_dict[j] = {
-            'Predicted PIRAD Score': pirads_list[i] + 2.0,
-            'csPCa risk': cspca_risk_list[i],
-            'Top left coordinate of top 5 patches(x,y,z)': coords_list[i],
+            "Predicted PIRAD Score": pirads_list[i] + 2.0,
+            "csPCa risk": cspca_risk_list[i],
+            "Top left coordinate of top 5 patches(x,y,z)": coords_list[i],
         }
-    with open(os.path.join(args.output_dir, "results.json"), 'w') as f:
-        json.dump(output_dict, f, indent=4)
+
+    with open(os.path.join(args.output_dir, "results.json"), "w") as f:
+        json.dump(output_dict, f, indent=4)

run_pirads.py

@@ -1,32 +1,14 @@
 import argparse
-import collections.abc
 import os
 import shutil
 import time
 import yaml
 import sys
-import gdown
 import numpy as np
 import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-import torch.nn as nn
-import torch.nn.functional as F
-from monai.config import KeysCollection
-from monai.data import Dataset, load_decathlon_datalist, ITKReader, NumpyReader, PersistentDataset
-from monai.data.wsi_reader import WSIReader
-from monai.metrics import Cumulative, CumulativeAverage
-from monai.networks.nets import milmodel, resnet, MILModel
-
-from sklearn.metrics import cohen_kappa_score
-from torch.cuda.amp import GradScaler, autocast
-from torch.utils.data.dataloader import default_collate
 from torch.utils.tensorboard import SummaryWriter
 from monai.utils import set_determinism
-import matplotlib.pyplot as plt
-
 import wandb
-import math
 import logging
 from pathlib import Path
 from src.data.data_loader import get_dataloader
@@ -37,13 +19,9 @@ from src.utils import save_pirads_checkpoint, setup_logging
 
 def main_worker(args):
     if args.device == torch.device("cuda"):
-        torch.cuda.set_device(args.gpu)  # use this default device (same as args.device if not distributed)
         torch.backends.cudnn.benchmark = True
 
-    model = MILModel_3D(
-        num_classes=args.num_classes,  
-        mil_mode=args.mil_mode 
-    )
+    model = MILModel_3D(num_classes=args.num_classes, mil_mode=args.mil_mode)
     start_epoch = 0
     best_acc = 0.0
     if args.checkpoint is not None:
@@ -54,41 +32,54 @@ def main_worker(args):
             start_epoch = checkpoint["epoch"]
         if "best_acc" in checkpoint:
             best_acc = checkpoint["best_acc"]
-        logging.info("=> loaded checkpoint %s (epoch %d) (bestacc %f)",args.checkpoint, start_epoch, best_acc)
+        logging.info(
+            "=> loaded checkpoint %s (epoch %d) (bestacc %f)",
+            args.checkpoint,
+            start_epoch,
+            best_acc,
+        )
     cache_dir_ = os.path.join(args.logdir, "cache")
     model.to(args.device)
     params = model.parameters()
-    if args.mode == 'train':
-        train_loader = get_dataloader(args, split=args.mode)
+    if args.mode == "train":
+        train_loader = get_dataloader(args, split="train")
         valid_loader = get_dataloader(args, split="test")
-        logging.info("Dataset training:", str(len(train_loader.dataset)), "test:", str(len(valid_loader.dataset)))
-    
+        logging.info(
+            f"Dataset training: {len(train_loader.dataset)}, test: {len(valid_loader.dataset)}"
+        )
+
         if args.mil_mode in ["att_trans", "att_trans_pyramid"]:
             params = [
-                {"params": list(model.attention.parameters()) + list(model.myfc.parameters()) + list(model.net.parameters())},
+                {
+                    "params": list(model.attention.parameters())
+                    + list(model.myfc.parameters())
+                    + list(model.net.parameters())
+                },
                 {"params": list(model.transformer.parameters()), "lr": 6e-5, "weight_decay": 0.1},
             ]
 
         optimizer = torch.optim.AdamW(params, lr=args.optim_lr, weight_decay=args.weight_decay)
-        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs, eta_min=0)
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer, T_max=args.epochs, eta_min=0
+        )
         scaler = torch.amp.GradScaler(device=str(args.device), enabled=args.amp)
 
         if args.logdir is not None:
             writer = SummaryWriter(log_dir=args.logdir)
-            logging.info("Writing Tensorboard logs to ", writer.log_dir)
+            logging.info(f"Writing Tensorboard logs to {writer.log_dir}")
         else:
             writer = None
 
-
         # RUN TRAINING
         n_epochs = args.epochs
         val_loss_min = float("inf")
         epochs_no_improve = 0
         for epoch in range(start_epoch, n_epochs):
-
             logging.info(time.ctime(), "Epoch:", epoch)
             epoch_time = time.time()
-            train_loss, train_acc, train_att_loss, batch_norm = train_epoch(model, train_loader, optimizer, scaler=scaler, epoch=epoch, args=args)
+            train_loss, train_acc, train_att_loss, batch_norm = train_epoch(
+                model, train_loader, optimizer, scaler=scaler, epoch=epoch, args=args
+            )
             logging.info(
                 "Final training %d/%d loss: %.4f attention loss: %.4f acc: %.4f time %.2fs",
                 epoch,
@@ -98,13 +89,21 @@ def main_worker(args):
                 train_acc,
                 time.time() - epoch_time,
             )
-        
+
             if writer is not None:
                 writer.add_scalar("train_loss", train_loss, epoch)
                 writer.add_scalar("train_attention_loss", train_att_loss, epoch)
                 writer.add_scalar("train_acc", train_acc, epoch)
-            wandb.log({"Train Loss": train_loss, "Train Accuracy": train_acc, "Train Attention Loss": train_att_loss, "Batch Norm": batch_norm}, step=epoch)
-                
+            wandb.log(
+                {
+                    "Train Loss": train_loss,
+                    "Train Accuracy": train_acc,
+                    "Train Attention Loss": train_att_loss,
+                    "Batch Norm": batch_norm,
+                },
+                step=epoch,
+            )
+
             model_new_best = False
             val_acc = 0
             if (epoch + 1) % args.val_every == 0:
@@ -125,35 +124,39 @@ def main_worker(args):
                     writer.add_scalar("test_acc", val_acc, epoch)
                     writer.add_scalar("test_qwk", qwk, epoch)
 
-                    #val_acc = qwk
-                wandb.log({"Test Loss": val_loss, "Test Accuracy": val_acc,"Cohen Kappa": qwk}, step=epoch)
+                    # val_acc = qwk
+                wandb.log(
+                    {"Test Loss": val_loss, "Test Accuracy": val_acc, "Cohen Kappa": qwk},
+                    step=epoch,
+                )
                 if val_loss < val_loss_min:
                     logging.info("Loss (%.6f --> %.6f)", val_loss_min, val_loss)
                     val_loss_min = val_loss
                     model_new_best = True
 
             if args.logdir is not None:
-                save_pirads_checkpoint(model, epoch, args, best_acc=val_acc, filename=f"model_{epoch}.pt")
+                save_pirads_checkpoint(
+                    model, epoch, args, best_acc=val_acc, filename=f"model_{epoch}.pt"
+                )
                 if model_new_best:
-                    logging.info("Copying to model.pt new best model!!!!")
-                    shutil.copyfile(os.path.join(args.logdir, f"model_{epoch}.pt"), os.path.join(args.logdir, "model.pt"))
+                    logging.info("Copying to model.pt new best model")
+                    shutil.copyfile(
+                        os.path.join(args.logdir, f"model_{epoch}.pt"),
+                        os.path.join(args.logdir, "model.pt"),
+                    )
                     epochs_no_improve = 0
-                
+
                 else:
                     epochs_no_improve += 1
                     if epochs_no_improve == args.early_stop:
-                        logging.info('Early stopping!')
+                        logging.info("Early stopping!")
                         break
-                    
-
 
             scheduler.step()
 
         logging.info("ALL DONE")
-    
-    elif args.mode == 'test':
 
-        
+    elif args.mode == "test":
         kappa_list = []
         for seed in list(range(args.num_seeds)):
             set_determinism(seed=seed)
@@ -163,51 +166,73 @@ def main_worker(args):
             kappa_list.append(qwk)
             logging.info(f"Seed {seed}, QWK: {qwk}")
             if os.path.exists(cache_dir_):
-                logging.info("Removing cache directory ", cache_dir_)
+                logging.info(f"Removing cache directory {cache_dir_}")
                 shutil.rmtree(cache_dir_)
-                
+
         logging.info(f"Mean QWK over {args.num_seeds} seeds: {np.mean(kappa_list)}")
 
-    
     if os.path.exists(cache_dir_):
-        logging.info("Removing cache directory ", cache_dir_)
+        logging.info(f"Removing cache directory {cache_dir_}")
         shutil.rmtree(cache_dir_)
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(description="Multiple Instance Learning (MIL) for PIRADS Classification.")
-    parser.add_argument('--mode', type=str, choices=['train', 'test'], required=True, help='operation mode: train or infer')
-    parser.add_argument('--wandb', action='store_true', help='Add this flag to enable WandB logging')
-    parser.add_argument('--project_name', type=str, default='Classification_prostate', help='WandB project name')
-    parser.add_argument('--run_name', type=str, default='train_pirads', help='run name for WandB logging')
-    parser.add_argument('--config', type=str, help='path to YAML config file')
+    parser = argparse.ArgumentParser(
+        description="Multiple Instance Learning (MIL) for PIRADS Classification."
+    )
+    parser.add_argument(
+        "--mode",
+        type=str,
+        choices=["train", "test"],
+        required=True,
+        help="operation mode: train or infer",
+    )
+    parser.add_argument(
+        "--wandb", action="store_true", help="Add this flag to enable WandB logging"
+    )
     parser.add_argument(
-        "--project_dir", default=None, help="path to project firectory"
+        "--project_name", type=str, default="Classification_prostate", help="WandB project name"
     )
     parser.add_argument(
-        "--data_root", default=None, help="path to root folder of images"
+        "--run_name", type=str, default="train_pirads", help="run name for WandB logging"
     )
+    parser.add_argument("--config", type=str, help="path to YAML config file")
+    parser.add_argument("--project_dir", default=None, help="path to project firectory")
+    parser.add_argument("--data_root", default=None, help="path to root folder of images")
     parser.add_argument("--dataset_json", default=None, type=str, help="path to dataset json file")
     parser.add_argument("--num_classes", default=4, type=int, help="number of output classes")
-    parser.add_argument("--mil_mode", default="att_trans", help="MIL algorithm: choose either att_trans or att_pyramid")
     parser.add_argument(
-        "--tile_count", default=24, type=int, help="number of patches (instances) to extract from MRI input"
+        "--mil_mode",
+        default="att_trans",
+        help="MIL algorithm: choose either att_trans or att_pyramid",
     )
-    parser.add_argument("--tile_size", default=64, type=int, help="size of square patch (instance) in pixels")
-    parser.add_argument("--depth", default=3, type=int, help="number of slices in each 3D patch (instance)")
     parser.add_argument(
-        "--use_heatmap", action="store_true",
-        help="enable weak attention heatmap guided patch generation"
+        "--tile_count",
+        default=24,
+        type=int,
+        help="number of patches (instances) to extract from MRI input",
+    )
+    parser.add_argument(
+        "--tile_size", default=64, type=int, help="size of square patch (instance) in pixels"
+    )
+    parser.add_argument(
+        "--depth", default=3, type=int, help="number of slices in each 3D patch (instance)"
     )
     parser.add_argument(
-        "--no_heatmap", dest="use_heatmap", action="store_false",
-        help="disable heatmap"
+        "--use_heatmap",
+        action="store_true",
+        help="enable weak attention heatmap guided patch generation",
+    )
+    parser.add_argument(
+        "--no_heatmap", dest="use_heatmap", action="store_false", help="disable heatmap"
     )
     parser.set_defaults(use_heatmap=True)
     parser.add_argument("--workers", default=2, type=int, help="number of workers for data loading")
 
     parser.add_argument("--checkpoint", default=None, help="load existing checkpoint")
-    parser.add_argument("--epochs", "--max_epochs", default=50, type=int, help="number of training epochs")
+    parser.add_argument(
+        "--epochs", "--max_epochs", default=50, type=int, help="number of training epochs"
+    )
     parser.add_argument("--early_stop", default=40, type=int, help="early stopping criteria")
     parser.add_argument("--batch_size", default=4, type=int, help="number of MRI scans per batch")
     parser.add_argument("--optim_lr", default=3e-5, type=float, help="initial learning rate")
@@ -220,19 +245,27 @@ def parse_args():
         type=int,
         help="run validation after this number of epochs, default 1 to run every epoch",
     )
-    parser.add_argument("--dry_run", action="store_true", help="Run the script in dry-run mode (default: False)")
+    parser.add_argument(
+        "--dry_run", action="store_true", help="Run the script in dry-run mode (default: False)"
+    )
     args = parser.parse_args()
     if args.config:
-        with open(args.config, 'r') as config_file:
+        with open(args.config, "r") 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
+
+    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
+
     args.logdir = os.path.join(args.project_dir, "logs", args.run_name)
     os.makedirs(args.logdir, exist_ok=True)
     args.logfile = os.path.join(args.logdir, f"{args.run_name}.log")
@@ -243,15 +276,14 @@ if __name__ == "__main__":
         logging.info(f"{k} => {v}")
     logging.info("-----------------")
 
-    args.num_seeds = 10
     args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     if args.device == torch.device("cpu"):
         args.amp = False
     if args.dataset_json is None:
-        logging.error('Dataset JSON file not provided. Quitting.')
+        logging.error("Dataset JSON file not provided. Quitting.")
         sys.exit(1)
-    if args.checkpoint is None and args.mode == 'test':
-        logging.error('Model checkpoint path not provided. Quitting.')
+    if args.checkpoint is None and args.mode == "test":
+        logging.error("Model checkpoint path not provided. Quitting.")
         sys.exit(1)
 
     if args.dry_run:
@@ -261,8 +293,8 @@ if __name__ == "__main__":
         args.workers = 0
         args.num_seeds = 2
         args.wandb = False
-    
-    mode_wandb = "online" if args.wandb else "disabled"
+
+    mode_wandb = "online" if args.wandb and args.mode != "test" else "disabled"
 
     config_wandb = {
         "learning_rate": args.optim_lr,
@@ -271,13 +303,14 @@ if __name__ == "__main__":
         "patch size": args.tile_size,
         "patch count": args.tile_count,
     }
-    wandb.init(project=args.project_name, 
-               name=args.run_name,
-               dir=os.path.join(args.logdir, "wandb"),
-               config=config_wandb,
-               mode=mode_wandb)
-
+    wandb.init(
+        project=args.project_name,
+        name=args.run_name,
+        dir=os.path.join(args.logdir, "wandb"),
+        config=config_wandb,
+        mode=mode_wandb,
+    )
 
     main_worker(args)
-    
+
     wandb.finish()

src/data/custom_transforms.py

@@ -1,27 +1,25 @@
 import numpy as np
 import torch
-from typing import Union, Optional
-
+from typing import Union
 from monai.transforms import MapTransform
 from monai.config import DtypeLike, KeysCollection
-from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor
+from monai.config.type_definitions import NdarrayOrTensor
 from monai.data.meta_obj import get_track_meta
 from monai.transforms.transform import Transform
 from monai.transforms.utils import soft_clip
-from monai.transforms.utils_pytorch_numpy_unification import clip, percentile, where
+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, get_equivalent_dtype
+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 typing import Union, Sequence
 from collections.abc import Hashable, Mapping, Sequence
 
 
-
 class DilateAndSaveMaskd(MapTransform):
     """
     Custom transform to dilate binary mask and save a copy.
     """
+
     def __init__(self, keys, dilation_size=10, copy_key="original_mask"):
         super().__init__(keys)
         self.dilation_size = dilation_size
@@ -29,37 +27,61 @@ class DilateAndSaveMaskd(MapTransform):
 
     def __call__(self, data):
         d = dict(data)
-        
+
         for key in self.keys:
             mask = d[key].numpy() if isinstance(d[key], torch.Tensor) else d[key]
             mask = mask.squeeze(0)  # Remove channel dimension if present
 
             # Save a copy of the original mask
-            d[self.copy_key] = torch.tensor(mask, dtype=torch.float32).unsqueeze(0)  # Save to a new key
+            d[self.copy_key] = torch.tensor(mask, dtype=torch.float32).unsqueeze(
+                0
+            )  # Save to a new key
 
             # Apply binary dilation to the mask
             dilated_mask = binary_dilation(mask, iterations=self.dilation_size).astype(np.uint8)
 
             # Store the dilated mask
-            d[key] = torch.tensor(dilated_mask, dtype=torch.float32).unsqueeze(0)  # Add channel dimension back
+            d[key] = torch.tensor(dilated_mask, dtype=torch.float32).unsqueeze(
+                0
+            )  # Add channel dimension back
 
         return d
 
 
 class ClipMaskIntensityPercentiles(Transform):
+    """
+    Clip image intensity values based on percentiles computed from a masked region.
+    This transform clips the intensity range of an image to values between lower and upper
+    percentiles calculated only from voxels where the mask is positive. It supports both
+    hard clipping and soft (smooth) clipping via a sharpness factor.
+    Args:
+        lower: Lower percentile threshold in range [0, 100]. If None, no lower clipping applied.
+        upper: Upper percentile threshold in range [0, 100]. If None, no upper clipping applied.
+        sharpness_factor: If provided, applies soft clipping with this sharpness parameter.
+            Must be greater than 0. If None, applies hard clipping instead.
+        channel_wise: If True, applies clipping independently to each channel using the
+            corresponding channel's mask. If False, uses the same mask for all channels.
+        dtype: Output data type for the clipped image. Defaults to np.float32.
+    Raises:
+        ValueError: If both lower and upper are None, if percentiles are outside [0, 100],
+            if upper < lower, or if sharpness_factor <= 0.
+    Returns:
+        Clipped image with intensities adjusted based on masked percentiles.
+    Note:
+        Supports both torch.Tensor and numpy.ndarray inputs.
 
 
     backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
+    """
 
     def __init__(
         self,
         lower: Union[float, None],
         upper: Union[float, None],
-        sharpness_factor : Union[float, None] = None,
+        sharpness_factor: Union[float, None] = None,
         channel_wise: bool = False,
         dtype: DtypeLike = np.float32,
     ) -> None:
-
         if lower is None and upper is None:
             raise ValueError("lower or upper percentiles must be provided")
         if lower is not None and (lower < 0.0 or lower > 100.0):
@@ -71,7 +93,7 @@ class ClipMaskIntensityPercentiles(Transform):
         if sharpness_factor is not None and sharpness_factor <= 0:
             raise ValueError("sharpness_factor must be greater than 0")
 
-        #self.mask_data = mask_data
+        # self.mask_data = mask_data
         self.lower = lower
         self.upper = upper
         self.sharpness_factor = sharpness_factor
@@ -81,14 +103,26 @@ class ClipMaskIntensityPercentiles(Transform):
     def _clip(self, img: NdarrayOrTensor, mask_data: NdarrayOrTensor) -> NdarrayOrTensor:
         masked_img = img * (mask_data > 0)
         if self.sharpness_factor is not None:
-
-            lower_percentile = percentile(masked_img, self.lower) if self.lower is not None else None
-            upper_percentile = percentile(masked_img, self.upper) if self.upper is not None else None
-            img = soft_clip(img, self.sharpness_factor, lower_percentile, upper_percentile, self.dtype)
+            lower_percentile = (
+                percentile(masked_img, self.lower) if self.lower is not None else None
+            )
+            upper_percentile = (
+                percentile(masked_img, self.upper) if self.upper is not None else None
+            )
+            img = soft_clip(
+                img, self.sharpness_factor, lower_percentile, upper_percentile, self.dtype
+            )
         else:
-            
-            lower_percentile = percentile(masked_img, self.lower) if self.lower is not None else percentile(masked_img, 0)
-            upper_percentile = percentile(masked_img, self.upper) if self.upper is not None else percentile(masked_img, 100)
+            lower_percentile = (
+                percentile(masked_img, self.lower)
+                if self.lower is not None
+                else percentile(masked_img, 0)
+            )
+            upper_percentile = (
+                percentile(masked_img, self.upper)
+                if self.upper is not None
+                else percentile(masked_img, 100)
+            )
             img = clip(img, lower_percentile, upper_percentile)
 
         img = convert_to_tensor(img, track_meta=False)
@@ -102,7 +136,9 @@ class ClipMaskIntensityPercentiles(Transform):
         img_t = convert_to_tensor(img, track_meta=False)
         mask_t = convert_to_tensor(mask_data, track_meta=False)
         if self.channel_wise:
-            img_t = torch.stack([self._clip(img=d, mask_data=mask_t[e]) for e,d in enumerate(img_t)])  # type: ignore
+            img_t = torch.stack(
+                [self._clip(img=d, mask_data=mask_t[e]) for e, d in enumerate(img_t)]
+            )  # type: ignore
         else:
             img_t = self._clip(img=img_t, mask_data=mask_t)
 
@@ -110,7 +146,28 @@ class ClipMaskIntensityPercentiles(Transform):
 
         return img
 
+
 class ClipMaskIntensityPercentilesd(MapTransform):
+    """
+    Dictionary wrapper for ClipMaskIntensityPercentiles.
+    Args:
+        keys: Keys of the corresponding items to be transformed.
+        mask_key: Key to the mask data in the input dictionary used to compute percentiles. Only intensity values where the mask is positive will be considered.
+        lower: Lower percentile value (0-100) for clipping. If None, no lower clipping is applied.
+        upper: Upper percentile value (0-100) for clipping. If None, no upper clipping is applied.
+        sharpness_factor: Optional factor to enhance contrast after clipping. If None, no sharpness enhancement is applied.
+        channel_wise: If True, compute percentiles separately for each channel. If False, compute globally.
+        dtype: Data type of the output. Defaults to np.float32.
+        allow_missing_keys: If True, missing keys will not raise an error. Defaults to False.
+    Example:
+        >>> transform = ClipMaskIntensityPercentilesd(
+        ...     keys=["image"],
+        ...     mask_key="mask",
+        ...     lower=2,
+        ...     upper=98,
+        ...     sharpness_factor=1.0
+        ... )
+    """
 
     def __init__(
         self,
@@ -125,7 +182,11 @@ class ClipMaskIntensityPercentilesd(MapTransform):
     ) -> None:
         super().__init__(keys, allow_missing_keys)
         self.scaler = ClipMaskIntensityPercentiles(
-            lower=lower, upper=upper, sharpness_factor=sharpness_factor, channel_wise=channel_wise, dtype=dtype
+            lower=lower,
+            upper=upper,
+            sharpness_factor=sharpness_factor,
+            channel_wise=channel_wise,
+            dtype=dtype,
         )
         self.mask_key = mask_key
 
@@ -134,16 +195,32 @@ class ClipMaskIntensityPercentilesd(MapTransform):
         for key in self.key_iterator(d):
             d[key] = self.scaler(d[key], d[self.mask_key])
         return d
-    
-    
 
 
 class ElementwiseProductd(MapTransform):
+    """
+    A dictionary-based transform that computes the elementwise product of two arrays.
+    This transform multiplies two input arrays element-by-element and stores the result
+    in a specified output key.
+    Args:
+        keys: Collection of keys to select from the input dictionary. Must contain exactly
+            two keys whose corresponding values will be multiplied together.
+        output_key: Key in the output dictionary where the product result will be stored.
+    Returns:
+        Dictionary with the elementwise product stored at the output_key.
+    Example:
+        >>> transform = ElementwiseProductd(keys=["image1", "image2"], output_key="product")
+        >>> data = {"image1": np.array([1, 2, 3]), "image2": np.array([2, 3, 4])}
+        >>> result = transform(data)
+        >>> result["product"]
+        array([ 2,  6, 12])
+    """
+
     def __init__(self, keys: KeysCollection, output_key: str) -> None:
         super().__init__(keys)
         self.output_key = output_key
 
-    def __call__(self, data) -> NdarrayOrTensor:
+    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
         d = dict(data)
         d[self.output_key] = d[self.keys[0]] * d[self.keys[1]]
         return d
@@ -159,6 +236,7 @@ class CLAHEd(MapTransform):
         clip_limit (float): Threshold for contrast limiting. Default is 2.0.
         tile_grid_size (Union[tuple, Sequence[int]]): Size of grid for histogram equalization (default: (8,8)).
     """
+
     def __init__(
         self,
         keys: KeysCollection,
@@ -184,13 +262,13 @@ class CLAHEd(MapTransform):
 
             image_clahe = np.stack([clahe.apply(slice) for slice in image[0]])
 
-
             # Convert back to float in [0,1]
             processed_img = image_clahe.astype(np.float32) / 255.0
             reshaped_ = processed_img.reshape(1, *processed_img.shape)
             d[key] = torch.from_numpy(reshaped_).to(image_.device)
         return d
-    
+
+
 class NormalizeIntensity_custom(Transform):
     """
     Normalize input based on the `subtrahend` and `divisor`: `(img - subtrahend) / divisor`.
@@ -240,9 +318,11 @@ class NormalizeIntensity_custom(Transform):
         x = torch.std(x.float(), unbiased=False)
         return x.item() if x.numel() == 1 else x
 
-    def _normalize(self, img: NdarrayOrTensor, mask_data: NdarrayOrTensor, sub=None, div=None) -> NdarrayOrTensor:
+    def _normalize(
+        self, img: NdarrayOrTensor, mask_data: NdarrayOrTensor, sub=None, div=None
+    ) -> NdarrayOrTensor:
         img, *_ = convert_data_type(img, dtype=torch.float32)
-        '''
+        """
         if self.nonzero:
             slices = img != 0
             masked_img = img[slices]
@@ -251,7 +331,7 @@ class NormalizeIntensity_custom(Transform):
         else:
             slices = None
             masked_img = img
-        '''
+        """
         slices = None
         mask_data = mask_data.squeeze(0)
         slices_mask = mask_data > 0
@@ -288,9 +368,13 @@ class NormalizeIntensity_custom(Transform):
         dtype = self.dtype or img.dtype
         if self.channel_wise:
             if self.subtrahend is not None and len(self.subtrahend) != len(img):
-                raise ValueError(f"img has {len(img)} channels, but subtrahend has {len(self.subtrahend)} components.")
+                raise ValueError(
+                    f"img has {len(img)} channels, but subtrahend has {len(self.subtrahend)} components."
+                )
             if self.divisor is not None and len(self.divisor) != len(img):
-                raise ValueError(f"img has {len(img)} channels, but divisor has {len(self.divisor)} components.")
+                raise ValueError(
+                    f"img has {len(img)} channels, but divisor has {len(self.divisor)} components."
+                )
 
             if not img.dtype.is_floating_point:
                 img, *_ = convert_data_type(img, dtype=torch.float32)
@@ -308,21 +392,27 @@ class NormalizeIntensity_custom(Transform):
         out = convert_to_dst_type(img, img, dtype=dtype)[0]
         return out
 
+
 class NormalizeIntensity_customd(MapTransform):
     """
-    Dictionary-based wrapper of :py:class:`monai.transforms.NormalizeIntensity`.
-    This transform can normalize only non-zero values or entire image, and can also calculate
-    mean and std on each channel separately.
+    Dictionary-based wrapper of :class:`NormalizeIntensity_custom`.
+
+    The mean and standard deviation are calculated only from intensities which are
+    defined in the mask provided through ``mask_key``.
 
     Args:
         keys: keys of the corresponding items to be transformed.
-            See also: monai.transforms.MapTransform
-        subtrahend: the amount to subtract by (usually the mean)
-        divisor: the amount to divide by (usually the standard deviation)
+            See also: :py:class:`monai.transforms.MapTransform`
+        mask_key: key of the corresponding mask item to be used for calculating
+            statistics (mean and std).
+        subtrahend: the amount to subtract by (usually the mean). If None,
+            the mean is calculated from the masked region of the input image.
+        divisor: the amount to divide by (usually the standard deviation). If None,
+            the std is calculated from the masked region of the input image.
         nonzero: whether only normalize non-zero values.
         channel_wise: if True, calculate on each channel separately, otherwise, calculate on
-            the entire image directly. default to False.
-        dtype: output data type, if None, same as input image. defaults to float32.
+            the entire image directly. Defaults to False.
+        dtype: output data type, if None, same as input image. Defaults to float32.
         allow_missing_keys: don't raise exception if key is missing.
     """
 
@@ -332,19 +422,21 @@ class NormalizeIntensity_customd(MapTransform):
         self,
         keys: KeysCollection,
         mask_key: str,
-        subtrahend:Union[ NdarrayOrTensor, None] = None,
-        divisor: Union[ NdarrayOrTensor, None] = None,
+        subtrahend: Union[NdarrayOrTensor, None] = None,
+        divisor: Union[NdarrayOrTensor, None] = None,
         nonzero: bool = False,
         channel_wise: bool = False,
         dtype: DtypeLike = np.float32,
         allow_missing_keys: bool = False,
     ) -> None:
         super().__init__(keys, allow_missing_keys)
-        self.normalizer = NormalizeIntensity_custom(subtrahend, divisor, nonzero, channel_wise, dtype)
+        self.normalizer = NormalizeIntensity_custom(
+            subtrahend, divisor, nonzero, channel_wise, dtype
+        )
         self.mask_key = mask_key
 
     def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
         d = dict(data)
         for key in self.key_iterator(d):
             d[key] = self.normalizer(d[key], d[self.mask_key])
-        return d
+        return d

src/data/data_loader.py

@@ -1,43 +1,29 @@
-import argparse
 import os
 import numpy as np
-from monai.config import KeysCollection
-from monai.data import Dataset, load_decathlon_datalist, ITKReader, NumpyReader, PersistentDataset
+from monai.data import load_decathlon_datalist, ITKReader, PersistentDataset
 from monai.transforms import (
     Compose,
     LoadImaged,
-    MapTransform,
-    ScaleIntensityRanged,
-    SplitDimd,
     ToTensord,
-    ConcatItemsd, 
-    SelectItemsd,
-    EnsureChannelFirstd,
-    RepeatChanneld,
+    ConcatItemsd,
     DeleteItemsd,
     EnsureTyped,
-    ClipIntensityPercentilesd,
-    MaskIntensityd,
     RandCropByPosNegLabeld,
     NormalizeIntensityd,
-    SqueezeDimd,
-    ScaleIntensityd,
-    ScaleIntensityd,
     Transposed,
     RandWeightedCropd,
 )
 from .custom_transforms import (
-    NormalizeIntensity_customd, 
-    ClipMaskIntensityPercentilesd, 
+    NormalizeIntensity_customd,
+    ClipMaskIntensityPercentilesd,
     ElementwiseProductd,
 )
 import torch
 from torch.utils.data.dataloader import default_collate
-import matplotlib.pyplot as plt
 from typing import Literal
-import monai
 import collections.abc
 
+
 def list_data_collate(batch: collections.abc.Sequence):
     """
     Combine instances from a list of dicts into a single dict, by stacking them along first dim
@@ -51,55 +37,71 @@ def list_data_collate(batch: collections.abc.Sequence):
 
         if all("final_heatmap" in ix for ix in item):
             data["final_heatmap"] = torch.stack([ix["final_heatmap"] for ix in item], dim=0)
-            
+
         batch[i] = data
     return default_collate(batch)
 
 
-
 def data_transform(args):
-
     if args.use_heatmap:
         transform = Compose(
             [
-                LoadImaged(keys=["image", "mask","dwi", "adc", "heatmap"], reader=ITKReader(), ensure_channel_first=True, dtype=np.float32),
+                LoadImaged(
+                    keys=["image", "mask", "dwi", "adc", "heatmap"],
+                    reader=ITKReader(),
+                    ensure_channel_first=True,
+                    dtype=np.float32,
+                ),
                 ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"),
-                ConcatItemsd(keys=["image", "dwi", "adc"], name="image", dim=0),  # stacks to (3, H, W)
+                ConcatItemsd(
+                    keys=["image", "dwi", "adc"], name="image", dim=0
+                ),  # stacks to (3, H, W)
                 NormalizeIntensity_customd(keys=["image"], channel_wise=True, mask_key="mask"),
                 ElementwiseProductd(keys=["mask", "heatmap"], output_key="final_heatmap"),
-                RandWeightedCropd(keys=["image", "final_heatmap"],
-                                w_key="final_heatmap",
-                                spatial_size=(args.tile_size,args.tile_size,args.depth),
-                                num_samples=args.tile_count),
+                RandWeightedCropd(
+                    keys=["image", "final_heatmap"],
+                    w_key="final_heatmap",
+                    spatial_size=(args.tile_size, args.tile_size, args.depth),
+                    num_samples=args.tile_count,
+                ),
                 EnsureTyped(keys=["label"], dtype=torch.float32),
                 Transposed(keys=["image"], indices=(0, 3, 1, 2)),
-                DeleteItemsd(keys=['mask', 'dwi', 'adc', 'heatmap']),
+                DeleteItemsd(keys=["mask", "dwi", "adc", "heatmap"]),
                 ToTensord(keys=["image", "label", "final_heatmap"]),
             ]
         )
     else:
         transform = Compose(
             [
-                LoadImaged(keys=["image", "mask","dwi", "adc"], reader=ITKReader(), ensure_channel_first=True, dtype=np.float32),
+                LoadImaged(
+                    keys=["image", "mask", "dwi", "adc"],
+                    reader=ITKReader(),
+                    ensure_channel_first=True,
+                    dtype=np.float32,
+                ),
                 ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"),
-                ConcatItemsd(keys=["image", "dwi", "adc"], name="image", dim=0),  # stacks to (3, H, W)
+                ConcatItemsd(
+                    keys=["image", "dwi", "adc"], name="image", dim=0
+                ),  # stacks to (3, H, W)
                 NormalizeIntensityd(keys=["image"], channel_wise=True),
-                RandCropByPosNegLabeld(keys=["image"],
-                                label_key="mask",
-                                spatial_size=(args.tile_size,args.tile_size,args.depth),
-                                pos=1,
-                                neg=0,
-                                num_samples=args.tile_count),
+                RandCropByPosNegLabeld(
+                    keys=["image"],
+                    label_key="mask",
+                    spatial_size=(args.tile_size, args.tile_size, args.depth),
+                    pos=1,
+                    neg=0,
+                    num_samples=args.tile_count,
+                ),
                 EnsureTyped(keys=["label"], dtype=torch.float32),
                 Transposed(keys=["image"], indices=(0, 3, 1, 2)),
-                DeleteItemsd(keys=['mask', 'dwi', 'adc']),
+                DeleteItemsd(keys=["mask", "dwi", "adc"]),
                 ToTensord(keys=["image", "label"]),
             ]
         )
     return transform
 
-def get_dataloader(args, split: Literal["train", "test"]):
 
+def get_dataloader(args, split: Literal["train", "test"]):
     data_list = load_decathlon_datalist(
         data_list_file_path=args.dataset_json,
         data_list_key=split,
@@ -110,16 +112,17 @@ def get_dataloader(args, split: Literal["train", "test"]):
     cache_dir_ = os.path.join(args.logdir, "cache")
     os.makedirs(os.path.join(cache_dir_, split), exist_ok=True)
     transform = data_transform(args)
-    dataset = PersistentDataset(data=data_list, transform=transform, cache_dir= os.path.join(cache_dir_, split))
+    dataset = PersistentDataset(
+        data=data_list, transform=transform, cache_dir=os.path.join(cache_dir_, split)
+    )
     loader = torch.utils.data.DataLoader(
         dataset,
         batch_size=args.batch_size,
         shuffle=(split == "train"),
         num_workers=args.workers,
         pin_memory=True,
-        multiprocessing_context="spawn" if args.workers > 0 else None,
+        multiprocessing_context="fork" if args.workers > 0 else None,
         sampler=None,
         collate_fn=list_data_collate,
     )
     return loader
-

src/model/MIL.py

@@ -1,13 +1,10 @@
-
 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
+
 models, _ = optional_import("torchvision.models")
 
 
@@ -16,8 +13,7 @@ class MILModel_3D(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]`,
     where `B` is the batch_size of PyTorch Dataloader and `N` is the number of instances
-    extracted from every original image in the batch. A tutorial example is available at:
-    https://github.com/Project-MONAI/tutorials/tree/master/pathology/multiple_instance_learning.
+    extracted from every original image in the batch.
 
     Args:
         num_classes: number of output classes.
@@ -29,10 +25,9 @@ class MILModel_3D(nn.Module):
             - ``"att_trans"`` - transformer MIL https://arxiv.org/abs/2111.01556.
             - ``"att_trans_pyramid"`` - transformer pyramid MIL https://arxiv.org/abs/2111.01556.
 
-        pretrained: init backbone with pretrained weights, defaults to ``True``.
         backbone: Backbone classifier CNN (either ``None``, a ``nn.Module`` that returns features,
             or a string name of a torchvision model).
-            Defaults to ``None``, in which case ResNet50 is used.
+            Defaults to ``None``, in which case ResNet18 is used.
         backbone_num_features: Number of output features of the backbone CNN
             Defaults to ``None`` (necessary only when using a custom backbone)
         trans_blocks: number of the blocks in `TransformEncoder` layer.
@@ -63,7 +58,11 @@ class MILModel_3D(nn.Module):
         self.transformer: nn.Module | None = None
 
         if backbone is None:
-            net = resnet.resnet18(spatial_dims=3, n_input_channels=3, num_classes=5, )
+            net = resnet.resnet18(
+                spatial_dims=3,
+                n_input_channels=3,
+                num_classes=5,
+            )
             nfc = net.fc.in_features  # save the number of final features
             net.fc = torch.nn.Identity()  # remove final linear layer
 
@@ -72,7 +71,6 @@ class MILModel_3D(nn.Module):
             if mil_mode == "att_trans_pyramid":
                 # register hooks to capture outputs of intermediate layers
                 def forward_hook(layer_name):
-
                     def hook(module, input, output):
                         self.extra_outputs[layer_name] = output
 
@@ -105,31 +103,23 @@ class MILModel_3D(nn.Module):
             nfc = backbone_num_features
             net.fc = torch.nn.Identity()  # remove final linear layer
 
-            self.extra_outputs: dict[str, torch.Tensor] = {}
-
             if mil_mode == "att_trans_pyramid":
                 # register hooks to capture outputs of intermediate layers
-                def forward_hook(layer_name):
-
-                    def hook(module, input, output):
-                        self.extra_outputs[layer_name] = output
-
-                    return hook
-
-                net.layer1.register_forward_hook(forward_hook("layer1"))
-                net.layer2.register_forward_hook(forward_hook("layer2"))
-                net.layer3.register_forward_hook(forward_hook("layer3"))
-                net.layer4.register_forward_hook(forward_hook("layer4"))
+                raise ValueError(
+                    "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")
+                raise ValueError(
+                    "Number of endencoder features must be provided for a custom backbone model"
+                )
 
         else:
             raise ValueError("Unsupported backbone")
-        
+
         if backbone is not None and mil_mode not in ["mean", "max", "att", "att_trans"]:
             raise ValueError("Custom backbone is not supported for the mode:" + str(mil_mode))
-        
+
         if self.mil_mode in ["mean", "max"]:
             pass
         elif self.mil_mode == "att":
@@ -144,7 +134,8 @@ class MILModel_3D(nn.Module):
             transformer_list = nn.ModuleList(
                 [
                     nn.TransformerEncoder(
-                        nn.TransformerEncoderLayer(d_model=64, nhead=8, dropout=trans_dropout), num_layers=trans_blocks
+                        nn.TransformerEncoderLayer(d_model=64, nhead=8, dropout=trans_dropout),
+                        num_layers=trans_blocks,
                     ),
                     nn.Sequential(
                         nn.Linear(192, 64),
@@ -206,10 +197,26 @@ class MILModel_3D(nn.Module):
             x = self.myfc(x)
 
         elif self.mil_mode == "att_trans_pyramid" and self.transformer is not None:
-            l1 = torch.mean(self.extra_outputs["layer1"], dim=(2, 3, 4)).reshape(sh[0], sh[1], -1).permute(1, 0, 2)
-            l2 = torch.mean(self.extra_outputs["layer2"], dim=(2, 3, 4)).reshape(sh[0], sh[1], -1).permute(1, 0, 2)
-            l3 = torch.mean(self.extra_outputs["layer3"], dim=(2, 3, 4)).reshape(sh[0], sh[1], -1).permute(1, 0, 2)
-            l4 = torch.mean(self.extra_outputs["layer4"], dim=(2, 3, 4)).reshape(sh[0], sh[1], -1).permute(1, 0, 2)
+            l1 = (
+                torch.mean(self.extra_outputs["layer1"], dim=(2, 3, 4))
+                .reshape(sh[0], sh[1], -1)
+                .permute(1, 0, 2)
+            )
+            l2 = (
+                torch.mean(self.extra_outputs["layer2"], dim=(2, 3, 4))
+                .reshape(sh[0], sh[1], -1)
+                .permute(1, 0, 2)
+            )
+            l3 = (
+                torch.mean(self.extra_outputs["layer3"], dim=(2, 3, 4))
+                .reshape(sh[0], sh[1], -1)
+                .permute(1, 0, 2)
+            )
+            l4 = (
+                torch.mean(self.extra_outputs["layer4"], dim=(2, 3, 4))
+                .reshape(sh[0], sh[1], -1)
+                .permute(1, 0, 2)
+            )
 
             transformer_list = cast(nn.ModuleList, self.transformer)
 
@@ -242,7 +249,3 @@ class MILModel_3D(nn.Module):
             x = self.calc_head(x)
 
         return x
-
-
-
-

src/model/csPCa_model.py

@@ -1,32 +1,72 @@
-
 from __future__ import annotations
-
-from typing import cast
-
 import torch
 import torch.nn as nn
-
 from monai.utils.module import optional_import
-models, _ = optional_import("torchvision.models")
 
+models, _ = optional_import("torchvision.models")
 
 
 class SimpleNN(nn.Module):
+    """
+    A simple Multi-Layer Perceptron (MLP) for binary classification.
+
+    This network consists of two hidden layers with ReLU activation and a dropout layer,
+    followed by a final sigmoid activation for probability output.
+
+    Args:
+        input_dim (int): The number of input features.
+    """
+
     def __init__(self, input_dim):
         super(SimpleNN, self).__init__()
         self.net = nn.Sequential(
             nn.Linear(input_dim, 256),
             nn.ReLU(),
-            nn.Linear( 256,128),
+            nn.Linear(256, 128),
             nn.ReLU(),
-            nn.Dropout(p=0.3),    
+            nn.Dropout(p=0.3),
             nn.Linear(128, 1),
-            nn.Sigmoid()   # since binary classification
+            nn.Sigmoid(),  # since binary classification
         )
+
     def forward(self, x):
+        """
+        Forward pass of the classifier.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (Batch, input_dim).
+
+        Returns:
+            torch.Tensor: Output probabilities of shape (Batch, 1).
+        """
         return self.net(x)
 
+
 class csPCa_Model(nn.Module):
+    """
+    Clinically Significant Prostate Cancer (csPCa) risk prediction model using a MIL backbone.
+
+    This model repurposes a pre-trained Multiple Instance Learning (MIL) backbone (originally
+    designed for PI-RADS prediction) for binary csPCa risk assessment. It utilizes the
+    backbone's feature extractor, transformer, and attention mechanism to aggregate instance-level
+    features into a bag-level embedding.
+
+    The original fully connected classification head of the backbone is replaced by a
+    custom :class:`SimpleNN` head for the new task.
+
+    Args:
+        backbone (nn.Module): A pre-trained MIL model. The backbone must possess the
+            following attributes/sub-modules:
+            - ``net``: The CNN feature extractor.
+            - ``transformer``: A sequence modeling module.
+            - ``attention``: An attention mechanism for pooling.
+            - ``myfc``: The original fully connected layer (used to determine feature dimensions).
+
+    Attributes:
+        fc_cspca (SimpleNN): The new classification head for csPCa prediction.
+        backbone: The MIL based PI-RADS classifier.
+    """
+
     def __init__(self, backbone):
         super().__init__()
         self.backbone = backbone
@@ -47,4 +87,3 @@ class csPCa_Model(nn.Module):
 
         x = self.fc_cspca(x)
         return x
-

src/preprocessing/center_crop.py

@@ -8,10 +8,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#python scripts/center_crop.py --file_name path/to/t2_image --out_name cropped_t2
+# python scripts/center_crop.py --file_name path/to/t2_image --out_name cropped_t2
 
 
-#import argparse
+# import argparse
 from typing import Union
 
 import SimpleITK as sitk  # noqa N813
@@ -41,7 +41,9 @@ def crop(image: sitk.Image, margin: Union[int, float], interpolator=sitk.sitkLin
 
     # calculate new origin and new image size
     if all([isinstance(m, float) for m in _flatten(margin)]):
-        assert all([m >= 0 and m < 0.5 for m in _flatten(margin)]), "margins must be between 0 and 0.5"
+        assert all([m >= 0 and m < 0.5 for m in _flatten(margin)]), (
+            "margins must be between 0 and 0.5"
+        )
         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
@@ -61,4 +63,3 @@ def crop(image: sitk.Image, margin: Union[int, float], interpolator=sitk.sitkLin
     ref_image.SetDirection(image.GetDirection())
 
     return sitk.Resample(image, ref_image, interpolator=interpolator)
-

src/preprocessing/generate_heatmap.py

@@ -1,76 +1,93 @@
-
 import os
 import numpy as np
 import nrrd
-import json
-import pandas as pd
-import json
-import SimpleITK as sitk
-import multiprocessing
-
+from tqdm import tqdm
 import logging
 
 
 def get_heatmap(args):
+    """
+    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
+    segmentation masks to create normalized heatmaps. It combines the DWI
+    and ADC heatmaps through element-wise multiplication.
+    Args:
+        args: An object containing the following attributes:
+            - t2_dir (str): Directory path containing T2 image files.
+            - dwi_dir (str): Directory path containing DWI image files.
+            - adc_dir (str): Directory path containing ADC image files.
+            - seg_dir (str): Directory path containing segmentation mask files.
+            - output_dir (str): Base output directory where 'heatmaps/' subdirectory will be created.
+            - heatmapdir (str): Output directory for heatmap files (created by function).
+    Returns:
+        args: The modified args object with heatmapdir attribute set.
+    Raises:
+        FileNotFoundError: If input directories or files do not exist.
+        ValueError: If NRRD files cannot be read properly.
+    Notes:
+        - DWI heatmap is normalized as (dwi - min) / (max - min)
+        - ADC heatmap is normalized as (max - adc) / (max - min) (inverted)
+        - Final heatmap is re-normalized to [0, 1] range
+        - If all values in a mask region are identical, the heatmap is skipped for that modality
+        - Output files are written in NRRD format with the same header as the input DWI file
+    """
 
     files = os.listdir(args.t2_dir)
-    args.heatmapdir = os.path.join(args.output_dir, 'heatmaps/')
+    args.heatmapdir = os.path.join(args.output_dir, "heatmaps/")
     os.makedirs(args.heatmapdir, exist_ok=True)
-    for file in files:
-
+    logging.info("Starting heatmap generation")
+    for file in tqdm(files):
         bool_dwi = False
         bool_adc = False
         mask, _ = nrrd.read(os.path.join(args.seg_dir, file))
-        dwi, _ = nrrd.read(os.path.join(args.dwi_dir, file))  
-        adc, _ = nrrd.read(os.path.join(args.adc_dir, file))
-
+        dwi, header_dwi = nrrd.read(os.path.join(args.dwi_dir, file))
+        adc, header_adc = nrrd.read(os.path.join(args.adc_dir, file))
         nonzero_vals_dwi = dwi[mask > 0]
+
         if len(nonzero_vals_dwi) > 0:
             min_val = nonzero_vals_dwi.min()
             max_val = nonzero_vals_dwi.max()
             heatmap_dwi = np.zeros_like(dwi, dtype=np.float32)
-            
+
             if min_val != max_val:
                 heatmap_dwi = (dwi - min_val) / (max_val - min_val)
-                masked_heatmap_dwi = np.where(mask > 0, heatmap_dwi, heatmap_dwi[mask>0].min())
+                masked_heatmap_dwi = np.where(mask > 0, heatmap_dwi, heatmap_dwi[mask > 0].min())
             else:
                 bool_dwi = True
-                
+
         else:
             bool_dwi = True
 
         nonzero_vals_adc = adc[mask > 0]
+
         if len(nonzero_vals_adc) > 0:
             min_val = nonzero_vals_adc.min()
             max_val = nonzero_vals_adc.max()
             heatmap_adc = np.zeros_like(adc, dtype=np.float32)
-            
+
             if min_val != max_val:
                 heatmap_adc = (max_val - adc) / (max_val - min_val)
-                masked_heatmap_adc = np.where(mask > 0, heatmap_adc, heatmap_adc[mask>0].min())
+                masked_heatmap_adc = np.where(mask > 0, heatmap_adc, heatmap_adc[mask > 0].min())
             else:
                 bool_adc = True
 
         else:
             bool_adc = True
-            
 
-        if bool_dwi:
-            mix_mask = masked_heatmap_adc  
-        if bool_adc:
-            mix_mask = masked_heatmap_dwi 
         if not bool_dwi and not bool_adc:
             mix_mask = masked_heatmap_dwi * masked_heatmap_adc
+            write_header = header_dwi
+        elif bool_dwi:
+            mix_mask = masked_heatmap_adc
+            write_header = header_adc
+        elif bool_adc:
+            mix_mask = masked_heatmap_dwi
+            write_header = header_dwi
         else:
             mix_mask = np.ones_like(adc, dtype=np.float32)
+            write_header = header_dwi
 
         mix_mask = (mix_mask - mix_mask.min()) / (mix_mask.max() - mix_mask.min())
+        nrrd.write(os.path.join(args.heatmapdir, file), mix_mask, write_header)
 
-
-        nrrd.write(os.path.join(args.heatmapdir, file), mix_mask)
-    
     return args
-
-
-    
-    

src/preprocessing/histogram_match.py

@@ -1,16 +1,30 @@
-import SimpleITK as sitk 
 import os
-import numpy as np 
 import nrrd
-from tqdm import tqdm
-import pandas as pd
-import random
-import json
 from skimage import exposure
-import multiprocessing
 import logging
-def get_histmatched(data, ref_data, mask, ref_mask):
+import numpy as np
+from tqdm import tqdm
+
 
+def get_histmatched(
+    data: np.ndarray, ref_data: np.ndarray, mask: np.ndarray, ref_mask: np.ndarray
+) -> np.ndarray:
+    """
+    Perform histogram matching on source data using a reference image.
+    This function adjusts the histogram of the source image to match the
+    histogram of the reference image within masked regions. Only pixels
+    where the mask is greater than 0 are considered for matching.
+    Args:
+        data: Source image array to be histogram matched.
+        ref_data: Reference image array whose histogram will be used as target.
+        mask: Binary mask for source image indicating valid pixels (values > 0).
+        ref_mask: Binary mask for reference image indicating valid pixels (values > 0).
+    Returns:
+        Histogram-matched image with the same shape as input data.
+        Only pixels in masked regions are modified; unmasked pixels remain unchanged.
+    Example:
+        >>> matched = get_histmatched(source_img, reference_img, source_mask, ref_mask)
+    """
     source_pixels = data[mask > 0]
     ref_pixels = ref_data[ref_mask > 0]
     matched_pixels = exposure.match_histograms(source_pixels, ref_pixels)
@@ -19,36 +33,35 @@ def get_histmatched(data, ref_data, mask, ref_mask):
 
     return matched_img
 
-def histmatch(args):
 
+def histmatch(args):
     files = os.listdir(args.t2_dir)
-    
-    t2_histmatched_dir = os.path.join(args.output_dir, 't2_histmatched')
-    dwi_histmatched_dir = os.path.join(args.output_dir, 'DWI_histmatched')
-    adc_histmatched_dir = os.path.join(args.output_dir, 'ADC_histmatched')
+
+    t2_histmatched_dir = os.path.join(args.output_dir, "t2_histmatched")
+    dwi_histmatched_dir = os.path.join(args.output_dir, "DWI_histmatched")
+    adc_histmatched_dir = os.path.join(args.output_dir, "ADC_histmatched")
     os.makedirs(t2_histmatched_dir, exist_ok=True)
     os.makedirs(dwi_histmatched_dir, exist_ok=True)
     os.makedirs(adc_histmatched_dir, exist_ok=True)
     logging.info("Starting histogram matching")
-    for file in files:
 
+    for file in tqdm(files):
         t2_image, header_t2 = nrrd.read(os.path.join(args.t2_dir, file))
         dwi_image, header_dwi = nrrd.read(os.path.join(args.dwi_dir, file))
         adc_image, header_adc = nrrd.read(os.path.join(args.adc_dir, file))
 
-        ref_t2, _ = nrrd.read(os.path.join(args.project_dir, 'dataset', 't2_reference.nrrd'))
-        ref_dwi, _ = nrrd.read(os.path.join(args.project_dir, 'dataset', 'dwi_reference.nrrd'))
-        ref_adc , _ = nrrd.read(os.path.join(args.project_dir, 'dataset', 'adc_reference.nrrd'))
-
+        ref_t2, _ = nrrd.read(os.path.join(args.project_dir, "dataset", "t2_reference.nrrd"))
+        ref_dwi, _ = nrrd.read(os.path.join(args.project_dir, "dataset", "dwi_reference.nrrd"))
+        ref_adc, _ = nrrd.read(os.path.join(args.project_dir, "dataset", "adc_reference.nrrd"))
         prostate_mask, _ = nrrd.read(os.path.join(args.seg_dir, file))
-        ref_prostate_mask, _ = nrrd.read(os.path.join(args.project_dir, 'dataset', 'prostate_segmentation_reference.nrrd'))
+        ref_prostate_mask, _ = nrrd.read(
+            os.path.join(args.project_dir, "dataset", "prostate_segmentation_reference.nrrd")
+        )
 
         histmatched_t2 = get_histmatched(t2_image, ref_t2, prostate_mask, ref_prostate_mask)
         histmatched_dwi = get_histmatched(dwi_image, ref_dwi, prostate_mask, ref_prostate_mask)
         histmatched_adc = get_histmatched(adc_image, ref_adc, prostate_mask, ref_prostate_mask)
 
-
-
         nrrd.write(os.path.join(t2_histmatched_dir, file), histmatched_t2, header_t2)
         nrrd.write(os.path.join(dwi_histmatched_dir, file), histmatched_dwi, header_dwi)
         nrrd.write(os.path.join(adc_histmatched_dir, file), histmatched_adc, header_adc)
@@ -58,5 +71,3 @@ def histmatch(args):
         args.adc_dir = adc_histmatched_dir
 
     return args
-
-

src/preprocessing/prostate_mask.py

@@ -1,65 +1,63 @@
 import os
-from typing import Union
-import SimpleITK as sitk
 import numpy as np
 import nrrd
-import matplotlib.pyplot as plt
 from tqdm import tqdm
-from AIAH_utility.viewer import BasicViewer, ListViewer
-from PIL import Image
-import monai
 from monai.bundle import ConfigParser
-from monai.config import print_config
 import torch
-import sys
-import os
-import nibabel as nib
-import shutil
 
-from tqdm import trange, tqdm
 
-from monai.data import DataLoader, Dataset, TestTimeAugmentation, create_test_image_2d
-from monai.losses import DiceLoss
-from monai.metrics import DiceMetric
-from monai.networks.nets import UNet
 from monai.transforms import (
-    Activationsd,
-    AsDiscreted,
     Compose,
-    CropForegroundd,
-    DivisiblePadd,
-    Invertd,
     LoadImaged,
     ScaleIntensityd,
-    RandRotated,
-    RandRotate,
-    InvertibleTransform,
-    RandFlipd,
-    Activations,
-    AsDiscrete,
     NormalizeIntensityd,
 )
 from monai.utils import set_determinism
 from monai.transforms import (
-    Resize,
     EnsureChannelFirstd,
     Orientationd,
     Spacingd,
     EnsureTyped,
 )
-import nrrd
-
-set_determinism(43)
 from monai.data import MetaTensor
-import SimpleITK as sitk
-import pandas as pd
 import logging
+
+set_determinism(43)
+
+
 def get_segmask(args):
+    """
+    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.
+    It applies preprocessing transformations, runs the segmentation model, and saves the
+    predicted masks. Post-processing is applied to retain only the top 10 slices with
+    the highest non-zero voxel counts.
+    Args:
+        args: An arguments object containing:
+            - output_dir (str): Base output directory where segmentation masks will be saved
+            - project_dir (str): Root project directory containing model config and checkpoint
+            - t2_dir (str): Directory containing input T2-weighted MRI images in NRRD format
+    Returns:
+        args: The updated arguments object with seg_dir added, pointing to the
+              prostate_mask subdirectory within output_dir
+    Raises:
+        FileNotFoundError: If the model checkpoint or config file is not found
+        RuntimeError: If CUDA operations fail on GPU
+    Notes:
+        - Automatically selects GPU (CUDA) if available, otherwise uses CPU
+        - Applies MONAI transformations: loading, orientation (RAS), spacing (0.5mm isotropic),
+          intensity scaling and normalization
+        - Post-processing filters predictions to top 10 slices by non-zero voxel density
+        - Output masks are saved in NRRD format preserving original image headers
+    """
 
     args.seg_dir = os.path.join(args.output_dir, "prostate_mask")
     os.makedirs(args.seg_dir, exist_ok=True)
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    torch.cuda.empty_cache()
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
     model_config_file = os.path.join(args.project_dir, "config", "inference.json")
     model_config = ConfigParser()
     model_config.read_config(model_config_file)
@@ -67,24 +65,16 @@ def get_segmask(args):
     model_config["dataset_dir"] = args.t2_dir
     files = os.listdir(args.t2_dir)
     model_config["datalist"] = [os.path.join(args.t2_dir, f) for f in files]
-    
-
     checkpoint = os.path.join(
         args.project_dir,
         "models",
         "prostate_segmentation_model.pt",
     )
-    preprocessing = model_config.get_parsed_content("preprocessing")
     model = model_config.get_parsed_content("network_def").to(device)
     inferer = model_config.get_parsed_content("inferer")
-    postprocessing = model_config.get_parsed_content("postprocessing")
-    dataloader = model_config.get_parsed_content("dataloader")
     model.load_state_dict(torch.load(checkpoint, map_location=device))
     model.eval()
 
-    torch.cuda.empty_cache()
-    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-
     keys = "image"
     transform = Compose(
         [
@@ -99,8 +89,8 @@ def get_segmask(args):
     )
     logging.info("Starting prostate segmentation")
     for file in tqdm(files):
-
         data = {"image": os.path.join(args.t2_dir, file)}
+        _, header_t2 = nrrd.read(data["image"])
         transformed_data = transform(data)
         a = transformed_data
         with torch.no_grad():
@@ -114,15 +104,12 @@ def get_segmask(args):
         temp = transform.inverse(transformed_data)
         pred_temp = temp["image"][0].numpy()
         pred_nrrd = np.round(pred_temp)
-        
-        nonzero_counts = np.count_nonzero(pred_nrrd, axis=(0,1))
+
+        nonzero_counts = np.count_nonzero(pred_nrrd, axis=(0, 1))
         top_slices = np.argsort(nonzero_counts)[-10:]
         output_ = np.zeros_like(pred_nrrd)
-        output_[:,:,top_slices] = pred_nrrd[:,:,top_slices]
-
-        nrrd.write(os.path.join(args.seg_dir, file), output_)
-
-        return args
+        output_[:, :, top_slices] = pred_nrrd[:, :, top_slices]
 
+        nrrd.write(os.path.join(args.seg_dir, file), output_, header_t2)
 
-    
+    return args

src/preprocessing/register_and_crop.py

@@ -1,25 +1,42 @@
-import SimpleITK as sitk 
+import SimpleITK as sitk
 import os
-import numpy as np 
-import nrrd
 from tqdm import tqdm
-import pandas as pd
 from picai_prep.preprocessing import PreprocessingSettings, Sample
-import multiprocessing
 from .center_crop import crop
 import logging
+
+
 def register_files(args):
+    """
+    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
+    new spacing of (0.4, 0.4, 3.0) mm, preprocesses them using the Sample class, and crops
+    them with specified margins. The processed images are saved to new output directories.
+    Args:
+        args: An argument object containing:
+            - t2_dir (str): Directory path containing T2 weighted images
+            - dwi_dir (str): Directory path containing DWI (Diffusion Weighted Imaging) images
+            - adc_dir (str): Directory path containing ADC (Apparent Diffusion Coefficient) images
+            - output_dir (str): Directory path where registered images will be saved
+            - margin (float): Margin in mm to crop from x and y dimensions
+    Returns:
+        args: Updated argument object with modified directory paths pointing to the
+              registered image directories (t2_registered, DWI_registered, ADC_registered)
+    Raises:
+        FileNotFoundError: If input directories do not exist or files cannot be read
+        RuntimeError: If image preprocessing or cropping fails
+    """
+
     files = os.listdir(args.t2_dir)
-    new_spacing = (0.4, 0.4, 3.0) 
-    t2_registered_dir = os.path.join(args.output_dir, 't2_registered')
-    dwi_registered_dir = os.path.join(args.output_dir, 'DWI_registered')
-    adc_registered_dir = os.path.join(args.output_dir, 'ADC_registered')
+    new_spacing = (0.4, 0.4, 3.0)
+    t2_registered_dir = os.path.join(args.output_dir, "t2_registered")
+    dwi_registered_dir = os.path.join(args.output_dir, "DWI_registered")
+    adc_registered_dir = os.path.join(args.output_dir, "ADC_registered")
     os.makedirs(t2_registered_dir, exist_ok=True)
     os.makedirs(dwi_registered_dir, exist_ok=True)
     os.makedirs(adc_registered_dir, exist_ok=True)
     logging.info("Starting registration and cropping")
     for file in tqdm(files):
-            
         t2_image = sitk.ReadImage(os.path.join(args.t2_dir, file))
         dwi_image = sitk.ReadImage(os.path.join(args.dwi_dir, file))
         adc_image = sitk.ReadImage(os.path.join(args.adc_dir, file))
@@ -30,38 +47,37 @@ def register_files(args):
             int(round(osz * ospc / nspc))
             for osz, ospc, nspc in zip(original_size, original_spacing, new_spacing)
         ]
-    
+
         images_to_preprocess = {}
-        images_to_preprocess['t2'] = t2_image
-        images_to_preprocess['hbv'] = dwi_image
-        images_to_preprocess['adc'] = adc_image
+        images_to_preprocess["t2"] = t2_image
+        images_to_preprocess["hbv"] = dwi_image
+        images_to_preprocess["adc"] = adc_image
 
         pat_case = Sample(
             scans=[
-                images_to_preprocess.get('t2'),
-                images_to_preprocess.get('hbv'),
-                images_to_preprocess.get('adc'),
+                images_to_preprocess.get("t2"),
+                images_to_preprocess.get("hbv"),
+                images_to_preprocess.get("adc"),
             ],
-            settings=PreprocessingSettings(spacing=[3.0,0.4,0.4], matrix_size=[new_size[2],new_size[1],new_size[0]]),
+            settings=PreprocessingSettings(
+                spacing=[3.0, 0.4, 0.4], matrix_size=[new_size[2], new_size[1], new_size[0]]
+            ),
         )
         pat_case.preprocess()
-    
-        t2_post = pat_case.__dict__['scans'][0]
-        dwi_post = pat_case.__dict__['scans'][1]
-        adc_post = pat_case.__dict__['scans'][2]
+
+        t2_post = pat_case.__dict__["scans"][0]
+        dwi_post = pat_case.__dict__["scans"][1]
+        adc_post = pat_case.__dict__["scans"][2]
         cropped_t2 = crop(t2_post, [args.margin, args.margin, 0.0])
         cropped_dwi = crop(dwi_post, [args.margin, args.margin, 0.0])
         cropped_adc = crop(adc_post, [args.margin, args.margin, 0.0])
 
-
-
         sitk.WriteImage(cropped_t2, os.path.join(t2_registered_dir, file))
         sitk.WriteImage(cropped_dwi, os.path.join(dwi_registered_dir, file))
         sitk.WriteImage(cropped_adc, os.path.join(adc_registered_dir, file))
 
-        args.t2_dir = t2_registered_dir
-        args.dwi_dir = dwi_registered_dir
-        args.adc_dir = adc_registered_dir
-
-        return args
+    args.t2_dir = t2_registered_dir
+    args.dwi_dir = dwi_registered_dir
+    args.adc_dir = adc_registered_dir
 
+    return args

src/train/train_cspca.py

@@ -1,91 +1,18 @@
-import argparse
-import collections.abc
-import os
-import shutil
-import time
-import yaml
-from scipy.stats import pearsonr
-import gdown
-import numpy as np
 import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp 
 import torch.nn as nn
-import torch.nn.functional as F
-from monai.config import KeysCollection
-from monai.data import Dataset, load_decathlon_datalist, ITKReader, NumpyReader, PersistentDataset
-from monai.data.wsi_reader import WSIReader
 from monai.metrics import Cumulative, CumulativeAverage
-from monai.networks.nets import milmodel, resnet, MILModel
-from monai.transforms import (
-    Compose,
-    GridPatchd,
-    LoadImaged,
-    MapTransform,
-    RandFlipd,
-    RandGridPatchd,
-    RandRotate90d,
-    ScaleIntensityRanged,
-    SplitDimd,
-    ToTensord,
-    ConcatItemsd, 
-    SelectItemsd,
-    EnsureChannelFirstd,
-    RepeatChanneld,
-    DeleteItemsd,
-    EnsureTyped,
-    ClipIntensityPercentilesd,
-    MaskIntensityd,
-    HistogramNormalized,
-    RandBiasFieldd,
-    RandCropByPosNegLabeld,
-    NormalizeIntensityd,
-    SqueezeDimd,
-    CropForegroundd,
-    ScaleIntensityd,
-    SpatialPadd,
-    CenterSpatialCropd,
-    ScaleIntensityd,
-    Transposed,
-    RandWeightedCropd,
-)
-from sklearn.metrics import cohen_kappa_score, roc_curve, confusion_matrix
-from torch.cuda.amp import GradScaler, autocast
-from torch.utils.data.dataloader import default_collate
-from torchvision.models.resnet import ResNet50_Weights
-import torch.optim as optim
-from torch.utils.data.distributed import DistributedSampler
-from torch.utils.tensorboard import SummaryWriter
-import matplotlib.pyplot as plt
-import matplotlib.patches as patches
-from tqdm import tqdm
-from sklearn.metrics import confusion_matrix, roc_auc_score
+from sklearn.metrics import confusion_matrix
 from sklearn.metrics import roc_auc_score
-from sklearn.preprocessing import label_binarize
-import numpy as np
-from AIAH_utility.viewer import BasicViewer
-from scipy.special import expit
-import nrrd
-import random
-from sklearn.metrics import roc_auc_score
-import SimpleITK as sitk 
-from AIAH_utility.viewer import BasicViewer
-import pandas as pd
-import json
-from sklearn.preprocessing import StandardScaler
-from torch.utils.data import DataLoader, TensorDataset, Dataset
-from sklearn.linear_model import LogisticRegression
-from sklearn.utils import resample
-import monai
+
 
 def train_epoch(cspca_model, loader, optimizer, epoch, args):
     cspca_model.train()
-    criterion = nn.BCELoss() 
+    criterion = nn.BCELoss()
     loss = 0.0
     run_loss = CumulativeAverage()
     TARGETS = Cumulative()
     PREDS = Cumulative()
-    
+
     for idx, 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)
@@ -108,9 +35,10 @@ def train_epoch(cspca_model, loader, optimizer, epoch, args):
 
     return loss_epoch, auc_epoch
 
+
 def val_epoch(cspca_model, loader, epoch, args):
     cspca_model.eval()
-    criterion = nn.BCELoss() 
+    criterion = nn.BCELoss()
     loss = 0.0
     run_loss = CumulativeAverage()
     TARGETS = Cumulative()
@@ -132,10 +60,15 @@ def val_epoch(cspca_model, loader, epoch, args):
     target_list = TARGETS.get_buffer().cpu().numpy()
     pred_list = PREDS.get_buffer().cpu().numpy()
     auc_epoch = roc_auc_score(target_list, pred_list)
-    y_pred_categoric = (pred_list >= 0.5)
+    y_pred_categoric = pred_list >= 0.5
     tn, fp, fn, tp = confusion_matrix(target_list, y_pred_categoric).ravel()
     sens_epoch = tp / (tp + fn)
     spec_epoch = tn / (tn + fp)
-    val_epoch_metric = {'epoch': epoch, 'loss': loss_epoch, 'auc': auc_epoch, 'sensitivity': sens_epoch, 'specificity': spec_epoch}
+    val_epoch_metric = {
+        "epoch": epoch,
+        "loss": loss_epoch,
+        "auc": auc_epoch,
+        "sensitivity": sens_epoch,
+        "specificity": spec_epoch,
+    }
     return val_epoch_metric
-

src/train/train_pirads.py

@@ -1,65 +1,11 @@
-import argparse
-import collections.abc
-import os
-import shutil
 import time
-import yaml
-import sys
-import gdown
 import numpy as np
 import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
 import torch.nn as nn
-import torch.nn.functional as F
-from monai.config import KeysCollection
-from monai.data import Dataset, load_decathlon_datalist, ITKReader, NumpyReader, PersistentDataset
-from monai.data.wsi_reader import WSIReader
 from monai.metrics import Cumulative, CumulativeAverage
-from monai.networks.nets import milmodel, resnet, MILModel
-from monai.transforms import (
-    Compose,
-    GridPatchd,
-    LoadImaged,
-    MapTransform,
-    RandFlipd,
-    RandGridPatchd,
-    RandRotate90d,
-    ScaleIntensityRanged,
-    SplitDimd,
-    ToTensord,
-    ConcatItemsd, 
-    SelectItemsd,
-    EnsureChannelFirstd,
-    RepeatChanneld,
-    DeleteItemsd,
-    EnsureTyped,
-    ClipIntensityPercentilesd,
-    MaskIntensityd,
-    HistogramNormalized,
-    RandBiasFieldd,
-    RandCropByPosNegLabeld,
-    NormalizeIntensityd,
-    SqueezeDimd,
-    CropForegroundd,
-    ScaleIntensityd,
-    SpatialPadd,
-    CenterSpatialCropd,
-    ScaleIntensityd,
-    Transposed,
-    RandWeightedCropd,
-)
 from sklearn.metrics import cohen_kappa_score
-from torch.cuda.amp import GradScaler, autocast
-from torch.utils.data.dataloader import default_collate
-from torch.utils.tensorboard import SummaryWriter
-import matplotlib.pyplot as plt
-import wandb
-import math
 import logging
-from pathlib import Path
-from src.data.data_loader import get_dataloader
-from src.utils import save_pirads_checkpoint, setup_logging
+
 
 def get_lambda_att(epoch, max_lambda=2.0, warmup_epochs=10):
     if epoch < warmup_epochs:
@@ -67,26 +13,53 @@ def get_lambda_att(epoch, max_lambda=2.0, warmup_epochs=10):
     else:
         return max_lambda
 
+
 def get_attention_scores(data, target, heatmap, args):
+    """
+    Compute attention scores from heatmaps and shuffle data accordingly.
+    This function generates attention scores based on spatial heatmaps, applies
+    sharpening, and creates shuffled versions of the input data and attention
+    labels. For PI-RADS 2 (target < 1), uniform attention scores are assigned.
+    Args:
+        data (torch.Tensor): Input data tensor of shape (batch_size, num_patches, ...).
+        target (torch.Tensor): Target labels tensor of shape (batch_size,).
+        heatmap (torch.Tensor): Attention heatmap tensor corresponding to input patches.
+        args: Arguments object containing device specification.
+    Returns:
+        tuple: A tuple containing:
+            - att_labels (torch.Tensor): Sharpened and normalized attention scores
+              of shape (batch_size, num_patches), moved to args.device.
+            - shuffled_images (torch.Tensor): Randomly permuted data samples
+              of shape (batch_size, num_patches, ...), moved to args.device.
+    Note:
+        - Attention scores are computed by summing heatmap values across spatial dimensions.
+        - Data and attention labels are shuffled with the same permutation per sample.
+        - PI-RADS 2 samples receive uniform attention distribution.
+        - Attention scores are squared for sharpening and then normalized.
+    """
+
     attention_score = torch.zeros((data.shape[0], data.shape[1]))
     for i in range(data.shape[0]):
         sample = heatmap[i]
         heatmap_patches = sample.squeeze(1)
-        raw_scores = heatmap_patches.view(len(heatmap_patches), -1).sum(dim=1)  
-        attention_score[i] = raw_scores / raw_scores.sum() 
+        raw_scores = heatmap_patches.view(len(heatmap_patches), -1).sum(dim=1)
+        attention_score[i] = raw_scores / raw_scores.sum()
     shuffled_images = torch.empty_like(data).to(args.device)
     att_labels = torch.empty_like(attention_score).to(args.device)
-    for i in range(data.shape[0]):  
-        perm = torch.randperm(data.shape[1])  
+    for i in range(data.shape[0]):
+        perm = torch.randperm(data.shape[1])
         shuffled_images[i] = data[i, perm]
-        att_labels[i] = attention_score[i, perm] 
-        
-    att_labels[torch.argwhere(target < 1)] = torch.ones_like(att_labels[0]) / len(att_labels[0])# Setting attention scores for cases
-    att_labels = att_labels ** 2  # Sharpening
+        att_labels[i] = attention_score[i, perm]
+
+    att_labels[torch.argwhere(target < 1)] = torch.ones_like(att_labels[0]) / len(
+        att_labels[0]
+    )  # For PI-RADS 2, uniform scores across patches
+    att_labels = att_labels**2  # Sharpening
     att_labels = att_labels / att_labels.sum(dim=1, keepdim=True)
-    
+
     return att_labels, shuffled_images
 
+
 def train_epoch(model, loader, optimizer, scaler, epoch, args):
     """One train epoch over the dataset"""
     lambda_att = get_lambda_att(epoch, warmup_epochs=25)
@@ -104,13 +77,14 @@ def train_epoch(model, loader, optimizer, scaler, epoch, args):
     loss, acc = 0.0, 0.0
 
     for idx, batch_data in enumerate(loader):
-
         eps = 1e-8
         data = batch_data["image"].as_subclass(torch.Tensor)
         target = batch_data["label"].as_subclass(torch.Tensor).to(args.device)
         target = target.long()
         if args.use_heatmap:
-            att_labels, shuffled_images = get_attention_scores(data, target, batch_data['final_heatmap'], args)
+            att_labels, shuffled_images = get_attention_scores(
+                data, target, batch_data["final_heatmap"], args
+            )
             att_labels = att_labels + eps
         else:
             shuffled_images = data.to(args.device)
@@ -139,7 +113,7 @@ def train_epoch(model, loader, optimizer, scaler, epoch, args):
                 b = b + eps
                 att_preds = torch.softmax(b, dim=1)
                 attn_loss = 1 - att_criterion(att_preds, att_labels).mean()
-                loss = class_loss + (lambda_att*attn_loss)
+                loss = class_loss + (lambda_att * attn_loss)
             else:
                 loss = class_loss
                 attn_loss = torch.tensor(0.0)
@@ -150,14 +124,10 @@ def train_epoch(model, loader, optimizer, scaler, epoch, args):
         if not torch.isfinite(total_norm):
             logging.warning("Non-finite gradient norm detected, skipping batch.")
             optimizer.zero_grad()
+            scaler.update()
         else:
             scaler.step(optimizer)
             scaler.update()
-            shuffled_images = shuffled_images.to('cpu')
-            logits = logits.to('cpu')
-            logits_attn = logits_attn.to('cpu')
-            target = target.to('cpu')
-            
             batch_norm.append(total_norm)
             pred = torch.argmax(logits, dim=1)
             acc = (pred == target).sum() / len(pred)
@@ -171,15 +141,15 @@ def train_epoch(model, loader, optimizer, scaler, epoch, args):
                     args.epochs,
                     idx,
                     len(loader),
-                    loss,
-                    attn_loss,
+                    loss.item(),
+                    attn_loss.item(),
                     acc,
                     total_norm,
-                    time.time() - start_time
+                    time.time() - start_time,
                 )
             )
             start_time = time.time()
-    
+
     del data, target, shuffled_images, logits, logits_attn
     torch.cuda.empty_cache()
     batch_norm_epoch = batch_norm.aggregate()
@@ -189,9 +159,7 @@ def train_epoch(model, loader, optimizer, scaler, epoch, args):
     return loss_epoch, acc_epoch, attn_loss_epoch, batch_norm_epoch
 
 
-
 def val_epoch(model, loader, epoch, args):
-
     criterion = nn.CrossEntropyLoss()
 
     run_loss = CumulativeAverage()
@@ -204,17 +172,17 @@ def val_epoch(model, loader, epoch, args):
     model.eval()
     with torch.no_grad():
         for idx, 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)
             target = target.long()
-            with torch.cuda.amp.autocast(enabled=args.amp):
+
+            with torch.amp.autocast(device_type=str(args.device), enabled=args.amp):
                 logits = model(data)
                 loss = criterion(logits, target)
 
-            data = data.to('cpu')
-            target = target.to('cpu')
-            logits = logits.to('cpu')
+            data = data.to("cpu")
+            target = target.to("cpu")
+            logits = logits.to("cpu")
             pred = torch.argmax(logits, dim=1)
             acc = (pred == target).sum() / len(target)
 
@@ -228,7 +196,7 @@ def val_epoch(model, loader, epoch, args):
                 )
             )
             start_time = time.time()
-            
+
             del data, target, logits
             torch.cuda.empty_cache()
 
@@ -237,10 +205,5 @@ def val_epoch(model, loader, epoch, args):
         TARGETS = TARGETS.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.astype(np.float64), PREDS.astype(np.float64))
     return loss_epoch, acc_epoch, qwk
-
-
-
-
-

src/utils.py

@@ -1,97 +1,46 @@
-import argparse
 import os
-import shutil
-import time
-import yaml
 import sys
-import gdown
 import numpy as np
 import torch
-import torch.distributed as dist
-import torch.multiprocessing as mp
-import torch.nn as nn
-import torch.nn.functional as F
-from monai.config import KeysCollection
-from monai.metrics import Cumulative, CumulativeAverage
-from monai.networks.nets import milmodel, resnet, MILModel
 from monai.transforms import (
     Compose,
-    GridPatchd,
     LoadImaged,
-    MapTransform,
-    RandFlipd,
-    RandGridPatchd,
-    RandRotate90d,
-    ScaleIntensityRanged,
-    SplitDimd,
     ToTensord,
-    ConcatItemsd, 
-    SelectItemsd,
-    EnsureChannelFirstd,
-    RepeatChanneld,
-    DeleteItemsd,
     EnsureTyped,
-    ClipIntensityPercentilesd,
-    MaskIntensityd,
-    HistogramNormalized,
-    RandBiasFieldd,
-    RandCropByPosNegLabeld,
-    NormalizeIntensityd,
-    SqueezeDimd,
-    CropForegroundd,
-    ScaleIntensityd,
-    SpatialPadd,
-    CenterSpatialCropd,
-    ScaleIntensityd,
-    Transposed,
-    RandWeightedCropd,
 )
-from sklearn.metrics import cohen_kappa_score
-from torch.cuda.amp import GradScaler, autocast
-from torch.utils.data.dataloader import default_collate
-from torchvision.models.resnet import ResNet50_Weights
 from .data.custom_transforms import ClipMaskIntensityPercentilesd, NormalizeIntensity_customd
-from torch.utils.data.distributed import DistributedSampler
-from torch.utils.tensorboard import SummaryWriter
-import matplotlib.patches as patches
-
-import matplotlib.pyplot as plt
-
-import wandb
-import math
-from monai.data import Dataset, load_decathlon_datalist, ITKReader, NumpyReader, PersistentDataset
-
-from src.model.MIL import MILModel_3D
-from src.model.csPCa_model import csPCa_Model
-
+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):
 
-    """Save checkpoint"""
+def save_pirads_checkpoint(model, epoch, args, filename="model.pth", best_acc=0):
+    """Save checkpoint for the PI-RADS model"""
 
     state_dict = model.state_dict()
-
     save_dict = {"epoch": epoch, "best_acc": best_acc, "state_dict": state_dict}
-
     filename = os.path.join(args.logdir, filename)
     torch.save(save_dict, filename)
-    logging.info("Saving checkpoint", filename)
+    logging.info(f"Saving checkpoint {filename}")
+
 
 def save_cspca_checkpoint(model, val_metric, model_dir):
+    """Save checkpoint for the csPCa model"""
+
     state_dict = model.state_dict()
     save_dict = {
-        'epoch' : val_metric['epoch'],
-        'loss' : val_metric['loss'],
-        'auc' : val_metric['auc'],
-        'sensitivity' : val_metric['sensitivity'],
-        'specificity' : val_metric['specificity'],
-        'state' : val_metric['state'],
-        'state_dict' : state_dict,
+        "epoch": val_metric["epoch"],
+        "loss": val_metric["loss"],
+        "auc": val_metric["auc"],
+        "sensitivity": val_metric["sensitivity"],
+        "specificity": val_metric["specificity"],
+        "state": val_metric["state"],
+        "state_dict": state_dict,
     }
-    torch.save(save_dict, os.path.join(model_dir,f"cspca_model.pth"))
-    logging.info('Saving model with auc: ', str(val_metric['auc']))
+    torch.save(save_dict, os.path.join(model_dir, "cspca_model.pth"))
+    logging.info(f"Saving model with auc: {val_metric['auc']}")
+
 
 def get_metrics(metric_dict: dict):
     for metric_name, metric_list in metric_dict.items():
@@ -102,6 +51,7 @@ def get_metrics(metric_dict: dict):
         logging.info(f"Mean {metric_name}: {mean_metric:.3f}")
         logging.info(f"95% CI: ({lower:.3f}, {upper:.3f})")
 
+
 def setup_logging(log_file):
     log_file = Path(log_file)
     log_file.parent.mkdir(parents=True, exist_ok=True)
@@ -115,6 +65,7 @@ def setup_logging(log_file):
         ],
     )
 
+
 def validate_steps(steps):
     REQUIRES = {
         "get_segmentation_mask": ["register_and_crop"],
@@ -126,38 +77,63 @@ def validate_steps(steps):
         for req in required:
             if req not in steps[:i]:
                 logging.error(
-                    f"Step '{step}' requires '{req}' to be executed before it. "
-                    f"Given order: {steps}"
+                    f"Step '{step}' requires '{req}' to be executed before it. Given order: {steps}"
                 )
                 sys.exit(1)
 
-def get_patch_coordinate(patches_top_5, parent_image, args):
 
-    sample = np.array([i.transpose(1,2,0) for i in patches_top_5])
+def get_patch_coordinate(patches_top_5, parent_image):
+    """
+    Locate the coordinates of top-5 patches within a parent image.
+
+    This function searches for the spatial location of the first slice (j=0) of each
+    top-5 patch within the parent 3D image volume. It returns the top-left corner
+    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)
+                             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.
+        args: Configuration arguments (currently unused in the function).
+
+    Returns:
+        list: List of tuples (row, col, slice_idx) representing the top-left corner
+              coordinates of each found patch in the parent image. Returns empty list
+              if no patches are found.
+
+    Note:
+        - Only searches for the first slice (j=0) of each patch.
+        - Uses exhaustive 2D spatial matching within each slice of the parent image.
+        - Returns coordinates of the first match found for each patch.
+    """
+
+    sample = np.array([i.transpose(1, 2, 0) for i in patches_top_5])
     coords = []
     rows, h, w, slices = sample.shape
 
     for i in range(rows):
-        for j in range(slices):
-            if j == 0:
-                for k in range(parent_image.shape[2]):
-                    img_temp = parent_image[:, :, k]
-                    H, W = img_temp.shape
-                    h, w = sample[i, :, :, j].shape
-                    a,b = 0, 0  # Initialize a and b
-                    bool1 = False
-                    for l in range(H - h + 1):
-                        for m in range(W - w + 1):
-                            if np.array_equal(img_temp[l:l+h, m:m+w], sample[i, :, :, j]):
-                                a,b = l, m  # top-left corner
-                                coords.append((a,b,k))
-                                bool1 = True
-                                break
-                        if bool1:
-                            break
-                        
-                    if bool1:
-                        break
+        template = sample[i, :, :, 0].astype(np.float32)
+        found = False
+        for k in list(range(parent_image.shape[2])):
+            img_slice = parent_image[:, :, k].astype(np.float32)
+            res = cv2.matchTemplate(img_slice, template, cv2.TM_CCOEFF_NORMED)
+            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
+
+            if max_val >= 0.99:
+                x, y = max_loc  # OpenCV returns (col, row) -> (x, y)
+
+                # 2. Verification Step: Check if it's actually the correct patch
+                # This mimics your original np.array_equal strictness
+                candidate_patch = img_slice[y : y + h, x : x + w]
+
+                if np.allclose(candidate_patch, template, atol=1e-5):
+                    coords.append((y, x, k))  # Original code stored (row, col, slice)
+                    found = True
+                    break
+
+        if not found:
+            print("Patch not found")
 
     return coords
 
@@ -165,7 +141,12 @@ def get_patch_coordinate(patches_top_5, parent_image, args):
 def get_parent_image(temp_data_list, args):
     transform_image = Compose(
         [
-            LoadImaged(keys=["image", "mask"], reader=ITKReader(), ensure_channel_first=True, dtype=np.float32),
+            LoadImaged(
+                keys=["image", "mask"],
+                reader=ITKReader(),
+                ensure_channel_first=True,
+                dtype=np.float32,
+            ),
             ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"),
             NormalizeIntensity_customd(keys=["image"], mask_key="mask", channel_wise=True),
             EnsureTyped(keys=["label"], dtype=torch.float32),
@@ -173,9 +154,10 @@ def get_parent_image(temp_data_list, args):
         ]
     )
     dataset_image = Dataset(data=temp_data_list, transform=transform_image)
-    return dataset_image[0]['image'][0].numpy()
+    return dataset_image[0]["image"][0].numpy()
+
 
-'''
+"""
 def visualise_patches():
     sample = np.array([i.transpose(1,2,0) for i in patches_top_5])
     rows = len(patches_top_5)
@@ -223,4 +205,4 @@ def visualise_patches():
     plt.tight_layout()
     plt.show()
     a=1
-'''
+"""

tests/test_run.py

@@ -0,0 +1,106 @@
+import subprocess
+import sys
+from pathlib import Path
+
+
+def test_run_pirads_training():
+    """
+    Test that run_cspca.py runs without crashing using an existing YAML config.
+    """
+
+    # Path to your run_pirads.py script
+    repo_root = Path(__file__).parent.parent
+    script_path = repo_root / "run_pirads.py"
+
+    # Path to your existing config.yaml
+    config_path = repo_root / "config" / "config_pirads_train.yaml"  # adjust this path
+
+    # Make sure the file exists
+    assert config_path.exists(), f"Config file not found: {config_path}"
+
+    # Run the script with the config
+    result = subprocess.run(
+        [sys.executable, str(script_path), "--mode", "train", "--config", str(config_path), "--dry_run", "True" ],
+        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_pirads_inference():
+    """
+    Test that run_cspca.py runs without crashing using an existing YAML config.
+    """
+
+    # Path to your run_pirads.py script
+    repo_root = Path(__file__).parent.parent
+    script_path = repo_root / "run_pirads.py"
+
+    # Path to your existing config.yaml
+    config_path = repo_root / "config" / "config_pirads_test.yaml"  # adjust this path
+
+    # Make sure the file exists
+    assert config_path.exists(), f"Config file not found: {config_path}"
+
+    # Run the script with the config
+    result = subprocess.run(
+        [sys.executable, str(script_path), "--mode", "test", "--config", str(config_path), "--dry_run", "True" ],
+        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.
+    """
+
+    # Path to your run_cspca.py script
+    repo_root = Path(__file__).parent.parent
+    script_path = repo_root / "run_cspca.py"
+
+    # Path to your existing config.yaml
+    config_path = repo_root / "config" / "config_cspca_train.yaml"  # adjust this path
+
+    # Make sure the file exists
+    assert config_path.exists(), f"Config file not found: {config_path}"
+
+    # Run the script with the config
+    result = subprocess.run(
+        [sys.executable, str(script_path), "--mode", "train", "--config", str(config_path), "--dry_run", "True" ],
+        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.
+    """
+
+    # Path to your run_cspca.py script
+    repo_root = Path(__file__).parent.parent
+    script_path = repo_root / "run_cspca.py"
+
+    # Path to your existing config.yaml
+    config_path = repo_root / "config" / "config_cspca_test.yaml"  # adjust this path
+
+    # Make sure the file exists
+    assert config_path.exists(), f"Config file not found: {config_path}"
+
+    # Run the script with the config
+    result = subprocess.run(
+        [sys.executable, str(script_path), "--mode", "test", "--config", str(config_path), "--dry_run", "True" ],
+        capture_output=True,
+        text=True,
+    )
+
+    # Check that it ran without errors
+    assert result.returncode == 0, f"Script failed with:\n{result.stderr}"
+
+

tests/test_run_cspca.py

@@ -1,28 +0,0 @@
-import subprocess
-import sys
-from pathlib import Path
-
-def test_run_cspca_with_existing_config():
-    """
-    Test that run_cspca.py runs without crashing using an existing YAML config.
-    """
-
-    # Path to your run_cspca.py script
-    repo_root = Path(__file__).parent.parent
-    script_path = repo_root / "run_cspca.py"
-
-    # Path to your existing config.yaml
-    config_path = repo_root / "config" / "config_cspca_test.yaml"  # adjust this path
-
-    # Make sure the file exists
-    assert config_path.exists(), f"Config file not found: {config_path}"
-
-    # Run the script with the config
-    result = subprocess.run(
-        [sys.executable, str(script_path), "--mode","test", "--config", str(config_path)],
-        capture_output=True,
-        text=True,
-    )
-
-    # Check that it ran without errors
-    assert result.returncode == 0, f"Script failed with:\n{result.stderr}"

tests/test_run_pirads.py

@@ -1,28 +0,0 @@
-import subprocess
-import sys
-from pathlib import Path
-
-def test_run_cspca_with_existing_config():
-    """
-    Test that run_cspca.py runs without crashing using an existing YAML config.
-    """
-
-    # Path to your run_pirads.py script
-    repo_root = Path(__file__).parent.parent
-    script_path = repo_root / "run_pirads.py"
-
-    # Path to your existing config.yaml
-    config_path = repo_root / "config" / "config_pirads_test.yaml"  # adjust this path
-
-    # Make sure the file exists
-    assert config_path.exists(), f"Config file not found: {config_path}"
-
-    # Run the script with the config
-    result = subprocess.run(
-        [sys.executable, str(script_path), "--mode","test", "--config", str(config_path)],
-        capture_output=True,
-        text=True,
-    )
-
-    # Check that it ran without errors
-    assert result.returncode == 0, f"Script failed with:\n{result.stderr}"