1. Add customDatasetLoader 2. Add predict function in test.py

Author: tylin7111095022

UNETR/BTCV/config.py

@@ -0,0 +1,3 @@
+NIFTI_DATA_ROOT = 'data/images' # nifti image directory
+NIFTI_LABEL_ROOT = 'data/labels' # nifti label directory
+PREDICT_DATA_ROOT = 'data/predict' # predict image directory

UNETR/BTCV/dataset/customDataset.py

@@ -0,0 +1,120 @@
+import os
+from torch.utils.data import DataLoader
+from monai.data import Dataset
+import monai.transforms as transforms
+import torch
+
+from config import NIFTI_DATA_ROOT, NIFTI_LABEL_ROOT, PREDICT_DATA_ROOT
+
+def _get_collate_fn(isTrain:bool):
+    def collate_fn(batch):
+        '''collate function'''
+        images = []
+        labels = []
+        if isTrain:
+            for p in batch: # [ {"image": (C, H, W ,D), "label": (C, H, W ,D)} , ...]
+                for i in range(len(p)): # list, RandCropByPosNegLabeld will produce multiple samples
+                    images.append(p[i]['image'])
+                    labels.append(p[i]['label'])
+        else:
+            for p in batch:
+                images.append(p['image'])
+                labels.append(p['label'])
+
+        images = torch.stack(images, dim=0)
+        labels = torch.stack(labels, dim=0)
+        
+        return [torch.Tensor(images), torch.Tensor(labels)]
+    
+    return collate_fn
+
+def getDatasetLoader(args):
+    dataName = [d for d in os.listdir(NIFTI_LABEL_ROOT)]
+    dataDicts = [{"image": f"{os.path.join(NIFTI_DATA_ROOT, d)}", "label": f"{os.path.join(NIFTI_LABEL_ROOT, d)}"} for d in dataName]
+    trainDicts, valDicts = _splitList(dataDicts)
+
+    train_transform = transforms.Compose(
+        [
+            transforms.LoadImaged(keys=["image", "label"]),
+            transforms.EnsureChannelFirstd(keys=["image", "label"]),
+            transforms.Orientationd(keys=["image", "label"], axcodes="RAS"),
+            transforms.Spacingd(
+                keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z), mode=("bilinear", "nearest")
+            ),
+            transforms.ScaleIntensityRanged(
+                keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
+            ),
+            transforms.CropForegroundd(keys=["image", "label"], source_key="image", allow_smaller=True),
+            transforms.RandCropByPosNegLabeld(
+                keys=["image", "label"],
+                label_key="label",
+                spatial_size=(args.roi_x, args.roi_y, args.roi_z),
+                pos=1,
+                neg=1,
+                num_samples=4,
+                image_key="image",
+                image_threshold=0,
+            ),
+            transforms.RandFlipd(keys=["image", "label"], prob=args.RandFlipd_prob, spatial_axis=0),
+            transforms.RandFlipd(keys=["image", "label"], prob=args.RandFlipd_prob, spatial_axis=1),
+            transforms.RandFlipd(keys=["image", "label"], prob=args.RandFlipd_prob, spatial_axis=2),
+            transforms.RandRotate90d(keys=["image", "label"], prob=args.RandRotate90d_prob, max_k=3),
+            transforms.RandScaleIntensityd(keys="image", factors=0.1, prob=args.RandScaleIntensityd_prob),
+            transforms.RandShiftIntensityd(keys="image", offsets=0.1, prob=args.RandShiftIntensityd_prob),
+            transforms.ToTensord(keys=["image", "label"]),
+        ]
+    )
+
+    val_transform = transforms.Compose(
+            [
+                transforms.LoadImaged(keys=["image", "label"]),
+                transforms.EnsureChannelFirstd(keys=["image", "label"]),
+                transforms.Orientationd(keys=["image", "label"], axcodes="RAS"),
+                transforms.Spacingd(
+                    keys=["image", "label"], pixdim=(args.space_x, args.space_y, args.space_z), mode=("bilinear", "nearest")
+                ),
+                transforms.ScaleIntensityRanged(
+                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
+                ),
+                transforms.CropForegroundd(keys=["image", "label"], source_key="image", allow_smaller=True),
+                transforms.ToTensord(keys=["image", "label"]),
+            ]
+        )
+    
+    trainDataset = Dataset(data=trainDicts, transform=train_transform)
+    valDataset = Dataset(data=valDicts, transform=val_transform)
+    trainLoader = DataLoader(trainDataset,batch_size=args.batch_size,shuffle=True,num_workers=args.workers, collate_fn=_get_collate_fn(isTrain=True))
+    valLoader = DataLoader(valDataset,batch_size=args.batch_size,shuffle=False,num_workers=args.workers, collate_fn=_get_collate_fn(isTrain=False))
+    loader = [trainLoader, valLoader]
+
+    return loader
+
+def _splitList(l, trainRatio:float = 0.8):
+    totalNum = len(l)
+    splitIdx = int(totalNum * trainRatio)
+
+    return l[:splitIdx], l[splitIdx :]
+
+def getPredictLoader(args):
+    dataName = [d for d in os.listdir(PREDICT_DATA_ROOT)]
+    dataDicts = [{"image": f"{os.path.join(PREDICT_DATA_ROOT, d)}" } for d in dataName]
+
+    preTransform = transforms.Compose(
+            [
+                transforms.LoadImaged(keys=["image"]),
+                transforms.EnsureChannelFirstd(keys=["image"]),
+                transforms.Orientationd(keys=["image"], axcodes="RAS"),
+                transforms.Spacingd(
+                    keys=["image"], pixdim=(args.space_x, args.space_y, args.space_z), mode=("bilinear")
+                ),
+                transforms.ScaleIntensityRanged(
+                    keys=["image"], a_min=args.a_min, a_max=args.a_max, b_min=args.b_min, b_max=args.b_max, clip=True
+                ),
+                transforms.CropForegroundd(keys=["image"], source_key="image", allow_smaller=True),
+                # transforms.ToTensord(keys=["image"],track_meta=True), # This transformation will transform MetaTensor to Tensor
+            ]
+        )
+    valDataset = Dataset(data=dataDicts, transform=preTransform)
+    valLoader = DataLoader(valDataset,batch_size=args.batch_size,shuffle=False,num_workers=args.workers)
+
+    return valLoader, preTransform

