update viewer

app.py

@@ -4,7 +4,6 @@ import os
 import shutil
 import subprocess
 
-import matplotlib.patches as patches
 import matplotlib.pyplot as plt
 import nrrd
 import numpy as np
@@ -45,6 +44,119 @@ def load_nrrd(file_path):
     return data, header
 
 
+def display_slicer(
+    scan_paths, mask_path=None, heatmap_path=None, title="Scan Viewer", key_suffix=""
+):
+    """
+    Displays slicer with Multi-Background Support, Mask Overlay, and Attention Heatmap.
+
+    Args:
+        scan_paths: Dict of {Label: FilePath}. Example: {"T2W": "path/to/t2.nrrd"}
+        mask_path: Path to the segmentation mask (optional)
+        heatmap_path: Path to the 3D attention heatmap nrrd file (optional)
+    """
+    # 1. Layout: Image/Slider (Left) | Controls (Right)
+    c_viewer, c_controls = st.columns([3, 1.5])
+
+    # --- CONTROLS SECTION (Right Column) ---
+    with c_controls:
+        st.write(f"**{title} Controls**")
+
+        # A. Background Selection
+        available_scans = list(scan_paths.keys())
+        selected_scan_name = st.radio(
+            "Background Image", available_scans, index=0, key=f"bg_{key_suffix}"
+        )
+        current_file_path = scan_paths[selected_scan_name]
+
+        # B. Heatmap Controls
+        st.write("---")
+        show_heatmap = False
+        if heatmap_path and os.path.exists(heatmap_path):
+            show_heatmap = st.checkbox("Show Attention Heatmap", value=True, key=f"hm_{key_suffix}")
+            if show_heatmap:
+                hm_alpha = st.slider(
+                    "Heatmap Opacity", 0.1, 1.0, 0.5, 0.1, key=f"hm_a_{key_suffix}"
+                )
+                # Threshold to hide low-attention areas (assumes heatmap is normalized 0-1)
+                hm_thresh = st.slider(
+                    "Hide Values Below", 0.0, 1.0, 0.1, 0.05, key=f"hm_t_{key_suffix}"
+                )
+
+        # C. Mask Controls
+        st.write("---")
+        show_mask = False
+        if mask_path and os.path.exists(mask_path):
+            show_mask = st.checkbox("Show Mask Overlay", value=False, key=f"mk_{key_suffix}")
+
+    # --- VIEWER SECTION (Left Column) ---
+    with c_viewer:
+        if not os.path.exists(current_file_path):
+            st.error(f"File not found: {current_file_path}")
+            return
+
+        # Load the selected background image
+        data, _ = load_nrrd(current_file_path)
+
+        if len(data.shape) != 3:
+            st.warning("Data is not 3D.")
+            return
+
+        total_slices = data.shape[2]
+
+        # D. Slider Logic
+        start_slice = total_slices // 2
+
+        slice_idx = st.slider(
+            "Select Slice (Z-Axis)", 0, total_slices - 1, start_slice, key=f"sl_{key_suffix}"
+        )
+
+        # E. Plotting
+        img_slice = data[:, :, slice_idx]
+
+        # Normalize Image (0-1)
+        img_slice = img_slice.astype(float)
+
+        fig, ax = plt.subplots(figsize=(5, 5))
+        ax.imshow(img_slice, cmap="gray", origin="upper")
+
+        # 1. Overlay Heatmap
+        if show_heatmap:
+            h_data, _ = load_nrrd(heatmap_path)
+
+            if h_data.shape == data.shape:
+                h_slice = h_data[:, :, slice_idx].astype(float)
+
+                # Normalize the heatmap slice to 0-1 if it isn't already
+                max_val = np.max(h_slice)
+                if max_val > 0:
+                    h_slice = h_slice / max_val
+
+                # Mask out values below the user-defined threshold
+                h_overlay = np.ma.masked_where(h_slice < hm_thresh, h_slice)
+
+                # Overlay using 'jet', 'inferno', or 'hot' colormap
+                ax.imshow(h_overlay, cmap="jet", alpha=hm_alpha, origin="upper")
+            else:
+                ax.text(5, 5, "Heatmap shape mismatch", color="red", fontsize=8)
+
+        # 2. Overlay Mask
+        if show_mask:
+            m_data, _ = load_nrrd(mask_path)
+
+            if m_data.shape == data.shape:
+                mslice = m_data[:, :, slice_idx]
+                overlay = np.ma.masked_where(mslice == 0, mslice)
+                ax.imshow(overlay, cmap="Reds", alpha=0.5, origin="upper")
+            else:
+                # Placed slightly lower so it doesn't overlap with the heatmap warning
+                ax.text(5, 15, "Mask shape mismatch", color="red", fontsize=8)
+
+        ax.axis("off")
+        st.pyplot(fig, use_container_width=False)
+
+
+'''
 def display_slicer(scan_paths, mask_path=None, bboxes=None, title="Scan Viewer", key_suffix=""):
     """
     Displays slicer with Multi-Background Support, Mask Overlay, and Bounding Box Multiselect.
@@ -156,6 +268,8 @@ def display_slicer(scan_paths, mask_path=None, bboxes=None, title="Scan Viewer",
         ax.axis("off")
         st.pyplot(fig, use_container_width=False)
 
+'''
+
 
 @st.cache_resource
 def download_all_models():
