submint_model_test.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os\n",
    "from functools import partial\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.distributed as dist\n",
    "import torch.multiprocessing as mp\n",
    "import torch.nn.parallel\n",
    "import torch.utils.data.distributed\n",
    "from trainer import run_training\n",
    "from monai.inferers import sliding_window_inference\n",
    "from monai.losses import DiceCELoss, FocalLoss\n",
    "from monai.metrics import DiceMetric, MeanIoU\n",
    "from monai.transforms import Activations, AsDiscrete, Compose\n",
    "from monai.utils.enums import MetricReduction\n",
    "from monai.visualize import matshow3d\n",
    "\n",
    "from utils.myModel import MyModel, MyModel2d, MyModel3dunet, MyFlexibleUNet2d, MyFlexibleUNet2dLSTM, MyBasicUNetPlusPlus\n",
    "from utils.my_loss import CustomWeightedDiceCELoss\n",
    "\n",
    "from pathlib import Path\n",
    "from easydict import EasyDict as edict\n",
    "TEST_DIR = Path('/root/autodl-tmp/vesuvius-challenge-ink-detection/test')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = edict(RandFlipd_prob=0.2, RandRotate90d_prob=0.2, RandScaleIntensityd_prob=0.1, RandShiftIntensityd_prob=0.1, a_max=65535.0, a_min=0.0, amp=True, b_max=255.0, b_min=0.0, batch_size=1, cache_rate=1.0, checkpoint=None, data_dir='/root/autodl-tmp/vesuvius-challenge-ink-detection', debug=False, dist_backend='nccl', dist_url='tcp://127.0.0.1:23456', distributed=False, dropout_path_rate=0.0, dropout_rate=0.0, eff='b3', feature_size=48, gpu=0, in_channels=65, infer_overlap=0.5, json_list='/root/autodl-tmp/data_split/data_split.json', logdir='./runs/512_funetlstm_b3_16_sgd_continue', loss_mode='custom', loss_weight=(2.0, 1.0), lrschedule='cosine_anneal', max_epochs=2000, mid=28, model_mode='2dfunetlstm', momentum=0.99, noamp=False, norm_name='instance', normal=False, num_channel=16, num_samples=4, optim_lr=0.0001, optim_name='sgd', out_channels=1, pretrained_dir='./pretrained_models/', pretrained_model_name='512_funetlstm_b3_16_sgd_1000.pt', rank=0, reg_weight=1e-05, resume_ckpt=True, roi_x=512, roi_y=512, roi_z=16, save_checkpoint=True, smooth_dr=1e-06, smooth_nr=0.0, space_x=1.5, space_y=1.5, space_z=1.0, spatial_dims=3, sw_batch_size=4, test_mode=False, threshold=0.4, use_checkpoint=False, use_normal_dataset=False, use_ssl_pretrained=False, val_every=10, warmup_epochs=50, workers=0, world_size=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sampler(torch.utils.data.Sampler):\n",
    "    def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True, make_even=True):\n",
    "        if num_replicas is None:\n",
    "            if not torch.distributed.is_available():\n",
    "                raise RuntimeError(\n",
    "                    \"Requires distributed package to be available\")\n",
    "            num_replicas = torch.distributed.get_world_size()\n",
    "        if rank is None:\n",
    "            if not torch.distributed.is_available():\n",
    "                raise RuntimeError(\n",
    "                    \"Requires distributed package to be available\")\n",
    "            rank = torch.distributed.get_rank()\n",
    "        self.shuffle = shuffle\n",
    "        self.make_even = make_even\n",
    "        self.dataset = dataset\n",
    "        self.num_replicas = num_replicas\n",
    "        self.rank = rank\n",
    "        self.epoch = 0\n",
    "        self.num_samples = int(\n",
    "            math.ceil(len(self.dataset) * 1.0 / self.num_replicas))\n",
    "        self.total_size = self.num_samples * self.num_replicas\n",
    "        indices = list(range(len(self.dataset)))\n",
    "        self.valid_length = len(\n",
    "            indices[self.rank: self.total_size: self.num_replicas])\n",
    "\n",
    "    def __iter__(self):\n",
    "        if self.shuffle:\n",
    "            g = torch.Generator()\n",
    "            g.manual_seed(self.epoch)\n",
    "            indices = torch.randperm(len(self.dataset), generator=g).tolist()\n",
    "        else:\n",
    "            indices = list(range(len(self.dataset)))\n",
    "        if self.make_even:\n",
    "            if len(indices) < self.total_size:\n",
    "                if self.total_size - len(indices) < len(indices):\n",
    "                    indices += indices[: (self.total_size - len(indices))]\n",
    "                else:\n",
    "                    extra_ids = np.random.randint(low=0, high=len(\n",
    "                        indices), size=self.total_size - len(indices))\n",
    "                    indices += [indices[ids] for ids in extra_ids]\n",
    "            assert len(indices) == self.total_size\n",
    "        indices = indices[self.rank: self.total_size: self.num_replicas]\n",
    "        self.num_samples = len(indices)\n",
    "        return iter(indices)\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.num_samples\n",
    "\n",
    "    def set_epoch(self, epoch):\n",
    "        self.epoch = epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.utils import get_transforms\n",
    "from monai.data import load_decathlon_datalist\n",
    "from monai import data, transforms\n",
    "from utils.my_transform import *\n",
    "\n",
    "def get_loader(args):\n",
    "    data_dir = args.data_dir\n",
    "    datalist_json = os.path.join(data_dir, args.json_list)\n",
    "    test_transform = transforms.Compose(\n",
    "        [\n",
    "            transforms.LoadImaged(\n",
    "                keys=[\"image\", \"label\"], reader=\"NumpyReader\"),\n",
    "            Copyd(keys=[\"label\"],\n",
    "                    num_channel=args.num_channel),\n",
    "            transforms.AddChanneld(keys=[\"image\", 'label']),\n",
    "            # transforms.CropForegroundd(\n",
    "            #     keys=[\"image\", \"label\"], source_key=\"image\"),\n",
    "            # transforms.Orientationd(keys=[\"image\"], axcodes=\"RAS\"),\n",
    "            # change_channeld(keys=[\"image\", \"label\", 'inklabels']),\n",
    "            # transforms.Spacingd(keys=\"image\", pixdim=(args.space_x, args.space_y, args.space_z), mode=\"bilinear\"),\n",
    "            transforms.ScaleIntensityRanged(\n",
    "                keys=[\"image\"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True\n",
    "            ),\n",
    "            transforms.ToTensord(keys=[\"image\"]),\n",
    "        ]\n",
    "    )\n",
    "    val_files = load_decathlon_datalist(\n",
    "        datalist_json, True, \"testing\", base_dir=data_dir)\n",
    "    val_ds = data.Dataset(data=val_files, transform=test_transform)\n",
    "    val_sampler = Sampler(\n",
    "        val_ds, shuffle=False) if args.distributed else None\n",
    "    val_loader = data.DataLoader(\n",
    "        val_ds, batch_size=8, shuffle=False, num_workers=args.workers, sampler=val_sampler, pin_memory=True\n",
    "    )\n",
    "    loader = val_loader\n",
    "    return loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.utils import resample_3d, resample_2d\n",
    "\n",
    "def test(model_infer, val_loader, loss_func, args):\n",
    "    output_directory = Path(\"./outputs/\") / args.exp_name\n",
    "    output_directory.mkdir(parents=True, exist_ok=True)\n",
    "    with torch.no_grad():\n",
    "        dice_list_case = []\n",
    "        for i, batch in enumerate(val_loader):\n",
    "            val_inputs, val_labels = (batch[\"image\"].cuda(), batch[\"label\"].cuda())\n",
    "            print(type(val_labels))\n",
    "            print(val_labels.shape)\n",
    "            _, d, h, w = val_labels.shape\n",
    "            target_shape = (h, w)\n",
    "            img_name = batch[\"image_meta_dict\"][\"filename_or_obj\"][0].split(\"/\")[-1]\n",
    "            print(\"Inference on case {}\".format(img_name))\n",
    "            val_outputs = model_infer(val_inputs)\n",
    "            print(val_outputs.shape)\n",
    "            val_outputs = torch.softmax(val_outputs, 1).cpu()\n",
    "            val_outputs = np.array(val_outputs)\n",
    "            val_outputs = np.argmax(val_outputs, axis=1).astype(np.uint8)[0]\n",
    "            val_labels = val_labels.cpu()\n",
    "            val_labels = np.array(val_labels)[0, 0, :, :]\n",
    "            if args.model_mode == \"2dswin\":\n",
    "                val_outputs = resample_2d(val_outputs, target_shape)\n",
    "            elif args.model_mode == \"3dswin\":\n",
    "                val_outputs = resample_3d(val_outputs, target_shape)\n",
    "            else:\n",
    "                raise ValueError(\"model_mode should be ['3dswin', '2dswin', '3dunet', '2dunet']\")\n",
    "                \n",
    "            dice_list_sub = []\n",
    "            for i in [1]:\n",
    "                organ_Dice = loss_func(val_outputs == i, val_labels == i)\n",
    "                dice_list_sub.append(organ_Dice)\n",
    "            mean_dice = np.mean(dice_list_sub)\n",
    "            print(\"Mean Organ Dice: {}\".format(mean_dice))\n",
    "            dice_list_case.append(mean_dice)\n",
    "            np.save(\n",
    "                os.path.join(output_directory, img_name), val_outputs[:,:]\n",
    "            )\n",
    "        mean_loss = np.mean(dice_list_case)\n",
    "        print(\"Overall Mean Dice: {}\".format(mean_loss))\n",
    "    return mean_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total parameters count 8160661\n",
      "Lodaer test\n",
      "2023-04-24 12:39:09,082 - > collate dict key \"image\" out of 4 keys\n",
      "2023-04-24 12:39:09,088 - >> collate/stack a list of tensors\n",
      "2023-04-24 12:39:09,094 - >> E: stack expects each tensor to be equal size, but got [1, 65, 2727, 6330] at entry 0 and [1, 65, 5454, 6330] at entry 1, shape [(1, 65, 2727, 6330), (1, 65, 5454, 6330)] in collate([tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]]]])\n",
      "Metadata\n",
      "\tspatial_shape: [  65 2727 6330]\n",
      "\tspace: RAS\n",
      "\toriginal_channel_dim: no_channel\n",
      "\tfilename_or_obj: /root/autodl-tmp/vesuvius-challenge-ink-detection/test/a/surface_volume.npy\n",
      "\taffine: tensor([[1., 0., 0., 0.],\n",
      "        [0., 1., 0., 0.],\n",
      "        [0., 0., 1., 0.],\n",
      "        [0., 0., 0., 1.]], dtype=torch.float64)\n",
      "\n",
      "Applied operations\n",
      "[]\n",
      "Is batch?: False, tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]]]])\n",
      "Metadata\n",
      "\tspatial_shape: [  65 5454 6330]\n",
      "\tspace: RAS\n",
      "\toriginal_channel_dim: no_channel\n",
      "\tfilename_or_obj: /root/autodl-tmp/vesuvius-challenge-ink-detection/test/b/surface_volume.npy\n",
      "\taffine: tensor([[1., 0., 0., 0.],\n",
      "        [0., 1., 0., 0.],\n",
      "        [0., 0., 1., 0.],\n",
      "        [0., 0., 0., 1.]], dtype=torch.float64)\n",
      "\n",
      "Applied operations\n",
      "[]\n",
      "Is batch?: False])\n",
      "2023-04-24 12:39:09,095 - > collate dict key \"label\" out of 4 keys\n",
      "2023-04-24 12:39:09,100 - >> collate/stack a list of tensors\n",
      "2023-04-24 12:39:09,108 - >> E: stack expects each tensor to be equal size, but got [1, 16, 2727, 6330] at entry 0 and [1, 16, 5454, 6330] at entry 1, shape [(1, 16, 2727, 6330), (1, 16, 5454, 6330)] in collate([tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]]]])\n",
      "Metadata\n",
      "\tspatial_shape: [2727 6330]\n",
      "\tspace: RAS\n",
      "\toriginal_channel_dim: no_channel\n",
      "\tfilename_or_obj: /root/autodl-tmp/vesuvius-challenge-ink-detection/test/a/mask.npy\n",
      "\taffine: tensor([[1., 0., 0., 0.],\n",
      "        [0., 1., 0., 0.],\n",
      "        [0., 0., 1., 0.],\n",
      "        [0., 0., 0., 1.]], dtype=torch.float64)\n",
      "\n",
      "Applied operations\n",
      "[]\n",
      "Is batch?: False, tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "         [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]]]])\n",
      "Metadata\n",
      "\tspatial_shape: [5454 6330]\n",
      "\tspace: RAS\n",
      "\toriginal_channel_dim: no_channel\n",
      "\tfilename_or_obj: /root/autodl-tmp/vesuvius-challenge-ink-detection/test/b/mask.npy\n",
      "\taffine: tensor([[1., 0., 0., 0.],\n",
      "        [0., 1., 0., 0.],\n",
      "        [0., 0., 1., 0.],\n",
      "        [0., 0., 0., 1.]], dtype=torch.float64)\n",
      "\n",
      "Applied operations\n",
      "[]\n",
      "Is batch?: False])\n",
      "2023-04-24 12:39:09,109 - > collate dict key \"image_meta_dict\" out of 4 keys\n",
      "2023-04-24 12:39:09,111 - >> collate dict key \"spatial_shape\" out of 5 keys\n",
      "2023-04-24 12:39:09,111 - >>> collate/stack a list of numpy arrays\n",
      "2023-04-24 12:39:09,113 - >>> collate/stack a list of tensors\n",
      "2023-04-24 12:39:09,113 - >> collate dict key \"space\" out of 5 keys\n",
      "2023-04-24 12:39:09,114 - >> collate dict key \"original_channel_dim\" out of 5 keys\n",
      "2023-04-24 12:39:09,114 - >> collate dict key \"filename_or_obj\" out of 5 keys\n",
      "2023-04-24 12:39:09,115 - >> collate dict key \"affine\" out of 5 keys\n",
      "2023-04-24 12:39:09,118 - >>> collate/stack a list of tensors\n",
      "2023-04-24 12:39:09,119 - > collate dict key \"label_meta_dict\" out of 4 keys\n",
      "2023-04-24 12:39:09,122 - >> collate dict key \"spatial_shape\" out of 5 keys\n",
      "2023-04-24 12:39:09,123 - >>> collate/stack a list of numpy arrays\n",
      "2023-04-24 12:39:09,124 - >>> collate/stack a list of tensors\n",
      "2023-04-24 12:39:09,125 - >> collate dict key \"space\" out of 5 keys\n",
      "2023-04-24 12:39:09,126 - >> collate dict key \"original_channel_dim\" out of 5 keys\n",
      "2023-04-24 12:39:09,126 - >> collate dict key \"filename_or_obj\" out of 5 keys\n",
      "2023-04-24 12:39:09,127 - >> collate dict key \"affine\" out of 5 keys\n",
      "2023-04-24 12:39:09,129 - >>> collate/stack a list of tensors\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "stack expects each tensor to be equal size, but got [1, 65, 2727, 6330] at entry 0 and [1, 65, 5454, 6330] at entry 1\nCollate error on the key 'image' of dictionary data.\n\nMONAI hint: if your transforms intentionally create images of different shapes, creating your `DataLoader` with `collate_fn=pad_list_data_collate` might solve this problem (check its documentation).",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m~/miniconda3/lib/python3.8/site-packages/monai/data/utils.py\u001b[0m in \u001b[0;36mlist_data_collate\u001b[0;34m(batch)\u001b[0m\n\u001b[1;32m    461\u001b[0m                 \u001b[0mdata_for_batch\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0md\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0md\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 462\u001b[0;31m                 \u001b[0mret\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcollate_meta_tensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata_for_batch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    463\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.8/site-packages/monai/data/utils.py\u001b[0m in \u001b[0;36mcollate_meta_tensor\u001b[0;34m(batch)\u001b[0m\n\u001b[1;32m    426\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0melem_0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mMetaObj\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 427\u001b[0;31m         \u001b[0mcollated\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdefault_collate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    428\u001b[0m         \u001b[0mcollated\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmeta\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdefault_collate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmeta\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mTraceKeys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mNONE\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.8/site-packages/torch/utils/data/_utils/collate.py\u001b[0m in \u001b[0;36mdefault_collate\u001b[0;34m(batch)\u001b[0m\n\u001b[1;32m     55\u001b[0m             \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0melem\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnew\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstorage\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 56\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     57\u001b[0m     \u001b[0;32melif\u001b[0m \u001b[0melem_type\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__module__\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'numpy'\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0melem_type\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__name__\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0;34m'str_'\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.8/site-packages/monai/data/meta_tensor.py\u001b[0m in \u001b[0;36m__torch_function__\u001b[0;34m(cls, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m    248\u001b[0m             \u001b[0mkwargs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 249\u001b[0;31m         \u001b[0mret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0msuper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__torch_function__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtypes\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    250\u001b[0m         \u001b[0;31m# if `out` has been used as argument, metadata is not copied, nothing to do.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.8/site-packages/torch/_tensor.py\u001b[0m in \u001b[0;36m__torch_function__\u001b[0;34m(cls, func, types, args, kwargs)\u001b[0m\n\u001b[1;32m   1050\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0m_C\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mDisableTorchFunction\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1051\u001b[0;31m             \u001b[0mret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1052\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mfunc\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mget_default_nowrap_functions\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: stack expects each tensor to be equal size, but got [1, 65, 2727, 6330] at entry 0 and [1, 65, 5454, 6330] at entry 1",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_155287/335093777.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     85\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     86\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"Lodaer test\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 87\u001b[0;31m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mloader\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     88\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'image'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     89\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'label'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.8/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    519\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_sampler_iter\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    520\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_reset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 521\u001b[0;31m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_next_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    522\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_num_yielded\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    523\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dataset_kind\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0m_DatasetKind\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mIterable\u001b[0m \u001b[0;32mand\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.8/site-packages/torch/utils/data/dataloader.py\u001b[0m in \u001b[0;36m_next_data\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    559\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_next_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    560\u001b[0m         \u001b[0mindex\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_next_index\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 561\u001b[0;31m         \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dataset_fetcher\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfetch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# may raise StopIteration\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    562\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_pin_memory\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    563\u001b[0m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpin_memory\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/miniconda3/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py\u001b[0m in \u001b[0;36mfetch\u001b[0;34m(self, possibly_batched_index)\u001b[0m\n\u001b[1;32m     50\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     51\u001b[0m             \u001b[0mdata\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdataset\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mpossibly_batched_index\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 52\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcollate_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/miniconda3/lib/python3.8/site-packages/monai/data/utils.py\u001b[0m in \u001b[0;36mlist_data_collate\u001b[0;34m(batch)\u001b[0m\n\u001b[1;32m    475\u001b[0m             )\n\u001b[1;32m    476\u001b[0m         \u001b[0m_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdev_collate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 477\u001b[0;31m         \u001b[0;32mraise\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mre_str\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mre\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    478\u001b[0m     \u001b[0;32mexcept\u001b[0m \u001b[0mTypeError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mre\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    479\u001b[0m         \u001b[0mre_str\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mre\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: stack expects each tensor to be equal size, but got [1, 65, 2727, 6330] at entry 0 and [1, 65, 5454, 6330] at entry 1\nCollate error on the key 'image' of dictionary data.\n\nMONAI hint: if your transforms intentionally create images of different shapes, creating your `DataLoader` with `collate_fn=pad_list_data_collate` might solve this problem (check its documentation)."
     ]
    }
   ],
   "source": [
    "np.set_printoptions(formatter={\"float\": \"{: 0.3f}\".format}, suppress=True)\n",
    "\n",
    "torch.cuda.set_device(0)\n",
    "torch.backends.cudnn.benchmark = True\n",
    "args.test_mode = True\n",
    "loader = get_loader(args)\n",
    "\n",
    "inf_size = [args.roi_x, args.roi_y, args.roi_z]\n",
    "\n",
    "pretrained_dir = args.pretrained_dir\n",
    "if args.model_mode == \"3dswin\":\n",
    "    model = MyModel(img_size=(args.roi_x,args.roi_y,args.roi_y))\n",
    "elif args.model_mode == \"2dswin\":\n",
    "    model = MyModel2d(img_size=(args.roi_x,args.roi_y))\n",
    "elif args.model_mode == \"3dunet\":\n",
    "    model = MyModel3dunet()\n",
    "elif args.model_mode == \"2dfunet\":\n",
    "    model = MyFlexibleUNet2d(args)\n",
    "elif args.model_mode == \"2dfunetlstm\":\n",
    "    model = MyFlexibleUNet2dLSTM(args)\n",
    "elif args.model_mode == \"3dunet++\":\n",
    "    model = MyBasicUNetPlusPlus(args)\n",
    "else:\n",
    "    raise ValueError(\"model mode error\")\n",
    "\n",
    "\n",
    "model_dict = torch.load(os.path.join(pretrained_dir, args.pretrained_model_name))[\"state_dict\"]\n",
    "if args.model_mode in [\"2dswin\", \"3dunet\", \"2dfunet\", \"2dfunetlstm\", \"3dunet++\"]:\n",
    "    model.load_state_dict(model_dict)\n",
    "elif args.model_mode == \"3dswin\":\n",
    "    model.load_swin_ckpt(model_dict)\n",
    "else:\n",
    "    raise ValueError(\"model mode error\")\n",
    "\n",
    "if args.loss_mode == 'focalLoss':\n",
    "    loss = FocalLoss(weight=[10.0])\n",
    "elif args.loss_mode == 'squared_dice':\n",
    "    loss = DiceCELoss(squared_pred=True, smooth_nr=args.smooth_nr, smooth_dr=args.smooth_dr)\n",
    "elif args.loss_mode == 'DiceCELoss':\n",
    "    loss = DiceCELoss(include_background=True, sigmoid=True, ce_weight=torch.Tensor([ 10])) # Normally\n",
    "elif args.loss_mode == 'custom':\n",
    "    loss = CustomWeightedDiceCELoss(ink_weight=3.0, weight=args.loss_weight)\n",
    "    \n",
    "\n",
    "post_label = AsDiscrete(to_onehot=args.out_channels)\n",
    "post_pred = AsDiscrete(argmax=True, to_onehot=args.out_channels)\n",
    "dice_acc = DiceMetric(include_background=False, reduction=MetricReduction.MEAN, get_not_nans=True)\n",
    "miou_acc = MeanIoU(include_background=False, reduction=MetricReduction.MEAN, get_not_nans=True)\n",
    "# f_beta_acc = FBetaScore()\n",
    "\n",
    "if args.model_mode in [\"3dswin\", \"3dunet\", \"3dunet++\"]:\n",
    "    model_inferer = partial(\n",
    "        sliding_window_inference,\n",
    "        roi_size = (args.roi_x,args.roi_y,args.roi_z),\n",
    "        sw_batch_size = 8,\n",
    "        predictor = model,\n",
    "        overlap = 0.5,\n",
    "        progress = True,\n",
    "        padding_mode = \"reflect\", \n",
    "        device = \"cpu\", \n",
    "        sw_device = \"cuda\"\n",
    "    )\n",
    "elif args.model_mode in [\"2dswin\", \"2dfunet\", \"2dfunetlstm\"]:\n",
    "    model_inferer = partial(\n",
    "        sliding_window_inference,\n",
    "        roi_size = (args.roi_x,args.roi_y),\n",
    "        sw_batch_size = 8,\n",
    "        predictor = model,\n",
    "        overlap = 0.5,\n",
    "        progress = True,\n",
    "        padding_mode = \"reflect\", \n",
    "        device = \"cpu\", \n",
    "        sw_device = \"cuda\"\n",
    "    )     \n",
    "else:\n",
    "    raise ValueError(\"model mode error\")\n",
    "    \n",
    "pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "print(\"Total parameters count\", pytorch_total_params)\n",
    "\n",
    "best_acc = 0\n",
    "start_epoch = 0\n",
    "\n",
    "model.cuda(0)\n",
    "\n",
    "print(\"Lodaer test\")\n",
    "for i in loader:\n",
    "    print(i['image'].shape)\n",
    "    print(i['label'].shape)\n",
    "    print(torch.unique(i['image']))\n",
    "    print(torch.unique(i['label']))\n",
    "    break\n",
    "print(\"Pass Test\")\n",
    "print(args)\n",
    "test(model_inferer, loader, loss, args)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}