update inference

.gitignore

@@ -12,3 +12,7 @@ __pycache__/
 site/
 updated_segmentations.zip
 updated_segmentations
+temp_2.ipynb
+temp copy.ipynb
+tcia_dataset.ipynb
+check_tum_datatset.ipynb

config/config_run_inference.yaml

@@ -0,0 +1,9 @@
+t2_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/t2
+dwi_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/dwi
+adc_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/adc
+output_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed
+json_path: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/psa_val.json
+
+
+
+

preprocess_main.py

@@ -1,15 +1,17 @@
 import argparse
 import logging
 import os
+from argparse import Namespace
+from collections.abc import Callable
 from pathlib import Path
 
 import yaml
 
-from src.preprocessing.generate_heatmap import get_heatmap
 from src.preprocessing.clip_intensity import clip_adc
+from src.preprocessing.generate_heatmap import get_heatmap
 from src.preprocessing.prostate_mask import get_segmask
 from src.preprocessing.register_and_crop import register_files
-from src.utils import setup_logging, validate_steps
+from src.utils import setup_logging
 
 
 def parse_args():
@@ -50,7 +52,7 @@ if __name__ == "__main__":
     if args.project_dir is None:
         args.project_dir = Path(__file__).resolve().parent  # Set project directory
 
-    FUNCTIONS = {
+    FUNCTIONS: dict[str, Callable[[Namespace], Namespace]] = {
         "register_and_crop": register_files,
         "clip_adc": clip_adc,
         "get_segmentation_mask": get_segmask,
@@ -62,7 +64,7 @@ if __name__ == "__main__":
     logging.info("Starting preprocessing")
     if args.steps is None:
         args.steps = ["register_and_crop", "get_segmentation_mask", "clip_adc", "get_heatmap"]
-    #validate_steps(args.steps)
+    # validate_steps(args.steps)
     for step in args.steps:
         func = FUNCTIONS[step]
         args = func(args)

pyproject.toml

@@ -1,6 +1,7 @@
 [tool.ruff]
 line-length = 100
 target-version = "py39" # Ensures ruff uses 3.9 compatible syntax
+extend-exclude = ["*.ipynb"]
 
 [tool.ruff.lint]
 select = ["E", "W", "F", "I", "B", "N", "UP"]
@@ -22,6 +23,7 @@ skip-magic-trailing-comma = false
 ignore-names = ["sitk", "NormalizeIntensity_custom", "NormalizeIntensity_customd"]
 
 [tool.mypy]
+
 ignore_missing_imports = true
 disable_error_code = ["override", "import-untyped"]
 mypy_path = "."

run_cspca.py

@@ -1,10 +1,10 @@
 import argparse
+import json
 import logging
 import os
 import shutil
 import sys
 from pathlib import Path
-import json
 
 import torch
 import yaml
@@ -21,11 +21,11 @@ from src.utils import get_metrics, save_cspca_checkpoint, setup_logging
 def main_worker(args):
     mil_model = MILModel3D(num_classes=args.num_classes, mil_mode=args.mil_mode)
     cache_dir_path = Path(os.path.join(args.logdir, "cache"))
-    
+
     scaler = StandardScaler()
     with open(os.path.join(args.project_dir, "dataset", "PICAI_cspca_updated_with_psa.json")) as f:
         dataset_json = json.load(f)
-    train_clinical = [i['psa'] for i in dataset_json['test']]
+    train_clinical = [i["psa"] for i in dataset_json["train"]]
     _ = scaler.fit_transform(train_clinical)
     args.psa_mean = scaler.mean_.tolist()
     args.psa_std = scaler.scale_.tolist()
@@ -60,7 +60,9 @@ def main_worker(args):
             train_loss, train_attn_loss, train_auc = train_epoch(
                 cspca_model, train_loader, optimizer, epoch=epoch, args=args
             )
-            logging.info(f"EPOCH {epoch} TRAIN loss: {train_loss:.4f} TRAIN ATTN LOSS: {train_attn_loss:.4f} TRAIN AUC: {train_auc:.4f}")
+            logging.info(
+                f"EPOCH {epoch} TRAIN loss: {train_loss:.4f} TRAIN ATTN LOSS: {train_attn_loss:.4f} TRAIN AUC: {train_auc:.4f}"
+            )
             val_metric = val_epoch(cspca_model, valid_loader, epoch=epoch, args=args)
             logging.info(
                 f"EPOCH {epoch} VAL loss: {val_metric['loss']:.4f} AUC: {val_metric['auc']:.4f}"
@@ -110,7 +112,9 @@ def parse_args():
         required=True,
         help="Operation mode: train or infer",
     )
-    parser.add_argument("--run_name", type=str, default="default_cspca", help="run name for log file")
+    parser.add_argument(
+        "--run_name", type=str, default="default_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")

run_inference.py

@@ -2,52 +2,52 @@ import argparse
 import json
 import logging
 import os
+from argparse import Namespace
+from collections.abc import Callable
 from pathlib import Path
-import json
 
+import streamlit as st
 import torch
 import yaml
 from monai.data import Dataset
+from sklearn.preprocessing import StandardScaler
 
 from src.data.data_loader import data_transform, list_data_collate
 from src.model.cspca_model import CSPCAModel
 from src.model.mil import MILModel3D
+from src.preprocessing.clip_intensity import clip_adc
 from src.preprocessing.generate_heatmap import get_heatmap
-from src.preprocessing.histogram_match import histmatch
 from src.preprocessing.prostate_mask import get_segmask
 from src.preprocessing.register_and_crop import register_files
-from src.utils import get_parent_image, get_patch_coordinate, setup_logging
-import streamlit as st
+from src.utils import get_parent_image, get_patch_coordinate, get_prostate_volume, setup_logging
 
-@st.cache_resource 
-def load_pirads_model(num_classes, mil_mode, project_dir, device):
 
+@st.cache_resource
+def load_pirads_model(num_classes, mil_mode, project_dir, device):
     model = MILModel3D(num_classes=num_classes, mil_mode=mil_mode)
-    checkpoint = torch.load(
-        os.path.join(project_dir, "models", "pirads.pt"), map_location="cpu"
-    )
+    checkpoint = torch.load(os.path.join(project_dir, "models", "pirads.pt"), map_location="cpu")
     model.load_state_dict(checkpoint["state_dict"])
     model.to(device)
-    
+
     model.eval()
     return model
+
+
 @st.cache_resource
 def load_cspca_model(_pirads_model, project_dir, device):
-
     model = CSPCAModel(backbone=_pirads_model).to(device)
-    checkpt = torch.load(
-        os.path.join(project_dir, "models", "cspca_model.pth"), map_location="cpu"
-    )
+    checkpt = torch.load(os.path.join(project_dir, "models", "cspca_model.pth"), map_location="cpu")
     model.load_state_dict(checkpt["state_dict"])
     model = model.to(device)
-    
-    model.eval() 
+
+    model.eval()
     return model
 
 
 def parse_args():
     parser = argparse.ArgumentParser(description="File preprocessing")
     parser.add_argument("--config", type=str, help="Path to YAML config file")
+    parser.add_argument("--json_path", default=None, help="Path to JSON file containing PSA values")
     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")
@@ -60,8 +60,8 @@ def parse_args():
     parser.add_argument("--mil_mode", default="att_trans", type=str)
     parser.add_argument("--use_heatmap", default=True, type=bool)
     parser.add_argument("--use_psa", default=True, type=bool)
-    parser.add_argument("--tile_size", default=64, type=int)
-    parser.add_argument("--tile_count", default=24, type=int)
+    parser.add_argument("--tile_size", default=48, type=int)
+    parser.add_argument("--tile_count", default=40, type=int)
     parser.add_argument("--depth", default=3, type=int)
     parser.add_argument("--project_dir", default=None, help="Project directory")
 
@@ -78,9 +78,9 @@ if __name__ == "__main__":
     if args.project_dir is None:
         args.project_dir = Path(__file__).resolve().parent  # Set project directory
 
-    FUNCTIONS = {
+    FUNCTIONS: dict[str, Callable[[Namespace], Namespace]] = {
         "register_and_crop": register_files,
-        "histogram_match": histmatch,
+        "clip_adc": clip_adc,
         "get_segmentation_mask": get_segmask,
         "get_heatmap": get_heatmap,
     }
@@ -88,7 +88,7 @@ if __name__ == "__main__":
     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"]
+    steps = ["register_and_crop", "get_segmentation_mask", "clip_adc", "get_heatmap"]
     for step in steps:
         func = FUNCTIONS[step]
         args = func(args)
@@ -99,26 +99,23 @@ if __name__ == "__main__":
 
     logging.info("Loading PIRADS model")
     pirads_model = load_pirads_model(args.num_classes, args.mil_mode, args.project_dir, args.device)
-    '''
-    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 = load_cspca_model(pirads_model, args.project_dir, args.device)
-    '''
-    cspca_model = CSPCAModel(backbone=pirads_model).to(args.device)
-    checkpt = torch.load(
-        os.path.join(args.project_dir, "models", "cspca_model.pth"), map_location="cpu"
-    )
-    cspca_model.load_state_dict(checkpt["state_dict"])
-    cspca_model = cspca_model.to(args.device)
-    '''
-    transform = data_transform(args)
+
+    scaler = StandardScaler()
+    with open(os.path.join(args.project_dir, "dataset", "PICAI_cspca_updated_with_psa.json")) as f:
+        dataset_json = json.load(f)
+    train_clinical = [i["psa"] for i in dataset_json["train"]]
+    _ = scaler.fit_transform(train_clinical)
+    args.psa_mean = scaler.mean_.tolist()
+    args.psa_std = scaler.scale_.tolist()
+
+    transform = data_transform(args, split="test")
     files = os.listdir(args.t2_dir)
     args.data_list = []
+    with open(args.json_path) as f:
+        psa_data = json.load(f)
     for file in files:
         temp = {}
         temp["image"] = os.path.join(args.t2_dir, file)
@@ -127,6 +124,9 @@ if __name__ == "__main__":
         temp["heatmap"] = os.path.join(args.heatmapdir, file)
         temp["mask"] = os.path.join(args.seg_dir, file)
         temp["label"] = 0  # dummy label
+        temp["smooth_mask"] = os.path.join(args.smooth_seg_dir, file)
+        prostate_vol = get_prostate_volume(temp["mask"])
+        temp["psa"] = [psa_data[file.split(".nrrd")[0]], prostate_vol]
         args.data_list.append(temp)
 
     dataset = Dataset(data=args.data_list, transform=transform)
@@ -149,21 +149,21 @@ if __name__ == "__main__":
     with torch.no_grad():
         for _, batch_data in enumerate(loader):
             data = batch_data["image"].as_subclass(torch.Tensor).to(args.device)
+            psa_data = batch_data["psa"].as_subclass(torch.Tensor).to(args.device)
             logits = pirads_model(data)
             pirads_score = torch.argmax(logits, dim=1)
             pirads_list.append(pirads_score.item())
 
-            output = cspca_model(data)
+            output = cspca_model(x=data, psa_data=psa_data)
             output = output.squeeze(1)
+            output = torch.sigmoid(output)
             cspca_risk_list.append(output.item())
 
             sh = data.shape
             x = data.reshape(sh[0] * sh[1], sh[2], sh[3], sh[4], sh[5])
             x = cspca_model.backbone.net(x)
             x = x.reshape(sh[0], sh[1], -1)
-            x = x.permute(1, 0, 2)
             x = cspca_model.backbone.transformer(x)
-            x = x.permute(1, 0, 2)
             a = cspca_model.backbone.attention(x)
             a = torch.softmax(a, dim=1)
             a = a.view(-1)

run_pirads.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 import logging
 import os
 import shutil
@@ -11,9 +12,8 @@ import torch
 import wandb
 import yaml
 from monai.utils import set_determinism
-from torch.utils.tensorboard import SummaryWriter
 from sklearn.preprocessing import StandardScaler
-import json
+from torch.utils.tensorboard import SummaryWriter
 
 from src.data.data_loader import get_dataloader
 from src.model.mil import MILModel3D
@@ -45,15 +45,15 @@ def main_worker(args):
     cache_dir_ = os.path.join(args.logdir, "cache")
     model.to(args.device)
     params = model.parameters()
-    
+
     scaler = StandardScaler()
     with open(os.path.join(args.project_dir, "dataset", "PICAI_cspca_updated_with_psa.json")) as f:
         dataset_json = json.load(f)
-    train_clinical = [i['psa'] for i in dataset_json['test']]
+    train_clinical = [i["psa"] for i in dataset_json["train"]]
     _ = scaler.fit_transform(train_clinical)
     args.psa_mean = scaler.mean_.tolist()
     args.psa_std = scaler.scale_.tolist()
-    
+
     if args.mode == "train":
         train_loader = get_dataloader(args, split="train")
         valid_loader = get_dataloader(args, split="test")

src/data/custom_transforms.py

@@ -15,6 +15,7 @@ from monai.utils.enums import TransformBackends
 from monai.utils.type_conversion import convert_data_type, convert_to_dst_type, convert_to_tensor
 from scipy.ndimage import binary_dilation
 
+
 class LabelEncodeIntegerGraded(MapTransform):
     """
     Convert an integer label to encoded array representation of length num_classes,
@@ -42,13 +43,14 @@ class LabelEncodeIntegerGraded(MapTransform):
         for key in self.keys:
             label = int(d[key])
 
-            lz = np.zeros(self.num_classes , dtype=np.float32)
+            lz = np.zeros(self.num_classes, dtype=np.float32)
             lz[:label] = 1.0
             # alternative oneliner lz=(np.arange(self.num_classes)<int(label)).astype(np.float32) #same oneliner
             d[key] = lz
 
         return d
 
+
 class DilateAndSaveMaskd(MapTransform):
     """
     Custom transform to dilate binary mask and save a copy.
@@ -80,7 +82,8 @@ class DilateAndSaveMaskd(MapTransform):
             )  # Add channel dimension back
 
         return d
-    
+
+
 class NormalizePSAd(MapTransform):
     """
     Custom transform to dilate binary mask and save a copy.
@@ -88,7 +91,7 @@ class NormalizePSAd(MapTransform):
 
     def __init__(self, keys, mean, std):
         super().__init__(keys)
-        self.mean = torch.tensor(mean)       
+        self.mean = torch.tensor(mean)
         self.std = torch.tensor(std)
 
     def __call__(self, data):
@@ -386,7 +389,7 @@ class NormalizeIntensity_custom(Transform):
             masked_img = img
         """
         slices = None
-        #mask_data = mask_data.squeeze(0)
+        # mask_data = mask_data.squeeze(0)
         slices_mask = mask_data > 0
         masked_img = img[slices_mask]
 
@@ -440,6 +443,8 @@ class NormalizeIntensity_custom(Transform):
                     sub=self.subtrahend[i] if self.subtrahend is not None else None,
                     div=self.divisor[i] if self.divisor is not None else None,
                 )
+                if isinstance(norm_temp, np.ndarray):
+                    norm_temp = torch.from_numpy(norm_temp)
                 img[i] = norm_temp.squeeze(0)
         else:
             img = self._normalize(img, mask_data, self.subtrahend, self.divisor)

src/data/data_loader.py

@@ -11,14 +11,11 @@ from monai.transforms import (
     DeleteItemsd,
     EnsureTyped,
     LoadImaged,
-    NormalizeIntensityd,
     RandCropByPosNegLabeld,
-    RandWeightedCropd,
+    RandRotate90d,
     ToTensord,
     Transform,
     Transposed,
-    RandFlipd,
-    RandRotate90d,
 )
 from torch.utils.data.dataloader import default_collate
 
@@ -27,10 +24,7 @@ from .custom_transforms import (
     ElementwiseProductd,
     NormalizeIntensity_customd,
     NormalizePSAd,
-    LabelEncodeIntegerGraded,
-
 )
-from sklearn.preprocessing import StandardScaler
 
 
 def list_data_collate(batch: list):
@@ -51,18 +45,18 @@ def list_data_collate(batch: list):
         batch[i] = data
     return default_collate(batch)
 
+
 def data_transform(args: argparse.Namespace, split) -> Transform:
     if split == "train":
-
         transform = Compose(
             [
                 LoadImaged(
-                    keys=["image", "mask", "dwi", "adc", "heatmap","smooth_mask"],
+                    keys=["image", "mask", "dwi", "adc", "heatmap", "smooth_mask"],
                     reader="ITKReader",
                     ensure_channel_first=True,
                     dtype=np.float32,
                 ),
-                #LabelEncodeIntegerGraded(keys=["label"], num_classes=args.num_classes),
+                # LabelEncodeIntegerGraded(keys=["label"], num_classes=args.num_classes),
                 ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"),
                 ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"),
                 NormalizeIntensity_customd(keys=["image"], mask_key="mask"),
@@ -79,11 +73,16 @@ def data_transform(args: argparse.Namespace, split) -> Transform:
                     neg=0,
                     num_samples=args.tile_count,
                 ),
-                RandRotate90d(keys=["image", "final_heatmap", "smooth_mask"], prob=0.6, spatial_axes=(0, 1), max_k=3),
+                RandRotate90d(
+                    keys=["image", "final_heatmap", "smooth_mask"],
+                    prob=0.6,
+                    spatial_axes=(0, 1),
+                    max_k=3,
+                ),
                 NormalizePSAd(keys=["psa"], mean=args.psa_mean, std=args.psa_std),
                 EnsureTyped(keys=["label", "psa"], dtype=torch.float32),
                 Transposed(keys=["image"], indices=(0, 3, 1, 2)),
-                DeleteItemsd(keys=[ "dwi", "adc", "heatmap", "mask"]),
+                DeleteItemsd(keys=["dwi", "adc", "heatmap", "mask"]),
                 ToTensord(keys=["image", "label", "final_heatmap", "smooth_mask", "psa"]),
             ]
         )
@@ -91,12 +90,12 @@ def data_transform(args: argparse.Namespace, split) -> Transform:
         transform = Compose(
             [
                 LoadImaged(
-                    keys=["image", "mask", "dwi", "adc", "heatmap","smooth_mask"],
+                    keys=["image", "mask", "dwi", "adc", "heatmap", "smooth_mask"],
                     reader="ITKReader",
                     ensure_channel_first=True,
                     dtype=np.float32,
                 ),
-                #LabelEncodeIntegerGraded(keys=["label"], num_classes=args.num_classes),
+                # LabelEncodeIntegerGraded(keys=["label"], num_classes=args.num_classes),
                 ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"),
                 ClipMaskIntensityPercentilesd(keys=["dwi"], lower=0, upper=99.5, mask_key="mask"),
                 NormalizeIntensity_customd(keys=["image"], mask_key="mask"),
@@ -116,12 +115,14 @@ def data_transform(args: argparse.Namespace, split) -> Transform:
                 NormalizePSAd(keys=["psa"], mean=args.psa_mean, std=args.psa_std),
                 EnsureTyped(keys=["label", "psa"], dtype=torch.float32),
                 Transposed(keys=["image"], indices=(0, 3, 1, 2)),
-                DeleteItemsd(keys=[ "dwi", "adc", "heatmap", "mask"]),
+                DeleteItemsd(keys=["dwi", "adc", "heatmap", "mask"]),
                 ToTensord(keys=["image", "label", "final_heatmap", "smooth_mask", "psa"]),
             ]
         )
     return transform
-'''
+
+
+"""
 def data_transform(args: argparse.Namespace) -> Transform:
     if args.use_heatmap:
         if args.use_psa:
@@ -251,22 +252,18 @@ def data_transform(args: argparse.Namespace) -> Transform:
                 ]
             )
     return transform
-'''
+"""
 
 
 def get_dataloader(
     args: argparse.Namespace, split: Literal["train", "test"]
 ) -> torch.utils.data.DataLoader:
-
     data_list = load_decathlon_datalist(
         data_list_file_path=args.dataset_json,
         data_list_key=split,
         base_dir=args.data_root,
     )
-    data_list_updated = [
-    {**i, 'psa': i.get('psa', [0, 0])}
-    for i in data_list
-    ]
+    data_list_updated = [{**i, "psa": i.get("psa", [0, 0])} for i in data_list]
     cache_dir_ = os.path.join(args.logdir, "cache")
     os.makedirs(os.path.join(cache_dir_, split), exist_ok=True)
     transform = data_transform(args, split)

src/model/cspca_model.py

@@ -70,17 +70,17 @@ class CSPCAModel(nn.Module):
     def __init__(self, backbone: nn.Module) -> None:
         super().__init__()
         self.backbone = backbone
-        
+
         self.clinical_dim = 2
         self.projection_dim = 32
         self.clinical_projection = nn.Sequential(
             nn.Linear(self.clinical_dim, self.projection_dim),
             nn.ReLU(),
-            nn.BatchNorm1d(self.projection_dim) # Helps stabilize the merged scale
+            nn.BatchNorm1d(self.projection_dim),  # Helps stabilize the merged scale
         )
-        
+
         self.fc_dim = backbone.myfc.in_features
-        self.fc_cspca = SimpleNN(input_dim=self.fc_dim + self.projection_dim) 
+        self.fc_cspca = SimpleNN(input_dim=self.fc_dim + self.projection_dim)
 
     def forward(self, x, psa_data):
         sh = x.shape
@@ -91,7 +91,7 @@ class CSPCAModel(nn.Module):
         a = self.backbone.attention(x)
         a = torch.softmax(a, dim=1)
         x = torch.sum(x * a, dim=1)
-        
+
         psa_features = self.clinical_projection(psa_data)
         x = torch.cat((x, psa_features), dim=1)
 

src/model/mil.py

@@ -130,7 +130,9 @@ class MILModel3D(nn.Module):
             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, batch_first=True)
+            transformer = nn.TransformerEncoderLayer(
+                d_model=nfc, nhead=8, dropout=trans_dropout, batch_first=True
+            )
             self.transformer = nn.TransformerEncoder(transformer, num_layers=trans_blocks)
             self.attention = nn.Sequential(nn.Linear(nfc, 2048), nn.Tanh(), nn.Linear(2048, 1))
 
@@ -190,7 +192,6 @@ class MILModel3D(nn.Module):
             x = self.myfc(x)
 
         elif self.mil_mode == "att_trans" and self.transformer is not None:
-            
             x = self.transformer(x)
 
             a = self.attention(x)

src/preprocessing/clip_intensity.py

@@ -7,20 +7,18 @@ import numpy as np
 from tqdm import tqdm
 
 
-def clip_adc(args: argparse.Namespace, adc_min = 0.0, adc_max = 3500.0):
-
+def clip_adc(args: argparse.Namespace, adc_min=0.0, adc_max=3500.0):
     files = os.listdir(args.adc_dir)
     clip_adc_dir = os.path.join(args.output_dir, "ADC_clipped")
     os.makedirs(clip_adc_dir, exist_ok=True)
     logging.info("Starting clipping ADC")
 
     for file in tqdm(files):
-
         adc, header_adc = nrrd.read(os.path.join(args.adc_dir, file))
-    
+
         if np.percentile(adc, 99) < 100:
             adc = adc * 100
- 
+
         adc_clipped = np.clip(adc, adc_min, adc_max)
         adc_normalized = adc_clipped / adc_max
         nrrd.write(os.path.join(clip_adc_dir, file), adc_normalized, header_adc)

src/preprocessing/generate_heatmap.py

@@ -1,50 +1,18 @@
 import argparse
 import logging
 import os
-import SimpleITK as sitk
+import shutil
 
 import nrrd
 import numpy as np
+import SimpleITK as sitk
 from tqdm import tqdm
-from scipy.ndimage import binary_dilation
-import shutil
 
-'''
-def dilate_prostate_mask_xyz(mask_array, expand_mm=2.0, pixel_spacing_xy=0.4):
-    """
-    Dilates a 3D prostate mask only in the axial (X,Y) plane.
-    Specifically designed for arrays where Z is the LAST dimension (X, Y, Z).
-    
-    Args:
-        mask_array: 3D numpy array of the mask (shape: 232, 232, 23)
-        expand_mm: How many millimeters of safety margin you want to add
-        pixel_spacing_xy: The in-plane spacing (e.g., 0.4 mm)
-    """
-    # 1. Calculate how many pixels to expand radially
-    iterations = int(np.round(expand_mm / pixel_spacing_xy))
-    
-    # 2. Create the structuring element for (X, Y, Z)
-    structuring_element = np.zeros((3, 3, 3), dtype=bool)
-    
-    # We set the middle Z-index (index 1) to True across all X and Y.
-    # This forces the math to spread outward in X and Y, but never up/down in Z.
-    structuring_element[:, :, 1] = True  
-    
-    # 3. Apply the dilation
-    dilated_mask = binary_dilation(mask_array > 0, 
-                                   structure=structuring_element, 
-                                   iterations=iterations)
-    
-    # 4. Return as the original integer type
-    return dilated_mask.astype(mask_array.dtype)
-'''
 
 def smoothen_mask(args, sigma=1.0):
-    
     files = os.listdir(args.seg_dir)
-    
+
     for file in files:
-    
         mask = sitk.ReadImage(os.path.join(args.seg_dir, file))
         mask = sitk.Cast(mask, sitk.sitkFloat32)
 
@@ -54,12 +22,11 @@ def smoothen_mask(args, sigma=1.0):
         # Threshold back to binary
         smooth_mask = smoothed > 0.5
         smooth_mask = sitk.Cast(smooth_mask, sitk.sitkUInt8)
-        
+
         out_path = os.path.join(args.smooth_seg_dir_temp, file)
         sitk.WriteImage(smooth_mask, out_path)
 
 
-
 def get_heatmap(args: argparse.Namespace) -> argparse.Namespace:
     """
     Generate heatmaps from DWI (Diffusion Weighted Imaging) and ADC (Apparent Diffusion Coefficient) medical imaging data.
@@ -95,13 +62,13 @@ def get_heatmap(args: argparse.Namespace) -> argparse.Namespace:
     args.smooth_seg_dir = os.path.join(args.output_dir, "smooth_prostate_mask/")
     os.makedirs(args.smooth_seg_dir, exist_ok=True)
     smoothen_mask(args)
-    
+
     logging.info("Starting heatmap generation")
     for file in tqdm(files):
         bool_dwi = False
         bool_adc = False
         mask_temp, header_mask = nrrd.read(os.path.join(args.seg_dir, file))
-        #spacing = np.linalg.norm(header_mask['space directions'], axis=1)
+        # spacing = np.linalg.norm(header_mask['space directions'], axis=1)
         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_temp > 0]
@@ -109,12 +76,11 @@ def get_heatmap(args: argparse.Namespace) -> argparse.Namespace:
         mask = np.maximum(mask, mask_temp)
 
         if len(nonzero_vals_dwi) > 0:
-            #min_val = nonzero_vals_dwi.min()
-            #max_val = nonzero_vals_dwi.max()
+            # min_val = nonzero_vals_dwi.min()
+            # max_val = nonzero_vals_dwi.max()
             min_val, max_val = np.percentile(nonzero_vals_dwi, [1, 99])
             clipped_dwi = np.clip(dwi, min_val, max_val)
 
-
             heatmap_dwi = np.zeros_like(clipped_dwi, dtype=np.float32)
 
             if min_val != max_val:
@@ -135,13 +101,12 @@ def get_heatmap(args: argparse.Namespace) -> argparse.Namespace:
             heatmap_adc = (max_val - adc) / (max_val - min_val)
             masked_heatmap_adc = np.where(mask > 0, heatmap_adc, 0)
 
-
         else:
             bool_adc = True
 
         if not bool_dwi and not bool_adc:
-            mix_mask = (masked_heatmap_dwi * 0.3 ) + (masked_heatmap_adc * 0.7)
-            #mix_mask = (masked_heatmap_dwi**0.5) * (masked_heatmap_adc**2.0)
+            mix_mask = (masked_heatmap_dwi * 0.3) + (masked_heatmap_adc * 0.7)
+            # mix_mask = (masked_heatmap_dwi**0.5) * (masked_heatmap_adc**2.0)
             write_header = header_dwi
         elif bool_dwi:
             mix_mask = masked_heatmap_adc
@@ -150,8 +115,10 @@ def get_heatmap(args: argparse.Namespace) -> argparse.Namespace:
             mix_mask = np.ones_like(adc, dtype=np.float32)
             write_header = header_dwi
 
-        mix_mask = (mix_mask - mix_mask[mask>0].min()) / (mix_mask[mask>0].max() - mix_mask[mask>0].min())
-        mix_mask =  np.where(mask > 0, mix_mask, 0)
+        mix_mask = (mix_mask - mix_mask[mask > 0].min()) / (
+            mix_mask[mask > 0].max() - mix_mask[mask > 0].min()
+        )
+        mix_mask = np.where(mask > 0, mix_mask, 0)
         nrrd.write(os.path.join(args.heatmapdir, file), mix_mask, write_header)
         nrrd.write(os.path.join(args.smooth_seg_dir, file), mask, header_mask)
 

src/preprocessing/prostate_mask.py

@@ -22,10 +22,12 @@ from tqdm import tqdm
 
 set_determinism(43)
 
+
 def is_continuous(idx):
     idx = sorted(idx)
     return all(b - a == 1 for a, b in zip(idx, idx[1:]))
 
+
 def get_segmask(args: argparse.Namespace) -> argparse.Namespace:
     """
     Generate prostate segmentation masks using a pre-trained deep learning model.
@@ -110,8 +112,10 @@ def get_segmask(args: argparse.Namespace) -> argparse.Namespace:
         top_slices = np.argsort(nonzero_counts)[-10:]
         try:
             assert is_continuous(top_slices), "Top slices are not continuous"
-        except AssertionError as e:
-            logging.warning(f"Top slices for {file} are not continuous: {top_slices}. Proceeding with non-continuous slices.")
+        except AssertionError:
+            logging.warning(
+                f"Top slices for {file} are not continuous: {top_slices}. Proceeding with non-continuous slices."
+            )
         output_ = np.zeros_like(pred_nrrd)
         output_[:, :, top_slices] = pred_nrrd[:, :, top_slices]
 

src/preprocessing/register_and_crop.py

@@ -84,85 +84,3 @@ def register_files(args: argparse.Namespace) -> argparse.Namespace:
     args.adc_dir = adc_registered_dir
 
     return args
-'''
-
-def create_master_grid(image, target_spacing=[0.4, 0.4, 3.0]):
-    """
-    Resamples an image to the target spacing to create the Master Grid.
-    """
-    original_spacing = image.GetSpacing()
-    original_size = image.GetSize()
-    
-    # Calculate the new dimensions based on the target spacing
-    new_size = [
-        int(np.round(original_size[0] * (original_spacing[0] / target_spacing[0]))),
-        int(np.round(original_size[1] * (original_spacing[1] / target_spacing[1]))),
-        int(np.round(original_size[2] * (original_spacing[2] / target_spacing[2])))
-    ]
-    
-    # Setup the resampler
-    resampler = sitk.ResampleImageFilter()
-    resampler.SetSize(new_size)
-    resampler.SetOutputSpacing(target_spacing)
-    resampler.SetOutputOrigin(image.GetOrigin())
-    resampler.SetOutputDirection(image.GetDirection())
-    # B-Spline is excellent for T2W anatomical structure
-    resampler.SetInterpolator(sitk.sitkBSpline) 
-    
-    return resampler.Execute(image)
-
-def register_files(args: argparse.Namespace) -> argparse.Namespace:
-
-    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):
-        
-        t2w = sitk.ReadImage(os.path.join(args.t2_dir, file), sitk.sitkFloat32)
-        dwi = sitk.ReadImage(os.path.join(args.dwi_dir, file), sitk.sitkFloat32)
-        adc = sitk.ReadImage(os.path.join(args.adc_dir, file), sitk.sitkFloat32)
-
-        t2w_master = create_master_grid(t2w, target_spacing=new_spacing)
-
-        registration = sitk.ImageRegistrationMethod()
-        registration.SetMetricAsMattesMutualInformation(numberOfHistogramBins=50)
-        registration.SetMetricSamplingStrategy(registration.RANDOM)
-        registration.SetMetricSamplingPercentage(0.1)
-        registration.SetOptimizerAsGradientDescent(learningRate=1.0, 
-                                                    numberOfIterations=100, 
-                                                    convergenceMinimumValue=1e-6, 
-                                                    convergenceWindowSize=10)
-
-        initial_transform = sitk.CenteredTransformInitializer(
-            t2w_master, dwi, sitk.TranslationTransform(3), sitk.CenteredTransformInitializerFilter.GEOMETRY
-        )
-        registration.SetInitialTransform(initial_transform, inPlace=False)
-        final_transform = registration.Execute(t2w_master, dwi)
-
-        dwi_aligned = sitk.Resample(
-            dwi, t2w_master, final_transform, sitk.sitkLinear, 0.0, dwi.GetPixelID()
-        )
-        adc_aligned = sitk.Resample(
-            adc, t2w_master, final_transform, sitk.sitkLinear, 0.0, adc.GetPixelID()
-        )
-
-        cropped_t2 = crop(t2w_master, [args.margin, args.margin, 0.0])
-        cropped_dwi = crop(dwi_aligned, [args.margin, args.margin, 0.0])
-        cropped_adc = crop(adc_aligned, [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

@@ -1,8 +1,9 @@
+import argparse
+
 import torch
 import torch.nn as nn
 from monai.metrics import Cumulative, CumulativeAverage
 from sklearn.metrics import confusion_matrix, roc_auc_score
-import argparse
 
 
 def get_lambda_att(epoch: int, max_lambda: float = 2.0, warmup_epochs: int = 10) -> float:
@@ -10,7 +11,8 @@ def get_lambda_att(epoch: int, max_lambda: float = 2.0, warmup_epochs: int = 10)
         return (epoch / warmup_epochs) * max_lambda
     else:
         return max_lambda
-    
+
+
 def get_attention_scores(
     data: torch.Tensor,
     target: torch.Tensor,
@@ -68,10 +70,10 @@ def get_attention_scores(
 
     return att_labels, shuffled_images
 
+
 def train_epoch(cspca_model, loader, optimizer, epoch, args):
-    
     lambda_att = get_lambda_att(epoch, warmup_epochs=25)
-    
+
     cspca_model.train()
     criterion = nn.BCEWithLogitsLoss()
     att_criterion = nn.CosineSimilarity(dim=1, eps=1e-6)
@@ -87,7 +89,7 @@ def train_epoch(cspca_model, loader, optimizer, epoch, args):
         data = batch_data["image"].as_subclass(torch.Tensor).to(args.device)
         target = batch_data["label"].as_subclass(torch.Tensor).to(args.device)
         psa_data = batch_data["psa"].as_subclass(torch.Tensor).to(args.device)
-        
+
         if args.use_heatmap:
             att_labels, shuffled_images = get_attention_scores(
                 data, target, batch_data["final_heatmap"], batch_data["smooth_mask"], args
@@ -95,9 +97,9 @@ def train_epoch(cspca_model, loader, optimizer, epoch, args):
             att_labels = att_labels + eps
         else:
             shuffled_images = data.to(args.device)
-            
+
         optimizer.zero_grad()
-        output = cspca_model(x = shuffled_images, psa_data = psa_data)
+        output = cspca_model(x=shuffled_images, psa_data=psa_data)
         output = output.squeeze(1)
         class_loss = criterion(output, target)
         if args.use_heatmap:
@@ -117,8 +119,7 @@ def train_epoch(cspca_model, loader, optimizer, epoch, args):
         else:
             loss = class_loss
             attn_loss = torch.tensor(0.0)
-  
-                
+
         loss.backward()
         optimizer.step()
 
@@ -149,7 +150,7 @@ def val_epoch(cspca_model, loader, epoch, args):
             target = batch_data["label"].as_subclass(torch.Tensor).to(args.device)
             psa_data = batch_data["psa"].as_subclass(torch.Tensor).to(args.device)
 
-            output = cspca_model(x = data, psa_data = psa_data)
+            output = cspca_model(x=data, psa_data=psa_data)
             output = output.squeeze(1)
             loss = criterion(output, target)
 

src/train/train_pirads.py

@@ -152,7 +152,7 @@ def train_epoch(model, loader, optimizer, scaler, epoch, args):
             start_time = time.time()
 
         del data, target, shuffled_images, logits, logits_attn
-    #torch.cuda.empty_cache()
+    # torch.cuda.empty_cache()
     batch_norm_epoch = batch_norm.aggregate()
     attn_loss_epoch = run_att_loss.aggregate()
     loss_epoch = run_loss.aggregate()
@@ -197,7 +197,7 @@ def val_epoch(model, loader, epoch, args):
             start_time = time.time()
 
             del data, target, logits
-            #torch.cuda.empty_cache()
+            # torch.cuda.empty_cache()
 
         # Calculate QWK metric (Quadratic Weigted Kappa) https://en.wikipedia.org/wiki/Cohen%27s_kappa
         preds_cumulative = preds_cumulative.get_buffer().cpu().numpy()

src/utils.py

@@ -8,6 +8,7 @@ from typing import Any, Union
 import cv2
 import matplotlib.patches as patches
 import matplotlib.pyplot as plt
+import nrrd
 import numpy as np
 import torch
 from monai.data import Dataset
@@ -245,3 +246,19 @@ def visualise_patches(coords, image, tile_size=64, depth=3):
     plt.subplots_adjust(left=0.06)
     plt.tight_layout()
     plt.show()
+
+
+def get_prostate_volume(mask_path) -> np.ndarray:
+    mask_data, header = nrrd.read(mask_path)
+    space_directions = header.get("space directions")
+    spacing = np.array([np.linalg.norm(vec) for vec in space_directions if np.any(vec)])
+
+    if len(spacing) != 3:
+        raise ValueError(f"Expected 3 spatial dimensions, found {len(spacing)}")
+
+    voxel_volume_mm3 = np.prod(spacing)
+    voxel_count = np.sum(mask_data > 0)
+    true_volume_mm3 = voxel_count * voxel_volume_mm3
+    true_volume_cc = true_volume_mm3 / 1000.0  # Convert mm³ to cc (mL)
+
+    return true_volume_cc

tcia_dataset.ipynb

@@ -7,16 +7,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import json\n",
     "import os\n",
-    "import nrrd\n",
-    "import numpy as np\n",
-    "import nrrd\n",
-    "import SimpleITK as sitk\n",
-    "from AIAH_utility.viewer import BasicViewer\n",
-    "from tqdm import tqdm\n",
     "import re\n",
+    "\n",
+    "import numpy as np\n",
     "import pandas as pd\n",
-    "import json"
+    "import SimpleITK as sitk\n",
+    "from tqdm import tqdm"
    ]
   },
   {
@@ -52,15 +50,15 @@
     }
    ],
    "source": [
-    "data_dir = '/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/manifest-1777460989222/prostate_mri_us_biopsy'\n",
-    "out_dir = '/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/'\n",
+    "data_dir = \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/manifest-1777460989222/prostate_mri_us_biopsy\"\n",
+    "out_dir = \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/\"\n",
     "print(len(os.listdir(data_dir)))\n",
-    "'''\n",
+    "\"\"\"\n",
     "for file in os.listdir(data_dir):\n",
     "    if not os.path.isfile(os.path.join(data_dir, file)):\n",
     "        if len(os.listdir(os.path.join(data_dir, file))) > 1:\n",
     "            print(file)\n",
-    "'''  "
+    "\"\"\""
    ]
   },
   {
@@ -70,20 +68,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n",
-    "\n",
     "def identify_mri_sequence(dicom_dir):\n",
     "    # 1. Get all DICOM files in the folder\n",
     "    reader = sitk.ImageSeriesReader()\n",
     "    series_IDs = reader.GetGDCMSeriesIDs(dicom_dir)\n",
-    "    \n",
+    "\n",
     "    if not series_IDs:\n",
     "        print(\"No DICOM files found.\")\n",
     "        return None\n",
     "\n",
     "    # Get the filenames for the first series found\n",
     "    dicom_names = reader.GetGDCMSeriesFileNames(dicom_dir, series_IDs[0])\n",
-    "    first_file = dicom_names[0] # We only need to check the first slice's header\n",
+    "    first_file = dicom_names[0]  # We only need to check the first slice's header\n",
     "\n",
     "    # 2. Read only the header information (fast, doesn't load image pixels)\n",
     "    file_reader = sitk.ImageFileReader()\n",
@@ -97,7 +93,7 @@
     "    except RuntimeError:\n",
     "        series_description = \"unknown (tag missing)\"\n",
     "\n",
-    "    #print(f\"Raw DICOM Series Description: '{series_description}'\")\n",
+    "    # print(f\"Raw DICOM Series Description: '{series_description}'\")\n",
     "\n",
     "    # 4. Keyword matching to identify the sequence\n",
     "    sequence_type = \"Unclassified\"\n",
@@ -108,42 +104,40 @@
     "        sequence_type = \"t2\"\n",
     "    elif \"adc\" in desc_lower:\n",
     "        sequence_type = \"adc\"\n",
-    "    elif any(word in desc_lower for word in dwi_keywords) or re.search(r'b\\d+', desc_lower):\n",
+    "    elif any(word in desc_lower for word in dwi_keywords) or re.search(r\"b\\d+\", desc_lower):\n",
     "        sequence_type = \"dwi\"\n",
     "\n",
-    "        \n",
-    "    #print(f\"Identified Sequence Type: {sequence_type}\")\n",
+    "    # print(f\"Identified Sequence Type: {sequence_type}\")\n",
     "    return sequence_type, series_description\n",
     "\n",
     "\n",
-    "\n",
     "def dicom_to_nrrd(dicom_dir, output_filepath):\n",
     "    # 1. Initialize the ImageSeriesReader\n",
     "    reader = sitk.ImageSeriesReader()\n",
     "\n",
     "    # 2. Find the DICOM series within the folder\n",
-    "    # A single folder might contain multiple scans (e.g., a T1 and a T2). \n",
+    "    # A single folder might contain multiple scans (e.g., a T1 and a T2).\n",
     "    # This function groups them by their unique Series Instance UID.\n",
     "    series_IDs = reader.GetGDCMSeriesIDs(dicom_dir)\n",
-    "    \n",
+    "\n",
     "    if not series_IDs:\n",
     "        print(f\"Error: No DICOM series found in the directory {dicom_dir}.\")\n",
     "        return\n",
     "\n",
     "    # Assuming you want the first (or only) series in the folder\n",
     "    series_ID = series_IDs[0]\n",
-    "    \n",
+    "\n",
     "    # 3. Get the list of files belonging to this specific series.\n",
     "    # SimpleITK automatically sorts them by spatial position/instance number here.\n",
     "    dicom_names = reader.GetGDCMSeriesFileNames(dicom_dir, series_ID)\n",
     "    reader.SetFileNames(dicom_names)\n",
     "\n",
-    "    #print(f\"Found series {series_ID}\")\n",
-    "    #print(f\"Reading {len(dicom_names)} DICOM files...\")\n",
+    "    # print(f\"Found series {series_ID}\")\n",
+    "    # print(f\"Reading {len(dicom_names)} DICOM files...\")\n",
     "\n",
     "    # 4. Read the files into a single 3D image object\n",
     "    image = reader.Execute()\n",
-    "    '''\n",
+    "    \"\"\"\n",
     "    # (Optional) Print out the header info that SimpleITK extracted\n",
     "    print(\"\\nExtracted Header Information:\")\n",
     "    print(f\"Size (X, Y, Z): {image.GetSize()}\")\n",
@@ -153,13 +147,9 @@
     "\n",
     "    # 5. Write the 3D image to an NRRD file\n",
     "    # SimpleITK automatically translates the spatial metadata into the NRRD header format.\n",
-    "    '''\n",
+    "    \"\"\"\n",
     "    sitk.WriteImage(image, output_filepath)\n",
-    "    #print(f\"Successfully saved 3D volume to: {output_filepath}\")\n",
-    "\n",
-    "\n",
-    "\n",
-    "\n"
+    "    # print(f\"Successfully saved 3D volume to: {output_filepath}\")"
    ]
   },
   {
@@ -171,28 +161,28 @@
    "source": [
     "import logging\n",
     "\n",
-    "\n",
     "# Keep your basic logging config\n",
     "logging.basicConfig(\n",
-    "    filename='dicom_warnings.log',\n",
-    "    filemode='a',\n",
+    "    filename=\"dicom_warnings.log\",\n",
+    "    filemode=\"a\",\n",
     "    level=logging.WARNING,\n",
-    "    format='%(asctime)s - %(levelname)s - %(message)s'\n",
+    "    format=\"%(asctime)s - %(levelname)s - %(message)s\",\n",
     ")\n",
     "\n",
+    "\n",
     "class SimpleITKWarningTracker(sitk.LoggerBase):\n",
     "    def __init__(self, logger: logging.Logger = logging.getLogger(\"SimpleITK\")):\n",
     "        super().__init__()\n",
     "        self._logger = logger\n",
-    "        \n",
+    "\n",
     "        # --- NEW ADDITIONS ---\n",
     "        self.current_uid = None\n",
-    "        self.problematic_uids = set() # Using a set so we don't get duplicates\n",
+    "        self.problematic_uids = set()  # Using a set so we don't get duplicates\n",
     "\n",
     "    def DisplayWarningText(self, s):\n",
     "        # 1. Log to the file, but inject the UID so you know exactly which one failed\n",
     "        self._logger.warning(f\"[{self.current_uid}] {s.rstrip()}\")\n",
-    "        \n",
+    "\n",
     "        # 2. Capture the UID into our set\n",
     "        if self.current_uid is not None:\n",
     "            self.problematic_uids.add(self.current_uid)\n",
@@ -203,9 +193,14 @@
     "            self.problematic_uids.add(self.current_uid)\n",
     "\n",
     "    # Standard passthroughs for everything else\n",
-    "    def DisplayText(self, s): self._logger.info(s.rstrip())\n",
-    "    def DisplayGenericOutputText(self, s): self._logger.info(s.rstrip())\n",
-    "    def DisplayDebugText(self, s): self._logger.debug(s.rstrip())"
+    "    def DisplayText(self, s):\n",
+    "        self._logger.info(s.rstrip())\n",
+    "\n",
+    "    def DisplayGenericOutputText(self, s):\n",
+    "        self._logger.info(s.rstrip())\n",
+    "\n",
+    "    def DisplayDebugText(self, s):\n",
+    "        self._logger.debug(s.rstrip())"
    ]
   },
   {
@@ -224,26 +219,24 @@
     }
    ],
    "source": [
-    "\n",
-    "\n",
-    "\n",
     "def get_clean_file_set(folder_path):\n",
     "    \"\"\"\n",
-    "    Returns a set of filenames in a folder, \n",
+    "    Returns a set of filenames in a folder,\n",
     "    ignoring hidden files like .DS_Store.\n",
     "    \"\"\"\n",
     "    if not os.path.exists(folder_path):\n",
     "        print(f\"Warning: Folder not found -> {folder_path}\")\n",
     "        return set()\n",
-    "        \n",
-    "    return set(f for f in os.listdir(folder_path) if not f.startswith('.'))\n",
+    "\n",
+    "    return set(f for f in os.listdir(folder_path) if not f.startswith(\".\"))\n",
+    "\n",
     "\n",
     "def audit_mri_folders(t2_dir, dwi_dir, adc_dir):\n",
     "    # 1. Grab the sets of files from each folder\n",
     "    t2_files = get_clean_file_set(t2_dir)\n",
     "    dwi_files = get_clean_file_set(dwi_dir)\n",
     "    adc_files = get_clean_file_set(adc_dir)\n",
-    "    \n",
+    "\n",
     "    # 2. Check for absolute perfection\n",
     "    if t2_files == dwi_files == adc_files:\n",
     "        print(\"✅ SUCCESS: All three folders contain the exact same files.\")\n",
@@ -252,32 +245,36 @@
     "\n",
     "    # 3. If they don't match, figure out exactly what went wrong\n",
     "    print(\"❌ MISMATCH DETECTED: The folders do not have the same files.\\n\")\n",
-    "    \n",
+    "\n",
     "    # Create a master list of EVERY unique file found across all three folders\n",
     "    all_known_files = t2_files | dwi_files | adc_files\n",
     "    print(f\"Total unique cases found across all folders: {len(all_known_files)}\\n\")\n",
-    "    \n",
+    "\n",
     "    # Subtracting a folder's files from the master list reveals exactly what it is missing\n",
     "    missing_from_t2 = all_known_files - t2_files\n",
     "    missing_from_dwi = all_known_files - dwi_files\n",
     "    missing_from_adc = all_known_files - adc_files\n",
-    "    \n",
+    "\n",
     "    # 4. Print the detailed report\n",
     "    if missing_from_t2:\n",
     "        print(f\"Missing from T2 ({len(missing_from_t2)} files):\")\n",
-    "        for f in missing_from_t2: print(f\"  - {f}\")\n",
+    "        for f in missing_from_t2:\n",
+    "            print(f\"  - {f}\")\n",
     "        print()\n",
-    "        \n",
+    "\n",
     "    if missing_from_dwi:\n",
     "        print(f\"Missing from DWI ({len(missing_from_dwi)} files):\")\n",
-    "        for f in missing_from_dwi: print(f\"  - {f}\")\n",
+    "        for f in missing_from_dwi:\n",
+    "            print(f\"  - {f}\")\n",
     "        print()\n",
-    "        \n",
+    "\n",
     "    if missing_from_adc:\n",
     "        print(f\"Missing from ADC ({len(missing_from_adc)} files):\")\n",
-    "        for f in missing_from_adc: print(f\"  - {f}\")\n",
+    "        for f in missing_from_adc:\n",
+    "            print(f\"  - {f}\")\n",
     "        print()\n",
     "\n",
+    "\n",
     "# ==========================================\n",
     "# Run the auditor\n",
     "# ==========================================\n",
@@ -285,7 +282,7 @@
     "t2_folder = os.path.join(out_dir, \"t2\")\n",
     "adc_folder = os.path.join(out_dir, \"adc\")\n",
     "dwi_folder = os.path.join(out_dir, \"dwi\")\n",
-    "audit_mri_folders(t2_folder, dwi_folder, adc_folder)\n"
+    "audit_mri_folders(t2_folder, dwi_folder, adc_folder)"
    ]
   },
   {
@@ -323,44 +320,44 @@
     }
    ],
    "source": [
-    "\n",
     "def get_clean_file_set(folder_path):\n",
     "    if not os.path.exists(folder_path):\n",
     "        return set()\n",
-    "    return set(f for f in os.listdir(folder_path) if not f.startswith('.'))\n",
+    "    return set(f for f in os.listdir(folder_path) if not f.startswith(\".\"))\n",
+    "\n",
     "\n",
     "def remove_orphan_files(t2_dir, dwi_dir, adc_dir, dry_run=True):\n",
     "    # 1. Grab the sets of files from each folder\n",
     "    t2_files = get_clean_file_set(t2_dir)\n",
     "    dwi_files = get_clean_file_set(dwi_dir)\n",
     "    adc_files = get_clean_file_set(adc_dir)\n",
-    "    \n",
+    "\n",
     "    # 2. Find the \"Perfect Matches\" (files that exist in ALL three folders)\n",
     "    common_files = t2_files & dwi_files & adc_files\n",
-    "    \n",
+    "\n",
     "    # 3. Setup directories to check\n",
     "    directories = [\n",
     "        (\"T2\", t2_dir, t2_files),\n",
     "        (\"DWI\", dwi_dir, dwi_files),\n",
-    "        (\"ADC\", adc_dir, adc_files)\n",
+    "        (\"ADC\", adc_dir, adc_files),\n",
     "    ]\n",
-    "    \n",
+    "\n",
     "    total_deleted = 0\n",
-    "    \n",
+    "\n",
     "    print(f\"--- Running in {'DRY RUN (Safe)' if dry_run else 'ACTIVE DELETE'} mode ---\\n\")\n",
-    "    \n",
+    "\n",
     "    # 4. Loop through each folder and delete files that aren't in the common pool\n",
     "    for name, dir_path, files in directories:\n",
     "        orphans = files - common_files\n",
-    "        \n",
+    "\n",
     "        if not orphans:\n",
     "            print(f\"✅ {name} folder is already clean.\")\n",
     "            continue\n",
-    "            \n",
+    "\n",
     "        print(f\"🧹 Cleaning {name} folder ({len(orphans)} orphan files found)...\")\n",
     "        for orphan in orphans:\n",
     "            file_path = os.path.join(dir_path, orphan)\n",
-    "            \n",
+    "\n",
     "            if dry_run:\n",
     "                print(f\"   [Would Delete] -> {file_path}\")\n",
     "            else:\n",
@@ -370,16 +367,15 @@
     "                    total_deleted += 1\n",
     "                except Exception as e:\n",
     "                    print(f\"   [Error] -> Could not delete {file_path}: {e}\")\n",
-    "                    \n",
+    "\n",
     "    # 5. Final Summary\n",
-    "    print(\"\\n\" + \"=\"*40)\n",
+    "    print(\"\\n\" + \"=\" * 40)\n",
     "    if dry_run:\n",
     "        print(\"🛑 This was a DRY RUN. No files were actually deleted.\")\n",
     "        print(\"To permanently delete these files, change 'dry_run=False' in the script.\")\n",
     "    else:\n",
     "        print(f\"✅ Cleanup complete. {total_deleted} orphan files permanently deleted.\")\n",
-    "    print(\"=\"*40)\n",
-    "\n",
+    "    print(\"=\" * 40)\n",
     "\n",
     "\n",
     "# Run it once with True to see what will happen.\n",
@@ -521,28 +517,25 @@
    "source": [
     "import os\n",
     "\n",
+    "\n",
     "def delete_bad_uids(t2_dir, dwi_dir, adc_dir, bad_uids_list, dry_run=True):\n",
     "    # The folders we need to check\n",
-    "    directories = {\n",
-    "        \"T2\": t2_dir,\n",
-    "        \"DWI\": dwi_dir,\n",
-    "        \"ADC\": adc_dir\n",
-    "    }\n",
-    "    \n",
+    "    directories = {\"T2\": t2_dir, \"DWI\": dwi_dir, \"ADC\": adc_dir}\n",
+    "\n",
     "    total_deleted = 0\n",
     "    total_not_found = 0\n",
-    "    \n",
+    "\n",
     "    print(f\"--- Running in {'DRY RUN (Safe)' if dry_run else 'ACTIVE DELETE'} mode ---\\n\")\n",
-    "    \n",
+    "\n",
     "    # Loop through every bad UID in your list\n",
     "    for uid in bad_uids_list:\n",
     "        filename = f\"{uid}.nrrd\"\n",
     "        print(f\"Targeting: {filename}\")\n",
-    "        \n",
+    "\n",
     "        # Check all three folders for this specific file\n",
     "        for folder_name, folder_path in directories.items():\n",
     "            file_path = os.path.join(folder_path, filename)\n",
-    "            \n",
+    "\n",
     "            if os.path.exists(file_path):\n",
     "                if dry_run:\n",
     "                    print(f\"   [{folder_name}] -> [Would Delete]: {file_path}\")\n",
@@ -557,16 +550,15 @@
     "                # Optional: Uncomment the line below if you want to know when a file was already missing\n",
     "                # print(f\"   [{folder_name}] -> Not found (already gone)\")\n",
     "                total_not_found += 1\n",
-    "                \n",
+    "\n",
     "    # Final Summary\n",
-    "    print(\"\\n\" + \"=\"*40)\n",
+    "    print(\"\\n\" + \"=\" * 40)\n",
     "    if dry_run:\n",
     "        print(\"🛑 This was a DRY RUN. No files were actually deleted.\")\n",
     "        print(\"To permanently delete these files, change 'dry_run=False' in the script.\")\n",
     "    else:\n",
     "        print(f\"✅ Cleanup complete. {total_deleted} files permanently deleted.\")\n",
-    "    print(\"=\"*40)\n",
-    "\n",
+    "    print(\"=\" * 40)\n",
     "\n",
     "\n",
     "# Paste your list of bad UIDs here\n",
@@ -574,7 +566,7 @@
     "\n",
     "bad_uids = list(tracker.problematic_uids)\n",
     "\n",
-    "# Run it once with True to verify. \n",
+    "# Run it once with True to verify.\n",
     "# Change to False to actually delete the files.\n",
     "delete_bad_uids(t2_folder, dwi_folder, adc_folder, bad_uids, dry_run=False)"
    ]
@@ -654,19 +646,18 @@
     "            for seq in os.listdir(os.path.join(data_dir, file, st_uid)):\n",
     "                seq_path = os.path.join(data_dir, file, st_uid, seq)\n",
     "                seq_type, _ = identify_mri_sequence(seq_path)\n",
-    "                output_nrrd_path = os.path.join(out_dir,seq_type, f\"{st_uid}.nrrd\")       \n",
+    "                output_nrrd_path = os.path.join(out_dir, seq_type, f\"{st_uid}.nrrd\")\n",
     "\n",
     "                dicom_to_nrrd(seq_path, output_nrrd_path)\n",
     "                bool_ = True\n",
     "\n",
-    "                \n",
-    "                \n",
+    "\n",
     "bad_uids = list(tracker.problematic_uids)\n",
     "for uid in bad_uids:\n",
-    "    print(uid)      \n",
+    "    print(uid)\n",
     "\n",
     "\n",
-    "'''\n",
+    "\"\"\"\n",
     "\n",
     "st_uid = os.listdir(os.path.join(data_dir, file))[0]\n",
     "print(st_uid)\n",
@@ -675,7 +666,7 @@
     "print(seqs)\n",
     "seq = os.path.join(data_dir, file, st_uid, seqs[1])\n",
     "len(os.listdir(seq))\n",
-    "'''"
+    "\"\"\""
    ]
   },
   {
@@ -725,8 +716,10 @@
     }
    ],
    "source": [
-    "data_dir = '/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/'\n",
-    "meta_folder = \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/manifest-1777460989222/metadata\"\n",
+    "data_dir = \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/\"\n",
+    "meta_folder = (\n",
+    "    \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/manifest-1777460989222/metadata\"\n",
+    ")\n",
     "t2_dir = os.path.join(data_dir, \"t2\")\n",
     "adc_dir = os.path.join(data_dir, \"adc\")\n",
     "dwi_dir = os.path.join(data_dir, \"dwi\")\n",
@@ -742,8 +735,10 @@
    "source": [
     "st_id = os.listdir(t2_dir)\n",
     "\n",
-    "df_parent = pd.read_excel(os.path.join(meta_folder, \"Prostate-MRI-US-Biopsy-NBIA-manifest_v2_20231020-nbia-digest.xlsx\"))\n",
-    "df = pd.read_excel(os.path.join(meta_folder, \"TCIA-Biopsy-Data_2020-07-14.xlsx\"))\n"
+    "df_parent = pd.read_excel(\n",
+    "    os.path.join(meta_folder, \"Prostate-MRI-US-Biopsy-NBIA-manifest_v2_20231020-nbia-digest.xlsx\")\n",
+    ")\n",
+    "df = pd.read_excel(os.path.join(meta_folder, \"TCIA-Biopsy-Data_2020-07-14.xlsx\"))"
    ]
   },
   {
@@ -767,7 +762,6 @@
     "exclude = []\n",
     "test_list = []\n",
     "for file in st_id:\n",
-    "    \n",
     "    id = file.split(\".nrrd\")[0]\n",
     "    filtered_parent = df_parent[df_parent[\"Study Instance UID\"] == id]\n",
     "    patient_ids = filtered_parent[\"Patient ID\"].unique()\n",
@@ -775,23 +769,51 @@
     "        print(f\"Warning: Multiple patient IDs found for Study Instance UID {id}: {patient_ids}\")\n",
     "    else:\n",
     "        patient_id = patient_ids[0]\n",
-    "        \n",
-    "    t2_id = filtered_parent[filtered_parent[\"Series Description\"].str.contains(\"t2\", case=False, na=False)][\"Series Instance UID\"].iloc[0]\n",
+    "\n",
+    "    t2_id = filtered_parent[\n",
+    "        filtered_parent[\"Series Description\"].str.contains(\"t2\", case=False, na=False)\n",
+    "    ][\"Series Instance UID\"].iloc[0]\n",
     "    filtered_df = df[df[\"Series Instance UID (MRI)\"] == t2_id]\n",
-    "    assert filtered_df[\"Patient Number\"].unique()[0] == patient_id, f\"Mismatch: Patient ID from parent manifest ({patient_id}) does not match Patient Number in biopsy data ({filtered_df['Patient Number'].unique()[0]}) for Study Instance UID {id}\"\n",
+    "    assert filtered_df[\"Patient Number\"].unique()[0] == patient_id, (\n",
+    "        f\"Mismatch: Patient ID from parent manifest ({patient_id}) does not match Patient Number in biopsy data ({filtered_df['Patient Number'].unique()[0]}) for Study Instance UID {id}\"\n",
+    "    )\n",
     "    if filtered_df[\"Series Instance UID (US)\"].unique().shape[0] == 1:\n",
-    "        assert filtered_df[\"PSA (ng/mL)\"].unique().shape[0] == 1, f\"Expected 1 unique PSA value for Study Instance UID {id}, but found {filtered_df['PSA (ng/mL)'].unique().shape[0]}\"\n",
-    "        assert filtered_df[\"Prostate Volume (CC)\"].unique().shape[0] == 1, f\"Expected 1 unique Volume value for Study Instance UID {id}, but found {filtered_df['Volume (cc)'].unique().shape[0]}\"\n",
-    "        gg = (filtered_df['Primary Gleason'].fillna(0) + filtered_df['Secondary Gleason'].fillna(0)).max()\n",
+    "        assert filtered_df[\"PSA (ng/mL)\"].unique().shape[0] == 1, (\n",
+    "            f\"Expected 1 unique PSA value for Study Instance UID {id}, but found {filtered_df['PSA (ng/mL)'].unique().shape[0]}\"\n",
+    "        )\n",
+    "        assert filtered_df[\"Prostate Volume (CC)\"].unique().shape[0] == 1, (\n",
+    "            f\"Expected 1 unique Volume value for Study Instance UID {id}, but found {filtered_df['Volume (cc)'].unique().shape[0]}\"\n",
+    "        )\n",
+    "        gg = (\n",
+    "            filtered_df[\"Primary Gleason\"].fillna(0) + filtered_df[\"Secondary Gleason\"].fillna(0)\n",
+    "        ).max()\n",
     "        temp = {}\n",
     "        temp[\"image\"] = file\n",
-    "        temp[\"psa\"] = [filtered_df[\"PSA (ng/mL)\"].unique()[0], filtered_df[\"Prostate Volume (CC)\"].unique()[0] ]\n",
-    "        temp[\"dwi\"] = os.path.join(\"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed/DWI_registered\", file)\n",
-    "        temp[\"adc\"] = os.path.join(\"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed/ADC_clipped\", file)\n",
-    "        temp[\"heatmap\"] = os.path.join(\"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed/heatmaps\", file)\n",
-    "        temp[\"mask\"] = os.path.join(\"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed/prostate_mask\", file)\n",
-    "        temp[\"smooth_mask\"] = os.path.join(\"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed/smooth_prostate_mask\", file)\n",
-    "        temp[\"label\"] = 1.0 if gg>=7 else 0.0\n",
+    "        temp[\"psa\"] = [\n",
+    "            filtered_df[\"PSA (ng/mL)\"].unique()[0],\n",
+    "            filtered_df[\"Prostate Volume (CC)\"].unique()[0],\n",
+    "        ]\n",
+    "        temp[\"dwi\"] = os.path.join(\n",
+    "            \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed/DWI_registered\",\n",
+    "            file,\n",
+    "        )\n",
+    "        temp[\"adc\"] = os.path.join(\n",
+    "            \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed/ADC_clipped\",\n",
+    "            file,\n",
+    "        )\n",
+    "        temp[\"heatmap\"] = os.path.join(\n",
+    "            \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed/heatmaps\",\n",
+    "            file,\n",
+    "        )\n",
+    "        temp[\"mask\"] = os.path.join(\n",
+    "            \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed/prostate_mask\",\n",
+    "            file,\n",
+    "        )\n",
+    "        temp[\"smooth_mask\"] = os.path.join(\n",
+    "            \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed/smooth_prostate_mask\",\n",
+    "            file,\n",
+    "        )\n",
+    "        temp[\"label\"] = 1.0 if gg >= 7 else 0.0\n",
     "        test_list.append(temp)\n",
     "    else:\n",
     "        exclude.append(id)\n",
@@ -849,8 +871,8 @@
     }
    ],
    "source": [
-    "labs = [i['label'] for  i in test_list ]\n",
-    "np.unique(np.array(labs), return_counts=True)\n"
+    "labs = [i[\"label\"] for i in test_list]\n",
+    "np.unique(np.array(labs), return_counts=True)"
    ]
   },
   {

temp copy.ipynb

@@ -15,26 +15,25 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import argparse\n",
+    "import json\n",
     "import logging\n",
     "import os\n",
     "import shutil\n",
     "import sys\n",
     "from pathlib import Path\n",
-    "import json\n",
+    "from types import SimpleNamespace\n",
     "\n",
     "import torch\n",
     "import yaml\n",
     "from monai.utils import set_determinism\n",
     "from sklearn.preprocessing import StandardScaler\n",
+    "from tqdm import tqdm\n",
     "\n",
     "from src.data.data_loader import get_dataloader\n",
     "from src.model.cspca_model import CSPCAModel\n",
     "from src.model.mil import MILModel3D\n",
-    "from src.train.train_cspca import train_epoch, val_epoch\n",
-    "from src.utils import get_metrics, save_cspca_checkpoint, setup_logging\n",
-    "from types import SimpleNamespace\n",
-    "from tqdm import tqdm"
+    "from src.train.train_cspca import val_epoch\n",
+    "from src.utils import setup_logging"
    ]
   },
   {
@@ -52,28 +51,28 @@
     }
    ],
    "source": [
-    "with open(\"config/config_cspca_test.yaml\", \"r\") as f:\n",
+    "with open(\"config/config_cspca_test.yaml\") as f:\n",
     "    config = yaml.safe_load(f)\n",
     "\n",
     "args = SimpleNamespace(**config)\n",
     "print(args.data_root)\n",
     "\n",
-    "args.mode = 'test'\n",
+    "args.mode = \"test\"\n",
     "args.project_dir = None\n",
     "args.project_dir = Path.cwd()\n",
-    "args.run_name = 'test_dummy'\n",
+    "args.run_name = \"test_dummy\"\n",
     "args.checkpoint_pirads = None\n",
     "\n",
     "scaler = StandardScaler()\n",
     "with open(os.path.join(args.project_dir, \"dataset\", \"PICAI_cspca_updated_with_psa.json\")) as f:\n",
     "    dataset_json = json.load(f)\n",
-    "train_clinical = [i['psa'] for i in dataset_json['train']]\n",
+    "train_clinical = [i[\"psa\"] for i in dataset_json[\"train\"]]\n",
     "_ = scaler.fit_transform(train_clinical)\n",
     "args.psa_mean = scaler.mean_.tolist()\n",
     "args.psa_std = scaler.scale_.tolist()\n",
     "args.use_psa = True\n",
     "args.batch_size = 1\n",
-    "args.dataset_json = '/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/TCIA_test_data_updated_mask.json'\n"
+    "args.dataset_json = \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/TCIA_test_data_updated_mask.json\""
    ]
   },
   {
@@ -83,7 +82,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n",
     "args.logdir = os.path.join(args.project_dir, \"logs\", args.run_name)\n",
     "os.makedirs(args.logdir, exist_ok=True)\n",
     "args.logfile = os.path.join(args.logdir, f\"{args.run_name}.log\")\n",
@@ -136,7 +134,7 @@
     "        f\"csPCa Model loaded from {args.checkpoint_cspca} with AUC: {auc}, Sensitivity: {sens}, Specificity: {spec} on the test set.\"\n",
     "    )\n",
     "else:\n",
-    "    logging.info(f\"csPCa Model loaded from {args.checkpoint_cspca}.\")\n"
+    "    logging.info(f\"csPCa Model loaded from {args.checkpoint_cspca}.\")"
    ]
   },
   {
@@ -169,7 +167,7 @@
     "metrics_dict[\"specificity\"].append(test_metric[\"specificity\"])\n",
     "\n",
     "\n",
-    "import shutil\n",
+    "\n",
     "shutil.rmtree(cache_dir_path)\n",
     "\n",
     "metrics_dict"
@@ -206,19 +204,18 @@
     }
    ],
    "source": [
+    "import numpy as np\n",
     "import torch\n",
     "import torch.nn as nn\n",
     "from monai.metrics import Cumulative, CumulativeAverage\n",
     "from sklearn.metrics import confusion_matrix, roc_auc_score\n",
     "\n",
-    "from monai.utils import set_determinism\n",
-    "from sklearn.metrics import roc_curve\n",
-    "import numpy as np\n",
     "spec_list = []\n",
     "\n",
     "auc_list = []\n",
     "import os\n",
-    "os.environ['CUDA_LAUNCH_BLOCKING'] = \"1\"\n",
+    "\n",
+    "os.environ[\"CUDA_LAUNCH_BLOCKING\"] = \"1\"\n",
     "\n",
     "for st in list(range(10)):\n",
     "    set_determinism(seed=st)\n",
@@ -236,7 +233,7 @@
     "            target = batch_data[\"label\"].as_subclass(torch.Tensor).to(args.device)\n",
     "            psa_data = batch_data[\"psa\"].as_subclass(torch.Tensor).to(args.device)\n",
     "\n",
-    "            output = cspca_model(x = data, psa_data = psa_data)\n",
+    "            output = cspca_model(x=data, psa_data=psa_data)\n",
     "            output = output.squeeze(1)\n",
     "            loss = criterion(output, target)\n",
     "\n",
@@ -247,9 +244,9 @@
     "    loss_epoch = run_loss.aggregate()\n",
     "    target_list = targets_cumulative.get_buffer().cpu().numpy()\n",
     "    pred_list = preds_cumulative.get_buffer().cpu().numpy()\n",
-    "    #auc_epoch = roc_auc_score(target_list, pred_list)\n",
+    "    # auc_epoch = roc_auc_score(target_list, pred_list)\n",
     "    y_pred_categoric = pred_list >= 0.5\n",
-    "    '''\n",
+    "    \"\"\"\n",
     "    tn, fp, fn, tp = confusion_matrix(target_list, y_pred_categoric).ravel()\n",
     "    sens_epoch = tp / (tp + fn)\n",
     "    spec_epoch = tn / (tn + fp)\n",
@@ -261,7 +258,7 @@
     "        \"specificity\": spec_epoch,\n",
     "    }\n",
     "    return val_epoch_metric\n",
-    "    '''\n",
+    "    \"\"\"\n",
     "    shutil.rmtree(cache_dir_path)\n",
     "    auc_epoch = roc_auc_score(target_list, pred_list)\n",
     "    auc_list.append(auc_epoch)"
@@ -322,16 +319,17 @@
     }
    ],
    "source": [
-    "import numpy as np\n",
     "from scipy import stats\n",
-    "'''\n",
+    "\n",
+    "\"\"\"\n",
     "# Example values from 10 seeds\n",
     "auc_list = [0.81, 0.84, 0.79, 0.83, 0.82,\n",
     "            0.85, 0.80, 0.84, 0.83, 0.81]\n",
     "\n",
     "spec_list = [0.72, 0.75, 0.70, 0.74, 0.73,\n",
     "             0.76, 0.71, 0.75, 0.74, 0.72]\n",
-    "'''\n",
+    "\"\"\"\n",
+    "\n",
     "\n",
     "def mean_confidence_interval(values, confidence=0.95):\n",
     "    values = np.array(values)\n",
@@ -339,25 +337,19 @@
     "    mean = np.mean(values)\n",
     "    sem = stats.sem(values)  # standard error\n",
     "\n",
-    "    ci = stats.t.interval(\n",
-    "        confidence,\n",
-    "        df=len(values) - 1,\n",
-    "        loc=mean,\n",
-    "        scale=sem\n",
-    "    )\n",
+    "    ci = stats.t.interval(confidence, df=len(values) - 1, loc=mean, scale=sem)\n",
     "\n",
     "    return mean, ci\n",
     "\n",
     "\n",
     "auc_mean, auc_ci = mean_confidence_interval(auc_list)\n",
-    "#spec_mean, spec_ci = mean_confidence_interval(spec_list)\n",
+    "# spec_mean, spec_ci = mean_confidence_interval(spec_list)\n",
     "\n",
-    "print(f\"AUC: {auc_mean:.3f} \"\n",
-    "      f\"(95% CI: {auc_ci[0]:.3f} - {auc_ci[1]:.3f})\")\n",
-    "'''\n",
+    "print(f\"AUC: {auc_mean:.3f} (95% CI: {auc_ci[0]:.3f} - {auc_ci[1]:.3f})\")\n",
+    "\"\"\"\n",
     "print(f\"Specificity: {spec_mean:.3f} \"\n",
     "      f\"(95% CI: {spec_ci[0]:.3f} - {spec_ci[1]:.3f})\")\n",
-    "'''"
+    "\"\"\""
    ]
   },
   {
@@ -663,7 +655,7 @@
    ],
    "source": [
     "for t, p in zip(target_list, pred_list):\n",
-    "   print(f\"target: {t} | pred: {p:.4f}\")"
+    "    print(f\"target: {t} | pred: {p:.4f}\")"
    ]
   },
   {
@@ -703,14 +695,13 @@
     "pred_cat = [1.0 if i > 0.5 else 0.0 for i in pred_list]\n",
     "from sklearn.metrics import (\n",
     "    accuracy_score,\n",
+    "    balanced_accuracy_score,\n",
+    "    classification_report,\n",
+    "    f1_score,\n",
     "    precision_score,\n",
     "    recall_score,\n",
-    "    f1_score,\n",
-    "    balanced_accuracy_score,\n",
-    "    confusion_matrix,\n",
-    "    classification_report\n",
     ")\n",
-    "import numpy as np\n",
+    "\n",
     "pred_cat = np.array(pred_cat)\n",
     "# Basic metrics\n",
     "acc = accuracy_score(target_list, pred_cat)\n",
@@ -834,11 +825,11 @@
     }
    ],
    "source": [
-    "updated_masks = os.listdir('updated_segmentations/')\n",
+    "updated_masks = os.listdir(\"updated_segmentations/\")\n",
     "c = 0\n",
     "for i in false_negative_ids[0]:\n",
-    "    if test_data['test'][i]['image'] in updated_masks:\n",
-    "        print(i)\n"
+    "    if test_data[\"test\"][i][\"image\"] in updated_masks:\n",
+    "        print(i)"
    ]
   },
   {
@@ -896,10 +887,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n",
     "import json\n",
+    "\n",
     "import nrrd\n",
-    "with open('/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/TCIA_test_data.json', 'r') as f:\n",
+    "\n",
+    "with open(\n",
+    "    \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/TCIA_test_data.json\",\n",
+    ") as f:\n",
     "    test_data = json.load(f)"
    ]
   },
