add scripts

.gitignore

@@ -0,0 +1,2 @@
+logs/
+models/

config/config_cspca_test.yaml

@@ -0,0 +1,17 @@
+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
+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 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
+
+
+
+

config/config_cspca_train.yaml

@@ -0,0 +1,19 @@
+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
+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 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
+batch_size: !!int 8
+optim_lr: !!float 2e-4
+
+
+
+
+

config/config_pirads_test.yaml

@@ -0,0 +1,18 @@
+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
+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
+amp: !!bool True
+dry_run: !!bool True
+
+
+
+

config/config_pirads_train.yaml

@@ -0,0 +1,22 @@
+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
+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
+batch_size: !!int 8
+optim_lr: !!float 2e-4
+weight_decay: !!float 1e-5
+amp: !!bool True
+wandb: !!bool True
+dry_run: !!bool True
+
+
+
+

config/config_preprocess.yaml

@@ -0,0 +1,11 @@
+t2_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/data_temp/t2
+dwi_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/data_temp/dwi
+adc_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/data_temp/adc
+output_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/data_temp/processed
+project_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/new_folder
+
+
+
+
+
+

config/inference.json

@@ -0,0 +1,153 @@
+{
+    "imports": [
+        "$import pandas as pd",
+        "$import os"
+    ],
+    "bundle_root": "/workspace/data/prostate_mri_anatomy",
+    "output_dir": "$@bundle_root + '/eval'",
+    "dataset_dir": "/workspace/data/prostate158/prostate158_train/",
+    "datalist": "$list(@dataset_dir + pd.read_csv(@dataset_dir + 'valid.csv').t2)",
+    "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')",
+    "network_def": {
+        "_target_": "UNet",
+        "spatial_dims": 3,
+        "in_channels": 1,
+        "out_channels": 3,
+        "channels": [
+            16,
+            32,
+            64,
+            128,
+            256,
+            512
+        ],
+        "strides": [
+            2,
+            2,
+            2,
+            2,
+            2
+        ],
+        "num_res_units": 4,
+        "norm": "batch",
+        "act": "prelu",
+        "dropout": 0.15
+    },
+    "network": "$@network_def.to(@device)",
+    "preprocessing": {
+        "_target_": "Compose",
+        "transforms": [
+            {
+                "_target_": "LoadImaged",
+                "keys": "image"
+            },
+            {
+                "_target_": "EnsureChannelFirstd",
+                "keys": "image"
+            },
+            {
+                "_target_": "Orientationd",
+                "keys": "image",
+                "axcodes": "RAS"
+            },
+            {
+                "_target_": "Spacingd",
+                "keys": "image",
+                "pixdim": [
+                    0.5,
+                    0.5,
+                    0.5
+                ],
+                "mode": "bilinear"
+            },
+            {
+                "_target_": "ScaleIntensityd",
+                "keys": "image",
+                "minv": 0,
+                "maxv": 1
+            },
+            {
+                "_target_": "NormalizeIntensityd",
+                "keys": "image"
+            },
+            {
+                "_target_": "EnsureTyped",
+                "keys": "image"
+            }
+        ]
+    },
+    "dataset": {
+        "_target_": "Dataset",
+        "data": "$[{'image': i} for i in @datalist]",
+        "transform": "@preprocessing"
+    },
+    "dataloader": {
+        "_target_": "DataLoader",
+        "dataset": "@dataset",
+        "batch_size": 1,
+        "shuffle": false,
+        "num_workers": 4
+    },
+    "inferer": {
+        "_target_": "SlidingWindowInferer",
+        "roi_size": [
+            96,
+            96,
+            96
+        ],
+        "sw_batch_size": 4,
+        "overlap": 0.5
+    },
+    "postprocessing": {
+        "_target_": "Compose",
+        "transforms": [
+            {
+                "_target_": "AsDiscreted",
+                "keys": "pred",
+                "argmax": true
+            },
+            {
+                "_target_": "KeepLargestConnectedComponentd",
+                "keys": "pred",
+                "applied_labels": [
+                    1,
+                    2
+                ]
+            },
+            {
+                "_target_": "SaveImaged",
+                "keys": "pred",
+                "resample": false,
+                "meta_keys": "pred_meta_dict",
+                "output_dir": "@output_dir"
+            }
+        ]
+    },
+    "handlers": [
+        {
+            "_target_": "CheckpointLoader",
+            "load_path": "$@bundle_root + '/models/model.pt'",
+            "load_dict": {
+                "model": "@network"
+            }
+        },
+        {
+            "_target_": "StatsHandler",
+            "iteration_log": false
+        }
+    ],
+    "evaluator": {
+        "_target_": "SupervisedEvaluator",
+        "device": "@device",
+        "val_data_loader": "@dataloader",
+        "network": "@network",
+        "inferer": "@inferer",
+        "postprocessing": "@postprocessing",
+        "val_handlers": "@handlers",
+        "amp": true
+    },
+    "evaluating": [
+        "$setattr(torch.backends.cudnn, 'benchmark', True)",
+        "$@evaluator.run()"
+    ]
+}

