swin unetr的3D语义分割

基于monai库。其实我不是很喜欢这种，可扩展性太差了，除非说你想快速在自己的数据集上出结果。但是它的transform可以对3d医学图像增强操作，比torch的transform强一点，因为它的数据增强输入是（x,y,z）h,w,d格式的，我还没有试过单独用它的transform来结合torch训练。

前提

pip install monai

目录结构

train.py

from nets.swin_model import GetSwinUnetr
import torch
from utils.dataloaderd import GetDataLoader
from monai.losses import DiceCELoss
from monai.metrics import DiceMetric
from monai.inferers import sliding_window_inference
from tqdm import tqdm
import os
import matplotlib.pyplot as plt
from torch import nn

os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"


from monai.data import (
    ThreadDataLoader,
    CacheDataset,
    load_decathlon_datalist,
    decollate_batch,
    set_track_meta,
)

from monai.transforms import (AsDiscrete,)


def validation(epoch_iterator_val):
    model.eval()
    with torch.no_grad():
        for batch in epoch_iterator_val:
            val_inputs, val_labels = (batch["image"].cuda(), batch["label"].cuda())
            with torch.cuda.amp.autocast():
                val_outputs = sliding_window_inference(val_inputs, (cutSize[0], cutSize[1], cutSize[2]), 1, model)
            val_labels_list = decollate_batch(val_labels)
            val_labels_convert = [post_label(val_label_tensor) for val_label_tensor in val_labels_list]
            val_outputs_list = decollate_batch(val_outputs)
            val_output_convert = [post_pred(val_pred_tensor) for val_pred_tensor in val_outputs_list]
            dice_metric(y_pred=val_output_convert, y=val_labels_convert)
            epoch_iterator_val.set_description("Validate (%d / %d Steps)" % (global_step, 10.0))
        mean_dice_val = dice_metric.aggregate().item()
        dice_metric.reset()
    return mean_dice_val

def train(global_step, train_loader, dice_val_best, global_step_best):
    model.train()
    epoch_loss = 0
    step = 0
    epoch_iterator = tqdm(train_loader, desc="Training (X / X Steps) (loss=X.X)", dynamic_ncols=True)
    for step, batch in enumerate(epoch_iterator):
        step  = 1
        x, y = (batch["image"].cuda(), batch["label"].cuda())
        # with torch.cuda.amp.autocast():
        logit_map = model(x)
        loss = loss_function(logit_map, y)
        # scaler.scale(loss).backward()
        epoch_loss  = loss.item()
        # scaler.unscale_(optimizer)
        # scaler.step(optimizer)
        # scaler.update()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        epoch_iterator.set_description(f"Training ({global_step} / {max_iterations} Steps) (loss={loss:2.5f})")
        if (global_step % eval_num == 0 and global_step != 0) or global_step == max_iterations:
            epoch_iterator_val = tqdm(val_loader, desc="Validate (X / X Steps) (dice=X.X)", dynamic_ncols=True)
            dice_val = validation(epoch_iterator_val)
            epoch_loss /= step
            epoch_loss_values.append(epoch_loss)
            metric_values.append(dice_val)
            if dice_val > dice_val_best:
                dice_val_best = dice_val
                global_step_best = global_step
                torch.save(model.state_dict(), os.path.join(root_dir, "best_metric_model.pth"))
                print(
                    "Model Was Saved ! Current Best Avg. Dice: {} Current Avg. Dice: {}".format(dice_val_best, dice_val)
                )
            else:
                print(
                    "Model Was Not Saved ! Current Best Avg. Dice: {} Current Avg. Dice: {}".format(
                        dice_val_best, dice_val
                    )
                )
        global_step  = 1

    return global_step, dice_val_best, global_step_best


