Add salient patch info

run_inference.py

@@ -34,7 +34,7 @@ from pathlib import Path
 from src.model.MIL import MILModel_3D
 from src.model.csPCa_model import csPCa_Model
 from src.data.data_loader import get_dataloader
-from src.utils import save_cspca_checkpoint, get_metrics, setup_logging, save_pirads_checkpoint
+from src.utils import save_cspca_checkpoint, get_metrics, setup_logging, save_pirads_checkpoint, get_parent_image, get_patch_coordinate
 from src.train import train_cspca, train_pirads
 import SimpleITK as sitk 
 
@@ -57,6 +57,7 @@ import argparse
 import yaml 
 from src.data.data_loader import data_transform, list_data_collate
 from monai.data import Dataset, load_decathlon_datalist, ITKReader, NumpyReader, PersistentDataset
+import json
 
 def parse_args():
 
@@ -120,7 +121,7 @@ if __name__ == "__main__":
 
     transform = data_transform(args)
     files = os.listdir(args.t2_dir)
-    data_list = []
+    args.data_list = []
     for file in files:
         temp = {}
         temp['image'] = os.path.join(args.t2_dir, file)
@@ -129,9 +130,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
-        data_list.append(temp)
+        args.data_list.append(temp)
 
-    dataset = Dataset(data=data_list, transform=transform)
+    dataset = Dataset(data=args.data_list, transform=transform)
     loader = torch.utils.data.DataLoader(
         dataset,
         batch_size=1,
@@ -147,6 +148,7 @@ if __name__ == "__main__":
     pirads_model.eval()
     cspca_risk_list = []
     cspca_model.eval()
+    top5_patches = []
     with torch.no_grad():
         for idx, batch_data in enumerate(loader):
             data = batch_data["image"].as_subclass(torch.Tensor).to(args.device)
@@ -158,5 +160,34 @@ if __name__ == "__main__":
             output = output.squeeze(1)
             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)
+            top5_values, top5_indices = torch.topk(a, 5)
+
+            patches_top_5 = []
+            for i in range(5):
+                patch_temp = data[0, top5_indices.cpu().numpy()[i]][0].cpu().numpy()
+                patches_top_5.append(patch_temp)
+
+    parent_image = get_parent_image(args)
+
+    coords = get_patch_coordinate(patches_top_5, parent_image, args)
+
     for i,j in enumerate(files):
-        logging.info(f"File: {j}, PIRADS score: {pirads_list[i]}, csPCa risk score: {cspca_risk_list[i]:.4f}")
+        logging.info(f"File: {j}, PIRADS score: {pirads_list[i]}, csPCa risk score: {cspca_risk_list[i]:.4f}")
+
+    output_dict = {
+        'Predicted PIRAD Score': pirads_list[i] + 2.0,
+        'csPCa risk': cspca_risk_list[i],
+        'Top left coordinate of top 5 patches(x,y,z)': coords,
+    }
+    with open(os.path.join(args.output_dir, "results.json"), 'w') as f:
+        json.dump(output_dict, f, indent=4)

src/utils.py

@@ -14,19 +14,52 @@ import torch.nn.functional as F
 from monai.config import KeysCollection
 from monai.metrics import Cumulative, CumulativeAverage
 from monai.networks.nets import milmodel, resnet, MILModel
-
+from monai.transforms import (
+    Compose,
+    GridPatchd,
+    LoadImaged,
+    MapTransform,
+    RandFlipd,
+    RandGridPatchd,
+    RandRotate90d,
+    ScaleIntensityRanged,
+    SplitDimd,
+    ToTensord,
+    ConcatItemsd, 
+    SelectItemsd,
+    EnsureChannelFirstd,
+    RepeatChanneld,
+    DeleteItemsd,
+    EnsureTyped,
+    ClipIntensityPercentilesd,
+    MaskIntensityd,
+    HistogramNormalized,
+    RandBiasFieldd,
+    RandCropByPosNegLabeld,
+    NormalizeIntensityd,
+    SqueezeDimd,
+    CropForegroundd,
+    ScaleIntensityd,
+    SpatialPadd,
+    CenterSpatialCropd,
+    ScaleIntensityd,
+    Transposed,
+    RandWeightedCropd,
+)
 from sklearn.metrics import cohen_kappa_score
 from torch.cuda.amp import GradScaler, autocast
 from torch.utils.data.dataloader import default_collate
 from torchvision.models.resnet import ResNet50_Weights
-
+from .data.custom_transforms import ClipMaskIntensityPercentilesd, NormalizeIntensity_customd
 from torch.utils.data.distributed import DistributedSampler
 from torch.utils.tensorboard import SummaryWriter
+import matplotlib.patches as patches
 
 import matplotlib.pyplot as plt
 
 import wandb
 import math