UNETR/BTCV/main.py

@@ -12,53 +12,55 @@
 import argparse
 import os
 from functools import partial
-
 import numpy as np
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 import torch.nn.parallel
 import torch.utils.data.distributed
-from networks.unetr import UNETR
-from optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR
-from trainer import run_training
-from utils.data_utils import get_loader
 
 from monai.inferers import sliding_window_inference
 from monai.losses import DiceCELoss, DiceLoss
 from monai.metrics import DiceMetric
 from monai.transforms import Activations, AsDiscrete, Compose
 from monai.utils.enums import MetricReduction
 
+from networks.unetr import UNETR
+from optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR
+from trainer import run_training
+from utils.data_utils import get_loader
+from dataset.customDataset import getDatasetLoader
+
 parser = argparse.ArgumentParser(description="UNETR segmentation pipeline")
 parser.add_argument("--checkpoint", default=None, help="start training from saved checkpoint")
 parser.add_argument("--logdir", default="test", type=str, help="directory to save the tensorboard logs")
 parser.add_argument(
     "--pretrained_dir", default="./pretrained_models/", type=str, help="pretrained checkpoint directory"
 )
-parser.add_argument("--data_dir", default="/dataset/dataset0/", type=str, help="dataset directory")
+parser.add_argument("--btcv", action="store_true", help="Use BTCV dataset")
+parser.add_argument("--data_dir", default="./dataset/dataset0/", type=str, help="dataset directory")
 parser.add_argument("--json_list", default="dataset_0.json", type=str, help="dataset json file")
 parser.add_argument(
     "--pretrained_model_name", default="UNETR_model_best_acc.pth", type=str, help="pretrained model name"
 )
-parser.add_argument("--save_checkpoint", action="store_true", help="save checkpoint during training")
-parser.add_argument("--max_epochs", default=5000, type=int, help="max number of training epochs")
+parser.add_argument("--save_checkpoint", action="store_true", default=True, help="save checkpoint during training")
+parser.add_argument("--max_epochs", default=100, type=int, help="max number of training epochs")
 parser.add_argument("--batch_size", default=1, type=int, help="number of batch size")
 parser.add_argument("--sw_batch_size", default=1, type=int, help="number of sliding window batch size")
 parser.add_argument("--optim_lr", default=1e-4, type=float, help="optimization learning rate")
 parser.add_argument("--optim_name", default="adamw", type=str, help="optimization algorithm")
 parser.add_argument("--reg_weight", default=1e-5, type=float, help="regularization weight")
 parser.add_argument("--momentum", default=0.99, type=float, help="momentum")
 parser.add_argument("--noamp", action="store_true", help="do NOT use amp for training")
-parser.add_argument("--val_every", default=100, type=int, help="validation frequency")
+parser.add_argument("--val_every", default=10, type=int, help="validation frequency")
 parser.add_argument("--distributed", action="store_true", help="start distributed training")
 parser.add_argument("--world_size", default=1, type=int, help="number of nodes for distributed training")
 parser.add_argument("--rank", default=0, type=int, help="node rank for distributed training")
 parser.add_argument("--dist-url", default="tcp://127.0.0.1:23456", type=str, help="distributed url")
 parser.add_argument("--dist-backend", default="nccl", type=str, help="distributed backend")
 parser.add_argument("--workers", default=8, type=int, help="number of workers")
 parser.add_argument("--model_name", default="unetr", type=str, help="model name")