@@ -936,18 +930,19 @@
     }
    ],
    "source": [
-    "\n",
-    "t2_dir_proc = '/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed/t2_registered'\n",
+    "t2_dir_proc = (\n",
+    "    \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed/t2_registered\"\n",
+    ")\n",
     "\n",
     "id = 195\n",
     "\n",
     "\n",
-    "file = test_data['test'][id]['image']\n",
+    "file = test_data[\"test\"][id][\"image\"]\n",
     "\n",
-    "t2_proc, _ = nrrd.read(os.path.join(t2_dir_proc,file))\n",
-    "dwi_proc, _ = nrrd.read(test_data['test'][id]['dwi'])\n",
-    "adc_proc, _ = nrrd.read(test_data['test'][id]['adc'])#\n",
-    "heats, _ = nrrd.read(test_data['test'][id]['heatmap'])\n",
+    "t2_proc, _ = nrrd.read(os.path.join(t2_dir_proc, file))\n",
+    "dwi_proc, _ = nrrd.read(test_data[\"test\"][id][\"dwi\"])\n",
+    "adc_proc, _ = nrrd.read(test_data[\"test\"][id][\"adc\"])  #\n",
+    "heats, _ = nrrd.read(test_data[\"test\"][id][\"heatmap\"])\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "\n",
@@ -969,12 +964,12 @@
     "else:\n",
     "    # 3. Create the dynamic grid based on how many valid slices were found\n",
     "    fig, axes = plt.subplots(nrows=num_slices, ncols=3, figsize=(15, 5 * num_slices))\n",
-    "    \n",
-    "    # If there is only exactly 1 slice with a heatmap, axes is 1D. \n",
+    "\n",
+    "    # If there is only exactly 1 slice with a heatmap, axes is 1D.\n",
     "    # We force it to 2D so our loop doesn't break.\n",
     "    if num_slices == 1:\n",
     "        axes = np.expand_dims(axes, axis=0)\n",
-    "    \n",
+    "\n",
     "    # 4. Plot the columns for the filtered slices\n",
     "    for i, z in enumerate(valid_slices):\n",
     "        # Extract the 2D slices\n",
@@ -982,30 +977,30 @@
     "        adc_slice = adc_proc[:, :, z].T\n",
     "        dwi_slice = dwi_proc[:, :, z].T\n",
     "        heat_slice = heats[:, :, z].T\n",
-    "        \n",
+    "\n",
     "        # Mask the background zeroes of the heatmap so it doesn't tint the healthy tissue\n",
     "        masked_heat = np.ma.masked_where(heat_slice <= 0.01, heat_slice)\n",
-    "    \n",
+    "\n",
     "        # --- Column 1: T2 + Heatmap ---\n",
-    "        axes[i, 0].imshow(t2_slice, cmap='gray')\n",
-    "        axes[i, 0].imshow(masked_heat, cmap='jet', alpha=0.5)\n",
-    "        axes[i, 0].set_title(f'T2 + Heatmap (Slice {z})', fontsize=14)\n",
-    "        axes[i, 0].axis('off')\n",
-    "    \n",
+    "        axes[i, 0].imshow(t2_slice, cmap=\"gray\")\n",
+    "        axes[i, 0].imshow(masked_heat, cmap=\"jet\", alpha=0.5)\n",
+    "        axes[i, 0].set_title(f\"T2 + Heatmap (Slice {z})\", fontsize=14)\n",
+    "        axes[i, 0].axis(\"off\")\n",
+    "\n",
     "        # --- Column 2: ADC + Heatmap ---\n",
-    "        axes[i, 1].imshow(adc_slice, cmap='gray')\n",
-    "        axes[i, 1].imshow(masked_heat, cmap='jet', alpha=0.5)\n",
-    "        axes[i, 1].set_title(f'ADC + Heatmap (Slice {z})', fontsize=14)\n",
-    "        axes[i, 1].axis('off')\n",
-    "    \n",
+    "        axes[i, 1].imshow(adc_slice, cmap=\"gray\")\n",
+    "        axes[i, 1].imshow(masked_heat, cmap=\"jet\", alpha=0.5)\n",
+    "        axes[i, 1].set_title(f\"ADC + Heatmap (Slice {z})\", fontsize=14)\n",
+    "        axes[i, 1].axis(\"off\")\n",
+    "\n",
     "        # --- Column 3: DWI (No Heatmap) ---\n",
-    "        axes[i, 2].imshow(dwi_slice, cmap='gray')\n",
-    "        axes[i, 2].set_title(f'DWI (Slice {z})', fontsize=14)\n",
-    "        axes[i, 2].axis('off')\n",
-    "    \n",
+    "        axes[i, 2].imshow(dwi_slice, cmap=\"gray\")\n",
+    "        axes[i, 2].set_title(f\"DWI (Slice {z})\", fontsize=14)\n",
+    "        axes[i, 2].axis(\"off\")\n",
+    "\n",
     "    plt.tight_layout()\n",
     "    plt.show()\n",
-    "print(file)\n"
+    "print(file)"
    ]
   },
   {
@@ -1017,10 +1012,10 @@
    "source": [
     "import shutil\n",
     "\n",
-    "in_folder = '/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed'\n",
-    "out_folder = '/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/keno_files'\n",
-    "for i in ['t2_registered', 'ADC_clipped', 'DWI_registered', 'smooth_prostate_mask']:\n",
-    "    shutil.copy(os.path.join(in_folder, i, file), os.path.join(out_folder, i.split('_')[0], file))"
+    "in_folder = \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/nrrd_files/processed\"\n",
+    "out_folder = \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/TCIA_prostate/keno_files\"\n",
+    "for i in [\"t2_registered\", \"ADC_clipped\", \"DWI_registered\", \"smooth_prostate_mask\"]:\n",
+    "    shutil.copy(os.path.join(in_folder, i, file), os.path.join(out_folder, i.split(\"_\")[0], file))"
    ]
   },
   {
@@ -1038,26 +1033,24 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import argparse\n",
     "import logging\n",
     "import os\n",
     "import shutil\n",
     "import sys\n",
     "import time\n",
     "from pathlib import Path\n",
+    "from types import SimpleNamespace\n",
     "\n",
     "import numpy as np\n",
     "import torch\n",
     "import wandb\n",
     "import yaml\n",
     "from monai.utils import set_determinism\n",
-    "from torch.utils.tensorboard import SummaryWriter\n",
     "\n",
     "from src.data.data_loader import get_dataloader\n",
     "from src.model.mil import MILModel3D\n",
-    "from src.train.train_pirads import train_epoch, val_epoch\n",
-    "from src.utils import save_pirads_checkpoint, setup_logging\n",
-    "from types import SimpleNamespace\n"
+    "from src.train.train_pirads import val_epoch\n",
+    "from src.utils import setup_logging"
    ]
   },
   {
@@ -1069,22 +1062,22 @@
    "source": [
     "from types import SimpleNamespace\n",
     "\n",
-    "with open(\"config/config_pirads_test.yaml\", \"r\") as f:\n",
+    "with open(\"config/config_pirads_test.yaml\") as f:\n",
     "    config = yaml.safe_load(f)\n",
     "\n",
     "\n",
     "args = SimpleNamespace(**config)\n",
     "print(args.data_root)\n",
     "\n",
-    "args.mode = 'test'\n",
+    "args.mode = \"test\"\n",
     "args.project_dir = Path.cwd()\n",
-    "args.run_name = 'test_dummy_pirads'\n",
+    "args.run_name = \"test_dummy_pirads\"\n",
     "args.checkpoint_pirads = None\n",
     "args.wandb = False\n",
     "args.optim_lr = 2e-5\n",
     "args.batch_size = 8\n",
     "args.epochs = 1\n",
-    "args.project_name = 'Dummy'\n",
+    "args.project_name = \"Dummy\"\n",
     "args.workers = 0\n",
     "args.use_psa = False"
    ]
@@ -1096,8 +1089,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n",
-    "\n",
     "args.logdir = os.path.join(args.project_dir, \"logs\", args.run_name)\n",
     "os.makedirs(args.logdir, exist_ok=True)\n",
     "args.logfile = os.path.join(args.logdir, f\"{args.run_name}.log\")\n",
@@ -1133,7 +1124,7 @@
     "    dir=os.path.join(args.logdir, \"wandb\"),\n",
     "    config=config_wandb,\n",
     "    mode=mode_wandb,\n",
-    ")\n"
+    ")"
    ]
   },
   {
@@ -1197,15 +1188,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import argparse\n",
     "import logging\n",
-    "import time\n",
-    "from tqdm import tqdm\n",
+    "\n",
     "import numpy as np\n",
     "import torch\n",
     "import torch.nn as nn\n",
     "from monai.metrics import Cumulative, CumulativeAverage\n",
     "from sklearn.metrics import cohen_kappa_score\n",
+    "from tqdm import tqdm\n",
+    "\n",
     "valid_loader = get_dataloader(args, split=args.mode)\n",
     "\n",
     "criterion = nn.CrossEntropyLoss()\n",
@@ -1230,7 +1221,6 @@
     "        with torch.amp.autocast(device_type=str(args.device), enabled=args.amp):\n",
     "            logits = model(data)\n",
     "\n",
-    "\n",
     "        data = data.to(\"cpu\")\n",
     "        target = target.to(\"cpu\")\n",
     "        logits = logits.to(\"cpu\")\n",
@@ -1242,7 +1232,6 @@
     "        targets_cumulative.extend(target.detach().cpu())\n",
     "        targets_cumulative_bin.extend(target_bin.detach().cpu())\n",
     "\n",
-    "\n",
     "        del data, target, logits\n",
     "        torch.cuda.empty_cache()\n",
     "\n",
@@ -1253,7 +1242,7 @@
     "preds_updated = preds_cumulative + 2\n",
     "\n",
     "if os.path.exists(cache_dir_):\n",
-    "    shutil.rmtree(cache_dir_)\n"
+    "    shutil.rmtree(cache_dir_)"
    ]
   },
   {
@@ -1263,9 +1252,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "qwk = cohen_kappa_score(\n",
-    "    targets_cumulative.astype(np.float64), preds_updated.astype(np.float64)\n",
-    ")\n",
+    "qwk = cohen_kappa_score(targets_cumulative.astype(np.float64), preds_updated.astype(np.float64))\n",
     "qwk"
    ]
   },