+from monai.data import Dataset, load_decathlon_datalist, ITKReader, NumpyReader, PersistentDataset
 
 from src.model.MIL import MILModel_3D
 from src.model.csPCa_model import csPCa_Model
@@ -96,4 +129,98 @@ def validate_steps(steps):
                     f"Step '{step}' requires '{req}' to be executed before it. "
                     f"Given order: {steps}"
                 )
-                sys.exit(1)
+                sys.exit(1)
+
+def get_patch_coordinate(patches_top_5, parent_image, args):
+
+    sample = np.array([i.transpose(1,2,0) for i in patches_top_5])
+    coords = []
+    rows, h, w, slices = sample.shape
+
+    for i in range(rows):
+        for j in range(slices):
+            if j == 0:
+                for k in range(parent_image.shape[2]):
+                    img_temp = parent_image[:, :, k]
+                    H, W = img_temp.shape
+                    h, w = sample[i, :, :, j].shape
+                    a,b = 0, 0  # Initialize a and b
+                    bool1 = False
+                    for l in range(H - h + 1):
+                        for m in range(W - w + 1):
+                            if np.array_equal(img_temp[l:l+h, m:m+w], sample[i, :, :, j]):
+                                a,b = l, m  # top-left corner
+                                coords.append((a,b,k))
+                                bool1 = True
+                                break
+                        if bool1:
+                            break
+                        
+                    if bool1:
+                        break
+
+    return coords
+
+
+def get_parent_image(args):
+    transform_image = Compose(
+        [
+            LoadImaged(keys=["image", "mask"], reader=ITKReader(), ensure_channel_first=True, dtype=np.float32),
+            ClipMaskIntensityPercentilesd(keys=["image"], lower=0, upper=99.5, mask_key="mask"),
+            NormalizeIntensity_customd(keys=["image"], mask_key="mask", channel_wise=True),
+            EnsureTyped(keys=["label"], dtype=torch.float32),
+            ToTensord(keys=["image", "label"]),
+        ]
+    )
+    dataset_image = Dataset(data=args.data_list, transform=transform_image)
+    return dataset_image[0]['image'][0].numpy()
+
+'''
+def visualise_patches():
+    sample = np.array([i.transpose(1,2,0) for i in patches_top_5])
+    rows = len(patches_top_5)
+    img = sample[0]
+    coords = []
+    rows, h, w, slices = sample.shape
+
+    fig, axes = plt.subplots(nrows=rows, ncols=slices, figsize=(slices * 3, rows * 3))
+
+    for i in range(rows):
+        for j in range(slices):
+            ax = axes[i, j]
+            
+            if j == 0:
+            
+                for k in range(parent_image.shape[2]):
+                    img_temp = parent_image[:, :, k]
+                    H, W = img_temp.shape
+                    h, w = sample[i, :, :, j].shape
+                    a,b = 0, 0  # Initialize a and b
+                    bool1 = False
+                    for l in range(H - h + 1):
+                        for m in range(W - w + 1):
+                            if np.array_equal(img_temp[l:l+h, m:m+w], sample[i, :, :, j]):
+                                a,b = l, m  # top-left corner
+                                coords.append((a,b,k))
+                                bool1 = True
+                                break
+                        if bool1:
+                            break
+                        
+                    if bool1:
+                        break
+
+                
+
+
+            ax.imshow(parent_image[:, :, k+j], cmap='gray')
+            rect = patches.Rectangle((b, a), args.tile_size, args.tile_size,
+                                    linewidth=2, edgecolor='red', facecolor='none')
+            ax.add_patch(rect)
+            ax.axis('off')
+        
+
+    plt.tight_layout()
+    plt.show()
+    a=1
+'''

temp.ipynb

@@ -2,10 +2,19 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 1,
    "id": "cec738fb",
    "metadata": {},
-   "outputs": [],
+   "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 os\n",
@@ -70,7 +79,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "id": "c91a5802",
    "metadata": {},
    "outputs": [
@@ -78,9 +87,9 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "  0%|          | 0/1 [00:02<?, ?it/s]\n",
+      "  0%|          | 0/1 [00:03<?, ?it/s]\n",
       "You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
-      "  0%|          | 0/1 [00:07<?, ?it/s]\n"
+      "  0%|          | 0/1 [00:05<?, ?it/s]\n"
      ]
     }
    ],
