pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install -r environment.yml -i https://pypi.tuna.tsinghua.edu.cn/simple
conda env create -f environment.yml
conda activate swin_unetr
# 导入必要的库和模块
import os
import json
import torch
from monai.data import CacheDataset, DataLoader, decollate_batch
from monai.transforms import (
Compose,
LoadImaged,
EnsureChannelFirstd,
EnsureTyped,
Orientationd,
Spacingd,
RandSpatialCropd,
RandFlipd,
NormalizeIntensityd,
RandScaleIntensityd,
RandShiftIntensityd,
Activationsd,
AsDiscreted,
Invertd,
MapTransform,
Resized
)
# 定义一个自定义转换，将单通道标签转换为多通道格式
class ConvertToMultiChannelBasedOnBrats25Classesd(MapTransform):
"""
根据 Brats25 标签的类别，将单通道标签转换为多通道格式。
- 标签 1：坏死和非增强肿瘤核心
- 标签 2：肿瘤周围水肿
- 标签 3：GD 增强肿瘤
- 标签 4：其他
输出通道对应：
- 通道 0：标签 1
- 通道 1：标签 2
- 通道 2：标签 3
- 通道 3：标签 4
"""
def __call__(self, data):
d = dict(data)
for key in self.keys:
result = []
# 标签 1：坏死和非增强肿瘤核心
result.append(d[key] == 1)
# 标签 2：肿瘤周围水肿
result.append(d[key] == 2)
# 标签 3：GD 增强肿瘤
result.append(d[key] == 3)
# 标签 4：其他
result.append(d[key] == 4)
d[key] = torch.stack(result, axis=0).float()
return d
# 定义一个函数，用于读取数据列表并转换路径
def datafold_read(datalist, basedir, key="training"):
"""
读取数据列表，并将相对路径转换为绝对路径。
:param datalist: 数据列表文件路径
:param basedir: 数据基目录
:param key: 数据类型（训练、验证或测试）
:return: 转换后的数据列表
"""
with open(datalist) as f:
json_data = json.load(f)
json_data = json_data[key]
for d in json_data:
for k, v in d.items():
if isinstance(d[k], list):
d[k] = [os.path.join(basedir, iv) for iv in d[k]]
elif isinstance(d[k], str):
d[k] = os.path.join(basedir, d[k]) if len(d[k]) > 0 else d[k]
return json_data
# 定义一个函数，用于获取训练和验证数据加载器
def get_loader(args):
"""
获取训练和验证数据加载器。
:param args: 参数对象，包含数据目录、数据列表文件等信息
:return: 训练和验证数据加载器
"""
data_dir = args.data_dir
datalist_json = args.json_list
# 获取训练和验证数据
train_files = datafold_read(datalist=datalist_json, basedir=data_dir, key="training")
validation_files = datafold_read(datalist=datalist_json, basedir=data_dir, key="validation")
# 定义训练和验证转换
train_transform = Compose(
[
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys="image"),
EnsureTyped(keys=["image", "label"]),
ConvertToMultiChannelBasedOnBrats25Classesd(keys="label"),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(keys=["image", "label"], pixdim=(1.0, 1.0, 1.0), mode=("bilinear", "nearest")),
RandSpatialCropd(keys=["image", "label"], roi_size=[args.roi_x, args.roi_y, args.roi_z], random_size=False),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
RandScaleIntensityd(keys="image", factors=0.1, prob=1.0),
RandShiftIntensityd(keys="image", offsets=0.1, prob=1.0),
]
)
val_transform = Compose(
[
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys="image"),
EnsureTyped(keys=["image", "label"]),
ConvertToMultiChannelBasedOnBrats25Classesd(keys="label"),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(keys=["image", "label"], pixdim=(1.0, 1.0, 1.0), mode=("bilinear", "nearest")),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
]
)
# 使用 CacheDataset 并设置缓存比例为 1
train_ds = CacheDataset(
data=train_files,
transform=train_transform,
cache_rate=0.1,
num_workers=8
)
val_ds = CacheDataset(
data=validation_files,
transform=val_transform,
cache_rate=0.1,
num_workers=8
)
# 创建数据加载器并启用 persistent_workers
train_loader = DataLoader(
train_ds,
batch_size=args.batch_size,
shuffle=True,
num_workers=4, # 设置为4个worker
persistent_workers=True
)
val_loader = DataLoader(
val_ds,
batch_size=args.batch_size,
shuffle=False,
num_workers=4, # 设置为4个worker
persistent_workers=True
)
print(f"Number of training files: {len(train_files)}")
print(f"Number of validation files: {len(validation_files)}")
return train_loader, val_loader
def get_test_loader(args):
"""
获取测试数据加载器。
:param args: 参数对象，包含数据目录、数据列表文件等信息
:return: 测试数据加载器
"""
data_dir = args.data_dir
datalist_json = args.json_list
test_files = datafold_read(datalist=datalist_json, basedir=data_dir, key="test")
test_transform = Compose(
[
LoadImaged(keys=["image", "label"]),
EnsureChannelFirstd(keys="image"),
EnsureTyped(keys=["image", "label"]),
ConvertToMultiChannelBasedOnBrats25Classesd(keys="label"),
Orientationd(keys=["image", "label"], axcodes="RAS"),
Spacingd(keys=["image", "label"], pixdim=(1.0, 1.0, 1.0), mode=("bilinear", "nearest")),
NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
]
)
test_ds = CacheDataset(
data=test_files,
transform=test_transform,
cache_rate=1.0,
num_workers=8
)
test_loader = DataLoader(
test_ds,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
persistent_workers=True
)
return test_loader
def get_post_transforms():
"""
获取后处理转换。
:return: 后处理转换
"""
return Compose(
[
Activationsd(keys="pred", softmax=True),
AsDiscreted(keys="pred", threshold=0.5),
]
)
# 导入必要的库和模块
import argparse
import os
import torch
from torch.cuda.amp import GradScaler
from torch.optim.lr_scheduler import CosineAnnealingLR
import time
from monai.losses import DiceLoss, DiceCELoss
from monai.inferers import sliding_window_inference
from monai.metrics import DiceMetric
from monai.networks.nets import SwinUNETR
import matplotlib.pyplot as plt
from monai.transforms import Compose, EnsureTyped, Activations, AsDiscrete, MapTransform
from monai.data import decollate_batch
from dataloader import get_loader, get_post_transforms
# 定义命令行参数解析器
parser = argparse.ArgumentParser(description="Training")
parser.add_argument("--data_dir", default="/data/coding/TEST", type=str, help="dataset directory")
parser.add_argument("--json_list", default="/data/coding/research-contributions-main/UNETR/BTCV/dataset_monai_cleaned.json", type=str, help="dataset json file")
parser.add_argument("--fold", default=0, type=int, help="fold number")
parser.add_argument("--batch_size", default=1, type=int, help="batch size")
parser.add_argument("--num_workers", default=4, type=int, help="number of workers")
parser.add_argument("--roi_x", default=128, type=int, help="roi size in x direction")
parser.add_argument("--roi_y", default=128, type=int, help="roi size in y direction")
parser.add_argument("--roi_z", default=128, type=int, help="roi size in z direction")
parser.add_argument("--in_channels", default=4, type=int, help="number of input channels")
parser.add_argument("--out_channels", default=4, type=int, help="number of output channels")
parser.add_argument("--optimizer", default="adamw", type=str, help="optimizer to use (adamw or adam)")
parser.add_argument("--loss", default="dicece", type=str, help="loss function to use (dicece or diceloss)")
parser.add_argument("--lr", default=2e-4, type=float, help="learning rate")
parser.add_argument("--weight_decay", default=1e-5, type=float, help="weight decay for optimizer")
parser.add_argument("--save_dir", default="/data/coding/TEST", type=str, help="directory to save checkpoints and plots")
parser.add_argument("--save_checkpoint_interval", default=10, type=int, help="interval to save checkpoints")
parser.add_argument("--max_epochs", default=100, type=int, help="maximum number of training epochs")
args = parser.parse_args()
# 定义训练函数
def train(args):
"""
执行模型训练过程，包括数据加载、模型初始化、训练循环、验证和指标计算。
"""
global model, optimizer, lr_scheduler, scaler, dice_metric, dice_metric_batch, post_trans, best_metric, best_metric_epoch, best_metrics_epochs_and_time, epoch_loss_values, metric_values, VAL_AMP, metric_values_tc, metric_values_wt, metric_values_et, metric_values_other
# 初始化变量和对象
max_epochs = args.max_epochs
val_interval = 5
VAL_AMP = False
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 定义 SwinUNETR 模型
model = SwinUNETR(
img_size=(args.roi_x, args.roi_y, args.roi_z),
in_channels=args.in_channels,
out_channels=args.out_channels,
feature_size=48,
use_checkpoint=False,
).to(device)
# 初始化损失函数
if args.loss == "dicece":
loss_function = DiceCELoss(include_background=True, softmax=True)
else:
loss_function = DiceLoss(smooth_nr=0, smooth_dr=1e-5, squared_pred=True, to_onehot_y=False, softmax=True)
# 初始化优化器
if args.optimizer == "adamw":
optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
else:
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
# 定义余弦退火学习率调度器
lr_scheduler = CosineAnnealingLR(optimizer, T_max=max_epochs)
from torch.amp import GradScaler
scaler = GradScaler('cuda')
torch.backends.cudnn.benchmark = True
dice_metric = DiceMetric(include_background=True, reduction="mean")
dice_metric_batch = DiceMetric(include_background=True, reduction="mean_batch")
post_trans = Compose([Activations(softmax=True), AsDiscrete(to_onehot=None, threshold=0.5)])
best_metric = -1
best_metric_epoch = -1
best_metrics_epochs_and_time = [[], [], []]
epoch_loss_values = []
metric_values = []
metric_values_tc = []
metric_values_wt = []
metric_values_et = []
metric_values_other = []
train_loader, val_loader = get_loader(args)
total_start = time.time()
for epoch in range(max_epochs):
epoch_start = time.time()
print("-" * 10)
print(f"epoch {epoch + 1}/{max_epochs}")
model.train()
epoch_loss = 0
step = 0
for batch_data in train_loader:
step_start = time.time()
step += 1
inputs, labels = (
batch_data["image"].to(device),
batch_data["label"].to(device),
)
optimizer.zero_grad()
with torch.autocast("cuda"):
outputs = model(inputs)
loss = loss_function(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
epoch_loss += loss.item()
print(
f"{step}/{len(train_loader.dataset) // train_loader.batch_size} "
f"train_loss: {loss.item():.4f} "
f"step time: {(time.time() - step_start):.4f}"
)
lr_scheduler.step()
epoch_loss /= step
epoch_loss_values.append(epoch_loss)
print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")
if (epoch + 1) % val_interval == 0:
model.eval()
with torch.no_grad():
all_iou = []
all_dice = []
for val_data in val_loader:
val_inputs, val_labels = (
val_data["image"].to(device),
val_data["label"].to(device),
)
val_outputs = sliding_window_inference(
inputs=val_inputs,
roi_size=(args.roi_x, args.roi_y, args.roi_z),
sw_batch_size=1,
predictor=model,
overlap=0.5,
)
val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]
dice_metric(y_pred=val_outputs, y=val_labels)
dice_metric_batch(y_pred=val_outputs, y=val_labels)
# 计算自定义的 IOU 和 Dice
val_outputs = torch.cat(val_outputs, dim=0)
val_labels = torch.cat([val_label.unsqueeze(0) for val_label in val_labels], dim=0)
iou, dice = calculate_metrics(val_outputs, val_labels)
all_iou.append(iou)
all_dice.append(dice)
# 计算平均 IOU 和 Dice
avg_iou = torch.tensor(all_iou).mean(dim=0)
avg_dice = torch.tensor(all_dice).mean(dim=0)
# 处理类别不平衡问题
first_non_zero = None
for i in range(avg_iou.shape[0]):
if avg_iou[i] != 0:
first_non_zero = i
break
if first_non_zero is not None:
for i in range(first_non_zero, avg_iou.shape[0]):
if avg_iou[i] == 0:
avg_iou[i] = 1.0
avg_dice[i] = 1.0
metric = dice_metric.aggregate().item()
metric_values.append(metric)
metric_batch = dice_metric_batch.aggregate()
metric_tc = metric_batch[0].item()
metric_values_tc.append(metric_tc)
metric_wt = metric_batch[1].item()
metric_values_wt.append(metric_wt)
metric_et = metric_batch[2].item()
metric_values_et.append(metric_et)
metric_other = metric_batch[3].item()
metric_values_other.append(metric_other)
dice_metric.reset()
dice_metric_batch.reset()
if metric > best_metric:
best_metric = metric
best_metric_epoch = epoch + 1
best_metrics_epochs_and_time[0].append(best_metric)
best_metrics_epochs_and_time[1].append(best_metric_epoch)
best_metrics_epochs_and_time[2].append(time.time() - total_start)
torch.save(
model.state_dict(),
os.path.join(args.save_dir, "best_metric_model.pth"),
)
print("saved new best metric model")
if (epoch + 1) % args.save_checkpoint_interval == 0:
torch.save(
{
"epoch": epoch + 1,
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"lr_scheduler_state_dict": lr_scheduler.state_dict(),
"best_metric": best_metric,
"best_metric_epoch": best_metric_epoch,
},
os.path.join(args.save_dir, f"checkpoint_epoch_{epoch + 1}.pth"),
)
print(f"saved checkpoint for epoch {epoch + 1}")
print(
f"current epoch: {epoch + 1} current mean dice: {metric:.4f} "
f"tc: {metric_tc:.4f} wt: {metric_wt:.4f} et: {metric_et:.4f} other: {metric_other:.4f} "
f"\nbest mean dice: {best_metric:.4f} at epoch: {best_metric_epoch}"
)
print(f"custom iou metric: tensor({avg_iou.tolist()})")
print(f"custom Dice metric: tensor({avg_dice.tolist()})")
print(f"time consuming of epoch {epoch + 1} is: {(time.time() - epoch_start):.4f}")
total_time = time.time() - total_start
print(f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}, total time: {total_time}.")
torch.save(
model.state_dict(),
os.path.join(args.save_dir, "final_model.pth"),
)
print("saved final model")
training_metrics = {
"epoch_loss_values": epoch_loss_values,
"metric_values": metric_values,
"metric_values_tc": metric_values_tc,
"metric_values_wt": metric_values_wt,
"metric_values_et": metric_values_et,
"metric_values_other": metric_values_other,
}
torch.save(
training_metrics,
os.path.join(args.save_dir, "training_metrics.pth"),
)
print("saved training metrics")
try:
plt.figure("train", (12, 6))
plt.subplot(1, 2, 1)
plt.title("Epoch Average Loss")
x = [i + 1 for i in range(len(epoch_loss_values))]
y = epoch_loss_values
plt.xlabel("epoch")
plt.plot(x, y, color="red")
plt.subplot(1, 2, 2)
plt.title("Val Mean Dice")
x = [val_interval * (i + 1) for i in range(len(metric_values))]
y = metric_values
plt.xlabel("epoch")
plt.plot(x, y, color="green")
plt.savefig(os.path.join(args.save_dir, "training_curves.png"))
plt.close()
plt.figure("train", (18, 6))
plt.subplot(1, 3, 1)
plt.title("Val Mean Dice TC")
x = [val_interval * (i + 1) for i in range(len(metric_values_tc))]
y = metric_values_tc
plt.xlabel("epoch")
plt.plot(x, y, color="blue")
plt.subplot(1, 3, 2)
plt.title("Val Mean Dice WT")
x = [val_interval * (i + 1) for i in range(len(metric_values_wt))]
y = metric_values_wt
plt.xlabel("epoch")
plt.plot(x, y, color="brown")
plt.subplot(1, 3, 3)
plt.title("Val Mean Dice ET")
x = [val_interval * (i + 1) for i in range(len(metric_values_et))]
y = metric_values_et
plt.xlabel("epoch")
plt.plot(x, y, color="purple")
plt.savefig(os.path.join(args.save_dir, "validation_curves.png"))
plt.close()
except Exception as e:
print(f"Warning: Failed to save plots. Error: {e}")
# 定义指标计算函数
def calculate_metrics(y_pred, y):
"""
计算每个标签的 IOU 和 Dice 指标。
"""
epsilon = 1e-6
iou = []
dice = []
for i in range(y.shape[1]):
pred = y_pred[:, i, :, :, :]
label = y[:, i, :, :, :]
intersection = torch.sum(pred * label)
union = torch.sum(pred) + torch.sum(label)
iou.append((intersection / (union - intersection + epsilon)).item())
dice.append((2. * intersection / (torch.sum(pred) + torch.sum(label) + epsilon)).item())
return iou, dice
if __name__ == "__main__":
train(args)
# 导入必要的库和模块
import argparse
import os
import torch
from monai.inferers import sliding_window_inference
from monai.metrics import DiceMetric
from monai.networks.nets import SwinUNETR
from monai.transforms import Compose, EnsureTyped, from_engine
from dataloader import get_test_loader, get_post_transforms
# 定义一个函数，用于评估模型
def evaluate(args):
"""
评估模型。
:param args: 参数对象，包含数据目录、数据列表文件等信息
"""
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 初始化模型
model = SwinUNETR(
img_size=(args.roi_x, args.roi_y, args.roi_z),
in_channels=args.in_channels,
out_channels=args.out_channels,
feature_size=48,
use_checkpoint=False,
).to(device)
# 加载模型权重
model.load_state_dict(torch.load(os.path.join(args.data_dir, "best_metric_model.pth"), map_location=device))
model.eval()
# 获取测试数据加载器
test_loader = get_test_loader(args)
# 初始化Dice指标
dice_metric = DiceMetric(include_background=True, reduction="mean")
dice_metric_batch = DiceMetric(include_background=True, reduction="mean_batch")
# 获取后处理转换
post_trans = get_post_transforms()
# 定义推理函数
def inference(input):
return sliding_window_inference(
inputs=input,
roi_size=(args.roi_x, args.roi_y, args.roi_z),
sw_batch_size=1,
predictor=model,
overlap=0.5,
)
# 执行推理和评估
with torch.no_grad():
for test_data in test_loader:
test_inputs = test_data["image"].to(device)
test_data["pred"] = inference(test_inputs)
test_data = [post_trans(i) for i in decollate_batch(test_data)]
test_outputs, test_labels = from_engine(["pred", "label"])(test_data)
dice_metric(y_pred=test_outputs, y=test_labels)
dice_metric_batch(y_pred=test_outputs, y=test_labels)
metric_org = dice_metric.aggregate().item()
metric_batch_org = dice_metric_batch.aggregate()
dice_metric.reset()
dice_metric_batch.reset()
# 假设 metric_batch_org 是一个长度为4的列表，分别对应四个类别的指标
metric_tc = metric_batch_org[0].item() # 类别1
metric_wt = metric_batch_org[1].item() # 类别2
metric_et = metric_batch_org[2].item() # 类别3
metric_other = metric_batch_org[3].item() # 类别4
print("Metric on original image spacing: ", metric_org)
print(f"metric_tc: {metric_tc:.4f}")
print(f"metric_wt: {metric_wt:.4f}")
print(f"metric_et: {metric_et:.4f}")
print(f"metric_other: {metric_other:.4f}")
if __name__ == "__main__":
# 定义命令行参数解析器
parser = argparse.ArgumentParser(description="Evaluation")
parser.add_argument("--data_dir", default="/data/coding/TEST", type=str, help="dataset directory")
parser.add_argument("--json_list", default="/data/coding/research-contributions-main/UNETR/BTCV/dataset_monai_cleaned.json", type=str, help="dataset json file")
parser.add_argument("--fold", default=0, type=int, help="fold number")
parser.add_argument("--batch_size", default=1, type=int, help="batch size")
parser.add_argument("--num_workers", default=4, type=int, help="number of workers")
parser.add_argument("--roi_x", default=128, type=int, help="roi size in x direction")
parser.add_argument("--roi_y", default=128, type=int, help="roi size in y direction")
parser.add_argument("--roi_z", default=128, type=int, help="roi size in z direction")
parser.add_argument("--in_channels", default=4, type=int, help="number of input channels")
parser.add_argument("--out_channels", default=4, type=int, help="number of output channels")
args = parser.parse_args()
evaluate(args)