@@ -1288,10 +1275,7 @@
    "source": [
     "import pandas as pd\n",
     "\n",
-    "df = pd.DataFrame({\n",
-    "    \"target\": targets_cumulative,\n",
-    "    \"pred\": preds_updated\n",
-    "})\n",
+    "df = pd.DataFrame({\"target\": targets_cumulative, \"pred\": preds_updated})\n",
     "\n",
     "# Group by class and show summary stats\n",
     "\n",
@@ -1309,6 +1293,7 @@
    "outputs": [],
    "source": [
     "from sklearn.metrics import roc_auc_score\n",
+    "\n",
     "roc_auc_score(targets_cumulative_bin, preds_cumulative)"
    ]
   },
@@ -1321,28 +1306,23 @@
    "source": [
     "import pandas as pd\n",
     "\n",
-    "df = pd.DataFrame({\n",
-    "    'pirads': preds_updated,\n",
-    "    'cspca': targets_cumulative_bin\n",
-    "})\n",
-    "import seaborn as sns\n",
+    "df = pd.DataFrame({\"pirads\": preds_updated, \"cspca\": targets_cumulative_bin})\n",
     "import matplotlib.pyplot as plt\n",
-    "\n",
     "import pandas as pd\n",
+    "import seaborn as sns\n",
     "\n",
-    "ct = pd.crosstab(df['pirads'], df['cspca'])\n",
+    "ct = pd.crosstab(df[\"pirads\"], df[\"cspca\"])\n",
     "\n",
     "# Add totals\n",
     "ct_with_totals = ct.copy()\n",
-    "ct_with_totals['Total'] = ct_with_totals.sum(axis=1)   # row totals\n",
-    "ct_with_totals.loc['Total'] = ct_with_totals.sum(axis=0)  # column totals\n",
-    "import pandas as pd\n",
-    "import seaborn as sns\n",
+    "ct_with_totals[\"Total\"] = ct_with_totals.sum(axis=1)  # row totals\n",
+    "ct_with_totals.loc[\"Total\"] = ct_with_totals.sum(axis=0)  # column totals\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
+    "import pandas as pd\n",
     "\n",
     "# Base crosstab\n",
-    "ct = pd.crosstab(df['pirads'], df['cspca'])\n",
+    "ct = pd.crosstab(df[\"pirads\"], df[\"cspca\"])\n",
     "\n",
     "# Row totals\n",
     "row_totals = ct.sum(axis=1)\n",
@@ -1351,50 +1331,43 @@
     "col_totals = ct.sum(axis=0)\n",
     "\n",
     "# ---- Plot ----\n",
-    "fig, ax = plt.subplots(figsize=(7,5))\n",
+    "fig, ax = plt.subplots(figsize=(7, 5))\n",
     "\n",
     "# Heatmap ONLY for actual data\n",
-    "sns.heatmap(\n",
-    "    ct,\n",
-    "    annot=True,\n",
-    "    fmt='d',\n",
-    "    cmap='Blues',\n",
-    "    cbar=False,\n",
-    "    ax=ax\n",
-    ")\n",
+    "sns.heatmap(ct, annot=True, fmt=\"d\", cmap=\"Blues\", cbar=False, ax=ax)\n",
     "\n",
     "# ---- Add row totals (right side, no color) ----\n",
     "for i, val in enumerate(row_totals):\n",
     "    ax.text(\n",
-    "        ct.shape[1] + 0.3,   # position to the right\n",
+    "        ct.shape[1] + 0.3,  # position to the right\n",
     "        i + 0.5,\n",
     "        str(val),\n",
-    "        va='center'\n",
+    "        va=\"center\",\n",
     "    )\n",
     "\n",
     "# Label for row totals\n",
-    "ax.text(ct.shape[1] + 0.3, -0.3, 'Total', ha='center')\n",
+    "ax.text(ct.shape[1] + 0.3, -0.3, \"Total\", ha=\"center\")\n",
     "\n",
     "# ---- Add column totals (below x-axis) ----\n",
     "for j, val in enumerate(col_totals):\n",
     "    ax.text(\n",
     "        j + 0.5,\n",
-    "        ct.shape[0] + 0.3,   # position below\n",
+    "        ct.shape[0] + 0.3,  # position below\n",
     "        str(val),\n",
-    "        ha='center'\n",
+    "        ha=\"center\",\n",
     "    )\n",
     "\n",
     "# Label for column totals\n",
-    "ax.text(-0.5, ct.shape[0] + 0.3, 'Total', va='center')\n",
+    "ax.text(-0.5, ct.shape[0] + 0.3, \"Total\", va=\"center\")\n",
     "\n",
     "# ---- Adjust limits so text is visible ----\n",
     "ax.set_xlim(0, ct.shape[1] + 1)\n",
     "ax.set_ylim(ct.shape[0] + 1, 0)\n",
     "\n",
     "# Titles and labels\n",
-    "ax.set_title('PI-RADS vs csPCa')\n",
-    "ax.set_xlabel('csPCa')\n",
-    "ax.set_ylabel('PI-RADS')\n",
+    "ax.set_title(\"PI-RADS vs csPCa\")\n",
+    "ax.set_xlabel(\"csPCa\")\n",
+    "ax.set_ylabel(\"PI-RADS\")\n",
     "\n",
     "plt.tight_layout()\n",
     "plt.show()"
@@ -1407,18 +1380,19 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "pirads_pred = [1.0 if i >=3 else 0.0 for i in preds_updated]\n",
+    "pirads_pred = [1.0 if i >= 3 else 0.0 for i in preds_updated]\n",
     "pirads_pred\n",
+    "import numpy as np\n",
     "from sklearn.metrics import (\n",
     "    accuracy_score,\n",
-    "    precision_score,\n",
-    "    recall_score,\n",
-    "    f1_score,\n",
     "    balanced_accuracy_score,\n",
+    "    classification_report,\n",
     "    confusion_matrix,\n",
-    "    classification_report\n",
+    "    f1_score,\n",
+    "    precision_score,\n",
+    "    recall_score,\n",
     ")\n",
-    "import numpy as np\n",
+    "\n",
     "pred_cat = np.array(pirads_pred)\n",
     "target_list = targets_cumulative_bin\n",
     "# Basic metrics\n",
@@ -1466,7 +1440,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import argparse\n",
+    "import json\n",
     "import logging\n",
     "import os\n",
     "import shutil\n",
@@ -1479,14 +1453,12 @@
     "import wandb\n",
     "import yaml\n",
     "from monai.utils import set_determinism\n",
-    "from torch.utils.tensorboard import SummaryWriter\n",
     "from sklearn.preprocessing import StandardScaler\n",
-    "import json\n",
     "\n",
     "from src.data.data_loader import get_dataloader\n",
     "from src.model.mil import MILModel3D\n",
-    "from src.train.train_pirads import train_epoch, val_epoch\n",
-    "from src.utils import save_pirads_checkpoint, setup_logging"
+    "from src.train.train_pirads import val_epoch\n",
+    "from src.utils import setup_logging"
    ]
   },
   {
@@ -1506,22 +1478,22 @@
    "source": [
     "from types import SimpleNamespace\n",
     "\n",
-    "with open(\"config/config_pirads_train.yaml\", \"r\") as f:\n",
+    "with open(\"config/config_pirads_train.yaml\") as f:\n",
     "    config = yaml.safe_load(f)\n",
     "\n",
     "\n",
     "args = SimpleNamespace(**config)\n",
     "print(args.data_root)\n",
     "\n",
-    "args.mode = 'rain'\n",
+    "args.mode = \"rain\"\n",
     "args.project_dir = Path.cwd()\n",
-    "args.run_name = 'test_dummy_pirads'\n",
+    "args.run_name = \"test_dummy_pirads\"\n",
     "args.checkpoint = None\n",
     "args.wandb = False\n",
     "args.optim_lr = 2e-5\n",
     "args.batch_size = 8\n",
     "args.epochs = 1\n",
-    "args.project_name = 'Dummy'\n",
+    "args.project_name = \"Dummy\"\n",
     "args.workers = 0\n",
     "args.use_psa = False"
    ]
@@ -1542,8 +1514,6 @@
     }
    ],
    "source": [
-    "\n",
-    "\n",
     "args.logdir = os.path.join(args.project_dir, \"logs\", args.run_name)\n",
     "os.makedirs(args.logdir, exist_ok=True)\n",
     "args.logfile = os.path.join(args.logdir, f\"{args.run_name}.log\")\n",
@@ -1579,7 +1549,7 @@
     "    dir=os.path.join(args.logdir, \"wandb\"),\n",
     "    config=config_wandb,\n",
     "    mode=mode_wandb,\n",
-    ")\n"
+    ")"
    ]
   },
   {
@@ -1616,7 +1586,7 @@
     "scaler = StandardScaler()\n",
     "with open(os.path.join(args.project_dir, \"dataset\", \"PICAI_cspca_updated_with_psa.json\")) as f:\n",
     "    dataset_json = json.load(f)\n",
-    "train_clinical = [i['psa'] for i in dataset_json['train']]\n",
+    "train_clinical = [i[\"psa\"] for i in dataset_json[\"train\"]]\n",
     "_ = scaler.fit_transform(train_clinical)\n",
     "args.psa_mean = scaler.mean_.tolist()\n",
     "args.psa_std = scaler.scale_.tolist()"
@@ -1631,9 +1601,7 @@
    "source": [
     "train_loader = get_dataloader(args, split=\"train\")\n",
     "valid_loader = get_dataloader(args, split=\"test\")\n",
-    "logging.info(\n",
-    "    f\"Dataset training: {len(train_loader.dataset)}, test: {len(valid_loader.dataset)}\"\n",
-    ")\n",
+    "logging.info(f\"Dataset training: {len(train_loader.dataset)}, test: {len(valid_loader.dataset)}\")\n",
     "\n",
     "if args.mil_mode in [\"att_trans\", \"att_trans_pyramid\"]:\n",
     "    params = [\n",
@@ -1641,15 +1609,13 @@
     "            \"params\": list(model.attention.parameters())\n",
     "            + list(model.myfc.parameters())\n",
     "            + list(model.net.parameters())\n",
-    "            #+ list(model.feature_proj.parameters()),\n",
+    "            # + list(model.feature_proj.parameters()),\n",
     "        },\n",
     "        {\"params\": list(model.transformer.parameters()), \"lr\": 6e-5, \"weight_decay\": 0.1},\n",
     "    ]\n",
     "\n",
     "optimizer = torch.optim.AdamW(params, lr=args.optim_lr, weight_decay=args.weight_decay)\n",
-    "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n",
-    "    optimizer, T_max=args.epochs, eta_min=0\n",
-    ")\n",
+    "scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs, eta_min=0)\n",
     "scaler = torch.amp.GradScaler(device=str(args.device), enabled=args.amp)"
    ]
   },