if __name__ == '__main__':
    data_dir = "/home/data/hablee_data_dir/aorta2023/"
    root_dir = "/home/hbli/pythonFiles3/c_swinUnetr_torch/"
    split_json = "dataset_0.json"
    cutSize = [64,64,32]
    numClasses = 3
    batchSize = 2
    train_loader,val_loader,val_ds = GetDataLoader(
        data_dir=data_dir,split_json=split_json,
        cut_size=cutSize,batch_size=batchSize)

    model = GetSwinUnetr(cut_size=cutSize,num_classes=numClasses)
    weight = torch.load("./best_metric_model.pth")
    model.load_state_dict(weight)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    gpu_counts = torch.cuda.device_count()
    if gpu_counts>1:
        batch_size = batchSize * gpu_counts
        model = nn.DataParallel(model)
    
    model.to(device=device)
    
    # torch.backends.cudnn.benchmark = True
    loss_function = DiceCELoss(to_onehot_y=True, softmax=True)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)
    # scaler = torch.cuda.amp.GradScaler()

    max_iterations = 30000
    eval_num = 500
    post_label = AsDiscrete(to_onehot=numClasses)
    post_pred = AsDiscrete(argmax=True, to_onehot=numClasses)
    dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
    global_step = 0
    dice_val_best = 0.0
    global_step_best = 0
    epoch_loss_values = []
    metric_values = []
    while global_step < max_iterations:
        global_step, dice_val_best, global_step_best = train(global_step, train_loader, dice_val_best, global_step_best)

    print(f"train completed, best_metric: {dice_val_best:.4f} " f"at iteration: {global_step_best}")

    plt.figure("train", (12, 6))
    plt.subplot(1, 2, 1)
    plt.title("Iteration Average Loss")
    x = [eval_num * (i   1) for i in range(len(epoch_loss_values))]
    y = epoch_loss_values
    plt.xlabel("Iteration")
    plt.plot(x, y)
    plt.subplot(1, 2, 2)
    plt.title("Val Mean Dice")
    x = [eval_num * (i   1) for i in range(len(metric_values))]
    y = metric_values
    plt.xlabel("Iteration")
    plt.plot(x, y)
    plt.show()
    plt.savefig('loss.jpg')


    case_num = 4
    slice_map = {
    "72.nii.gz": 50,
    "97.nii.gz": 50,
    "82.nii.gz": 50,
    "153.nii.gz": 50,
    "54.nii.gz": 50,
    "104.nii.gz": 50
    }
    model.load_state_dict(torch.load(os.path.join(root_dir, "best_metric_model.pth")))
    model.eval()
    with torch.no_grad():
        img_name = os.path.split(val_ds[case_num]["image"].meta["filename_or_obj"])[1]
        img = val_ds[case_num]["image"]
        label = val_ds[case_num]["label"]
        val_inputs = torch.unsqueeze(img, 1).cuda()
        val_labels = torch.unsqueeze(label, 1).cuda()
        val_outputs = sliding_window_inference(val_inputs, (cutSize[0], cutSize[1], cutSize[2]), 1, model, overlap=0.8)
        plt.figure("check", (18, 6))
        plt.subplot(1, 3, 1)
        plt.title("image")
        plt.imshow(val_inputs.cpu().numpy()[0, 0, :, :, slice_map[img_name]], cmap="gray")
        plt.subplot(1, 3, 2)
        plt.title("label")
        plt.imshow(val_labels.cpu().numpy()[0, 0, :, :, slice_map[img_name]])
        plt.subplot(1, 3, 3)
        plt.title("output")
        plt.imshow(torch.argmax(val_outputs, dim=1).detach().cpu()[0, :, :, slice_map[img_name]])
        plt.show()

    plt.savefig("result.jpg")
    

nets/swin_model.py

# import os
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
from monai.networks.nets.swin_unetr import SwinUNETR

# import torch
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def GetSwinUnetr(cut_size,num_classes):
    model = SwinUNETR(
        img_size=(cut_size[0], cut_size[1], cut_size[2]),
        in_channels=1,
        out_channels=num_classes,
        feature_size=48,
        use_checkpoint=True,)

    return model

utils/dataloaderd.py

import os
import shutil
import tempfile

import matplotlib.pyplot as plt
from tqdm import tqdm
import torch

