SynthRAD2023 CBCT冠军论文代码复现

原创

已于 2024-06-11 11:45:49 修改 · 1.3k 阅读

12 ·

CC 4.0 BY-SA版权

文章标签：

#算法 #图像处理 #神经网络 #深度学习 #pytorch

于 2024-06-11 11:28:05 首次发布

1、仅学习使用，有侵权行为请务必联系我，论文题目和作者大佬贴在这

2、因为是很多尝试的网络里面的一个，粗劣复现，代码提及没有复现的地方，且代码为节选粘贴，请仅作参考

3、文字为两个训练集的联合训练，本文仅以一个为例

4、数据集：

SynthRAD2023 Grand Challenge validation dataset: synthetizing computed tomography for radiotherapy (zenodo.org)

网络部分

网络是使用segmentation_models_pytorch实现，github

import segmentation_models_pytorch as smp
def load_net(args,net=None):
    if net==None:
        net = args.net
    if net == 'unetpp':
        # model = UnetPlusPlus(num_classes=3, deep_supervision=deep_supervision)
        model = smp.UnetPlusPlus(
            encoder_name="vgg16",        # 选择解码器, 例如 mobilenet_v2 或 efficientnet-b7
            encoder_weights="imagenet",     # 使用预先训练的权重imagenet进行解码器初始化
            in_channels=args.chan_out,                  # 模型输入通道（1个用于灰度图像，3个用于RGB等）
            classes=args.chan_out,                      # 模型输出通道（数据集所分的类别总数）
            )
    print("====> model is loaded")
    return model

2.5D dataset

本文仅展示大致2.5D内容，切割和归一化请自行调整

# 请按自己喜好创建pytorch 结构
from dataset_tool.dataset_nii import  SliceData_nii
train_set = SliceData_nii(datapath,patch_size=64)
#arg. 内容可以自行调整，测试可以全部为1
training_data_loader = DataLoader(dataset=train_set, num_workers=args.threads, batch_size=args.batchSize, shuffle=True)

import SimpleITK as sitk
import torch.utils.data as data
import numpy as np
import os
import pathlib

class SliceData_nii(data.Dataset):
    """
    A PyTorch Dataset template
    """

    def __init__(self, root, sample_rate=1, patch_size=None, add_noise_in=False):
        """
        Args:
            root (pathlib.Path): Path to the dataset.
            sample_rate (float, optional): A float between 0 and 1. This controls what fraction
                of the volumes should be loaded.
        """
        self.patch_size=patch_size
        self.add_noise_in=add_noise_in
        self.root=root
        self.example=[]
        namelist=os.listdir(root)
        for casename in namelist:
            ctpath=os