@@ -1660,7 +1626,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import argparse\n",
     "import logging\n",
     "import time\n",
     "\n",
@@ -1694,10 +1659,10 @@
     "        # Classification Loss\n",
     "        logits_attn = model(shuffled_images, no_head=True)\n",
     "        x = logits_attn.to(torch.float32)\n",
-    "        #x = x.permute(1, 0, 2)\n",
-    "        #x = model.feature_proj(x)\n",
+    "        # x = x.permute(1, 0, 2)\n",
+    "        # x = model.feature_proj(x)\n",
     "        x = model.transformer(x)\n",
-    "        #x = x.permute(1, 0, 2)\n",
+    "        # x = x.permute(1, 0, 2)\n",
     "        a = model.attention(x)\n",
     "        a = torch.softmax(a, dim=1)\n",
     "        x = torch.sum(x * a, dim=1)\n",
@@ -1940,7 +1905,7 @@
     }
    ],
    "source": [
-    "(pred_pirads == target.sum(dim=1)).sum()/len(pred_pirads)"
+    "(pred_pirads == target.sum(dim=1)).sum() / len(pred_pirads)"
    ]
   },
   {
@@ -2000,11 +1965,250 @@
    "outputs": [],
    "source": []
   },
+  {
+   "cell_type": "markdown",
+   "id": "fb57c919",
+   "metadata": {},
+   "source": [
+    "### Run Inference Amend"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "id": "df58821b",
    "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "If you have questions or suggestions, feel free to open an issue at https://github.com/DIAGNijmegen/picai_prep\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "import argparse\n",
+    "import json\n",
+    "import logging\n",
+    "import os\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import streamlit as st\n",
+    "import torch\n",
+    "import yaml\n",
+    "from monai.data import Dataset\n",
+    "\n",
+    "from src.data.data_loader import data_transform, list_data_collate\n",
+    "from src.model.cspca_model import CSPCAModel\n",
+    "from src.model.mil import MILModel3D\n",
+    "from src.preprocessing.clip_intensity import clip_adc\n",
+    "from src.preprocessing.generate_heatmap import get_heatmap\n",
+    "from src.preprocessing.prostate_mask import get_segmask\n",
+    "from src.preprocessing.register_and_crop import register_files\n",
+    "from src.utils import get_parent_image, get_patch_coordinate, setup_logging, get_prostate_volume\n",
+    "from sklearn.preprocessing import StandardScaler\n",
+    "\n",
+    "from argparse import Namespace\n",
+    "from collections.abc import Callable"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "01f61063",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "SyntaxError",
+     "evalue": "invalid syntax (2184847560.py, line 1)",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;36m  Cell \u001b[0;32mIn[22], line 1\u001b[0;36m\u001b[0m\n\u001b[0;31m    mask_patches.shapefrom argparse import Namespace\u001b[0m\n\u001b[0m                           ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m invalid syntax\n"
+     ]
+    }
+   ],
+   "source": [
+    "from argparse import Namespace\n",
+    "\n",
+    "args = Namespace(\n",
+    "    config=\"config.yaml\",\n",
+    "    json_path=\"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/psa_val.json\", \n",
+    "    t2_dir=\"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/t2\",\n",
+    "    dwi_dir=\"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/dwi\",\n",
+    "    adc_dir=\"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/adc\",\n",
+    "    seg_dir=\"/path/to/seg\",\n",
+    "    output_dir=\"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed\",\n",
+    "    margin=0.2,\n",
+    "    num_classes=4,\n",
+    "    mil_mode=\"att_trans\",\n",
+    "    use_heatmap=True,\n",
+    "    use_psa=True,\n",
+    "    tile_size=48,\n",
+    "    tile_count=40,\n",
+    "    depth=3,\n",
+    "    project_dir=\"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "e0117759",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 1/1 [00:02<00:00,  2.79s/it]\n",
+      "You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "100%|██████████| 1/1 [00:07<00:00,  7.29s/it]\n",
+      "100%|██████████| 1/1 [00:00<00:00, 10.55it/s]\n",
+      "100%|██████████| 1/1 [00:00<00:00,  4.04it/s]\n"
+     ]
+    }
+   ],
+   "source": [
+    "FUNCTIONS: dict[str, Callable[[Namespace], Namespace]] = {\n",
+    "    \"register_and_crop\": register_files,\n",
+    "    \"clip_adc\": clip_adc,\n",
+    "    \"get_segmentation_mask\": get_segmask,\n",
+    "    \"get_heatmap\": get_heatmap,\n",
+    "}\n",
+    "args.logfile = os.path.join(args.output_dir, \"inference.log\")\n",
+    "setup_logging(args.logfile)\n",
+    "logging.info(\"Starting preprocessing\")\n",
+    "steps = [\"register_and_crop\", \"get_segmentation_mask\", \"clip_adc\", \"get_heatmap\"]\n",
+    "for step in steps:\n",
+    "    #if step == 'get_segmentation_mask':\n",
+    "    #    args.seg_dir = os.path.join(args.output_dir, \"prostate_mask\")\n",
+    "    #    continue\n",
+    "    func = FUNCTIONS[step]\n",
+    "    \n",
+    "    args = func(args)\n",
+    "\n",
+    "logging.info(\"Preprocessing completed.\")\n",
+    "args.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "29a500c4",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "config: config.yaml\n",
+      "json_path: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/psa_val.json\n",
+      "t2_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed/t2_registered\n",
+      "dwi_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed/DWI_registered\n",
+      "adc_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed/ADC_clipped\n",
+      "seg_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed/prostate_mask\n",
+      "output_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed\n",
+      "margin: 0.2\n",
+      "num_classes: 4\n",
+      "mil_mode: att_trans\n",
+      "use_heatmap: True\n",
+      "use_psa: True\n",
+      "tile_size: 48\n",
+      "tile_count: 40\n",
+      "depth: 3\n",
+      "project_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate\n",
+      "logfile: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed/inference.log\n",
+      "heatmapdir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed/heatmaps/\n",
+      "smooth_seg_dir_temp: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed/smooth_prostate_mask_temp/\n",
+      "smooth_seg_dir: /sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed/smooth_prostate_mask/\n",
+      "device: cpu\n"
+     ]
+    }
+   ],
+   "source": [
+    "for key, value in vars(args).items():\n",
+    "    print(f\"{key}: {value}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "cf6aff0e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "scaler = StandardScaler()\n",
+    "with open(os.path.join(args.project_dir, \"dataset\", \"PICAI_cspca_updated_with_psa.json\")) as f:\n",
+    "    dataset_json = json.load(f)\n",
+    "train_clinical = [i[\"psa\"] for i in dataset_json[\"train\"]]\n",
+    "_ = scaler.fit_transform(train_clinical)\n",
+    "args.psa_mean = scaler.mean_.tolist()\n",
+    "args.psa_std = scaler.scale_.tolist()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "ca8c1554",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "transform = data_transform(args, split='test')\n",
+    "files = os.listdir(args.t2_dir)\n",
+    "args.data_list = []\n",
+    "with open(args.json_path, 'r') as f:\n",
+    "    psa_data = json.load(f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "b4398c65",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for file in files:\n",
+    "    temp = {}\n",
+    "    temp[\"image\"] = os.path.join(args.t2_dir, file)\n",
+    "    temp[\"dwi\"] = os.path.join(args.dwi_dir, file)\n",
+    "    temp[\"adc\"] = os.path.join(args.adc_dir, file)\n",
+    "    temp[\"heatmap\"] = os.path.join(args.heatmapdir, file)\n",
+    "    temp[\"mask\"] = os.path.join(args.seg_dir, file)\n",
+    "    temp[\"label\"] = 0  # dummy label\n",
+    "    temp[\"smooth_mask\"] = os.path.join(args.smooth_seg_dir, file)\n",
+    "    prostate_vol = get_prostate_volume(temp[\"mask\"])\n",
+    "    temp[\"psa\"] = [psa_data[file.split('.nrrd')[0]] , prostate_vol]\n",
+    "    args.data_list.append(temp)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "d2b5fa61",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "dataset = Dataset(data=args.data_list, transform=transform)\n",
+    "loader = torch.utils.data.DataLoader(\n",
+    "    dataset,\n",
+    "    batch_size=1,\n",
+    "    shuffle=False,\n",
+    "    num_workers=0,\n",
+    "    pin_memory=True,\n",
+    "    multiprocessing_context=None,\n",
+    "    sampler=None,\n",
+    "    collate_fn=list_data_collate,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "824263f5",
+   "metadata": {},
    "outputs": [],
    "source": []
   }

temp_2.ipynb

@@ -7,10 +7,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import pandas as pd\n",
-    "import numpy as np\n",
     "import json\n",
-    "import os"
+    "import os\n",
+    "\n",
+    "import numpy as np\n",
+    "import pandas as pd"
    ]
   },
   {
@@ -20,10 +21,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open('/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/test_data_updated_with_psa.json', 'r') as f:\n",
+    "with open(\n",
+    "    \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/test_data_updated_with_psa.json\",\n",
+    ") as f:\n",
     "    data = json.load(f)\n",
-    "    \n",
-    "neg_label = [i for i in data['test'] if i['label'] == 0]"
+    "\n",
+    "neg_label = [i for i in data[\"test\"] if i[\"label\"] == 0]"
    ]
   },
   {
@@ -34,8 +37,8 @@
    "outputs": [],
    "source": [
     "print(len(neg_label))\n",
-    "pirads_label = np.array([i['pirads'] for i in neg_label])\n",
-    "print(np.unique(pirads_label, return_counts=True))\n"
+    "pirads_label = np.array([i[\"pirads\"] for i in neg_label])\n",
+    "print(np.unique(pirads_label, return_counts=True))"
    ]
   },
   {
@@ -53,7 +56,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "data = pd.read_csv(\"/sc-projects/sc-proj-cc06-ag-ki-radiologie/pirad_model_test_PICAI/marksheet.csv\")\n",
+    "data = pd.read_csv(\n",
+    "    \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/pirad_model_test_PICAI/marksheet.csv\"\n",
+    ")\n",
     "data.head()"
    ]
   },