dataset/adc_reference.nrrd

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46821ddb4373198d98b5877ec69a702718702cbaae6b9ef00b6b5ad235cf3f3e
+size 3815961

dataset/dwi_reference.nrrd

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a964835a0c8c016b162a2eaa56456f9704290b3a4128aafeb9b01805166ca7b9
+size 3815961

dataset/prostate_segmentation_reference.nrrd

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:95984795649662ae405077d334744ac2e9d9fb8db7864c4ea99791a706ccee19
+size 13434

dataset/t2_reference.nrrd

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:561af922d7461870fcc1e84134d659f7611a8ae73d1a0681fde79808f1cd99a9
+size 3815961

job_scripts/train_cspca.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+#SBATCH --job-name=cspca_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/git_updated/Prostate-Foundation/MIL/logs/%x/log.o%j      # File name for standard output
+#SBATCH --error=/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/git_updated/Prostate-Foundation/MIL/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/git_updated/Prostate-Foundation"
+
+
+srun python -u $RUNDIR/MIL/new_folder/run_cspca.py --mode train --config $RUNDIR/MIL/new_folder/config/config_cspca_train.yaml

preprocess_main.py

@@ -0,0 +1,68 @@
+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 
+
+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
+        required=True,
+        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("--project_dir", default=None, help="Project directory")
+    
+    args = parser.parse_args()
+    if args.config:
+        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,
+    }
+
+    args.logfile = os.path.join(args.output_dir, f"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) 

run_cspca.py

@@ -0,0 +1,220 @@
+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
+
+def main_worker(args):
+
+    mil_model = MILModel_3D(
+        num_classes=args.num_classes,  
+        mil_mode=args.mil_mode 
+    ).to(args.device)
+    cache_dir_path = Path(os.path.join(args.logdir, "cache"))
+
+    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':[]}
+        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 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)
+
+            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}")
+                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']):
+                        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)
+            if cache_dir_path.exists() and cache_dir_path.is_dir():
+                shutil.rmtree(cache_dir_path)
+
+        get_metrics(metrics_dict)
+
+    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 = 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.")
+        else:
+            logging.info(f"csPCa Model loaded from {args.checkpoint_cspca}.")
+        
+        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'])
+
+            if cache_dir_path.exists() and cache_dir_path.is_dir():
+                shutil.rmtree(cache_dir_path)
+
+        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.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"
+    )
+    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"
+    )
+    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("--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("--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(
+        "--val_every",
+        "--val_interval",
+        default=1,
+        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("--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")
+    args = parser.parse_args()
+    if args.config:
+        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()
+    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.')
+        sys.exit(1)
+    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.')
+        sys.exit(1)
+    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    if args.device == torch.device("cuda"):
+        torch.backends.cudnn.benchmark = True
+
+    if args.dry_run:
+        logging.info("Dry run mode enabled.")
+        args.epochs = 2
+        args.batch_size = 2
+        args.workers = 0
+        args.num_seeds = 2
+        args.wandb = False
+
+
+    main_worker(args)

run_pirads.py

@@ -0,0 +1,283 @@
+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
+from src.train.train_pirads import train_epoch, val_epoch
+from src.model.MIL import MILModel_3D
+from src.utils import save_pirads_checkpoint, setup_logging
+
+
+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 
+    )
+    start_epoch = 0
+    best_acc = 0.0
+    if args.checkpoint is not None:
+        checkpoint = torch.load(args.checkpoint, map_location="cpu")
+        model.load_state_dict(checkpoint["state_dict"])
+
+        if "epoch" in checkpoint:
+            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)
+    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)
+        valid_loader = get_dataloader(args, split="test")
+        logging.info("Dataset training:", str(len(train_loader.dataset)), "test:", str(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.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)
+        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)
+        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)
+            logging.info(
+                "Final training %d/%d loss: %.4f attention loss: %.4f acc: %.4f time %.2fs",
+                epoch,
+                n_epochs - 1,
+                train_loss,
+                train_att_loss,
+                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)
+                
+            model_new_best = False
+            val_acc = 0
+            if (epoch + 1) % args.val_every == 0:
+                epoch_time = time.time()
+                val_loss, val_acc, qwk = val_epoch(model, valid_loader, epoch=epoch, args=args)
+
+                logging.info(
+                    "Final test %d/%d loss: %.4f acc: %.4f qwk: %.4f time %.2fs",
+                    epoch,
+                    n_epochs - 1,
+                    val_loss,
+                    val_acc,
+                    qwk,
+                    time.time() - epoch_time,
+                )
+                if writer is not None:
+                    writer.add_scalar("test_loss", val_loss, epoch)
+                    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)
+                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")
+                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"))
+                    epochs_no_improve = 0
+                
+                else:
+                    epochs_no_improve += 1
+                    if epochs_no_improve == args.early_stop:
+                        logging.info('Early stopping!')
+                        break
+                    
+
+
+            scheduler.step()
+
+        logging.info("ALL DONE")
+    
+    elif args.mode == 'test':
+
+        
+        kappa_list = []
+        for seed in list(range(args.num_seeds)):
+            set_determinism(seed=seed)
+            valid_loader = get_dataloader(args, split=args.mode)
+            logging.info("test:", str(len(valid_loader.dataset)))
+            val_loss, val_acc, qwk = val_epoch(model, valid_loader, epoch=0, args=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_)
+                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_)
+        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.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"
+    )
+    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"
+    )
+    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("--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")
+    parser.add_argument("--weight_decay", default=0, type=float, help="optimizer weight decay")
+    parser.add_argument("--amp", action="store_true", help="use AMP, recommended")
+    parser.add_argument(
+        "--val_every",
+        "--val_interval",
+        default=1,
+        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)")
+    args = parser.parse_args()
+    if args.config:
+        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()
+    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("-----------------")
+
+    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.')
+        sys.exit(1)
+    if args.checkpoint is None and args.mode == 'test':
+        logging.error('Model checkpoint path not provided. Quitting.')
+        sys.exit(1)
+
+    if args.dry_run:
+        logging.info("Dry run mode enabled.")
+        args.epochs = 2
+        args.batch_size = 2
+        args.workers = 0
+        args.num_seeds = 2
+        args.wandb = False
+    
+    mode_wandb = "online" if args.wandb else "disabled"
+
+    config_wandb = {
+        "learning_rate": args.optim_lr,
+        "batch_size": args.batch_size,
+        "epochs": args.epochs,
+        "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)
+
+
+    main_worker(args)
+    
+    wandb.finish()

src/data/custom_transforms.py

@@ -0,0 +1,350 @@
+import numpy as np
+import torch
+from typing import Union, Optional
+
+from monai.transforms import MapTransform
+from monai.config import DtypeLike, KeysCollection
+from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor
+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.utils.enums import TransformBackends
+from monai.utils.type_conversion import convert_data_type, convert_to_dst_type, convert_to_tensor, get_equivalent_dtype
+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
+        self.copy_key = copy_key
+
+    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
+
+            # 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
+
+        return d
+
+
+class ClipMaskIntensityPercentiles(Transform):
+
+
+    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
+
+    def __init__(
+        self,
+        lower: Union[float, None],
+        upper: Union[float, 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):
+            raise ValueError("Percentiles must be in the range [0, 100]")
+        if upper is not None and (upper < 0.0 or upper > 100.0):
+            raise ValueError("Percentiles must be in the range [0, 100]")
+        if upper is not None and lower is not None and upper < lower:
+            raise ValueError("upper must be greater than or equal to lower")
+        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.lower = lower
+        self.upper = upper
+        self.sharpness_factor = sharpness_factor
+        self.channel_wise = channel_wise
+        self.dtype = dtype
+
+    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)
+        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)
+            img = clip(img, lower_percentile, upper_percentile)
+
+        img = convert_to_tensor(img, track_meta=False)
+        return img
+
+    def __call__(self, img: NdarrayOrTensor, mask_data: NdarrayOrTensor) -> NdarrayOrTensor:
+        """
+        Apply the transform to `img`.
+        """
+        img = convert_to_tensor(img, track_meta=get_track_meta())
+        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
+        else:
+            img_t = self._clip(img=img_t, mask_data=mask_t)
+
+        img = convert_to_dst_type(img_t, dst=img)[0]
+
+        return img
+
+class ClipMaskIntensityPercentilesd(MapTransform):
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        mask_key: str,
+        lower: Union[float, None],
+        upper: Union[float, None],
+        sharpness_factor: Union[float, None] = None,
+        channel_wise: bool = False,
+        dtype: DtypeLike = np.float32,
+        allow_missing_keys: bool = False,
+    ) -> None:
+        super().__init__(keys, allow_missing_keys)
+        self.scaler = ClipMaskIntensityPercentiles(
+            lower=lower, upper=upper, sharpness_factor=sharpness_factor, channel_wise=channel_wise, dtype=dtype
+        )
+        self.mask_key = mask_key
+
+    def __call__(self, data: dict) -> dict:
+        d = dict(data)
+        for key in self.key_iterator(d):
+            d[key] = self.scaler(d[key], d[self.mask_key])
+        return d
+    
+    
+
+
+class ElementwiseProductd(MapTransform):
+    def __init__(self, keys: KeysCollection, output_key: str) -> None:
+        super().__init__(keys)
+        self.output_key = output_key
+
+    def __call__(self, data) -> NdarrayOrTensor:
+        d = dict(data)
+        d[self.output_key] = d[self.keys[0]] * d[self.keys[1]]
+        return d
+
+
+class CLAHEd(MapTransform):
+    """
+    Apply CLAHE (Contrast Limited Adaptive Histogram Equalization) to images in a data dictionary.
+    Works on 2D images or 3D volumes (applied slice-by-slice).
+
+    Args:
+        keys (KeysCollection): Keys of the items to be transformed.
+        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,
+        clip_limit: float = 2.0,
+        tile_grid_size: Union[tuple, Sequence[int]] = (8, 8),
+    ) -> None:
+        super().__init__(keys)
+        self.clip_limit = clip_limit
+        self.tile_grid_size = tile_grid_size
+
+    def __call__(self, data):
+        d = dict(data)
+        for key in self.keys:
+            image_ = d[key]
+
+            image = image_.cpu().numpy()
+
+            if image.dtype != np.uint8:
+                image = image.astype(np.uint8)
+
+            clahe = cv2.createCLAHE(clipLimit=self.clip_limit, tileGridSize=self.tile_grid_size)
+            # Handle 2D images or process 3D images slice-by-slice.
+
+            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`.
+    Use calculated mean or std value of the input image if no `subtrahend` or `divisor` provided.
+    This transform can normalize only non-zero values or entire image, and can also calculate
+    mean and std on each channel separately.
+    When `channel_wise` is True, the first dimension of `subtrahend` and `divisor` should
+    be the number of image channels if they are not None.
+    If the input is not of floating point type, it will be converted to float32
+
+    Args:
+        subtrahend: the amount to subtract by (usually the mean).
+        divisor: the amount to divide by (usually the standard deviation).
+        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.
+    """
+
+    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
+
+    def __init__(
+        self,
+        subtrahend: Union[Sequence, NdarrayOrTensor, None] = None,
+        divisor: Union[Sequence, NdarrayOrTensor, None] = None,
+        nonzero: bool = False,
+        channel_wise: bool = False,
+        dtype: DtypeLike = np.float32,
+    ) -> None:
+        self.subtrahend = subtrahend
+        self.divisor = divisor
+        self.nonzero = nonzero
+        self.channel_wise = channel_wise
+        self.dtype = dtype
+
+    @staticmethod
+    def _mean(x):
+        if isinstance(x, np.ndarray):
+            return np.mean(x)
+        x = torch.mean(x.float())
+        return x.item() if x.numel() == 1 else x
+
+    @staticmethod
+    def _std(x):
+        if isinstance(x, np.ndarray):
+            return np.std(x)
+        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:
+        img, *_ = convert_data_type(img, dtype=torch.float32)
+        '''
+        if self.nonzero:
+            slices = img != 0
+            masked_img = img[slices]
+            if not slices.any():
+                return img
+        else:
+            slices = None
+            masked_img = img
+        '''
+        slices = None
+        mask_data = mask_data.squeeze(0)
+        slices_mask = mask_data > 0
+        masked_img = img[slices_mask]
+
+        _sub = sub if sub is not None else self._mean(masked_img)
+        if isinstance(_sub, (torch.Tensor, np.ndarray)):
+            _sub, *_ = convert_to_dst_type(_sub, img)
+            if slices is not None:
+                _sub = _sub[slices]
+
+        _div = div if div is not None else self._std(masked_img)
+        if np.isscalar(_div):
+            if _div == 0.0:
+                _div = 1.0
+        elif isinstance(_div, (torch.Tensor, np.ndarray)):
+            _div, *_ = convert_to_dst_type(_div, img)
+            if slices is not None:
+                _div = _div[slices]
+            _div[_div == 0.0] = 1.0
+
+        if slices is not None:
+            img[slices] = (masked_img - _sub) / _div
+        else:
+            img = (img - _sub) / _div
+        return img
+
+    def __call__(self, img: NdarrayOrTensor, mask_data: NdarrayOrTensor) -> NdarrayOrTensor:
+        """
+        Apply the transform to `img`, assuming `img` is a channel-first array if `self.channel_wise` is True,
+        """
+        img = convert_to_tensor(img, track_meta=get_track_meta())
+        mask_data = convert_to_tensor(mask_data, track_meta=get_track_meta())
+        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.")
+            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.")
+
+            if not img.dtype.is_floating_point:
+                img, *_ = convert_data_type(img, dtype=torch.float32)
+
+            for i, d in enumerate(img):
+                img[i] = self._normalize(  # type: ignore
+                    d,
+                    mask_data,
+                    sub=self.subtrahend[i] if self.subtrahend is not None else None,
+                    div=self.divisor[i] if self.divisor is not None else None,
+                )
+        else:
+            img = self._normalize(img, mask_data, self.subtrahend, self.divisor)
+
+        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.
+
+    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)
+        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.
+        allow_missing_keys: don't raise exception if key is missing.
+    """
+
+    backend = NormalizeIntensity_custom.backend
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        mask_key: str,
+        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.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

src/data/data_loader.py

@@ -0,0 +1,125 @@
+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.transforms import (
+    Compose,
+    LoadImaged,
+    MapTransform,
+    ScaleIntensityRanged,
+    SplitDimd,
+    ToTensord,
+    ConcatItemsd, 
+    SelectItemsd,
+    EnsureChannelFirstd,
+    RepeatChanneld,
+    DeleteItemsd,
+    EnsureTyped,
+    ClipIntensityPercentilesd,
+    MaskIntensityd,
+    RandCropByPosNegLabeld,
+    NormalizeIntensityd,
+    SqueezeDimd,
+    ScaleIntensityd,
+    ScaleIntensityd,
+    Transposed,
+    RandWeightedCropd,
+)
+from .custom_transforms import (
+    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
+    [{'image' : 3xHxW}, {'image' : 3xHxW}, {'image' : 3xHxW}...] - > {'image' : Nx3xHxW}
+    followed by the default collate which will form a batch BxNx3xHxW
+    """
+
+    for i, item in enumerate(batch):
+        data = item[0]
+        data["image"] = torch.stack([ix["image"] for ix in item], dim=0)
+
+        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),
+                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)
+                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),
+                EnsureTyped(keys=["label"], dtype=torch.float32),
+                Transposed(keys=["image"], indices=(0, 3, 1, 2)),
+                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),
+                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)
+                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),
+                EnsureTyped(keys=["label"], dtype=torch.float32),
+                Transposed(keys=["image"], indices=(0, 3, 1, 2)),
+                DeleteItemsd(keys=['mask', 'dwi', 'adc']),
+                ToTensord(keys=["image", "label"]),
+            ]
+        )
+    return transform
+
+def get_dataloader(args, split: Literal["train", "test"]):
+
+    data_list = load_decathlon_datalist(
+        data_list_file_path=args.dataset_json,
+        data_list_key=split,
+        base_dir=args.data_root,
+    )
+    if args.dry_run:
+        data_list = data_list[:8]  # Use only 8 samples for dry run
+    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))
+    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,
+        sampler=None,
+        collate_fn=list_data_collate,
+    )
+    return loader
+