from monai.losses import DiceCELoss
from monai.inferers import sliding_window_inference
from monai.transforms import (
    AsDiscrete,
    Compose,
    CropForegroundd,
    LoadImaged,
    Orientationd,
    RandFlipd,
    RandCropByPosNegLabeld,
    RandShiftIntensityd,
    ScaleIntensityRanged,
    Spacingd,
    RandRotate90d,
    EnsureTyped,
)

from monai.metrics import DiceMetric
from monai.networks.nets import SwinUNETR
from monai.config import print_config

from monai.data import (
    DataLoader,
    CacheDataset,
    load_decathlon_datalist,
    decollate_batch,
    set_track_meta,
)


num_samples = 4
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")



val_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"], ensure_channel_first=True),
        ScaleIntensityRanged(keys=["image"], a_min=-175, a_max=250, b_min=0.0, b_max=1.0, clip=True),
        CropForegroundd(keys=["image", "label"], source_key="image"),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(1.5, 1.5, 2.0),
            mode=("bilinear", "nearest"),
        ),
        EnsureTyped(keys=["image", "label"], device=device, track_meta=True),
    ]
)


def GetDataLoader(data_dir,split_json,cut_size,batch_size):
    train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"], ensure_channel_first=True),
        ScaleIntensityRanged(
            keys=["image"],
            a_min=-175,
            a_max=250,
            b_min=0.0,
            b_max=1.0,
            clip=True,
        ),
        CropForegroundd(keys=["image", "label"], source_key="image"),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(1.5, 1.5, 2.0),
            mode=("bilinear", "nearest"),
        ),
        EnsureTyped(keys=["image", "label"], device=device, track_meta=False),
        RandCropByPosNegLabeld(
            keys=["image", "label"],
            label_key="label",
            spatial_size=(cut_size[0], cut_size[1], cut_size[2]),
            pos=1,
            neg=1,
            num_samples=num_samples,
            image_key="image",
            image_threshold=0,
        ),
        RandFlipd(
            keys=["image", "label"],
            spatial_axis=[0],
            prob=0.10,
        ),
        RandFlipd(
            keys=["image", "label"],
            spatial_axis=[1],
            prob=0.10,
        ),
        RandFlipd(
            keys=["image", "label"],
            spatial_axis=[2],
            prob=0.10,
        ),
        RandRotate90d(
            keys=["image", "label"],
            prob=0.10,
            max_k=3,
        ),
        RandShiftIntensityd(
            keys=["image"],
            offsets=0.10,
            prob=0.50,
        ),
    ]
    )

    datasets = data_dir   split_json
    datalist = load_decathlon_datalist(datasets, True, "training")
    val_files = load_decathlon_datalist(datasets, True, "validation")
    train_ds = CacheDataset(
        data=datalist,
        transform=train_transforms,
        cache_num=24,
        cache_rate=1.0,
        num_workers=0,
    )
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=False)
    val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_num=6, cache_rate=1.0, num_workers=0)
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=False)
    
    return train_loader,val_loader,val_ds


# if __name__ == '__main__':
#     data_dir = "/home/data/hablee_data_dir/aorta2023/"
#     split_json = "dataset_0.json"