@@ -64,40 +69,45 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open('/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/PICAI_cspca_updated.json', 'r') as f:\n",
+    "with open(\n",
+    "    \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/PICAI_cspca_updated.json\",\n",
+    ") as f:\n",
     "    json_data = json.load(f)\n",
     "\n",
-    "new_train =[]    \n",
-    "for i in json_data['train']:\n",
-    "    p_id = i['image'].split('_')[0]\n",
-    "    st_id = i['image'].split('_')[1].split('.')[0]\n",
+    "new_train = []\n",
+    "for i in json_data[\"train\"]:\n",
+    "    p_id = i[\"image\"].split(\"_\")[0]\n",
+    "    st_id = i[\"image\"].split(\"_\")[1].split(\".\")[0]\n",
     "    row = data[(data[\"patient_id\"] == int(p_id)) & (data[\"study_id\"] == int(st_id))]\n",
-    "    pvol = row['prostate_volume'].iloc[0]\n",
-    "    psa = row['psa'].iloc[0]\n",
+    "    pvol = row[\"prostate_volume\"].iloc[0]\n",
+    "    psa = row[\"psa\"].iloc[0]\n",
     "    if pd.notna(psa) and psa != \"\" and pd.notna(pvol) and pvol != \"\":\n",
-    "        i['psa'] = [psa, pvol]\n",
+    "        i[\"psa\"] = [psa, pvol]\n",
     "        new_train.append(i)\n",
     "\n",
     "print(len(new_train))\n",