@@ -571,6 +685,7 @@ if st.session_state.inference_done:
         dwi_vis_path = None
         adc_vis_path = None
         mask_vis_path = None
+        att_map_path = None
 
         t2_vis_dir = os.path.join(OUTPUT_DIR, "t2_registered")
         if os.path.exists(t2_vis_dir) and len(os.listdir(t2_vis_dir)) > 0:
@@ -595,6 +710,8 @@ if st.session_state.inference_done:
         else:
             print("No mask dir")
 
+        att_map_path = os.path.join(OUTPUT_DIR, "attention_weights.nrrd")
+
         roi_bbox = None
         if "coords" in st.session_state:
             detected_boxes = []
@@ -614,7 +731,7 @@ if st.session_state.inference_done:
             display_slicer(
                 scan_paths=scan_dict,  # <--- Pass the Dict here
                 mask_path=mask_vis_path,
-                bboxes=detected_boxes,
+                heatmap_path=att_map_path,
                 title="Salient Patch Viewer",
                 key_suffix="main_viz",
             )
@@ -622,7 +739,7 @@ if st.session_state.inference_done:
             display_slicer(
                 scan_paths=scan_dict,  # <--- Pass the Dict here
                 mask_path=mask_vis_path,
-                bboxes=None,
+                heatmap_path=None,
                 title="Salient Patch Viewer",
                 key_suffix="main_viz",
             )

pirads_inference.ipynb