-parser.add_argument("--pos_embed", default="perceptron", type=str, help="type of position embedding")
+parser.add_argument("--pos_embed", default="learnable", type=str, help="type of position embedding")
 parser.add_argument("--norm_name", default="instance", type=str, help="normalization layer type in decoder")
 parser.add_argument("--num_heads", default=12, type=int, help="number of attention heads in ViT encoder")
 parser.add_argument("--mlp_dim", default=3072, type=int, help="mlp dimention in ViT encoder")
@@ -73,12 +75,12 @@
 parser.add_argument("--a_max", default=250.0, type=float, help="a_max in ScaleIntensityRanged")
 parser.add_argument("--b_min", default=0.0, type=float, help="b_min in ScaleIntensityRanged")
 parser.add_argument("--b_max", default=1.0, type=float, help="b_max in ScaleIntensityRanged")
-parser.add_argument("--space_x", default=1.5, type=float, help="spacing in x direction")
-parser.add_argument("--space_y", default=1.5, type=float, help="spacing in y direction")
-parser.add_argument("--space_z", default=2.0, type=float, help="spacing in z direction")
-parser.add_argument("--roi_x", default=96, type=int, help="roi size in x direction")
-parser.add_argument("--roi_y", default=96, type=int, help="roi size in y direction")
-parser.add_argument("--roi_z", default=96, type=int, help="roi size in z direction")
+parser.add_argument("--space_x", default=1.0, type=float, help="spacing in x direction")
+parser.add_argument("--space_y", default=1.0, type=float, help="spacing in y direction")
+parser.add_argument("--space_z", default=1.0, type=float, help="spacing in z direction")
+parser.add_argument("--roi_x", default=64, type=int, help="roi size in x direction")
+parser.add_argument("--roi_y", default=64, type=int, help="roi size in y direction")
+parser.add_argument("--roi_z", default=64, type=int, help="roi size in z direction")
 parser.add_argument("--dropout_rate", default=0.0, type=float, help="dropout rate")
 parser.add_argument("--RandFlipd_prob", default=0.2, type=float, help="RandFlipd aug probability")
 parser.add_argument("--RandRotate90d_prob", default=0.2, type=float, help="RandRotate90d aug probability")
@@ -102,10 +104,9 @@ def main():
         print("Found total gpus", args.ngpus_per_node)
         args.world_size = args.ngpus_per_node * args.world_size
         mp.spawn(main_worker, nprocs=args.ngpus_per_node, args=(args,))
-    else:
+    else:   
         main_worker(gpu=0, args=args)
 
-
 def main_worker(gpu, args):
     if args.distributed:
         torch.multiprocessing.set_start_method("fork", force=True)
@@ -119,7 +120,8 @@ def main_worker(gpu, args):
     torch.cuda.set_device(args.gpu)
     torch.backends.cudnn.benchmark = True
     args.test_mode = False
-    loader = get_loader(args)
+    loader = get_loader(args) if args.btcv else getDatasetLoader(args)
+    
     print(args.rank, " gpu", args.gpu)
     if args.rank == 0:
         print("Batch size is:", args.batch_size, "epochs", args.max_epochs)
@@ -157,8 +159,8 @@ def main_worker(gpu, args):
     dice_loss = DiceCELoss(
         to_onehot_y=True, softmax=True, squared_pred=True, smooth_nr=args.smooth_nr, smooth_dr=args.smooth_dr
     )
-    post_label = AsDiscrete(to_onehot=True, n_classes=args.out_channels)
-    post_pred = AsDiscrete(argmax=True, to_onehot=True, n_classes=args.out_channels)
+    post_label = AsDiscrete(to_onehot=args.out_channels)
+    post_pred = AsDiscrete(argmax=True, to_onehot=args.out_channels)
     dice_acc = DiceMetric(include_background=True, reduction=MetricReduction.MEAN, get_not_nans=True)
     model_inferer = partial(
         sliding_window_inference,
@@ -235,6 +237,5 @@ def main_worker(gpu, args):
     )
     return accuracy
 
-
 if __name__ == "__main__":
     main()

UNETR/BTCV/networks/unetr.py

@@ -73,7 +73,7 @@ def __init__(
         if hidden_size % num_heads != 0:
             raise AssertionError("hidden size should be divisible by num_heads.")
 
-        if pos_embed not in ["conv", "perceptron"]:
+        if pos_embed not in ['sincos', 'learnable', 'none']:
             raise KeyError(f"Position embedding layer of type {pos_embed} is not supported.")
 
         self.num_layers = 12
@@ -93,7 +93,7 @@ def __init__(
             mlp_dim=mlp_dim,
             num_layers=self.num_layers,
             num_heads=num_heads,
-            pos_embed=pos_embed,
+            pos_embed_type=pos_embed,
             classification=self.classification,
             dropout_rate=dropout_rate,
         )

UNETR/BTCV/requirements.txt

@@ -1,6 +1,6 @@
-torch==1.9.1
-monai==0.7.0
-nibabel==3.1.1
-tqdm==4.59.0
-einops==0.3.0
-tensorboardX==2.1
+monai==1.5.0
+numpy==2.3.2
+opencv_python
+simpleitk==2.5.2
+tensorboardx==2.6.4
+torch