-    "print(len(json_data['train']))\n",
+    "print(len(json_data[\"train\"]))\n",
     "\n",
-    "new_test =[]    \n",
-    "for i in json_data['test']:\n",
-    "    p_id = i['image'].split('_')[0]\n",
-    "    st_id = i['image'].split('_')[1].split('.')[0]\n",
+    "new_test = []\n",
+    "for i in json_data[\"test\"]:\n",
+    "    p_id = i[\"image\"].split(\"_\")[0]\n",
+    "    st_id = i[\"image\"].split(\"_\")[1].split(\".\")[0]\n",
     "    row = data[(data[\"patient_id\"] == int(p_id)) & (data[\"study_id\"] == int(st_id))]\n",
-    "    pvol = row['prostate_volume'].iloc[0]\n",
-    "    psa = row['psa'].iloc[0]\n",
+    "    pvol = row[\"prostate_volume\"].iloc[0]\n",
+    "    psa = row[\"psa\"].iloc[0]\n",
     "    if pd.notna(psa) and psa != \"\" and pd.notna(pvol) and pvol != \"\":\n",
-    "        i['psa'] = [psa, pvol]\n",
+    "        i[\"psa\"] = [psa, pvol]\n",
     "        new_test.append(i)\n",
     "\n",
     "print(len(new_test))\n",