@@ -104,36 +113,23 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
-   "id": "f32be6e9",
+   "execution_count": 4,
+   "id": "8b5d382e",
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'margin': 0.2,\n",
-       " 't2_dir': '/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed/t2_histmatched',\n",
-       " 'dwi_dir': '/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed/DWI_histmatched',\n",
-       " 'adc_dir': '/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed/ADC_histmatched',\n",
-       " 'output_dir': '/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed',\n",
-       " 'project_dir': '/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate',\n",
-       " 'seg_dir': '/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed/prostate_mask',\n",
-       " 'heatmapdir': '/sc-scratch/sc-scratch-cc06-ag-ki-radiologie/prostate_foundation/WSAttention-Prostate/datatemp/processed/heatmaps/'}"
-      ]
-     },
-     "execution_count": 6,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
-    "vars(args)"
+    "args.num_classes = 4\n",
+    "args.mil_mode = \"att_trans\"\n",
+    "args.use_heatmap = True\n",
+    "args.tile_size = 64\n",
+    "args.tile_count = 24\n",
+    "args.depth = 3\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
-   "id": "118c2549",
+   "execution_count": 5,
+   "id": "4cf061ec",
    "metadata": {},
    "outputs": [
     {
@@ -141,155 +137,23 @@
      "output_type": "stream",
      "text": [
       "enable_nested_tensor is True, but self.use_nested_tensor is False because encoder_layer.self_attn.batch_first was not True(use batch_first for better inference performance)\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",
       "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"
      ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "MILModel_3D(\n",
-       "  (attention): Sequential(\n",
-       "    (0): Linear(in_features=512, out_features=2048, bias=True)\n",
-       "    (1): Tanh()\n",
-       "    (2): Linear(in_features=2048, out_features=1, bias=True)\n",
-       "  )\n",
-       "  (transformer): TransformerEncoder(\n",
-       "    (layers): ModuleList(\n",
-       "      (0-3): 4 x TransformerEncoderLayer(\n",
-       "        (self_attn): MultiheadAttention(\n",
-       "          (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n",
-       "        )\n",
-       "        (linear1): Linear(in_features=512, out_features=2048, bias=True)\n",
-       "        (dropout): Dropout(p=0.0, inplace=False)\n",
-       "        (linear2): Linear(in_features=2048, out_features=512, bias=True)\n",
-       "        (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
-       "        (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
-       "        (dropout1): Dropout(p=0.0, inplace=False)\n",
-       "        (dropout2): Dropout(p=0.0, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "  )\n",
-       "  (myfc): Linear(in_features=512, out_features=4, bias=True)\n",
-       "  (net): ResNet(\n",
-       "    (conv1): Conv3d(3, 64, kernel_size=(7, 7, 7), stride=(1, 1, 1), padding=(3, 3, 3), bias=False)\n",
-       "    (bn1): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "    (act): ReLU(inplace=True)\n",
-       "    (maxpool): MaxPool3d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
-       "    (layer1): Sequential(\n",
-       "      (0): ResNetBlock(\n",
-       "        (conv1): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)\n",
-       "        (bn1): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        (act): ReLU(inplace=True)\n",
-       "        (conv2): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)\n",
-       "        (bn2): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "      )\n",
-       "      (1): ResNetBlock(\n",
-       "        (conv1): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)\n",
-       "        (bn1): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        (act): ReLU(inplace=True)\n",
-       "        (conv2): Conv3d(64, 64, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)\n",
-       "        (bn2): BatchNorm3d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "      )\n",
-       "    )\n",
-       "    (layer2): Sequential(\n",
-       "      (0): ResNetBlock(\n",
-       "        (conv1): Conv3d(64, 128, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1), bias=False)\n",
-       "        (bn1): BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        (act): ReLU(inplace=True)\n",
-       "        (conv2): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)\n",
-       "        (bn2): BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        (downsample): Sequential(\n",
-       "          (0): Conv3d(64, 128, kernel_size=(1, 1, 1), stride=(2, 2, 2))\n",
-       "          (1): BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        )\n",
-       "      )\n",
-       "      (1): ResNetBlock(\n",
-       "        (conv1): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)\n",
-       "        (bn1): BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        (act): ReLU(inplace=True)\n",
-       "        (conv2): Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)\n",
-       "        (bn2): BatchNorm3d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "      )\n",
-       "    )\n",
-       "    (layer3): Sequential(\n",
-       "      (0): ResNetBlock(\n",
-       "        (conv1): Conv3d(128, 256, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1), bias=False)\n",
-       "        (bn1): BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        (act): ReLU(inplace=True)\n",
-       "        (conv2): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)\n",
-       "        (bn2): BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        (downsample): Sequential(\n",
-       "          (0): Conv3d(128, 256, kernel_size=(1, 1, 1), stride=(2, 2, 2))\n",
-       "          (1): BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        )\n",
-       "      )\n",
-       "      (1): ResNetBlock(\n",
-       "        (conv1): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)\n",
-       "        (bn1): BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        (act): ReLU(inplace=True)\n",
-       "        (conv2): Conv3d(256, 256, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)\n",
-       "        (bn2): BatchNorm3d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "      )\n",
-       "    )\n",
-       "    (layer4): Sequential(\n",
-       "      (0): ResNetBlock(\n",
-       "        (conv1): Conv3d(256, 512, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1), bias=False)\n",
-       "        (bn1): BatchNorm3d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        (act): ReLU(inplace=True)\n",
-       "        (conv2): Conv3d(512, 512, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)\n",
-       "        (bn2): BatchNorm3d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        (downsample): Sequential(\n",
-       "          (0): Conv3d(256, 512, kernel_size=(1, 1, 1), stride=(2, 2, 2))\n",
-       "          (1): BatchNorm3d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        )\n",
-       "      )\n",
-       "      (1): ResNetBlock(\n",
-       "        (conv1): Conv3d(512, 512, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)\n",
-       "        (bn1): BatchNorm3d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "        (act): ReLU(inplace=True)\n",
-       "        (conv2): Conv3d(512, 512, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), bias=False)\n",
-       "        (bn2): BatchNorm3d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
-       "      )\n",
-       "    )\n",
-       "    (avgpool): AdaptiveAvgPool3d(output_size=(1, 1, 1))\n",
-       "    (fc): Identity()\n",
-       "  )\n",
-       ")"
-      ]
-     },
-     "execution_count": 8,
-     "metadata": {},
-     "output_type": "execute_result"
     }
    ],
    "source": [
     "args.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
-    "args.num_classes = 4\n",
-    "args.mil_mode = 'att_trans'\n",
+    "\n",
+    "\n",
     "pirads_model = MILModel_3D(\n",
     "    num_classes=args.num_classes,  \n",
     "    mil_mode=args.mil_mode \n",
     ")\n",
     "pirads_checkpoint = torch.load(os.path.join(args.project_dir, 'models', 'pirads.pt'), map_location=\"cpu\")\n",
     "pirads_model.load_state_dict(pirads_checkpoint[\"state_dict\"])\n",
-    "pirads_model.to(args.device)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 31,
-   "id": "01467cae",
-   "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": [
+    "pirads_model.to(args.device)\n",
+    "\n",
     "cspca_model = csPCa_Model(backbone=pirads_model).to(args.device)\n",
     "checkpt = torch.load(os.path.join(args.project_dir, 'models', 'cspca_model.pth'), map_location=\"cpu\")\n",
     "cspca_model.load_state_dict(checkpt['state_dict'])\n",
@@ -298,18 +162,41 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "id": "aa7850f3",
+   "execution_count": 6,
+   "id": "fac15515",
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n"
+    "transform = data_transform(args)\n",
+    "files = os.listdir(args.t2_dir)\n",
+    "data_list = []\n",
+    "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",
+    "    data_list.append(temp)\n",
+    "\n",
+    "dataset = Dataset(data=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": 33,
-   "id": "bd8884f8",
+   "execution_count": 7,
+   "id": "eb80047b",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -317,6 +204,7 @@
     "pirads_model.eval()\n",
     "cspca_risk_list = []\n",
     "cspca_model.eval()\n",
+    "top5_patches = []\n",
     "with torch.no_grad():\n",
     "    for idx, batch_data in enumerate(loader):\n",
     "        data = batch_data[\"image\"].as_subclass(torch.Tensor).to(args.device)\n",
@@ -326,13 +214,371 @@
     "\n",
     "        output = cspca_model(data)\n",
     "        output = output.squeeze(1)\n",
-    "        cspca_risk_list.append(output.item())"
+    "        cspca_risk_list.append(output.item())\n",
+    "\n",
+    "        sh = data.shape\n",
+    "        x = data.reshape(sh[0] * sh[1], sh[2], sh[3], sh[4], sh[5])\n",
+    "        x = cspca_model.backbone.net(x)\n",
+    "        x = x.reshape(sh[0], sh[1], -1)\n",
+    "        x = x.permute(1, 0, 2)\n",
+    "        x = cspca_model.backbone.transformer(x)\n",
+    "        x = x.permute(1, 0, 2)\n",
+    "        a = cspca_model.backbone.attention(x)\n",
+    "        a = torch.softmax(a, dim=1)\n",
+    "        a = a.view(-1)\n",
+    "        top5_values, top5_indices = torch.topk(a, 5)\n",
+    "\n",
+    "        patches_top_5 = []\n",
+    "        for i in range(5):\n",
+    "            patch_temp = data[0, top5_indices.cpu().numpy()[i]][0].cpu().numpy()\n",
+    "            patches_top_5.append(patch_temp)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "dbcfc97f",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[array([[[-0.43333182, -0.52243334, -1.4245864 , ..., -1.7698549 ,\n",
+       "          -1.6027895 , -1.257521  ],\n",
+       "         [-0.7786003 , -1.2797965 , -2.0817103 , ..., -2.0371597 ,\n",
+       "          -2.0817103 , -2.059435  ],\n",
+       "         [-1.4802749 , -1.658478  , -1.7809926 , ..., -1.4691372 ,\n",
+       "          -2.1596742 , -1.9591957 ],\n",
+       "         ...,\n",
+       "         [ 0.8697783 ,  0.34630668,  0.6358867 , ...,  0.624749  ,\n",
+       "           0.24606745,  0.03445129],\n",
+       "         [ 0.44654593,  0.23492976,  0.7806767 , ...,  1.2373221 ,\n",
+       "           0.2906182 , -1.0124918 ],\n",
+       "         [ 0.3017559 , -0.2551287 ,  0.27948052, ...,  1.1036698 ,\n",
+       "          -0.16602719, -1.2463834 ]],\n",
+       " \n",
+       "        [[-0.84542644, -1.3800358 , -1.4357241 , ..., -2.5383558 ,\n",
+       "          -2.0817103 , -0.99021643],\n",
+       "         [-0.4444695 , -1.3466226 , -1.8255434 , ..., -1.335485  ,\n",
+       "          -1.6362027 , -1.4579996 ],\n",
+       "         [-1.2129703 , -1.7921304 , -2.0371597 , ...,  0.14582822,\n",
+       "          -0.04351256, -0.82315105],\n",
+       "         ...,\n",
+       "         [-0.4556072 ,  1.4378005 ,  1.5603153 , ...,  0.2906182 ,\n",
+       "          -0.85656416, -1.1906949 ],\n",
+       "         [-1.2352457 , -0.35536796,  1.1816336 , ..., -0.39991874,\n",
+       "          -1.3800358 , -1.2129703 ],\n",
+       "         [-1.2909342 , -1.1127311 ,  0.5467852 , ..., -0.9568034 ,\n",
+       "          -1.4023111 , -0.9011149 ]],\n",
+       " \n",
+       "        [[-2.1596742 , -1.914645  , -1.7809926 , ..., -1.0124918 ,\n",
+       "           0.03445129,  0.44654593],\n",
+       "         [-2.0037465 , -1.8923696 , -1.8700942 , ..., -1.2129703 ,\n",
+       "          -1.1572819 , -0.50015795],\n",
+       "         [-1.8478189 , -1.9369203 , -1.9814711 , ..., -0.65608567,\n",
+       "          -1.2352457 , -1.4134488 ],\n",
+       "         ...,\n",
+       "         [-0.6338103 , -1.0124918 , -0.16602719, ...,  0.41313285,\n",
+       "           0.13469052, -0.80087566],\n",
+       "         [-0.7451872 , -1.3577603 , -0.2996795 , ...,  0.34630668,\n",
+       "           0.68043745,  0.45768362],\n",
+       "         [ 0.0121759 , -0.7674626 , -0.33309257, ...,  0.09013975,\n",
+       "           0.46882132,  1.0034306 ]]], dtype=float32),\n",
+       " array([[[-0.64494795, -0.789738  , -0.5447087 , ...,  0.0233136 ,\n",
+       "           0.14582822, -0.35536796],\n",
+       "         [-0.8120134 , -0.7340495 , -0.42219412, ..., -0.31081718,\n",
+       "           0.46882132,  0.15696591],\n",
+       "         [-0.08806333,  0.03445129,  0.12355283, ..., -0.2774041 ,\n",
+       "           0.73612595,  0.7249882 ],\n",
+       "         ...,\n",
+       "         [ 0.09013975,  0.0233136 , -0.24399103, ...,  0.34630668,\n",
+       "           0.914329  ,  0.6358867 ],\n",
+       "         [ 0.2906182 ,  0.335169  ,  0.624749  , ...,  0.24606745,\n",
+       "           0.9254667 ,  0.9922929 ],\n",
+       "         [ 0.44654593,  0.70271283,  1.0479814 , ..., -0.09920102,\n",
+       "           0.37971976,  0.70271283]],\n",
+       " \n",
+       "        [[-0.23285334, -0.01009948, -0.1326141 , ...,  1.1593583 ,\n",
+       "           1.5603153 ,  1.5603153 ],\n",
+       "         [-0.23285334,  0.13469052,  0.1792413 , ...,  0.98115516,\n",
+       "           1.5603153 ,  1.5603153 ],\n",
+       "         [-0.36650565, -0.01009948,  0.190379  , ...,  0.8697783 ,\n",
+       "           1.5603153 ,  1.5603153 ],\n",
+       "         ...,\n",
+       "         [-0.16602719, -0.06578795,  0.190379  , ..., -1.4357241 ,\n",
+       "          -1.368898  , -1.6027895 ],\n",
+       "         [-0.2662664 , -0.35536796,  0.190379  , ..., -1.513688  ,\n",
+       "          -1.3800358 , -1.6362027 ],\n",
+       "         [-0.789738  , -0.7786003 , -0.17716487, ..., -1.7141665 ,\n",
+       "          -1.5693765 , -1.9591957 ]],\n",
+       " \n",
+       "        [[-0.12147641, -0.01009948,  0.0233136 , ...,  0.769539  ,\n",
+       "           0.8140898 ,  0.2906182 ],\n",
+       "         [-0.35536796, -0.2774041 , -0.16602719, ...,  0.55792284,\n",
+       "           0.68043745,  0.335169  ],\n",
+       "         [-0.18830256,  0.1681036 ,  0.7918144 , ...,  0.70271283,\n",
+       "           0.4910967 ,  0.10127745],\n",
+       "         ...,\n",
+       "         [-0.97907877, -0.97907877, -0.9011149 , ..., -0.43333182,\n",
+       "          -0.5447087 , -0.1437518 ],\n",
+       "         [-1.1238688 , -1.1461442 , -1.0347673 , ..., -0.7117741 ,\n",
+       "          -0.12147641,  0.479959  ],\n",
+       "         [-0.8788395 , -0.62267256, -0.9122526 , ..., -0.92339027,\n",
+       "          -0.42219412,  0.22379206]]], dtype=float32),\n",
+       " array([[[ 2.7948052e-01,  1.0368437e+00,  1.5603153e+00, ...,\n",
+       "          -1.2797965e+00, -8.3428878e-01, -2.4399103e-01],\n",
+       "         [ 6.7864366e-02,  3.4630668e-01,  1.1704960e+00, ...,\n",
+       "          -7.7860028e-01, -8.3428878e-01, -3.1081718e-01],\n",
+       "         [ 2.6834285e-01,  9.0139754e-02,  2.3492976e-01, ...,\n",
+       "          -6.3381028e-01, -9.5680338e-01, -4.4446951e-01],\n",
+       "         ...,\n",
+       "         [-1.3688980e+00, -1.2241080e+00, -9.4566566e-01, ...,\n",
+       "          -1.0570426e+00, -2.6626641e-01,  1.2707351e+00],\n",
+       "         [-1.0793180e+00, -8.8997722e-01, -1.1015934e+00, ...,\n",
+       "          -7.6746261e-01, -5.4650255e-02,  1.0257059e+00],\n",
+       "         [-7.3404950e-01, -6.5608567e-01, -1.0124918e+00, ...,\n",
+       "          -5.3357106e-01,  2.6834285e-01,  8.5864055e-01]],\n",
+       " \n",
+       "        [[ 1.1593583e+00,  1.5603153e+00,  1.5603153e+00, ...,\n",
+       "          -4.5560721e-01, -1.1033872e-01,  8.6977828e-01],\n",
+       "         [ 8.4750289e-01,  1.5603153e+00,  1.5603153e+00, ...,\n",
+       "          -1.3020718e+00, -1.0459049e+00,  6.3588673e-01],\n",
+       "         [ 7.9002060e-02,  1.3469052e-01, -3.1081718e-01, ...,\n",
+       "          -1.3688980e+00, -1.0570426e+00,  4.9109671e-01],\n",
+       "         ...,\n",
+       "         [ 1.5603153e+00,  6.4702439e-01, -4.3512560e-02, ...,\n",
+       "           5.2450979e-01,  8.9205366e-01, -2.9967949e-01],\n",
+       "         [ 1.2930106e+00,  6.1361134e-01, -4.6674490e-01, ...,\n",
+       "           2.7948052e-01,  8.4750289e-01,  7.8067672e-01],\n",
+       "         [ 1.0382106e-03,  2.5720516e-01, -4.1105643e-01, ...,\n",
+       "           6.2474900e-01,  1.5603153e+00,  1.5603153e+00]],\n",
+       " \n",
+       "        [[-9.2339027e-01, -1.0236295e+00, -1.0347673e+00, ...,\n",
+       "           2.6834285e-01, -1.9944026e-01, -1.1033872e-01],\n",
+       "         [-8.2315105e-01, -1.1127311e+00, -9.7907877e-01, ...,\n",
+       "           1.2818729e+00,  3.2403129e-01,  6.7864366e-02],\n",
+       "         [-8.0087566e-01, -9.7907877e-01, -8.0087566e-01, ...,\n",
+       "           1.5046268e+00,  6.1361134e-01, -1.7716487e-01],\n",
+       "         ...,\n",
+       "         [ 1.3932499e+00,  6.8043745e-01,  2.3492976e-01, ...,\n",
+       "          -1.8144057e+00, -1.2129703e+00,  2.3313597e-02],\n",
+       "         [ 1.5603153e+00,  1.5603153e+00,  1.5603153e+00, ...,\n",
+       "          -1.7587173e+00, -1.1572819e+00,  5.6906056e-01],\n",
+       "         [ 1.5603153e+00,  1.5603153e+00,  1.5603153e+00, ...,\n",
+       "          -1.8366811e+00, -8.7883949e-01,  1.3598367e+00]]], dtype=float32),\n",
+       " array([[[-1.4245864 , -1.3466226 , -1.079318  , ..., -0.36650565,\n",
+       "          -0.66722333, -1.079318  ],\n",
+       "         [-1.1350064 , -1.2352457 , -1.4357241 , ..., -0.01009948,\n",
+       "          -0.41105643, -0.7563249 ],\n",
+       "         [-1.0459049 , -0.84542644, -1.368898  , ..., -0.22171564,\n",
+       "          -0.4778826 , -0.82315105],\n",
+       "         ...,\n",
+       "         [-2.2599134 , -2.1151235 , -1.6139272 , ...,  0.5245098 ,\n",
+       "           0.41313285,  0.37971976],\n",
+       "         [-2.560631  , -2.204225  , -1.7921304 , ...,  0.10127745,\n",
+       "           0.3128936 ,  0.37971976],\n",
+       "         [-2.2933266 , -2.1596742 , -2.0705726 , ...,  0.0233136 ,\n",
+       "           0.3017559 ,  0.4910967 ]],\n",
+       " \n",
+       "        [[-1.0570426 , -1.3911734 , -1.658478  , ...,  0.769539  ,\n",
+       "           0.7806767 ,  0.95887977],\n",
+       "         [-1.6362027 , -1.4134488 , -1.5693765 , ...,  1.0145682 ,\n",
+       "           1.1816336 ,  1.1370829 ],\n",
+       "         [-1.914645  , -1.1127311 , -1.224108  , ...,  0.68043745,\n",
+       "           0.9922929 ,  0.85864055],\n",
+       "         ...,\n",
+       "         [-2.1262612 , -2.026022  , -1.7921304 , ...,  0.4910967 ,\n",
+       "           0.9031913 ,  1.0702567 ],\n",
+       "         [-2.3712904 , -2.5494936 , -2.304464  , ...,  0.70271283,\n",
+       "           0.914329  ,  1.0479814 ],\n",
+       "         [-2.2710512 , -2.3267395 , -2.4715295 , ...,  1.0145682 ,\n",
+       "           0.7806767 ,  0.7472636 ]],\n",
+       " \n",
+       "        [[-0.94566566, -0.6003972 , -0.85656416, ...,  1.2373221 ,\n",
+       "           1.5603153 ,  1.5603153 ],\n",
+       "         [-1.1795572 , -0.38878104, -0.37764335, ...,  1.5603153 ,\n",
+       "           1.5603153 ,  1.5603153 ],\n",
+       "         [-1.335485  , -1.3577603 , -0.32195488, ...,  1.5603153 ,\n",
+       "           1.5603153 ,  1.5603153 ],\n",
+       "         ...,\n",
+       "         [-1.224108  , -1.1795572 , -1.4914126 , ..., -0.04351256,\n",
+       "          -0.5112957 , -0.64494795],\n",
+       "         [-1.3132095 , -1.8589565 , -1.7587173 , ..., -0.03237487,\n",
+       "          -0.42219412, -0.1437518 ],\n",
+       "         [-1.6473403 , -2.4381166 , -2.1262612 , ...,  0.37971976,\n",
+       "           0.34630668,  0.46882132]]], dtype=float32),\n",
+       " array([[[-0.84542644, -0.8120134 , -0.6894987 , ...,  0.46882132,\n",
+       "          -0.01009948, -0.39991874],\n",
+       "         [-0.65608567, -0.8120134 , -0.6338103 , ...,  0.59133595,\n",
+       "           0.04558898, -0.5112957 ],\n",
+       "         [-0.96794105, -0.8342888 , -0.35536796, ...,  0.5801982 ,\n",
+       "           0.13469052, -0.37764335],\n",
+       "         ...,\n",
+       "         [-0.09920102, -0.11033872, -1.224108  , ..., -0.92339027,\n",
+       "          -1.2686588 , -0.85656416],\n",
+       "         [ 0.0121759 , -0.01009948, -1.0236295 , ..., -1.2129703 ,\n",
+       "          -1.4023111 , -0.7340495 ],\n",
+       "         [ 0.23492976, -0.07692564, -0.70063645, ..., -0.70063645,\n",
+       "          -0.934528  , -1.0124918 ]],\n",
+       " \n",
+       "        [[-1.6139272 , -0.97907877, -0.7340495 , ..., -0.8342888 ,\n",
+       "          -0.6894987 , -1.0347673 ],\n",
+       "         [-1.0459049 , -0.5892595 , -0.789738  , ..., -1.257521  ,\n",
+       "          -1.0124918 , -1.4357241 ],\n",
+       "         [-0.42219412, -0.5892595 , -1.4134488 , ..., -1.5025504 ,\n",
+       "          -1.2463834 , -1.4802749 ],\n",
+       "         ...,\n",
+       "         [ 0.10127745,  0.40199515,  0.13469052, ...,  0.25720516,\n",
+       "           0.55792284,  0.12355283],\n",
+       "         [-0.2885418 , -0.2551287 ,  0.41313285, ...,  0.46882132,\n",
+       "          -0.33309257, -1.2797965 ],\n",
+       "         [-0.2996795 , -0.64494795,  0.9366044 , ...,  0.56906056,\n",
+       "          -0.41105643, -1.3466226 ]],\n",
+       " \n",
+       "        [[-1.1238688 , -1.4134488 , -1.7364419 , ..., -1.1350064 ,\n",
+       "          -1.6918911 , -2.3156018 ],\n",
+       "         [-1.4914126 , -1.3020718 , -0.99021643, ..., -1.658478  ,\n",
+       "          -1.7698549 , -1.8812319 ],\n",
+       "         [-1.5582387 , -0.85656416, -0.43333182, ..., -1.6027895 ,\n",
+       "          -1.914645  , -1.7698549 ],\n",
+       "         ...,\n",
+       "         [ 0.6358867 ,  0.9031913 ,  1.0702567 , ...,  0.11241514,\n",
+       "           0.07900206,  0.34630668],\n",
+       "         [ 0.70271283,  1.1593583 ,  0.9254667 , ...,  0.335169  ,\n",
+       "           0.41313285,  0.23492976],\n",
+       "         [ 0.35744438,  1.1482205 ,  0.8697783 , ..., -0.2662664 ,\n",
+       "           0.56906056,  0.624749  ]]], dtype=float32)]"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "patches_top_5"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "4edb20e7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import argparse\n",
+    "import os\n",
+    "import shutil\n",
+    "import time\n",
+    "import yaml\n",
+    "import sys\n",
+    "import gdown\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "import torch.distributed as dist\n",
+    "import torch.multiprocessing as mp\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from monai.config import KeysCollection\n",
+    "from monai.metrics import Cumulative, CumulativeAverage\n",
+    "from monai.networks.nets import milmodel, resnet, MILModel\n",
+    "from monai.transforms import (\n",
+    "    Compose,\n",
+    "    GridPatchd,\n",
+    "    LoadImaged,\n",
+    "    MapTransform,\n",
+    "    RandFlipd,\n",
+    "    RandGridPatchd,\n",
+    "    RandRotate90d,\n",
+    "    ScaleIntensityRanged,\n",
+    "    SplitDimd,\n",
+    "    ToTensord,\n",
+    "    ConcatItemsd, \n",
+    "    SelectItemsd,\n",
+    "    EnsureChannelFirstd,\n",
+    "    RepeatChanneld,\n",
+    "    DeleteItemsd,\n",
+    "    EnsureTyped,\n",
+    "    ClipIntensityPercentilesd,\n",
+    "    MaskIntensityd,\n",
+    "    HistogramNormalized,\n",
+    "    RandBiasFieldd,\n",
+    "    RandCropByPosNegLabeld,\n",
+    "    NormalizeIntensityd,\n",
+    "    SqueezeDimd,\n",
+    "    CropForegroundd,\n",
+    "    ScaleIntensityd,\n",
+    "    SpatialPadd,\n",
+    "    CenterSpatialCropd,\n",
+    "    ScaleIntensityd,\n",
+    "    Transposed,\n",
+    "    RandWeightedCropd,\n",
+    ")\n",
+    "from sklearn.metrics import cohen_kappa_score\n",
+    "from torch.cuda.amp import GradScaler, autocast\n",
+    "from torch.utils.data.dataloader import default_collate\n",
+    "from torchvision.models.resnet import ResNet50_Weights\n",
+    "from src.data.custom_transforms import ClipMaskIntensityPercentilesd, NormalizeIntensity_customd\n",
+    "from torch.utils.data.distributed import DistributedSampler\n",
+    "from torch.utils.tensorboard import SummaryWriter\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import wandb\n",
+    "import math\n",
+    "from monai.data import Dataset, load_decathlon_datalist, ITKReader, NumpyReader, PersistentDataset\n",
+    "\n",
+    "from src.model.MIL import MILModel_3D\n",
+    "from src.model.csPCa_model import csPCa_Model\n",
+    "\n",
+    "import logging\n",
+    "from pathlib import Path"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "e42cc132",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "transform_image = Compose(\n",
+    "    [\n",
+    "        LoadImaged(keys=[\"image\", \"mask\"], reader=ITKReader(), ensure_channel_first=True, dtype=np.float32),\n",
+    "        ClipMaskIntensityPercentilesd(keys=[\"image\"], lower=0, upper=99.5, mask_key=\"mask\"),\n",
+    "        NormalizeIntensity_customd(keys=[\"image\"], mask_key=\"mask\", channel_wise=True),\n",
+    "        EnsureTyped(keys=[\"label\"], dtype=torch.float32),\n",
+    "        ToTensord(keys=[\"image\", \"label\"]),\n",
+    "    ]\n",
+    ")\n",
+    "dataset_image = Dataset(data=data_list, transform=transform_image)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "bcdddd9e",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(270, 270, 28)"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dataset_image[0]['image'][0].numpy().shape"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4cf061ec",
+   "id": "56072a2b",
    "metadata": {},
    "outputs": [],
    "source": []