#     datasets = data_dir   split_json
#     datalist = load_decathlon_datalist(datasets, True, "training")
#     val_files = load_decathlon_datalist(datasets, True, "validation")
#     train_ds = CacheDataset(
#         data=datalist,
#         transform=train_transforms,
#         cache_num=24,
#         cache_rate=1.0,
#         num_workers=4,
#     )
#     train_loader = DataLoader(train_ds, batch_size=1, shuffle=True, num_workers=4, pin_memory=True)
#     val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_num=6, cache_rate=1.0, num_workers=4)
#     val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)

    # slice_map = {
    # "72.nii.gz": 50,
    # "97.nii.gz": 50,
    # "82.nii.gz": 50,
    # "153.nii.gz": 50,
    # "54.nii.gz": 50,
    # "104.nii.gz": 50
    # }

    # case_num = 2
    # img_name = os.path.split(val_ds[case_num]["image"].meta["filename_or_obj"])[1]
    # img = val_ds[case_num]["image"]
    # label = val_ds[case_num]["label"]
    # img_shape = img.shape
    # label_shape = label.shape
    # print(f"image shape: {img_shape}, label shape: {label_shape}")
    # plt.figure("image", (18, 6))
    # plt.subplot(1, 2, 1)
    # plt.title("image")
    # plt.imshow(img[0, :, :, slice_map[img_name]].detach().cpu(), cmap="gray")
    # plt.subplot(1, 2, 2)
    # plt.title("label")
    # plt.imshow(label[0, :, :, slice_map[img_name]].detach().cpu())
    # plt.show()
    # plt.savefig("a.jpg")

就这几个文件就可以训练了。

数据

数据都是.nii.gz格式的

dataset_0.json