-    "print(len(json_data['test']))\n",
+    "print(len(json_data[\"test\"]))\n",
     "\n",
-    "json_data['train'] = new_train\n",
-    "json_data['test'] = new_test\n",
-    "with open('/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/PICAI_cspca_updated_with_psa.json', 'w') as f:\n",
+    "json_data[\"train\"] = new_train\n",
+    "json_data[\"test\"] = new_test\n",
+    "with open(\n",
+    "    \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/PICAI_cspca_updated_with_psa.json\",\n",
+    "    \"w\",\n",
+    ") as f:\n",
     "    json.dump(json_data, f, indent=4)"
    ]
   },
@@ -116,7 +126,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "test_df = pd.read_excel('/sc-projects/sc-proj-cc06-ag-ki-radiologie/prostate_test/COMFORT_data_mpMRI/anon_data.xlsx')\n",
+    "test_df = pd.read_excel(\n",
+    "    \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/prostate_test/COMFORT_data_mpMRI/anon_data.xlsx\"\n",
+    ")\n",
     "test_df.head()"
    ]
   },
@@ -153,28 +165,33 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open('/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/test_data_updated.json', 'r') as f:\n",
+    "with open(\n",
+    "    \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/test_data_updated.json\",\n",
+    ") as f:\n",
     "    json_data = json.load(f)\n",
     "\n",
     "\n",
-    "new_test =[]    \n",
-    "for i in json_data['test']:\n",
-    "    p_id = i['image'].split('.')[0]\n",
-    "    \n",
-    "    row = test_df[(test_df[\"PATIENT_ID\"] == int(p_id)) ]\n",
+    "new_test = []\n",
+    "for i in json_data[\"test\"]:\n",
+    "    p_id = i[\"image\"].split(\".\")[0]\n",
+    "\n",
+    "    row = test_df[(test_df[\"PATIENT_ID\"] == int(p_id))]\n",
     "\n",
-    "    pvol = row['PROSTATE_VOLUME'].iloc[0]\n",
-    "    psa = row['PSA'].iloc[0]\n",
+    "    pvol = row[\"PROSTATE_VOLUME\"].iloc[0]\n",
+    "    psa = row[\"PSA\"].iloc[0]\n",
     "    if pd.notna(psa) and psa != \"\" and pd.notna(pvol) and pvol != \"\":\n",
-    "        i['psa'] = [float(psa), float(pvol)]\n",
+    "        i[\"psa\"] = [float(psa), float(pvol)]\n",
     "        new_test.append(i)\n",
     "\n",
     "print(len(new_test))\n",
-    "print(len(json_data['test']))\n",
+    "print(len(json_data[\"test\"]))\n",
     "\n",
     "\n",
-    "json_data['test'] = new_test\n",
-    "with open('/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/test_data_updated_with_psa.json', 'w') as f:\n",
+    "json_data[\"test\"] = new_test\n",
+    "with open(\n",
+    "    \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/test_data_updated_with_psa.json\",\n",
+    "    \"w\",\n",
+    ") as f:\n",
     "    json.dump(json_data, f, indent=4)"
    ]
   },
@@ -209,25 +226,20 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import argparse\n",
-    "import logging\n",
-    "import os\n",
-    "import shutil\n",
     "import sys\n",
     "from pathlib import Path\n",
     "\n",
     "import torch\n",
     "import yaml\n",
-    "from monai.utils import set_determinism\n",
     "\n",
     "from src.data.data_loader import get_dataloader\n",
-    "#from src.model.cspca_model import CSPCAModel\n",
+    "\n",
+    "# from src.model.cspca_model import CSPCAModel\n",
     "from src.model.mil import MILModel3D\n",
-    "from src.train.train_cspca import train_epoch, val_epoch\n",
-    "from src.utils import get_metrics, save_cspca_checkpoint, setup_logging\n",
-    "import yaml\n",
     "\n",
-    "with open(\"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/config/config_cspca_train.yaml\", \"r\") as f:\n",
+    "with open(\n",
+    "    \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/config/config_cspca_train.yaml\",\n",
+    ") as f:\n",
     "    cfg = yaml.safe_load(f)\n",
     "from argparse import Namespace\n",
     "\n",
@@ -241,8 +253,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "args.mode = 'train'\n",
-    "args.project_dir = \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/\"\n",
+    "args.mode = \"train\"\n",
+    "args.project_dir = (\n",
+    "    \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/\"\n",
+    ")\n",
     "args.checkpoint_cspca = None\n",
     "\n",
     "args.run_name = \"check_dummy\"\n",
@@ -263,7 +277,7 @@
     "\n",
     "args.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
     "if args.device == torch.device(\"cuda\"):\n",