@@ -0,0 +1,397 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "dfc57266",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import argparse\n",
+    "import json\n",
+    "import logging\n",
+    "import os\n",
+    "import shutil\n",
+    "import sys\n",
+    "import time\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import wandb\n",
+    "import yaml\n",
+    "from monai.utils import set_determinism\n",
+    "from sklearn.preprocessing import StandardScaler\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 monai.metrics import Cumulative, CumulativeAverage\n",
+    "from sklearn.metrics import cohen_kappa_score\n",
+    "from tqdm import tqdm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "b1264f17",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import yaml\n",
+    "from argparse import Namespace\n",
+    "\n",
+    "# 1. Read the YAML file into a dictionary\n",
+    "with open('/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/config/config_pirads_test.yaml', 'r') as file:\n",
+    "    config_dict = yaml.safe_load(file)\n",
+    "\n",
+    "# 2. Convert the dictionary directly to a Namespace\n",
+    "args = Namespace(**config_dict)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "6b12b122",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "args.project_dir = '/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate'\n",
+    "args.run_name = 'pirads_inf'\n",
+    "args.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "args.batch_size = 4\n",
+    "args.workers = 0\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",
+    "setup_logging(args.logfile)\n",
+    "if args.device == torch.device(\"cpu\"):\n",
+    "    args.amp = False"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "bd3292b2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "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"
+     ]
+    }
+   ],
+   "source": [
+    "if args.device == torch.device(\"cuda\"):\n",
+    "    torch.backends.cudnn.benchmark = True\n",
+    "\n",
+    "model = MILModel3D(num_classes=args.num_classes, mil_mode=args.mil_mode)\n",
+    "start_epoch = 0\n",
+    "best_acc = 0.0\n",
+    "checkpoint = torch.load(args.checkpoint, map_location=\"cpu\")\n",
+    "model.load_state_dict(checkpoint[\"state_dict\"])\n",
+    "\n",
+    "if \"epoch\" in checkpoint:\n",
+    "    start_epoch = checkpoint[\"epoch\"]\n",
+    "if \"best_acc\" in checkpoint:\n",
+    "    best_acc = checkpoint[\"best_acc\"]\n",
+    "logging.info(\n",
+    "    \"=> loaded checkpoint %s (epoch %d) (bestacc %f)\",\n",
+    "    args.checkpoint,\n",
+    "    start_epoch,\n",
+    "    best_acc,\n",
+    ")\n",
+    "cache_dir_ = os.path.join(args.logdir, \"cache\")\n",
+    "model.to(args.device)\n",
+    "params = model.parameters()\n",
+    "\n",
+    "scaler = StandardScaler()\n",
+    "with open(os.path.join(args.project_dir, \"dataset\", \"PICAI_cspca_updated_with_psa_updated_vol.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": 5,
+   "id": "e7ffa937",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "0it [00:00, ?it/s]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",
+      "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",
+      "204it [10:13,  3.01s/it]\n"
+     ]
+    }
+   ],
+   "source": [
+    "args.num_seeds = 1\n",
+    "args.epochs = 0\n",
+    "seed = 56\n",
+    "set_determinism(seed=seed)\n",
+    "valid_loader = get_dataloader(args, split='test')\n",
+    "\n",
+    "preds_cumulative = Cumulative()\n",
+    "targets_cumulative = Cumulative()\n",
+    "\n",
+    "model.eval()\n",
+    "with torch.no_grad():\n",
+    "    for idx, batch_data in tqdm(enumerate(valid_loader)):\n",
+    "        data = batch_data[\"image\"].as_subclass(torch.Tensor).to(args.device)\n",
+    "        target = batch_data[\"pirads\"].as_subclass(torch.Tensor).to(args.device)\n",
+    "        target = target.long()\n",
+    "\n",
+    "        with torch.amp.autocast(device_type=str(args.device), enabled=args.amp):\n",
+    "            logits = model(data)\n",
+    "\n",
+    "        data = data.to(\"cpu\")\n",
+    "        target = target.to(\"cpu\")\n",
+    "        logits = logits.to(\"cpu\")\n",
+    "        pred = torch.argmax(logits, dim=1)\n",
+    "\n",
+    "        preds_cumulative.extend(pred.detach().cpu())\n",
+    "        targets_cumulative.extend(target.detach().cpu())\n",
+    "\n",
+    "\n",
+    "        del data, target, logits\n",
+    "        # torch.cuda.empty_cache()\n",
+    "\n",
+    "    # Calculate QWK metric (Quadratic Weigted Kappa) https://en.wikipedia.org/wiki/Cohen%27s_kappa\n",
+    "    preds_cumulative = preds_cumulative.get_buffer().cpu().numpy()\n",
+    "    targets_cumulative = targets_cumulative.get_buffer().cpu().numpy()\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "if os.path.exists(cache_dir_):\n",
+    "\n",
+    "    shutil.rmtree(cache_dir_)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "5a544d4b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "np.save('pred.npy', preds_cumulative)\n",
+    "np.save('target.npy', targets_cumulative)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "c854d497",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pred = preds_cumulative\n",
+    "targets = targets_cumulative"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "f1a42cbe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "targets_clipped = np.clip(targets, a_min = 2, a_max = 5)\n",
+    "target = targets_clipped - 2\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "868e3173",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0.23748254253920043"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "qwk = cohen_kappa_score(\n",
+    "    target.astype(np.float64), pred.astype(np.float64)\n",
+    ")\n",
+    "qwk"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "4a496dc8",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0.2053401834331301"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "pred_mod = np.clip(pred, a_min = 0, a_max = 2)\n",
+    "target_mod = np.clip(target, a_min = 0, a_max = 2)\n",
+    "qwk = cohen_kappa_score(\n",
+    "    target_mod.astype(np.float64), pred_mod.astype(np.float64)\n",
+    ")\n",
+    "qwk"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "50e3195b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Counts for predicted classes [0, 1, 2, 3]:\n",
+      "[ 27  88 113  14]\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "preds_for_class_0 = pred[target == 0]\n",
+    "\n",
+    "\n",
+    "distribution = np.bincount(preds_for_class_0, minlength=4)\n",
+    "\n",
+    "print(\"Counts for predicted classes [0, 1, 2, 3]:\")\n",
+    "print(distribution)\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "b812ce5c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Counts for predicted classes [0, 1, 2, 3]:\n",
+      "[ 6 19 25 12]\n"
+     ]
+    }
+   ],
+   "source": [
+    "preds_for_class_0 = pred[target == 1]\n",
+    "\n",
+    "\n",
+    "distribution = np.bincount(preds_for_class_0, minlength=4)\n",
+    "\n",
+    "print(\"Counts for predicted classes [0, 1, 2, 3]:\")\n",
+    "print(distribution)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "46c55552",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Counts for predicted classes [0, 1, 2, 3]:\n",
+      "[ 12  34 129  49]\n"
+     ]
+    }
+   ],
+   "source": [
+    "preds_for_class_0 = pred[target == 2]\n",
+    "\n",
+    "\n",
+    "distribution = np.bincount(preds_for_class_0, minlength=4)\n",
+    "\n",
+    "print(\"Counts for predicted classes [0, 1, 2, 3]:\")\n",
+    "print(distribution)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "73283018",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Counts for predicted classes [0, 1, 2, 3]:\n",
+      "[  3  16  87 182]\n"
+     ]
+    }
+   ],
+   "source": [
+    "preds_for_class_0 = pred[target == 3]\n",
+    "\n",
+    "\n",
+    "distribution = np.bincount(preds_for_class_0, minlength=4)\n",
+    "\n",
+    "print(\"Counts for predicted classes [0, 1, 2, 3]:\")\n",
+    "print(distribution)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9cc17b5a",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "foundation",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.21"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

run_inference.py

@@ -6,6 +6,7 @@ from argparse import Namespace
 from collections.abc import Callable
 from pathlib import Path
 
+import nrrd
 import streamlit as st
 import torch
 import yaml
@@ -19,7 +20,13 @@ 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 get_parent_image, get_patch_coordinate, get_prostate_volume, setup_logging
+from src.utils import (
+    create_additive_heatmap,
+    get_parent_image,
+    get_patch_coordinate,
+    get_prostate_volume,
+    setup_logging,
+)
 
 
 @st.cache_resource
@@ -105,7 +112,7 @@ if __name__ == "__main__":
 
     scaler = StandardScaler()
     with open(
-        os.path.join(args.project_dir, "dataset", "PICAI_cspca_updated_with_psa_updated_vol.json")
+        os.path.join(args.project_dir, "dataset", "cspca_train_tcia.json.json")
     ) as f:
         dataset_json = json.load(f)
     train_clinical = [i["psa"] for i in dataset_json["train"]]
@@ -169,10 +176,10 @@ if __name__ == "__main__":
             a = cspca_model.backbone.attention(x)
             a = torch.softmax(a, dim=1)
             a = a.view(-1)
-            top5_values, top5_indices = torch.topk(a, 5)
+            top5_values, top5_indices = torch.topk(a, args.tile_count)
 
             patches_top_5 = []
-            for i in range(5):
+            for i in range(args.tile_count):
                 patch_temp = data[0, top5_indices.cpu().numpy()[i]][0].cpu().numpy()
                 patches_top_5.append(patch_temp)
             patches_top_5_list.append(patches_top_5)
@@ -193,8 +200,18 @@ if __name__ == "__main__":
             "Predicted PIRAD Score": pirads_list[i] + 2.0,
             "csPCa risk": cspca_risk_list[i],
             "Prostate Volume": args.data_list[i]["psa"][1],
-            "Top left coordinate of top 5 patches(x,y,z)": coords_list[i],
+            "Top left coordinates of the patches(x,y,z)": coords_list[i],
         }
 
