BraTS2025(UNETR)

pip install   -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

pip install "monai>=0.8.1,<1.0"

# ...existing code...
train_transform = transforms.Compose(
    [
        transforms.LoadImaged(keys=["image", "label"]),
        transforms.AddChanneld(keys=["label"]),  # 只对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"),
        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.AddChanneld(keys=["label"]),  # 只对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"),
        transforms.ToTensord(keys=["image", "label"]),
    ]
)
# ...existing code...

# ...existing code...
parser.add_argument("--in_channels", default=4, type=int, help="number of input channels")  # 改为4
parser.add_argument("--out_channels", default=5, type=int, help="number of output channels")  # 改为5
# ...existing code...

def main_worker(gpu, args):
    # ...existing code...
    if (args.model_name is None) or args.model_name == "unetr":
        model = UNETR(
            in_channels=args.in_channels,      # 已自动适配
            out_channels=args.out_channels,    # 已自动适配
            img_size=(args.roi_x, args.roi_y, args.roi_z),
            feature_size=args.feature_size,
            hidden_size=args.hidden_size,
            mlp_dim=args.mlp_dim,
            num_heads=args.num_heads,
            pos_embed=args.pos_embed,
            norm_name=args.norm_name,
            conv_block=True,
            res_block=True,
            dropout_rate=args.dropout_rate,
        )
    # ...existing code...
    post_label = AsDiscrete(to_onehot=True, n_classes=args.out_channels)  # 已自动适配
    post_pred = AsDiscrete(argmax=True, to_onehot=True, n_classes=args.out_channels)  # 已自动适配
    # ...existing code...

# ...existing code...
parser.add_argument("--in_channels", default=4, type=int, help="number of input channels")  # 改为4
parser.add_argument("--out_channels", default=5, type=int, help="number of output channels")  # 改为5
# ...existing code...

def main():
    # ...existing code...
    with torch.no_grad():
        dice_list_case = []
        for i, batch in enumerate(val_loader):
            val_inputs, val_labels = (batch["image"].cuda(), batch["label"].cuda())
            img_name = batch["image_meta_dict"]["filename_or_obj"][0].split("/")[-1]
            print("Inference on case {}".format(img_name))
            val_outputs = sliding_window_inference(val_inputs, (96, 96, 96), 4, model, overlap=args.infer_overlap)
            val_outputs = torch.softmax(val_outputs, 1).cpu().numpy()
            val_outputs = np.argmax(val_outputs, axis=1).astype(np.uint8)
            val_labels = val_labels.cpu().numpy()[:, 0, :, :, :]
            dice_list_sub = []
            for i in range(1, args.out_channels):  # 改为类别数
                organ_Dice = dice(val_outputs[0] == i, val_labels[0] == i)
                dice_list_sub.append(organ_Dice)
            mean_dice = np.mean(dice_list_sub)
            print("Mean Organ Dice: {}".format(mean_dice))
            dice_list_case.append(mean_dice)
        print("Overall Mean Dice: {}".format(np.mean(dice_list_case)))
# ...existing code...

import json

import nibabel as nib

import os

  

json_path = "/data/coding/research-contributions-main/UNETR/BTCV/dataset_monai.json"

  

with open(json_path, "r") as f:

    data = json.load(f)

  

# 兼容不同json结构

if isinstance(data, dict) and "training" in data:

    items = data["training"]

elif isinstance(data, list):

    items = data

else:

    raise RuntimeError("Unknown dataset_monai.json structure")

  

for item in items:

    label_path = item["label"] if isinstance(item, dict) else item["label"]

    if not os.path.isabs(label_path):

        label_path = os.path.join(os.path.dirname(json_path), label_path)

    if not os.path.exists(label_path):

        print(f"Label file not found: {label_path}")

        continue

    img = nib.load(label_path)

    arr = img.get_fdata()

    if arr.max() == 0:

        print(f"All-zero label: {label_path}")

def train_epoch(model, loader, optimizer, scaler, epoch, loss_func, args):
    model.train()
    start_time = time.time()
    run_loss = AverageMeter()
    for idx, batch_data in enumerate(loader):
        batch_start = time.time()
        try:
            if isinstance(batch_data, list):
                data, target = batch_data
            else:
                data, target = batch_data["image"], batch_data["label"]
            data, target = data.cuda(args.rank), target.cuda(args.rank)
            for param in model.parameters():
                param.grad = None
            with autocast(enabled=args.amp):
                logits = model(data)
                loss = loss_func(logits, target)
            if args.amp:
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
            else:
                loss.backward()
                optimizer.step()
            if args.distributed:
                loss_list = distributed_all_gather([loss], out_numpy=True, is_valid=idx < loader.sampler.valid_length)
                run_loss.update(
                    np.mean(np.mean(np.stack(loss_list, axis=0), axis=0), axis=0), n=args.batch_size * args.world_size
                )
            else:
                run_loss.update(loss.item(), n=args.batch_size)
            if args.rank == 0:
                print(
                    "Epoch {}/{} {}/{}".format(epoch, args.max_epochs, idx, len(loader)),
                    "loss: {:.4f}".format(run_loss.avg),
                    "time {:.2f}s".format(time.time() - start_time),
                )
            batch_end = time.time()
            print(f"Epoch {epoch} Batch {idx} time: {batch_end - batch_start:.2f}s")
            start_time = time.time()
        except RuntimeError as e:
            if "all zero" in str(e) or "high <= 0" in str(e):
                print(f"Skip batch {idx} due to all-zero label patch or sampling error.")
                continue
            else:
                raise e
    for param in model.parameters():
        param.grad = None
    return run_loss.avg