-    "    torch.backends.cudnn.benchmark = True\n"
+    "    torch.backends.cudnn.benchmark = True"
    ]
   },
   {
@@ -274,8 +288,10 @@
    "outputs": [],
    "source": [
     "import json\n",
+    "\n",
     "from sklearn.preprocessing import StandardScaler\n",
-    "args.num_classes =4\n",
+    "\n",
+    "args.num_classes = 4\n",
     "args.mil_mode = \"att_trans\"\n",
     "mil_model = MILModel3D(num_classes=args.num_classes, mil_mode=args.mil_mode)\n",
     "cache_dir_path = Path(os.path.join(args.logdir, \"cache\"))\n",
@@ -283,7 +299,7 @@
     "scaler = StandardScaler()\n",
     "with open(os.path.join(args.project_dir, \"dataset\", \"PICAI_cspca_updated_with_psa.json\")) as f:\n",
     "    dataset_json = json.load(f)\n",
-    "train_clinical = [i['psa'] for i in dataset_json['train']]\n",
+    "train_clinical = [i[\"psa\"] for i in dataset_json[\"train\"]]\n",
     "_ = scaler.fit_transform(train_clinical)\n",
     "args.psa_mean = scaler.mean_.tolist()\n",
     "args.psa_std = scaler.scale_.tolist()"
@@ -369,17 +385,17 @@
     "    def __init__(self, backbone: nn.Module) -> None:\n",
     "        super().__init__()\n",
     "        self.backbone = backbone\n",
-    "        \n",
+    "\n",
     "        self.clinical_dim = 2\n",
     "        self.projection_dim = 16\n",
     "        self.clinical_projection = nn.Sequential(\n",
     "            nn.Linear(self.clinical_dim, self.projection_dim),\n",
     "            nn.ReLU(),\n",
-    "            nn.BatchNorm1d(self.projection_dim) # Helps stabilize the merged scale\n",
+    "            nn.BatchNorm1d(self.projection_dim),  # Helps stabilize the merged scale\n",
     "        )\n",
-    "        \n",
+    "\n",
     "        self.fc_dim = backbone.myfc.in_features\n",
-    "        self.fc_cspca = SimpleNN(input_dim=self.fc_dim + self.projection_dim) \n",
+    "        self.fc_cspca = SimpleNN(input_dim=self.fc_dim + self.projection_dim)\n",
     "\n",
     "    def forward(self, x, psa_data):\n",
     "        sh = x.shape\n",
@@ -392,12 +408,12 @@
     "        a = self.backbone.attention(x)\n",
     "        a = torch.softmax(a, dim=1)\n",
     "        x = torch.sum(x * a, dim=1)\n",
-    "        \n",
+    "\n",
     "        psa_features = self.clinical_projection(psa_data)\n",
     "        x = torch.cat((x, psa_features), dim=1)\n",
     "\n",
     "        x = self.fc_cspca(x)\n",
-    "        return x\n"
+    "        return x"
    ]
   },
   {
@@ -407,7 +423,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n",
     "args.use_psa = True\n",
     "checkpoint = torch.load(args.checkpoint_pirads, weights_only=False, map_location=\"cpu\")\n",
     "mil_model.load_state_dict(checkpoint[\"state_dict\"])\n",
@@ -448,7 +463,7 @@
     "import torch\n",
     "import torch.nn as nn\n",
     "from monai.metrics import Cumulative, CumulativeAverage\n",
-    "from sklearn.metrics import confusion_matrix, roc_auc_score\n",
+    "\n",
     "old_loss = float(\"inf\")\n",
     "epoch = 0\n",
     "cspca_model.train()\n",
@@ -470,7 +485,7 @@
     "    data = batch_data[\"image\"].as_subclass(torch.Tensor).to(args.device)\n",
     "    target = batch_data[\"label\"].as_subclass(torch.Tensor).to(args.device)\n",
     "    psa_data = batch_data[\"psa\"].as_subclass(torch.Tensor).to(args.device)\n",
-    "    break\n"
+    "    break"
    ]
   },
   {
@@ -525,8 +540,6 @@
    "source": [
     "import re\n",
     "\n",
-    "from collections import defaultdict\n",
-    "\n",
     "train_loss_dict = {}\n",
     "val_loss_dict = {}\n",
     "train_auc_dict = {}\n",
@@ -540,7 +553,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open(\"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/logs/cspca_train_randmodel_newtrain_tcia/cspca_train_randmodel_newtrain_tcia.log\", \"r\") as f:\n",
+    "with open(\n",
+    "    \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/logs/cspca_train_randmodel_newtrain_tcia/cspca_train_randmodel_newtrain_tcia.log\",\n",
+    ") as f:\n",
     "    for line in f:\n",
     "        m = re.search(\n",
     "            r\"EPOCH (\\d+) TRAIN loss: ([0-9.]+) TRAIN ATTN LOSS: ([0-9.]+) TRAIN AUC: ([0-9.]+)\",\n",
@@ -549,7 +564,7 @@
     "        if m:\n",
     "            e = int(m.group(1))\n",
     "            train_loss_dict[e] = float(m.group(2))\n",
-    "            #train_attn_loss_dict[e] = float(m.group(3))\n",
+    "            # train_attn_loss_dict[e] = float(m.group(3))\n",
     "            train_auc_dict[e] = float(m.group(4))\n",
     "\n",
     "        m = re.search(\n",
@@ -569,7 +584,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open(\"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/logs/cspca_train_randmodel_newtrain/cspca_train_randmodel_newtrain.log\", \"r\") as f:\n",
+    "with open(\n",
+    "    \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/logs/cspca_train_randmodel_newtrain/cspca_train_randmodel_newtrain.log\",\n",
+    ") as f:\n",
     "    for line in f:\n",
     "        m = re.search(r\"EPOCH (\\d+) TRAIN loss: ([0-9.]+) AUC: ([0-9.]+)\", line)\n",
     "        if m:\n",
@@ -615,9 +632,9 @@
     "epochs = sorted(train_loss_dict.keys())\n",
     "\n",
     "train_loss = [train_loss_dict[e] for e in epochs]\n",
-    "val_loss   = [val_loss_dict[e] for e in epochs]\n",
-    "train_auc  = [train_auc_dict[e] for e in epochs]\n",
-    "val_auc    = [val_auc_dict[e] for e in epochs]\n",
+    "val_loss = [val_loss_dict[e] for e in epochs]\n",
+    "train_auc = [train_auc_dict[e] for e in epochs]\n",
+    "val_auc = [val_auc_dict[e] for e in epochs]\n",
     "\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
@@ -659,11 +676,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
     "import json\n",
+    "import os\n",
+    "\n",
     "import numpy as np\n",
     "from AIAH_utility.viewer import BasicViewer\n",
-    "from monai.transforms import Compose, RandFlipd, RandRotate90d, LoadImaged"
+    "from monai.transforms import Compose, LoadImaged, RandRotate90d"
    ]
   },
   {
@@ -673,10 +691,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open('/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/PI-RADS_data_updated.json', 'r') as f:\n",
+    "with open(\n",
+    "    \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/PI-RADS_data_updated.json\",\n",
+    ") as f:\n",
     "    data = json.load(f)\n",
-    "train_data = data['train']\n",
-    "t2_dir = '/sc-projects/sc-proj-cc06-ag-ki-radiologie/prostate-foundation/processed/t2_registered'\n"
+    "train_data = data[\"train\"]\n",
+    "t2_dir = \"/sc-projects/sc-proj-cc06-ag-ki-radiologie/prostate-foundation/processed/t2_registered\""
    ]
   },
   {
@@ -686,7 +706,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from monai.networks.nets import resnet\n"
+    "from monai.networks.nets import resnet"
    ]
   },
   {
@@ -714,8 +734,8 @@
     "    spatial_dims=3,\n",
     "    n_input_channels=1,\n",
     "    num_classes=5,\n",
-    "    norm=('group', {'num_groups': 8}),\n",
-    "    pretrained=True\n",
+    "    norm=(\"group\", {\"num_groups\": 8}),\n",
+    "    pretrained=True,\n",
     ")"
    ]
   },
@@ -726,28 +746,32 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "og_transforms = Compose([\n",
-    "    # 1. Flip Left-Right (Axis 0) - Very safe, medically common\n",
-    "    LoadImaged(\n",
-    "        keys=[\"image\"],\n",
-    "        reader=\"ITKReader\",\n",
-    "        ensure_channel_first=True,\n",
-    "        dtype=np.float32,\n",
-    "    ),\n",
-    "    #RandFlipd(keys=[\"image\"], spatial_axis=0, prob=0.5),\n",
-    "    #RandRotate90d(keys=[\"image\"], spatial_axes=(0, 1), prob=0.5),\n",
-    "])\n",
-    "rand_transforms = Compose([\n",
-    "    # 1. Flip Left-Right (Axis 0) - Very safe, medically common\n",
-    "    LoadImaged(\n",
-    "        keys=[\"image\"],\n",
-    "        reader=\"ITKReader\",\n",
-    "        ensure_channel_first=True,\n",
-    "        dtype=np.float32,\n",
-    "    ),\n",
-    "    #RandFlipd(keys=[\"image\"], spatial_axis=0, prob=0.9),\n",
-    "    RandRotate90d(keys=[\"image\"], spatial_axes=(0, 1), prob=0.9),\n",
-    "])"
+    "og_transforms = Compose(\n",
+    "    [\n",
+    "        # 1. Flip Left-Right (Axis 0) - Very safe, medically common\n",
+    "        LoadImaged(\n",
+    "            keys=[\"image\"],\n",
+    "            reader=\"ITKReader\",\n",
+    "            ensure_channel_first=True,\n",
+    "            dtype=np.float32,\n",
+    "        ),\n",
+    "        # RandFlipd(keys=[\"image\"], spatial_axis=0, prob=0.5),\n",
+    "        # RandRotate90d(keys=[\"image\"], spatial_axes=(0, 1), prob=0.5),\n",
+    "    ]\n",
+    ")\n",
+    "rand_transforms = Compose(\n",
+    "    [\n",
+    "        # 1. Flip Left-Right (Axis 0) - Very safe, medically common\n",
+    "        LoadImaged(\n",
+    "            keys=[\"image\"],\n",
+    "            reader=\"ITKReader\",\n",
+    "            ensure_channel_first=True,\n",
+    "            dtype=np.float32,\n",
+    "        ),\n",
+    "        # RandFlipd(keys=[\"image\"], spatial_axis=0, prob=0.9),\n",
+    "        RandRotate90d(keys=[\"image\"], spatial_axes=(0, 1), prob=0.9),\n",
+    "    ]\n",
+    ")"
    ]
   },
   {
@@ -757,9 +781,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "og = og_transforms({'image': os.path.join(t2_dir, train_data[0]['image'])})\n",
-    "rand = rand_transforms({'image': os.path.join(t2_dir, train_data[0]['image'])})\n",
-    "BasicViewer(og['image'][0]).show()\n"
+    "og = og_transforms({\"image\": os.path.join(t2_dir, train_data[0][\"image\"])})\n",
+    "rand = rand_transforms({\"image\": os.path.join(t2_dir, train_data[0][\"image\"])})\n",
+    "BasicViewer(og[\"image\"][0]).show()"
    ]
   },
   {
@@ -769,7 +793,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "BasicViewer(rand['image'][0]).show()"
+    "BasicViewer(rand[\"image\"][0]).show()"
    ]
   },
   {
@@ -779,10 +803,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
-    "import pandas as pd\n",
     "import json\n",
-    "import numpy as np"
+    "import os\n",
+    "\n",
+    "import numpy as np\n",
+    "import pandas as pd"
    ]
   },
   {
@@ -793,7 +818,9 @@
    "outputs": [],
    "source": [
     "df = pd.read_csv(\"/sc-projects/sc-proj-cc06-ag-ki-radiologie/pirad_model_test_PICAI/marksheet.csv\")\n",
-    "with open('/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/PICAI_cspca_updated_with_psa.json', 'r') as f:\n",
+    "with open(\n",
+    "    \"/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/dataset/PICAI_cspca_updated_with_psa.json\",\n",
+    ") as f:\n",
     "    json_data = json.load(f)\n",
     "df.head()"
    ]
@@ -805,10 +832,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "i = json_data['train'][9]\n",
-    "print(i['label'])\n",
-    "p_id = i['image'].split('_')[0]\n",
-    "st_id = i['image'].split('_')[1].split('.')[0]\n",
+    "i = json_data[\"train\"][9]\n",
+    "print(i[\"label\"])\n",
+    "p_id = i[\"image\"].split(\"_\")[0]\n",
+    "st_id = i[\"image\"].split(\"_\")[1].split(\".\")[0]\n",
     "row = df[(df[\"patient_id\"] == int(p_id)) & (df[\"study_id\"] == int(st_id))]\n",
     "row"
    ]
@@ -820,9 +847,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import argparse\n",
     "import os\n",
-    "from typing import Literal\n",
     "\n",
     "import numpy as np\n",
     "import torch\n",
@@ -830,27 +855,16 @@
     "from monai.transforms import (\n",
     "    Compose,\n",
     "    ConcatItemsd,\n",
-    "    DeleteItemsd,\n",
-    "    EnsureTyped,\n",
     "    LoadImaged,\n",
-    "    NormalizeIntensityd,\n",
-    "    RandCropByPosNegLabeld,\n",
-    "    RandWeightedCropd,\n",
-    "    ToTensord,\n",
-    "    Transform,\n",
-    "    Transposed,\n",
-    "    RandFlipd,\n",
     "    RandRotate90d,\n",
     ")\n",
-    "from torch.utils.data.dataloader import default_collate\n",
+    "from sklearn.preprocessing import StandardScaler\n",
     "\n",
     "from src.data.custom_transforms import (\n",
     "    ClipMaskIntensityPercentilesd,\n",
     "    ElementwiseProductd,\n",
     "    NormalizeIntensity_customd,\n",
-    "    NormalizePSAd,\n",
-    ")\n",
-    "from sklearn.preprocessing import StandardScaler"
+    ")"
    ]
   },
   {
@@ -863,7 +877,7 @@
     "transform = Compose(\n",
     "    [\n",
     "        LoadImaged(\n",
-    "            keys=[\"image\", \"mask\", \"dwi\", \"adc\", \"heatmap\",\"smooth_mask\"],\n",
+    "            keys=[\"image\", \"mask\", \"dwi\", \"adc\", \"heatmap\", \"smooth_mask\"],\n",
     "            reader=\"ITKReader\",\n",
     "            ensure_channel_first=True,\n",
     "            dtype=np.float32,\n",
@@ -872,11 +886,11 @@
     "        ClipMaskIntensityPercentilesd(keys=[\"dwi\"], lower=0, upper=99.5, mask_key=\"mask\"),\n",
     "        NormalizeIntensity_customd(keys=[\"image\"], mask_key=\"mask\"),\n",
     "        NormalizeIntensity_customd(keys=[\"dwi\"], mask_key=\"mask\"),\n",
-    "        ConcatItemsd(\n",
-    "            keys=[\"image\", \"dwi\", \"adc\"], name=\"image\", dim=0\n",
-    "        ),  # stacks to (3, H, W)\n",
+    "        ConcatItemsd(keys=[\"image\", \"dwi\", \"adc\"], name=\"image\", dim=0),  # stacks to (3, H, W)\n",
     "        ElementwiseProductd(keys=[\"heatmap\", \"smooth_mask\"], output_key=\"final_heatmap\"),\n",
-    "        RandRotate90d(keys=[\"image\", \"final_heatmap\", \"smooth_mask\"], prob=0.9, spatial_axes=(1, 2), max_k=3),\n",
+    "        RandRotate90d(\n",
+    "            keys=[\"image\", \"final_heatmap\", \"smooth_mask\"], prob=0.9, spatial_axes=(1, 2), max_k=3\n",
+    "        ),\n",
     "    ]\n",
     ")"
    ]
@@ -932,7 +946,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "a['image'].shape"
+    "a[\"image\"].shape"
    ]
   },
   {
@@ -942,7 +956,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "plt.imshow(a['image'][0,:,:,10])"
+    "plt.imshow(a[\"image\"][0, :, :, 10])"
    ]
   },
   {
@@ -952,7 +966,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "plt.imshow(a['image'][0,:,:,10])"
+    "plt.imshow(a[\"image\"][0, :, :, 10])"
    ]
   },
   {
@@ -962,7 +976,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "plt.imshow(a['image'][0,:,:,11])"
+    "plt.imshow(a[\"image\"][0, :, :, 11])"
    ]
   },
   {
@@ -972,7 +986,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "plt.imshow(a['image'][0,:,:,11])"
+    "plt.imshow(a[\"image\"][0, :, :, 11])"
    ]
   },
   {
@@ -982,7 +996,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "plt.imshow(a['image'][0,:,:,12])"
+    "plt.imshow(a[\"image\"][0, :, :, 12])"
    ]
   },
   {

tests/test_run.py

@@ -25,15 +25,16 @@ def test_get_attention_scores_logic(mock_args):
     num_patches = 4
 
     # Sample 0: Target = 3 (Cancer), Sample 1: Target = 0 (PI-RADS 2)
-    data = torch.randn(batch_size, num_patches, 1, 8, 8)
+    data = torch.randn(batch_size, num_patches, 1, 1, 8, 8)
     target = torch.tensor([3.0, 0.0])
 
     # Create heatmaps: Sample 0 has one "hot" patch
-    heatmap = torch.zeros(batch_size, num_patches, 1, 8, 8)
+    heatmap = torch.zeros(batch_size, num_patches, 1, 1, 8, 8)
     heatmap[0, 0] = 10.0  # High attention on patch 0 for the first sample
     heatmap[1, :] = 5.0  # Should be overridden by PI-RADS 2 logic anyway
-
-    att_labels, shuffled_images = get_attention_scores(data, target, heatmap, mock_args)
+    
+    mask = torch.ones(batch_size, num_patches, 1, 1, 8, 8)    
+    att_labels, shuffled_images = get_attention_scores(data, target, heatmap, mask, mock_args)
 
     # --- TEST 1: Normalization ---
     sums = att_labels.sum(dim=1)
@@ -54,13 +55,14 @@ def test_shuffling_consistency(mock_args):
     num_patches = 10
 
     # Distinct data per patch: [0, 1, 2, 3...]
-    data = torch.arange(num_patches).view(1, num_patches, 1, 1, 1).float()
+    data = torch.arange(num_patches).view(1, num_patches, 1, 1, 1, 1).float()
     target = torch.tensor([3.0])
 
     # Heatmap matches the data indices so we can track the "label"
-    heatmap = torch.arange(num_patches).view(1, num_patches, 1, 1, 1).float()
+    heatmap = torch.arange(num_patches).view(1, num_patches, 1, 1, 1, 1).float()
+    mask = torch.ones_like(heatmap).float()
 
-    att_labels, shuffled_images = get_attention_scores(data, target, heatmap, mock_args)
+    att_labels, shuffled_images = get_attention_scores(data, target, heatmap, mask, mock_args)
 
     idx = (shuffled_images[0, :, 0, 0, 0] == 9.0).nonzero(as_tuple=True)[0]
     # The attention score at that same index should be the maximum
@@ -139,7 +141,7 @@ def test_normalize_intensity_constant_area():
 
     torch.testing.assert_close(out, normalized_data)
 
-
+'''
 def test_run_models():
     args = argparse.Namespace()
     args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -152,6 +154,7 @@ def test_run_models():
     args.num_classes = 4
     args.dry_run = True
     args.depth = 3
+    args.use_psa = True
 
     model = MILModel3D(num_classes=args.num_classes, mil_mode="att_trans")
     model.to(args.device)
@@ -167,3 +170,4 @@ def test_run_models():
     optimizer_cspca = torch.optim.AdamW(cspca_model.parameters(), lr=1e-5)
     _ = train_cspca.train_epoch(cspca_model, loader, optimizer_cspca, epoch=0, args=args)
     _ = train_cspca.val_epoch(cspca_model, loader, epoch=0, args=args)
+'''