{
    "description": "zdm fjm",
    "labels": {
        "0": "background",
        "1": "zdm",
        "2": "fjm"
    },
    "licence": "hablee",
    "modality": {
        "0": "CT"
    },
    "name": "zdm fjm",
    "numTest": 5,
    "numTraining": 100,
    "reference": "None",
    "release": "2023/02/14",
    "tensorImageSize": "3D",
    "test":[
        "images/100.nii.gz",
        "images/102.nii.gz",
        "images/103.nii.gz",
        "images/104.nii.gz",
        "images/105.nii.gz"
    ],
    "training":[
            {
            "image": "images/72.nii.gz",
            "label": "labels/72.nii.gz"
            },
            {
            "image": "images/97.nii.gz",
            "label": "labels/97.nii.gz"
            },
            {
            "image": "images/82.nii.gz",
            "label": "labels/82.nii.gz"
            },
            {
            "image": "images/153.nii.gz",
            "label": "labels/153.nii.gz"
            },
            {
            "image": "images/54.nii.gz",
            "label": "labels/54.nii.gz"
            },
            {
            "image": "images/104.nii.gz",
            "label": "labels/104.nii.gz"
            },
            {
            "image": "images/110.nii.gz",
            "label": "labels/110.nii.gz"
            },
            {
            "image": "images/89.nii.gz",
            "label": "labels/89.nii.gz"
            },
            {
            "image": "images/95.nii.gz",
            "label": "labels/95.nii.gz"
            },
            {
            "image": "images/129.nii.gz",
            "label": "labels/129.nii.gz"
            },
            {
            "image": "images/118.nii.gz",
            "label": "labels/118.nii.gz"
            },
            {
            "image": "images/141.nii.gz",
            "label": "labels/141.nii.gz"
            },
            {
            "image": "images/134.nii.gz",
            "label": "labels/134.nii.gz"
            },
            {
            "image": "images/146.nii.gz",
            "label": "labels/146.nii.gz"
            },
            {
            "image": "images/100.nii.gz",
            "label": "labels/100.nii.gz"
            },
            {
            "image": "images/142.nii.gz",
            "label": "labels/142.nii.gz"
            },
            {
            "image": "images/85.nii.gz",
            "label": "labels/85.nii.gz"
            },
            {
            "image": "images/154.nii.gz",
            "label": "labels/154.nii.gz"
            },
            {
            "image": "images/103.nii.gz",
            "label": "labels/103.nii.gz"
            },
            {
            "image": "images/59.nii.gz",
            "label": "labels/59.nii.gz"
            },
            {
            "image": "images/84.nii.gz",
            "label": "labels/84.nii.gz"
            },
            {
            "image": "images/124.nii.gz",
            "label": "labels/124.nii.gz"
            },
            {
            "image": "images/125.nii.gz",
            "label": "labels/125.nii.gz"
            },
            {
            "image": "images/58.nii.gz",
            "label": "labels/58.nii.gz"
            },
            {
            "image": "images/68.nii.gz",
            "label": "labels/68.nii.gz"
            },
            {
            "image": "images/81.nii.gz",
            "label": "labels/81.nii.gz"
            },
            {
            "image": "images/115.nii.gz",
            "label": "labels/115.nii.gz"
            },
            {
            "image": "images/77.nii.gz",
            "label": "labels/77.nii.gz"
            },
            {
            "image": "images/127.nii.gz",
            "label": "labels/127.nii.gz"
            },
            {
            "image": "images/131.nii.gz",
            "label": "labels/131.nii.gz"
            },
            {
            "image": "images/147.nii.gz",
            "label": "labels/147.nii.gz"
            },
            {
            "image": "images/73.nii.gz",
            "label": "labels/73.nii.gz"
            },
            {
            "image": "images/102.nii.gz",
            "label": "labels/102.nii.gz"
            },
            {
            "image": "images/66.nii.gz",
            "label": "labels/66.nii.gz"
            },
            {
            "image": "images/67.nii.gz",
            "label": "labels/67.nii.gz"
            },
            {
            "image": "images/135.nii.gz",
            "label": "labels/135.nii.gz"
            },
            {
            "image": "images/149.nii.gz",
            "label": "labels/149.nii.gz"
            },
            {
            "image": "images/48.nii.gz",
            "label": "labels/48.nii.gz"
            },
            {
            "image": "images/83.nii.gz",
            "label": "labels/83.nii.gz"
            },
            {
            "image": "images/145.nii.gz",
            "label": "labels/145.nii.gz"
            },
            {
            "image": "images/45.nii.gz",
            "label": "labels/45.nii.gz"
            },
            {
            "image": "images/61.nii.gz",
            "label": "labels/61.nii.gz"
            },
            {
            "image": "images/122.nii.gz",
            "label": "labels/122.nii.gz"
            },
            {
            "image": "images/96.nii.gz",
            "label": "labels/96.nii.gz"
            },
            {
            "image": "images/60.nii.gz",
            "label": "labels/60.nii.gz"
            },
            {
            "image": "images/144.nii.gz",
            "label": "labels/144.nii.gz"
            },
            {
            "image": "images/91.nii.gz",
            "label": "labels/91.nii.gz"
            },
            {
            "image": "images/111.nii.gz",
            "label": "labels/111.nii.gz"
            },
            {
            "image": "images/114.nii.gz",
            "label": "labels/114.nii.gz"
            },
            {
            "image": "images/90.nii.gz",
            "label": "labels/90.nii.gz"
            },
            {
            "image": "images/52.nii.gz",
            "label": "labels/52.nii.gz"
            },
            {
            "image": "images/132.nii.gz",
            "label": "labels/132.nii.gz"
            },
            {
            "image": "images/107.nii.gz",
            "label": "labels/107.nii.gz"
            },
            {
            "image": "images/109.nii.gz",
            "label": "labels/109.nii.gz"
            },
            {
            "image": "images/139.nii.gz",
            "label": "labels/139.nii.gz"
            },
            {
            "image": "images/143.nii.gz",
            "label": "labels/143.nii.gz"
            },
            {
            "image": "images/119.nii.gz",
            "label": "labels/119.nii.gz"
            },
            {
            "image": "images/55.nii.gz",
            "label": "labels/55.nii.gz"
            },
            {
            "image": "images/80.nii.gz",
            "label": "labels/80.nii.gz"
            },
            {
            "image": "images/53.nii.gz",
            "label": "labels/53.nii.gz"
            },
            {
            "image": "images/120.nii.gz",
            "label": "labels/120.nii.gz"
            },
            {
            "image": "images/65.nii.gz",
            "label": "labels/65.nii.gz"
            },
            {
            "image": "images/88.nii.gz",
            "label": "labels/88.nii.gz"
            },
            {
            "image": "images/47.nii.gz",
            "label": "labels/47.nii.gz"
            },
            {
            "image": "images/57.nii.gz",
            "label": "labels/57.nii.gz"
            },
            {
            "image": "images/130.nii.gz",
            "label": "labels/130.nii.gz"
            },
            {
            "image": "images/108.nii.gz",
            "label": "labels/108.nii.gz"
            },
            {
            "image": "images/151.nii.gz",
            "label": "labels/151.nii.gz"
            },
            {
            "image": "images/113.nii.gz",
            "label": "labels/113.nii.gz"
            },
            {
            "image": "images/71.nii.gz",
            "label": "labels/71.nii.gz"
            },
            {
            "image": "images/46.nii.gz",
            "label": "labels/46.nii.gz"
            },
            {
            "image": "images/105.nii.gz",
            "label": "labels/105.nii.gz"
            },
            {
            "image": "images/148.nii.gz",
            "label": "labels/148.nii.gz"
            },
            {
            "image": "images/112.nii.gz",
            "label": "labels/112.nii.gz"
            },
            {
            "image": "images/106.nii.gz",
            "label": "labels/106.nii.gz"
            },
            {
            "image": "images/49.nii.gz",
            "label": "labels/49.nii.gz"
            },
            {
            "image": "images/140.nii.gz",
            "label": "labels/140.nii.gz"
            },
            {
            "image": "images/92.nii.gz",
            "label": "labels/92.nii.gz"
            },
            {
            "image": "images/137.nii.gz",
            "label": "labels/137.nii.gz"
            },
            {
            "image": "images/74.nii.gz",
            "label": "labels/74.nii.gz"
            },
            {
            "image": "images/62.nii.gz",
            "label": "labels/62.nii.gz"
            },
            {
            "image": "images/99.nii.gz",
            "label": "labels/99.nii.gz"
            },
            {
            "image": "images/150.nii.gz",
            "label": "labels/150.nii.gz"
            },
            {
            "image": "images/75.nii.gz",
            "label": "labels/75.nii.gz"
            },
            {
            "image": "images/98.nii.gz",
            "label": "labels/98.nii.gz"
            },
            {
            "image": "images/86.nii.gz",
            "label": "labels/86.nii.gz"
            },
            {
            "image": "images/50.nii.gz",
            "label": "labels/50.nii.gz"
            },
            {
            "image": "images/93.nii.gz",
            "label": "labels/93.nii.gz"
            },
            {
            "image": "images/138.nii.gz",
            "label": "labels/138.nii.gz"
            },
            {
            "image": "images/126.nii.gz",
            "label": "labels/126.nii.gz"
            },
            {
            "image": "images/69.nii.gz",
            "label": "labels/69.nii.gz"
            },
            {
            "image": "images/64.nii.gz",
            "label": "labels/64.nii.gz"
            },
            {
            "image": "images/136.nii.gz",
            "label": "labels/136.nii.gz"
            },
            {
            "image": "images/51.nii.gz",
            "label": "labels/51.nii.gz"
            },
            {
            "image": "images/70.nii.gz",
            "label": "labels/70.nii.gz"
            },
            {
            "image": "images/56.nii.gz",
            "label": "labels/56.nii.gz"
            },
            {
            "image": "images/128.nii.gz",
            "label": "labels/128.nii.gz"
            },
            {
            "image": "images/76.nii.gz",
            "label": "labels/76.nii.gz"
            },
            {
            "image": "images/123.nii.gz",
            "label": "labels/123.nii.gz"
            },
            {
            "image": "images/152.nii.gz",
            "label": "labels/152.nii.gz"
            }
    ],
    "validation":[
        {
            "image": "images/72.nii.gz",
            "label": "labels/72.nii.gz"
            },
            {
            "image": "images/97.nii.gz",
            "label": "labels/97.nii.gz"
            },
            {
            "image": "images/82.nii.gz",
            "label": "labels/82.nii.gz"
            },
            {
            "image": "images/153.nii.gz",
            "label": "labels/153.nii.gz"
            },
            {
            "image": "images/54.nii.gz",
            "label": "labels/54.nii.gz"
            },
            {
            "image": "images/104.nii.gz",
            "label": "labels/104.nii.gz"
            }
    ]
}

这篇好文章是转载于：学新通技术网