src/model/MIL.py

@@ -0,0 +1,248 @@
+
+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")
+
+
+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.
+
+    Args:
+        num_classes: number of output classes.
+        mil_mode: MIL algorithm, available values (Defaults to ``"att"``):
+
+            - ``"mean"`` - average features from all instances, equivalent to pure CNN (non MIL).
+            - ``"max"`` - retain only the instance with the max probability for loss calculation.
+            - ``"att"`` - attention based MIL https://arxiv.org/abs/1802.04712.
+            - ``"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.
+        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.
+        trans_dropout: dropout rate in `TransformEncoder` layer.
+
+    """
+
+    def __init__(
+        self,
+        num_classes: int,
+        mil_mode: str = "att",
+        pretrained: bool = True,
+        backbone: str | nn.Module | None = None,
+        backbone_num_features: int | None = None,
+        trans_blocks: int = 4,
+        trans_dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+
+        if num_classes <= 0:
+            raise ValueError("Number of classes must be positive: " + str(num_classes))
+
+        if mil_mode.lower() not in ["mean", "max", "att", "att_trans", "att_trans_pyramid"]:
+            raise ValueError("Unsupported mil_mode: " + str(mil_mode))
+
+        self.mil_mode = mil_mode.lower()
+        self.attention = nn.Sequential()
+        self.transformer: nn.Module | None = None
+
+        if backbone is None:
+            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
+
+            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"))
+
+        elif isinstance(backbone, str):
+            # assume torchvision model string is provided
+            torch_model = getattr(models, backbone, None)
+            if torch_model is None:
+                raise ValueError("Unknown torch vision model" + str(backbone))
+            net = torch_model(pretrained=pretrained)
+
+            if getattr(net, "fc", None) is not None:
+                nfc = net.fc.in_features  # save the number of final features
+                net.fc = torch.nn.Identity()  # remove final linear layer
+            else:
+                raise ValueError(
+                    "Unable to detect FC layer for the torchvision model " + str(backbone),
+                    ". Please initialize the backbone model manually.",
+                )
+
+        elif isinstance(backbone, nn.Module):
+            # use a custom backbone
+            net = backbone
+            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"))
+
+            if backbone_num_features is None:
+                raise ValueError("Number of endencoder features must be provided for a custom backbone model")
+
+        else:
+            raise ValueError("Unsupported backbone")
+        
+        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":
+            self.attention = nn.Sequential(nn.Linear(nfc, 2048), nn.Tanh(), nn.Linear(2048, 1))
+
+        elif self.mil_mode == "att_trans":
+            transformer = nn.TransformerEncoderLayer(d_model=nfc, nhead=8, dropout=trans_dropout)
+            self.transformer = nn.TransformerEncoder(transformer, num_layers=trans_blocks)
+            self.attention = nn.Sequential(nn.Linear(nfc, 2048), nn.Tanh(), nn.Linear(2048, 1))
+
+        elif self.mil_mode == "att_trans_pyramid":
+            transformer_list = nn.ModuleList(
+                [
+                    nn.TransformerEncoder(
+                        nn.TransformerEncoderLayer(d_model=64, nhead=8, dropout=trans_dropout), num_layers=trans_blocks
+                    ),
+                    nn.Sequential(
+                        nn.Linear(192, 64),
+                        nn.TransformerEncoder(
+                            nn.TransformerEncoderLayer(d_model=64, nhead=8, dropout=trans_dropout),
+                            num_layers=trans_blocks,
+                        ),
+                    ),
+                    nn.Sequential(
+                        nn.Linear(320, 64),
+                        nn.TransformerEncoder(
+                            nn.TransformerEncoderLayer(d_model=64, nhead=8, dropout=trans_dropout),
+                            num_layers=trans_blocks,
+                        ),
+                    ),
+                    nn.TransformerEncoder(
+                        nn.TransformerEncoderLayer(d_model=576, nhead=8, dropout=trans_dropout),
+                        num_layers=trans_blocks,
+                    ),
+                ]
+            )
+            self.transformer = transformer_list
+            nfc = nfc + 64
+            self.attention = nn.Sequential(nn.Linear(nfc, 2048), nn.Tanh(), nn.Linear(2048, 1))
+
+        else:
+            raise ValueError("Unsupported mil_mode: " + str(mil_mode))
+
+        self.myfc = nn.Linear(nfc, num_classes)
+        self.net = net
+
+    def calc_head(self, x: torch.Tensor) -> torch.Tensor:
+        sh = x.shape
+
+        if self.mil_mode == "mean":
+            x = self.myfc(x)
+            x = torch.mean(x, dim=1)
+
+        elif self.mil_mode == "max":
+            x = self.myfc(x)
+            x, _ = torch.max(x, dim=1)
+
+        elif self.mil_mode == "att":
+            a = self.attention(x)
+            a = torch.softmax(a, dim=1)
+            x = torch.sum(x * a, dim=1)
+
+            x = self.myfc(x)
+
+        elif self.mil_mode == "att_trans" and self.transformer is not None:
+            x = x.permute(1, 0, 2)
+            x = self.transformer(x)
+            x = x.permute(1, 0, 2)
+
+            a = self.attention(x)
+            a = torch.softmax(a, dim=1)
+            x = torch.sum(x * a, dim=1)
+
+            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)
+
+            transformer_list = cast(nn.ModuleList, self.transformer)
+
+            x = transformer_list[0](l1)
+            x = transformer_list[1](torch.cat((x, l2), dim=2))
+            x = transformer_list[2](torch.cat((x, l3), dim=2))
+            x = transformer_list[3](torch.cat((x, l4), dim=2))
+
+            x = x.permute(1, 0, 2)
+
+            a = self.attention(x)
+            a = torch.softmax(a, dim=1)
+            x = torch.sum(x * a, dim=1)
+
+            x = self.myfc(x)
+
+        else:
+            raise ValueError("Wrong model mode" + str(self.mil_mode))
+
+        return x
+
+    def forward(self, x: torch.Tensor, no_head: bool = False) -> torch.Tensor:
+        sh = x.shape
+        x = x.reshape(sh[0] * sh[1], sh[2], sh[3], sh[4], sh[5])
+
+        x = self.net(x)
+        x = x.reshape(sh[0], sh[1], -1)
+
+        if not no_head:
+            x = self.calc_head(x)
+
+        return x
+
+
+
+

src/model/csPCa_model.py

@@ -0,0 +1,50 @@
+
+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")
+
+
+
+class SimpleNN(nn.Module):
+    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.ReLU(),
+            nn.Dropout(p=0.3),    
+            nn.Linear(128, 1),
+            nn.Sigmoid()   # since binary classification
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class csPCa_Model(nn.Module):
+    def __init__(self, backbone):
+        super().__init__()
+        self.backbone = backbone
+        self.fc_dim = backbone.myfc.in_features
+        self.fc_cspca = SimpleNN(input_dim=self.fc_dim)
+
+    def forward(self, x):
+        sh = x.shape
+        x = x.reshape(sh[0] * sh[1], sh[2], sh[3], sh[4], sh[5])
+        x = self.backbone.net(x)
+        x = x.reshape(sh[0], sh[1], -1)
+        x = x.permute(1, 0, 2)
+        x = self.backbone.transformer(x)
+        x = x.permute(1, 0, 2)
+        a = self.backbone.attention(x)
+        a = torch.softmax(a, dim=1)
+        x = torch.sum(x * a, dim=1)
+
+        x = self.fc_cspca(x)
+        return x
+

src/preprocessing/center_crop.py

@@ -0,0 +1,64 @@
+# Copyright 2020 - 2022 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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
+
+
+#import argparse
+from typing import Union
+
+import SimpleITK as sitk  # noqa N813
+
+
+def _flatten(t):
+    return [item for sublist in t for item in sublist]
+
+
+def crop(image: sitk.Image, margin: Union[int, float], interpolator=sitk.sitkLinear):
+    """
+    Crops a sitk.Image while retaining correct spacing. Negative margins will lead to zero padding
+
+    Args:
+        image:  a sitk.Image
+        margin: margins to crop. Single integer or float (percentage crop),
+                lists of int/float or nestes lists are supported.
+    """
+    if isinstance(margin, (list, tuple)):
+        assert len(margin) == 3, "expected margin to be of length 3"
+    else:
+        assert isinstance(margin, (int, float)), "expected margin to be a float value"
+        margin = [margin, margin, margin]
+
+    margin = [m if isinstance(m, (tuple, list)) else [m, m] for m in margin]
+    old_size = image.GetSize()
+
+    # 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"
+        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
+    else:
+        raise ValueError("Wrong format of margins.")
+
+    new_size = [sz - sum(c) for sz, c in zip(old_size, to_crop)]
+
+    # origin has Index (0,0,0)
+    # new origin has Index (to_crop[0][0], to_crop[2][0], to_crop[2][0])
+    new_origin = image.TransformIndexToPhysicalPoint([c[0] for c in to_crop])
+
+    # create reference plane to resample image
+    ref_image = sitk.Image(new_size, image.GetPixelIDValue())
+    ref_image.SetSpacing(image.GetSpacing())
+    ref_image.SetOrigin(new_origin)
+    ref_image.SetDirection(image.GetDirection())
+
+    return sitk.Resample(image, ref_image, interpolator=interpolator)
+

src/preprocessing/generate_heatmap.py

@@ -0,0 +1,76 @@
+
+import os
+import numpy as np
+import nrrd
+import json
+import pandas as pd
+import json
+import SimpleITK as sitk
+import multiprocessing
+
+import logging
+
+
+def get_heatmap(args):
+
+    files = os.listdir(args.t2_dir)
+    args.heatmapdir = os.path.join(args.output_dir, 'heatmaps/')
+    os.makedirs(args.heatmapdir, exist_ok=True)
+    for file in 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))
+
+        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())
+            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())
+            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
+        else:
+            mix_mask = np.ones_like(adc, dtype=np.float32)
+
+        mix_mask = (mix_mask - mix_mask.min()) / (mix_mask.max() - mix_mask.min())
+
+
+        nrrd.write(os.path.join(args.heatmapdir, file), mix_mask)
+    
+    return args
+
+
+    
+    

src/preprocessing/histogram_match.py

@@ -0,0 +1,62 @@
+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):
+
+    source_pixels = data[mask > 0]
+    ref_pixels = ref_data[ref_mask > 0]
+    matched_pixels = exposure.match_histograms(source_pixels, ref_pixels)
+    matched_img = data.copy()
+    matched_img[mask > 0] = matched_pixels
+
+    return matched_img
+
+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')
+    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:
+
+        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'))
+
+        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'))
+
+        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)
+
+        args.t2_dir = t2_histmatched_dir
+        args.dwi_dir = dwi_histmatched_dir
+        args.adc_dir = adc_histmatched_dir
+
+    return args
+
+

src/preprocessing/prostate_mask.py

@@ -0,0 +1,128 @@
+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
+def get_segmask(args):
+
+    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")
+    model_config_file = os.path.join(args.project_dir, "config", "inference.json")
+    model_config = ConfigParser()
+    model_config.read_config(model_config_file)
+    model_config["output_dir"] = args.seg_dir
+    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(
+        [
+            LoadImaged(keys=keys),
+            EnsureChannelFirstd(keys=keys),
+            Orientationd(keys=keys, axcodes="RAS"),
+            Spacingd(keys=keys, pixdim=[0.5, 0.5, 0.5], mode="bilinear"),
+            ScaleIntensityd(keys=keys, minv=0, maxv=1),
+            NormalizeIntensityd(keys=keys),
+            EnsureTyped(keys=keys),
+        ]
+    )
+    logging.info("Starting prostate segmentation")
+    for file in tqdm(files):
+
+        data = {"image": os.path.join(args.t2_dir, file)}
+        transformed_data = transform(data)
+        a = transformed_data
+        with torch.no_grad():
+            images = a["image"].reshape(1, *(a["image"].shape)).to(device)
+            data["pred"] = inferer(images, network=model)
+        pred_img = data["pred"].argmax(1).cpu()
+
+        model_output = {}
+        model_output["image"] = MetaTensor(pred_img, meta=transformed_data["image"].meta)
+        transformed_data["image"].data = model_output["image"].data
+        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))
+        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
+
+
+    

src/preprocessing/register_and_crop.py

@@ -0,0 +1,67 @@
+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):
+    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')
+    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))
+
+        original_spacing = t2_image.GetSpacing()
+        original_size = t2_image.GetSize()
+        new_size = [
+            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
+
+        pat_case = Sample(
+            scans=[
+                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]]),
+        )
+        pat_case.preprocess()
+    
+        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
+

src/train/train_cspca.py

@@ -0,0 +1,141 @@
+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 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() 
+    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)
+
+        optimizer.zero_grad()
+        output = cspca_model(data)
+        output = output.squeeze(1)
+        loss = criterion(output, target)
+        loss.backward()
+        optimizer.step()
+
+        TARGETS.extend(target.detach().cpu())
+        PREDS.extend(output.detach().cpu())
+        run_loss.append(loss.item())
+
+    loss_epoch = run_loss.aggregate()
+    target_list = TARGETS.get_buffer().cpu().numpy()
+    pred_list = PREDS.get_buffer().cpu().numpy()
+    auc_epoch = roc_auc_score(target_list, pred_list)
+
+    return loss_epoch, auc_epoch
+
+def val_epoch(cspca_model, loader, epoch, args):
+    cspca_model.eval()
+    criterion = nn.BCELoss() 
+    loss = 0.0
+    run_loss = CumulativeAverage()
+    TARGETS = Cumulative()
+    PREDS = Cumulative()
+    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)
+
+            output = cspca_model(data)
+            output = output.squeeze(1)
+            loss = criterion(output, target)
+
+            TARGETS.extend(target.detach().cpu())
+            PREDS.extend(output.detach().cpu())
+            run_loss.append(loss.item())
+
+    loss_epoch = run_loss.aggregate()
+    target_list = TARGETS.get_buffer().cpu().numpy()
+    pred_list = PREDS.get_buffer().cpu().numpy()
+    auc_epoch = roc_auc_score(target_list, pred_list)
+    y_pred_categoric = (pred_list >= 0.5)
+    tn, fp, fn, tp = confusion_matrix(target_list, y_pred_categoric).ravel()
+    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}
+    return val_epoch_metric
+