tes

Author: BitterMongo

CETUS/test/patient03.pkl

@@ -0,0 +1,18 @@
+class Config:
+    gpu = '0'
+    data_path = './CETUS/'
+    model = 'vm2'
+    result_dir = './Result'
+    lr = 1e-5
+    max_epochs = 201
+    sim_loss = 'mse'
+    alpha = 0.04
+    batch_size = 1
+    n_save_epoch = 5
+    model_dir = './Checkpoint/exp_'
+    log_dir = './Log'
+    saved_unet_name = 'unet_model.pth'
+    saved_tnet_name = 'tnet_model.pth'
+
+    # test时参数
+    checkpoint_path = "./Checkpoint/exp_1715837905.2272804/"

CETUS/train/patient01.pkl

@@ -0,0 +1,31 @@
+import os
+
+import numpy as np
+from torch.utils.data import Dataset
+import pickle
+
+
+class CETUS(Dataset):
+
+    def __init__(self, dataroot='./dataset_pkl/', split='train'):
+        # 初始化
+        self.target_file = dataroot + split + '/'
+        self.files = os.listdir(self.target_file)
+        self.data_len = len(self.files)
+
+    def __len__(self):
+        # 返回数据集的大小
+        return self.data_len
+
+    def __getitem__(self, index):
+        # 索引数据集中的某个数据，还可以对数据进行预处理
+        # 下标index参数是必须有的，名字任意
+
+        file_current = self.files[index]
+        f = open(self.target_file + file_current, 'rb')
+        data = pickle.load(f)
+        for key in data:
+            data[key] = data[key][np.newaxis, ...]
+        return data
+
+

CETUS/valid/patient16.pkl

@@ -0,0 +1,137 @@
+"""
+*Preliminary* pytorch implementation.
+
+Losses for VoxelMorph
+"""
+
+import math
+import torch
+import numpy as np
+from Model.config import Config as args
+import torch.nn.functional as F
+import pystrum.pynd.ndutils as nd
+
+
+def gradient_loss(s, penalty='l2'):
+    dy = torch.abs(s[:, :, 1:, :] - s[:, :, :-1, :])
+    dx = torch.abs(s[:, :, :, 1:] - s[:, :, :, :-1])
+
+    if penalty == 'l2':
+        dy = dy * dy
+        dx = dx * dx
+
+    d = torch.mean(dx) + torch.mean(dy)
+    return d / 2.0
+
+
+def mse_loss(x, y):
+    return torch.mean((x - y) ** 2)
+
+
+def compute_label_dice(pred, gt):
+    return DSC(gt == 255, pred == 255)
+
+
+def DSC(pred, target):
+    smooth = 1e-5
+    intersection = torch.mul(pred, target).sum()
+    return (2. * intersection + smooth) / (pred.sum() + target.sum() + smooth)
+
+
+def ncc_loss(I, J, win=None):
+    '''
+    输入大小是[B,C,D,W,H]格式的，在计算ncc时用卷积来实现指定窗口内求和
+    '''
+    device = f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu'
+    ndims = len(list(I.size())) - 2
+    assert ndims in [1, 2, 3], "volumes should be 1 to 3 dimensions. found: %d" % ndims
+    if win is None:
+        win = [9] * ndims
+    sum_filt = torch.ones([1, 1, *win]).to(device)
+    pad_no = math.floor(win[0] / 2)
+    stride = [1] * ndims
+    padding = [pad_no] * ndims
+    I_var, J_var, cross = compute_local_sums(I, J, sum_filt, stride, padding, win)
+    cc = cross * cross / (I_var * J_var + 1e-5)
+    return -1 * torch.mean(cc)
+
+
+def compute_local_sums(I, J, filt, stride, padding, win):
+    I2, J2, IJ = I * I, J * J, I * J
+    I_sum = F.conv3d(I, filt, stride=stride, padding=padding)
+    J_sum = F.conv3d(J, filt, stride=stride, padding=padding)
+    I2_sum = F.conv3d(I2, filt, stride=stride, padding=padding)
+    J2_sum = F.conv3d(J2, filt, stride=stride, padding=padding)
+    IJ_sum = F.conv3d(IJ, filt, stride=stride, padding=padding)
+    win_size = np.prod(win)
+    u_I = I_sum / win_size
+    u_J = J_sum / win_size
+    cross = IJ_sum - u_J * I_sum - u_I * J_sum + u_I * u_J * win_size
+    I_var = I2_sum - 2 * u_I * I_sum + u_I * u_I * win_size
+    J_var = J2_sum - 2 * u_J * J_sum + u_J * u_J * win_size
+    return I_var, J_var, cross
+
+
+def cc_loss(x, y):
+    # 根据互相关公式进行计算
+    dim = [2, 3, 4]
+    mean_x = torch.mean(x, dim, keepdim=True)
+    mean_y = torch.mean(y, dim, keepdim=True)
+    mean_x2 = torch.mean(x ** 2, dim, keepdim=True)
+    mean_y2 = torch.mean(y ** 2, dim, keepdim=True)
+    stddev_x = torch.sum(torch.sqrt(mean_x2 - mean_x ** 2), dim, keepdim=True)
+    stddev_y = torch.sum(torch.sqrt(mean_y2 - mean_y ** 2), dim, keepdim=True)
+    return -torch.mean((x - mean_x) * (y - mean_y) / (stddev_x * stddev_y))
+
+
+def jacobian_determinant(disp):
+    """
+    jacobian determinant of a displacement field.
+    NB: to compute the spatial gradients, we use np.gradient.
+
+    Parameters:
+        disp: 2D or 3D displacement field of size [*vol_shape, nb_dims],
+              where vol_shape is of len nb_dims
+
+    Returns:
+        jacobian determinant (scalar)
+    """
+
+    # check input
+    disp = disp.transpose(1, 2, 3, 0)
+    volshape = disp.shape[:-1]
+    nb_dims = len(volshape)
+    assert len(volshape) in (2, 3), 'flow has to be 2D or 3D'
+
+    # compute grid
+    grid_lst = nd.volsize2ndgrid(volshape)
+    grid = np.stack(grid_lst, len(volshape))
+
+    # compute gradients
+    J = np.gradient(disp + grid)
+
+    # 3D glow
+    if nb_dims == 3:
+        dx = J[0]
+        dy = J[1]
+        dz = J[2]
+
+        # compute jacobian components
+        Jdet0 = dx[..., 0] * (dy[..., 1] * dz[..., 2] - dy[..., 2] * dz[..., 1])
+        Jdet1 = dx[..., 1] * (dy[..., 0] * dz[..., 2] - dy[..., 2] * dz[..., 0])
+        Jdet2 = dx[..., 2] * (dy[..., 0] * dz[..., 1] - dy[..., 1] * dz[..., 0])
+
+        return Jdet0 - Jdet1 + Jdet2
+
+    else:  # must be 2
+
+        dfdx = J[0]
+        dfdy = J[1]
+
+        return dfdx[..., 0] * dfdy[..., 1] - dfdy[..., 0] * dfdx[..., 1]
+
+
+def charbonnier_loss(flow_diff, alpha=0.45, beta=1.0, epsilon=0.01):
+    normalization = int(flow_diff.numel())
+    error = torch.pow(torch.square(flow_diff * beta) + epsilon, alpha)
+    return torch.sum(error) / normalization