+    pmask, _ = nrrd.read(args.data_list[0]["smooth_mask"])
+    score_map = create_additive_heatmap(
+        coords_list[0],
+        top5_values.cpu().numpy(),
+        parent_image.shape,
+        (args.tile_size, args.tile_size, args.depth),
+        pmask,
+    )
+    nrrd.write(os.path.join(args.output_dir, "attention_weights.nrrd"), score_map)
+
     with open(os.path.join(args.output_dir, "results.json"), "w") as f:
         json.dump(output_dict, f, indent=4)

src/utils.py

@@ -18,6 +18,7 @@ from monai.transforms import (
     LoadImaged,
     ToTensord,
 )
+from scipy.ndimage import gaussian_filter
 
 from .data.custom_transforms import ClipMaskIntensityPercentilesd, NormalizeIntensity_customd
 
@@ -262,3 +263,34 @@ def get_prostate_volume(mask_path) -> np.ndarray:
     true_volume_cc = true_volume_mm3 / 1000.0  # Convert mm³ to cc (mL)
 
     return true_volume_cc
+
+
+def create_additive_heatmap(
+    patch_coords, attention_scores, volume_shape, patch_size, pmask, apply_blur=True
+):
+    """
+    Sums attention scores in overlapping regions.
+    """
+    heatmap = np.zeros(volume_shape, dtype=np.float32)
+    dz, dy, dx = patch_size
+    for (z, y, x), score in zip(patch_coords, attention_scores):
+        # Define boundaries starting from top-left corner
+        z_end = min(volume_shape[0], z + dz)
+        y_end = min(volume_shape[1], y + dy)
+        x_end = min(volume_shape[2], x + dx)
+
+        # Add scores to the region
+        heatmap[z:z_end, y:y_end, x:x_end] += score
+        """
+            heatmap[z:z_end, y:y_end, x:x_end] = np.maximum(
+                heatmap[z:z_end, y:y_end, x:x_end], score
+            )
+            """
+    heatmap_masked = heatmap  # * (pmask > 0)
+    if heatmap_masked.max() > 0:
+        heatmap_masked /= heatmap_masked.max()
+
+    if apply_blur:
+        heatmap = gaussian_filter(heatmap_masked, sigma=1.0)
+
+    return heatmap