# x: 128x128 resolution for 32 frames.
import torch
import torch.nn as nn
import torch.nn.functional as F

basic_dims = 8
num_modals = 1
patch_size = [2, 8, 8]


def conv_block_3d(in_dim, out_dim, activation):
    return nn.Sequential(
        nn.Conv3d(in_dim, out_dim, kernel_size=3, stride=1, padding=1),
        nn.BatchNorm3d(out_dim),
        activation,)


def conv_trans_block_3d(in_dim, out_dim, activation):
    return nn.Sequential(
        nn.ConvTranspose3d(in_dim, out_dim, kernel_size=3, stride=2, padding=1, output_padding=1),
        nn.BatchNorm3d(out_dim),
        activation,)


def max_pooling_3d():
    return nn.MaxPool3d(kernel_size=2, stride=2, padding=0)


def conv_block_2_3d(in_dim, out_dim, activation):
    return nn.Sequential(
        conv_block_3d(in_dim, out_dim, activation),
        nn.Conv3d(out_dim, out_dim, kernel_size=3, stride=1, padding=1),
        nn.BatchNorm3d(out_dim),)

def normalization(planes, norm='bn'):
    if norm == 'bn':
        m = nn.BatchNorm3d(planes)
    elif norm == 'gn':
        m = nn.GroupNorm(4, planes)
    elif norm == 'in':
        m = nn.InstanceNorm3d(planes)
    # elif norm == 'sync_bn':
    #     m = SynchronizedBatchNorm3d(planes)
    else:
        raise ValueError('normalization type {} is not supported'.format(norm))
    return m


class general_conv3d_prenorm(nn.Module):
    def __init__(self, in_ch, out_ch, k_size=3, stride=1, padding=1, pad_type='zeros', norm='in', is_training=True,
                 act_type='lrelu', relufactor=0.2):
        super(general_conv3d_prenorm, self).__init__()
        self.conv = nn.Conv3d(in_channels=in_ch, out_channels=out_ch, kernel_size=k_size, stride=stride,
                              padding=padding, padding_mode=pad_type, bias=True)

        self.norm = normalization(out_ch, norm=norm)
        if act_type == 'relu':
            self.activation = nn.ReLU(inplace=True)
        elif act_type == 'lrelu':
            self.activation = nn.LeakyReLU(negative_slope=relufactor, inplace=True)

    def forward(self, x):
        x = self.norm(x)
        x = self.activation(x)
        x = self.conv(x)
        return x


class fusion_prenorm(nn.Module):
    def __init__(self, in_channel=64, num_cls=4):
        super(fusion_prenorm, self).__init__()
        self.fusion_layer = nn.Sequential(
            general_conv3d_prenorm(in_channel * num_modals, in_channel, k_size=1, padding=0, stride=1),
            general_conv3d_prenorm(in_channel, in_channel, k_size=3, padding=1, stride=1),
            general_conv3d_prenorm(in_channel, in_channel, k_size=1, padding=0, stride=1))

    def forward(self, x):
        return self.fusion_layer(x)


class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()

        self.e1_c1 = nn.Conv3d(in_channels=2, out_channels=basic_dims, kernel_size=3, stride=1, padding=1,
                                   padding_mode='zeros', bias=True)
        self.e1_c2 = general_conv3d_prenorm(basic_dims, basic_dims, pad_type='zeros')
        self.e1_c3 = general_conv3d_prenorm(basic_dims, basic_dims, pad_type='zeros')

        self.e2_c1 = general_conv3d_prenorm(basic_dims, basic_dims * 2, stride=2, pad_type='zeros')
        self.e2_c2 = general_conv3d_prenorm(basic_dims * 2, basic_dims * 2, pad_type='zeros')
        self.e2_c3 = general_conv3d_prenorm(basic_dims * 2, basic_dims * 2, pad_type='zeros')

        self.e3_c1 = general_conv3d_prenorm(basic_dims * 2, basic_dims * 4, stride=2, pad_type='zeros')
        self.e3_c2 = general_conv3d_prenorm(basic_dims * 4, basic_dims * 4, pad_type='zeros')
        self.e3_c3 = general_conv3d_prenorm(basic_dims * 4, basic_dims * 4, pad_type='zeros')

        self.e4_c1 = general_conv3d_prenorm(basic_dims * 4, basic_dims * 8, stride=2, pad_type='zeros')
        self.e4_c2 = general_conv3d_prenorm(basic_dims * 8, basic_dims * 8, pad_type='zeros')
        self.e4_c3 = general_conv3d_prenorm(basic_dims * 8, basic_dims * 8, pad_type='zeros')

        self.e5_c1 = general_conv3d_prenorm(basic_dims * 8, basic_dims * 16, stride=2, pad_type='zeros')
        self.e5_c2 = general_conv3d_prenorm(basic_dims * 16, basic_dims * 16, pad_type='zeros')
        self.e5_c3 = general_conv3d_prenorm(basic_dims * 16, basic_dims * 16, pad_type='zeros')

    def forward(self, x0,x1):
        x1 = self.e1_c1(torch.cat([x0,x1],1))
        x1 = x1 + self.e1_c3(self.e1_c2(x1))

        x2 = self.e2_c1(x1)
        x2 = x2 + self.e2_c3(self.e2_c2(x2))

        x3 = self.e3_c1(x2)
        x3 = x3 + self.e3_c3(self.e3_c2(x3))

        x4 = self.e4_c1(x3)
        x4 = x4 + self.e4_c3(self.e4_c2(x4))

        x5 = self.e5_c1(x4)
        x5 = x5 + self.e5_c3(self.e5_c2(x5))

        return x1, x2, x3, x4, x5


class Decoder(nn.Module):
    def __init__(self, num_cls=1):
        super(Decoder, self).__init__()

        self.d4_c1 = general_conv3d_prenorm(basic_dims * 16, basic_dims * 8, pad_type='zeros')
        self.d4_c2 = general_conv3d_prenorm(basic_dims * 16, basic_dims * 8, pad_type='zeros')
        self.d4_out = general_conv3d_prenorm(basic_dims * 8, basic_dims * 8, k_size=1, padding=0, pad_type='zeros')

        self.d3_c1 = general_conv3d_prenorm(basic_dims * 8, basic_dims * 4, pad_type='zeros')
        self.d3_c2 = general_conv3d_prenorm(basic_dims * 8, basic_dims * 4, pad_type='zeros')
        self.d3_out = general_conv3d_prenorm(basic_dims * 4, basic_dims * 4, k_size=1, padding=0, pad_type='zeros')

        self.d2_c1 = general_conv3d_prenorm(basic_dims * 4, basic_dims * 2, pad_type='zeros')
        self.d2_c2 = general_conv3d_prenorm(basic_dims * 4, basic_dims * 2, pad_type='zeros')
        self.d2_out = general_conv3d_prenorm(basic_dims * 2, basic_dims * 2, k_size=1, padding=0, pad_type='zeros')

        self.d1_c1 = general_conv3d_prenorm(basic_dims * 2, basic_dims, pad_type='zeros')
        self.d1_c2 = general_conv3d_prenorm(basic_dims * 2, basic_dims, pad_type='zeros')
        self.d1_out = general_conv3d_prenorm(basic_dims, basic_dims, k_size=1, padding=0, pad_type='zeros')

        self.seg_d4 = nn.Conv3d(in_channels=basic_dims * 16, out_channels=num_cls, kernel_size=1, stride=1, padding=0,
                                bias=True)
        self.seg_d3 = nn.Conv3d(in_channels=basic_dims * 8, out_channels=num_cls, kernel_size=1, stride=1, padding=0,
                                bias=True)
        self.seg_d2 = nn.Conv3d(in_channels=basic_dims * 4, out_channels=num_cls, kernel_size=1, stride=1, padding=0,
                                bias=True)
        self.seg_d1 = nn.Conv3d(in_channels=basic_dims * 2, out_channels=num_cls, kernel_size=1, stride=1, padding=0,
                                bias=True)
        self.seg_layer = nn.Conv3d(in_channels=basic_dims, out_channels=num_cls, kernel_size=1, stride=1, padding=0,
                                   bias=True)
        self.softmax = nn.Softmax(dim=1)

        self.up2 = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)

        self.RFM5 = fusion_prenorm(in_channel=basic_dims * 16, num_cls=num_cls)
        self.RFM4 = fusion_prenorm(in_channel=basic_dims * 8, num_cls=num_cls)
        self.RFM3 = fusion_prenorm(in_channel=basic_dims * 4, num_cls=num_cls)
        self.RFM2 = fusion_prenorm(in_channel=basic_dims * 2, num_cls=num_cls)
        self.RFM1 = fusion_prenorm(in_channel=basic_dims * 1, num_cls=num_cls)

    def forward(self, x1, x2, x3, x4, x5):
        de_x5 = self.RFM5(x5)
        de_x5 = self.d4_c1(self.up2(de_x5))
        de_x4 = self.RFM4(x4)
        de_x4 = torch.cat((de_x4, de_x5), dim=1)
        de_x4 = self.d4_out(self.d4_c2(de_x4))
        de_x4 = self.d3_c1(self.up2(de_x4))

        de_x3 = self.RFM3(x3)
        de_x3 = torch.cat((de_x3, de_x4), dim=1)
        de_x3 = self.d3_out(self.d3_c2(de_x3))
        de_x3 = self.d2_c1(self.up2(de_x3))

        de_x2 = self.RFM2(x2)
        de_x2 = torch.cat((de_x2, de_x3), dim=1)
        de_x2 = self.d2_out(self.d2_c2(de_x2))
        de_x2 = self.d1_c1(self.up2(de_x2))

        de_x1 = self.RFM1(x1)
        de_x1 = torch.cat((de_x1, de_x2), dim=1)
        de_x1 = self.d1_out(self.d1_c2(de_x1))

        logits = self.seg_layer(de_x1)
        pred = torch.sigmoid(logits)

        return pred


class RUNet_recon(nn.Module):
    def __init__(self,):
        super(RUNet_recon, self).__init__()
        self.encoder = Encoder()
        self.decoder = Decoder(num_cls=1)

        for m in self.modules():
            if isinstance(m, nn.Conv3d):
                torch.nn.init.kaiming_normal_(m.weight)

    def forward(self, x0,x1):

        z1, z2, z3, z4, z5 = self.encoder(x0,x1)
        recon = self.decoder(z1, z2, z3, z4, z5)

        return recon


class RUNet_seg(nn.Module):
    def __init__(self, in_dim, out_dim, num_filters):
        super(RUNet_seg, self).__init__()
        
        self.in_dim = in_dim
        self.out_dim = out_dim
        self.num_filters = num_filters
        activation = nn.LeakyReLU(0.2, inplace=True)
        act1=nn.Sigmoid()

        
        # Down sampling
        self.down_1 = conv_block_2_3d(self.in_dim, self.num_filters, activation)
        self.res_1 = conv_block_3d(self.in_dim, self.num_filters, activation)
        self.pool_1 = max_pooling_3d()

        self.down_2 = conv_block_2_3d(self.num_filters, self.num_filters * 2, activation)
        self.res_2 = conv_block_3d(self.num_filters, self.num_filters * 2, activation)
        self.pool_2 = max_pooling_3d()

        self.down_3 = conv_block_2_3d(self.num_filters * 2, self.num_filters * 4, activation)
        self.res_3 = conv_block_3d(self.num_filters * 2, self.num_filters * 4, activation)
        self.pool_3 = max_pooling_3d()

        self.down_4 = conv_block_2_3d(self.num_filters * 4, self.num_filters * 8, activation)
        self.res_4 = conv_block_3d(self.num_filters * 4, self.num_filters * 8, activation)
        self.pool_4 = max_pooling_3d()
        
        # Bridge
        self.bridge = conv_block_2_3d(self.num_filters * 8, self.num_filters * 16, activation)
        self.res_bridge = conv_block_3d(self.num_filters * 8, self.num_filters * 16, activation)
        # Up sampling

        self.trans_2 = conv_trans_block_3d(self.num_filters * 16, self.num_filters * 16, activation)
        self.up_2 = conv_block_2_3d(self.num_filters * 24, self.num_filters * 8, activation)
        self.res_up2 = conv_block_3d(self.num_filters * 24, self.num_filters * 8, activation)

        self.trans_3 = conv_trans_block_3d(self.num_filters * 8, self.num_filters * 8, activation)
        self.up_3 = conv_block_2_3d(self.num_filters * 12, self.num_filters * 4, activation)
        self.res_up3 = conv_block_3d(self.num_filters * 12, self.num_filters * 4, activation)

        self.trans_4 = conv_trans_block_3d(self.num_filters * 4, self.num_filters * 4, activation)
        self.up_4 = conv_block_2_3d(self.num_filters * 6, self.num_filters * 2, activation)
        self.res_up4 = conv_block_3d(self.num_filters * 6, self.num_filters * 2, activation)

        self.trans_5 = conv_trans_block_3d(self.num_filters * 2, self.num_filters * 2, activation)
        self.up_5 = conv_block_2_3d(self.num_filters * 3, self.num_filters * 1, activation)

        # Output
        self.out = conv_block_3d(self.num_filters, out_dim, act1)
    
    def forward(self, x0,x1):
        # Down sampling
        x = torch.cat([x0,x1-x0],1)
        down_1 = self.down_1(x) # -> [1, 4, 128, 128, 128]
        res_1 = self.res_1(x)
        pool_1 = self.pool_1(down_1+res_1) # -> [1, 4, 64, 64, 64]
        
        down_2 = self.down_2(pool_1) # -> [1, 8, 64, 64, 64]
        res_2 = self.res_2(pool_1)
        pool_2 = self.pool_2(down_2+res_2) # -> [1, 8, 32, 32, 32]
        
        down_3 = self.down_3(pool_2) # -> [1, 16, 32, 32, 32]
        res_3 = self.res_3(pool_2)
        pool_3 = self.pool_3(down_3+res_3) # -> [1, 16, 16, 16, 16]
        
        down_4 = self.down_4(pool_3) # -> [1, 32, 16, 16, 16]
        res_4 = self.res_4(pool_3)
        pool_4 = self.pool_4(down_4+res_4) # -> [1, 32, 8, 8, 8]

        
        # Bridge
        bridge = self.bridge(pool_4) # -> [1, 128, 4, 4, 4]
        res_bridge = self.res_bridge(pool_4)
        
        trans_2 = self.trans_2(bridge+res_bridge) # -> [1, 64, 16, 16, 16]
        concat_2 = torch.cat([trans_2, down_4+res_4], dim=1) # -> [1, 96, 16, 16, 16]
        up_2 = self.up_2(concat_2) # -> [1, 32, 16, 16, 16]
        res_up_2 = self.res_up2(concat_2)
        
        trans_3 = self.trans_3(up_2+res_up_2) # -> [1, 32, 32, 32, 32]
        concat_3 = torch.cat([trans_3, down_3+res_3], dim=1) # -> [1, 48, 32, 32, 32]
        up_3 = self.up_3(concat_3) # -> [1, 16, 32, 32, 32]
        res_up_3 = self.res_up3(concat_3)
        
        trans_4 = self.trans_4(up_3+res_up_3) # -> [1, 16, 64, 64, 64]
        concat_4 = torch.cat([trans_4, down_2+res_2], dim=1) # -> [1, 24, 64, 64, 64]
        up_4 = self.up_4(concat_4) # -> [1, 8, 64, 64, 64]
        res_up_4 = self.res_up4(concat_4)
        
        trans_5 = self.trans_5(up_4+res_up_4) # -> [1, 8, 128, 128, 128]
        concat_5 = torch.cat([trans_5, down_1+res_1], dim=1) # -> [1, 12, 128, 128, 128]
        up_5 = self.up_5(concat_5) # -> [1, 4, 128, 128, 128]
        
        # Output
        out = self.out(up_5) # -> [1, 3, 128, 128, 128]
        return out


if __name__ == '__main__':

    a = RUNet_recon()
    z1 = torch.zeros([2,1,128,128,128])
    z2 = torch.zeros([2,1,128,128,128])

    s = a(z1,z2)
    print(s.shape)
Model/runet.py
ZhangJD-ong/Iterative-Cycle-consistent-Semi-supervised-Learning-for-fibroglandular-tissue-segmentation
import numpy as np
import torch, os
from torch.utils.data import Dataset, DataLoader
from glob import glob
import random
import SimpleITK as sitk

def min_max_normalization(img):
    out=(img - np.min(img))/(np.max(img) - np.min(img) + 0.000001 )
    return out


def normalization( img, lmin=1, rmax=None, dividend=None, quantile=1):
    newimg = img.copy()
    newimg = newimg.astype(np.float32)
    if quantile is not None:
        maxval = round(np.percentile(newimg, 100 - quantile))
        minval = round(np.percentile(newimg, quantile))
        newimg[newimg >= maxval] = maxval
        newimg[newimg <= minval] = minval

    if lmin is not None:
        newimg[newimg < lmin] = lmin
    if rmax is not None:
        newimg[newimg > rmax] = rmax

    minval = np.min(newimg)
    if dividend is None:
        maxval = np.max(newimg)
        newimg = (np.asarray(newimg).astype(np.float32) - minval) / (maxval - minval)
    else:
        newimg = (np.asarray(newimg).astype(np.float32) - minval) / dividend
    return newimg, minval, maxval


def normalization_fix(img, minval, maxval, lmin=1):
    newimg = img.copy()
    newimg = newimg.astype(np.float32)
    if lmin is not None:
        newimg[newimg < lmin] = lmin

    newimg = (np.asarray(newimg).astype(np.float32) - minval) / (maxval - minval)
    return newimg



def load(file):
    itkimage = sitk.ReadImage(file)
    image = sitk.GetArrayFromImage(itkimage)
    return image

def random_crop_3d(pre, pos, sub, gt, crop_size):
    cor_box = maskcor_extract_3d(gt)
    random_x_min, random_x_max = max(cor_box[0, 1] - crop_size[0], 0), min(cor_box[0, 0], pre.shape[0] - crop_size[0])
    random_y_min, random_y_max = max(cor_box[1, 1] - crop_size[1], 0), min(cor_box[1, 0], pre.shape[1] - crop_size[1])
    random_z_min, random_z_max = max(cor_box[2, 1] - crop_size[2], 0), min(cor_box[2, 0], pre.shape[2] - crop_size[2])
    if random_x_min > random_x_max:
        random_x_min, random_x_max = cor_box[0, 0], cor_box[0, 1] - crop_size[0]
    if random_y_min > random_y_max:
        random_y_min, random_y_max = cor_box[1, 0], cor_box[1, 1] - crop_size[1]
    if random_z_min > random_z_max:
        random_z_min, random_z_max = cor_box[2, 0], cor_box[2, 1] - crop_size[2]

    x_random = random.randint(random_x_min, random_x_max)
    y_random = random.randint(random_y_min, random_y_max)
    z_random = random.randint(random_z_min, random_z_max)

    pre_patch = pre[x_random:x_random + crop_size[0], y_random:y_random + crop_size[1],
                z_random:z_random + crop_size[2]]
    pos_patch = pos[x_random:x_random + crop_size[0], y_random:y_random + crop_size[1],
                z_random:z_random + crop_size[2]]
    sub_patch = sub[x_random:x_random + crop_size[0], y_random:y_random + crop_size[1],
                z_random:z_random + crop_size[2]]
    gt_patch = gt[x_random:x_random + crop_size[0], y_random:y_random + crop_size[1],
               z_random:z_random + crop_size[2]]

    return pre_patch, pos_patch, sub_patch, gt_patch


def maskcor_extract_3d(mask, padding=(5, 5, 5)):
    # mask_s = mask.shape
    p = np.where(mask > 0)
    a = np.zeros([3, 2], dtype=np.int)
    for i in range(3):
        s = p[i].min()
        e = p[i].max() + 1

        ss = s - padding[i]
        ee = e + padding[i]
        if ss < 0:
            ss = 0
        if ee > mask.shape[i]:
            ee = mask.shape[i]

        a[i, 0] = ss
        a[i, 1] = ee
    return a

class Img_DataSet(Dataset):
    def __init__(self, x0,x1, cut_param):
        self.x0 = x0
        self.x1 = x1
        

        self.ori_shape = self.x0.shape
        self.cut_param = cut_param

        self.x0 = self.padding_img(self.x0, self.cut_param)
        self.x0 = self.extract_ordered_overlap(self.x0, self.cut_param)
        self.x1 = self.padding_img(self.x1, self.cut_param)
        self.x1 = self.extract_ordered_overlap(self.x1, self.cut_param)

        self.new_shape = self.x0.shape
        
    def __getitem__(self, index):
        pre = self.x0[index]
        pos = self.x1[index]

        
        return torch.from_numpy(pre).type(torch.float32),torch.from_numpy(pos).type(torch.float32)
    
    def __len__(self):
        return len(self.x0)


    def padding_img(self, img, C):
        assert (len(img.shape) == 3)  # 3D array
        img_s, img_h, img_w = img.shape
        leftover_s = (img_s - C['patch_s']) % C['stride_s']
        leftover_h = (img_h - C['patch_h']) % C['stride_h']
        leftover_w = (img_w - C['patch_w']) % C['stride_w']
        if (leftover_s != 0):
            s = img_s + (C['stride_s'] - leftover_s)
        else:
            s = img_s

        if (leftover_h != 0):
            h = img_h + (C['stride_h'] - leftover_h)
        else:
            h = img_h

        if (leftover_w != 0):
            w = img_w + (C['stride_w'] - leftover_w)
        else:
            w = img_w

        tmp_full_imgs = np.zeros((s, h, w))
        tmp_full_imgs[:img_s, :img_h, 0:img_w] = img
        #print("Padded images shape: " + str(tmp_full_imgs.shape))
        return tmp_full_imgs

    # Divide all the full_imgs in pacthes
    def extract_ordered_overlap(self, img, C):
        assert (len(img.shape) == 3)  # 3D arrays
        img_s, img_h, img_w = img.shape
        assert ((img_h - C['patch_h']) % C['stride_h'] == 0
                and (img_w - C['patch_w']) % C['stride_w'] == 0
                and (img_s - C['patch_s']) % C['stride_s'] == 0)
        N_patches_s = (img_s - C['patch_s']) // C['stride_s'] + 1
        N_patches_h = (img_h - C['patch_h']) // C['stride_h'] + 1
        N_patches_w = (img_w - C['patch_w']) // C['stride_w'] + 1
        N_patches_img = N_patches_s * N_patches_h * N_patches_w
#        print("Patches number of the image:{} [s={} | h={} | w={}]"\
#               .format(N_patches_img, N_patches_s, N_patches_h, N_patches_w))
        patches = np.empty((N_patches_img, C['patch_s'], C['patch_h'], C['patch_w']))
        iter_tot = 0  # iter over the total number of patches (N_patches)
        for s in range(N_patches_s):  # loop over the full images
            for h in range(N_patches_h):
                for w in range(N_patches_w):
                    patch = img[s * C['stride_s'] : s * C['stride_s']+C['patch_s'],
                            h * C['stride_h']: h * C['stride_h']+C['patch_h'],
                            w * C['stride_w']: w * C['stride_w']+C['patch_w']]
                   
                    patches[iter_tot] = patch
                    iter_tot += 1  # total
        assert (iter_tot == N_patches_img)
        return patches  # array with all the full_imgs divided in patches


class Recompone_tool():
    def __init__(self, img_ori_shape, img_new_shape, Cut_para):
        self.result = None
        self.ori_shape = img_ori_shape
        self.new_shape = img_new_shape
        self.C = Cut_para

    def add_result(self, tensor):
        # tensor = tensor.detach().cpu() # shape: [N,class,s,h,w]
        # tensor_np = np.squeeze(tensor_np,axis=0)
        if self.result is not None:
            self.result = torch.cat((self.result, tensor), dim=0)
        else:
            self.result = tensor

    def recompone_overlap(self):
        """
        :param preds: output of model  shape：[N_patchs_img,3,patch_s,patch_h,patch_w]
        :return: result of recompone output shape: [3,img_s,img_h,img_w]
        """
        patch_s = self.result.shape[2]
        patch_h = self.result.shape[3]
        patch_w = self.result.shape[4]
        N_patches_s = (self.new_shape[0] - patch_s) // self.C['stride_s'] + 1
        N_patches_h = (self.new_shape[1] - patch_h) // self.C['stride_h'] + 1
        N_patches_w = (self.new_shape[2] - patch_w) // self.C['stride_w'] + 1

        N_patches_img = N_patches_s * N_patches_h * N_patches_w
        #print("N_patches_s/h/w:", N_patches_s, N_patches_h, N_patches_w)
        #print("N_patches_img: " + str(N_patches_img))
        assert (self.result.shape[0] == N_patches_img)

        full_prob = torch.zeros(( self.new_shape[0], self.new_shape[1],self.new_shape[2]))  # itialize to zero mega array with sum of Probabilities
        full_sum = torch.zeros((self.new_shape[0], self.new_shape[1], self.new_shape[2]))
        k = 0  # iterator over all the patches
        for s in range(N_patches_s):
            for h in range(N_patches_h):
                for w in range(N_patches_w):
                    #print(k,self.result[k].squeeze().sum())
                    full_prob[s * self.C['stride_s']:s * self.C['stride_s'] + patch_s,
                                 h * self.C['stride_h']:h  * self.C['stride_h'] + patch_h,
                                 w * self.C['stride_w']:w * self.C['stride_w'] + patch_w] += self.result[k].squeeze()
                    full_sum[s * self.C['stride_s']:s * self.C['stride_s'] + patch_s,
                                h * self.C['stride_h']:h * self.C['stride_h'] + patch_h,
                                w * self.C['stride_w']:w * self.C['stride_w'] + patch_w] += 1
                    k += 1
        assert (k == self.result.size(0))
        assert (torch.min(full_sum) >= 1.0)  # at least one
        final_avg = full_prob / full_sum
        # print(final_avg.size())
        img = final_avg[:self.ori_shape[0], :self.ori_shape[1], :self.ori_shape[2]]
        return img

def cal_newshape( img, C):
    assert (len(img.shape) == 3)  # 3D array
    img_s, img_h, img_w = img.shape
    leftover_s = (img_s - C['patch_s']) % C['stride_s']
    leftover_h = (img_h - C['patch_h']) % C['stride_h']
    leftover_w = (img_w - C['patch_w']) % C['stride_w']
    if (leftover_s != 0):
        s = img_s + (C['stride_s'] - leftover_s)
    else:
        s = img_s

    if (leftover_h != 0):
        h = img_h + (C['stride_h'] - leftover_h)
    else:
        h = img_h

    if (leftover_w != 0):
        w = img_w + (C['stride_w'] - leftover_w)
    else:
        w = img_w

    return np.zeros((s, h, w)).shape

def package_torch(pre_patch, pos_patch, sub_patch, gt_patch):
    pre_patch = torch.from_numpy(pre_patch[np.newaxis,np.newaxis,:])
    pos_patch = torch.from_numpy(pos_patch[np.newaxis,np.newaxis,:])
    sub_patch = torch.from_numpy(sub_patch[np.newaxis,np.newaxis,:])
    gt_patch = torch.from_numpy(gt_patch[np.newaxis,np.newaxis,:])
    return pre_patch, pos_patch, sub_patch, gt_patch


def label_coding(gt, label_num):
    new_gt = np.zeros([label_num, gt.shape[0], gt.shape[1], gt.shape[2]])
    a1 = np.zeros([gt.shape[0], gt.shape[1], gt.shape[2]])
    a2 = np.zeros([gt.shape[0], gt.shape[1], gt.shape[2]])

    if label_num == 2:
        ff1 = np.where(gt == 1)
        a1[ff1] = 1
        ff2 = np.where(gt == 2)
        a2[ff2] = 1
        new_gt[0, :] = a1
        new_gt[1, :] = a2

    return new_gt


def Test_all_Datasets(dataset_path,size):
    f = open(os.path.join(dataset_path, 'test.txt'))
    data_list = f.read().splitlines()
    print("The number of test samples is: ", len(data_list))
    for file in data_list:
        print("\nStart Evaluate: ", file)
        x1,x1_min,x1_max = normalization(load(os.path.join(dataset_path,'w_label',file,'P1_reg.nii.gz')))
        x1 = x1.astype(np.float32)
        x0 = normalization_fix(load(os.path.join(dataset_path,'w_label',file,'P0.nii.gz')),x1_min,x1_max)
        x0 = x0.astype(np.float32)
        gt = load(os.path.join(dataset_path,'w_label',file,'Tissues.nii.gz')).astype(np.float32)
        new_gt = np.zeros(gt.shape).astype(np.float32)
        ff = np.where(gt==2)
        new_gt[ff]=1
        breast = load(os.path.join(dataset_path,'w_label',file,'Breast.nii.gz')).astype(np.float32)
        breast_erode = sitk.GetArrayFromImage(sitk.BinaryErode(sitk.ReadImage(os.path.join(dataset_path,'w_label',file,'Breast.nii.gz')),(2,2,2))).astype(np.float32)

        itkimage = sitk.ReadImage(os.path.join(dataset_path,'w_label',file, 'P0.nii.gz'))
        original_shape = gt.shape
        new_shape = cal_newshape(gt,size)

        yield Img_DataSet(x0,x1,size),original_shape, new_shape,itkimage, file,new_gt,breast,breast_erode
            
dataset/dataset_lits_test.py
ZhangJD-ong/Iterative-Cycle-consistent-Semi-supervised-Learning-for-fibroglandular-tissue-segmentation
import random
import numpy as np
import SimpleITK as sitk
import os
from torch.utils.data import Dataset
  
class Lits_DataSet(Dataset):
    def __init__(self,root, size = (32,128,128)):
        self.root = root
        self.size = size
        f = open(os.path.join(self.root,'train_1-10.txt'))
        self.filename = f.read().splitlines()

    def __getitem__(self, index):

        files = self.filename[index]
        source, file = files.split('-')[0], files.split('-')[1]

        if source == 'w':
            x1,x1_min,x1_max = self.normalization(self.load(os.path.join(self.root,'w_label',file,'P1_reg.nii.gz')))
            x1 = x1.astype(np.float32)
            x0 = self.normalization_fix(self.load(os.path.join(self.root,'w_label',file,'P0.nii.gz')),x1_min,x1_max)
            x0 = x0.astype(np.float32)
            gt = self.load(os.path.join(self.root,'w_label',file,'Tissues.nii.gz')).astype(np.float32)
            breast = self.load(os.path.join(self.root,'w_label',file,'Breast.nii.gz')).astype(np.float32)
            ff = np.where(gt==2)
            label = np.zeros(gt.shape).astype(np.float32)
            label[ff]=1
            x0_patch,x1_patch,gt_patch = self.random_crop_3d(label,x0,x1,breast,self.size)
            label_index = np.ones(x0_patch.shape)
        elif source == 'wo':
            x1,x1_min,x1_max = self.normalization(self.load(os.path.join(self.root,'wo_label',file,'P1_reg.nii.gz')))
            x1 = x1.astype(np.float32)
            x0 = self.normalization_fix(self.load(os.path.join(self.root,'wo_label',file,'P0.nii.gz')),x1_min,x1_max)
            x0 = x0.astype(np.float32)
            label = np.zeros(x1.shape).astype(np.float32)
            breast = self.load(os.path.join(self.root,'wo_label',file,'Breast.nii.gz')).astype(np.float32)
            x0_patch,x1_patch,gt_patch = self.random_crop_3d(label,x0,x1,breast,self.size)
            label_index = np.zeros(x0_patch.shape)
        else:
            print('Error!')

        return x0_patch[np.newaxis,:],x1_patch[np.newaxis,:],gt_patch[np.newaxis,:],label_index[np.newaxis,:]

    def __len__(self):
        return len(self.filename)

    def random_crop_3d(self,gt,x0,x1,breast,crop_size):

        cor_box = self.maskcor_extract_3d(breast)
        random_x_min, random_x_max = max(cor_box[0, 0] - crop_size[0], 0), min(cor_box[0, 1],
                                                                               gt.shape[0] - crop_size[0])
        random_y_min, random_y_max = max(cor_box[1, 0] - crop_size[1], 0), min(cor_box[1, 1],
                                                                               gt.shape[1] - crop_size[1])
        random_z_min, random_z_max = max(cor_box[2, 0] - crop_size[2], 0), min(cor_box[2, 1],
                                                                               gt.shape[2] - crop_size[2])

        x_random = random.randint(random_x_min, random_x_max)
        y_random = random.randint(random_y_min, random_y_max)
        z_random = random.randint(random_z_min, random_z_max)

        gt_patch = gt[x_random:x_random + crop_size[0], y_random:y_random + crop_size[1], z_random:z_random + crop_size[2]]
        x0_patch = x0[x_random:x_random + crop_size[0], y_random:y_random + crop_size[1], z_random:z_random + crop_size[2]]
        x1_patch = x1[x_random:x_random + crop_size[0], y_random:y_random + crop_size[1], z_random:z_random + crop_size[2]]


        return x0_patch,x1_patch,gt_patch

    def label_coding(self,gt,label_num):
        new_gt = np.zeros([label_num+1,gt.shape[0],gt.shape[1],gt.shape[2]])
        a1 = np.zeros([gt.shape[0],gt.shape[1],gt.shape[2]])
        a2 = np.zeros([gt.shape[0], gt.shape[1], gt.shape[2]])
        a3 = np.ones([gt.shape[0], gt.shape[1], gt.shape[2]])

        if label_num ==2:
            ff1 = np.where(gt==1)
            a1[ff1] = 1
            a3[ff1] = 0
            ff2 = np.where(gt==2)
            a2[ff2] = 1
            a3[ff2] = 0
            new_gt[0,:] = a1
            new_gt[1,:] = a2
            new_gt[2,:] = a3


        return new_gt


    def normalization(self, img, lmin=1, rmax=None, dividend=None, quantile=1):
        newimg = img.copy()
        newimg = newimg.astype(np.float32)
        if quantile is not None:
            maxval = round(np.percentile(newimg, 100 - quantile))
            minval = round(np.percentile(newimg, quantile))
            newimg[newimg >= maxval] = maxval
            newimg[newimg <= minval] = minval

        if lmin is not None:
            newimg[newimg < lmin] = lmin
        if rmax is not None:
            newimg[newimg > rmax] = rmax

        minval = np.min(newimg)
        if dividend is None:
            maxval = np.max(newimg)
            newimg = (np.asarray(newimg).astype(np.float32) - minval) / (maxval - minval)
        else:
            newimg = (np.asarray(newimg).astype(np.float32) - minval) / dividend
        return newimg, minval, maxval


    def normalization_fix(self, img, minval, maxval, lmin=1):
        newimg = img.copy()
        newimg = newimg.astype(np.float32)
        if lmin is not None:
            newimg[newimg < lmin] = lmin

        newimg = (np.asarray(newimg).astype(np.float32) - minval) / (maxval - minval)
        return newimg

    def load(self,file):
        itkimage = sitk.ReadImage(file)
        image = sitk.GetArrayFromImage(itkimage)
        return image

    def maskcor_extract_3d(self,mask, padding=(0, 0, 0)):
        p = np.where(mask > 0)
        a = np.zeros([3, 2], dtype=np.int)
        for i in range(3):
            s = p[i].min()
            e = p[i].max() + 1

            ss = s - padding[i]
            ee = e + padding[i]
            if ss < 0:
                ss = 0
            if ee > mask.shape[i]:
                ee = mask.shape[i]

            a[i, 0] = ss
            a[i, 1] = ee
        return a







dataset/dataset_lits_train.py
ZhangJD-ong/Iterative-Cycle-consistent-Semi-supervised-Learning-for-fibroglandular-tissue-segmentation
import argparse
import os
from utils import util
import torch


class BaseOptions():
    """This class defines options used during both training and test time.

    It also implements several helper functions such as parsing, printing, and saving the options.
    It also gathers additional options defined in <modify_commandline_options> functions in both dataset class and model class.
    """

    def __init__(self):
        """Reset the class; indicates the class hasn't been initailized"""
        self.initialized = False

    def initialize(self, parser):
        """Define the common options that are used in both training and test."""
        # basic parameters
        
        parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
        parser.add_argument('--num_threads', default=4, type=int, help='# threads for loading data')
        
        parser.add_argument('--batch_size', type=int, default=6, help='input train batch size')
        parser.add_argument('--test_batch', type=int, default=6, help='input test batch size')
        parser.add_argument('--epoch', type=int, default=300, help='number of epochs')
        parser.add_argument('--step', type=int, default=50, help='number of epochs to adjust learning rate')
        parser.add_argument('--datapath', default = r'/hpc/data', help='path of the raw data')
        parser.add_argument('--lr', type=float, default=0.001, help='initial learning rate of net for adam')
        parser.add_argument('--model_save_fre', type=int, default=100, help='frequency of saving model')
        parser.add_argument('--weight', type=float, default=5, help='weight of self-reconstruction loss')
        parser.add_argument('--test_fre', type=int, default=600, help='frequency of testing the model')
        parser.add_argument('--patch_size', type=int, default=(32,128,128), help='the size of crop patch')
        parser.add_argument('--patch_stride', type=int, default=(32,64,64), help='the stride of patch')
        parser.add_argument('--gpu_ids', type=str, default='0', help='gpu ids: e.g. 0  0,1,2, 0,2. use -1 for CPU')
        parser.add_argument('--task_name', type=str, default='Semi-supervised', help='the current task name')
        parser.add_argument('--seg_pretrain_path', type=str, default='/hpc/data/Pretrained/Segmentation_model.pth', help='pretrained segmentation model path')
        parser.add_argument('--recon_pretrain_path', type=str, default='/hpc/data/Pretrained/Reconstruction_model.pth', help='pretrained reconstruction model path')
        self.initialized = True
        return parser

    def gather_options(self):
        """Initialize our parser with basic options(only once).
        Add additional model-specific and dataset-specific options.
        These options are defined in the <modify_commandline_options> function
        in model and dataset classes.
        """
        if not self.initialized:  # check if it has been initialized
            parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
            parser = self.initialize(parser)

        # get the basic options
        opt, _ = parser.parse_known_args()
        # save and return the parser
        self.parser = parser
        return parser.parse_args()

    def print_options(self, opt):
        """Print and save options

        It will print both current options and default values(if different).
        It will save options into a text file / [checkpoints_dir] / opt.txt
        """
        message = ''
        message += '----------------- Options ---------------\n'
        for k, v in sorted(vars(opt).items()):
            comment = ''
            default = self.parser.get_default(k)
            if v != default:
                comment = '\t[default: %s]' % str(default)
            message += '{:>25}: {:<30}{}\n'.format(str(k), str(v), comment)
        message += '----------------- End -------------------'
        print(message)

        # save to the disk
        expr_dir = os.path.join(opt.checkpoints_dir, 'model_parameter_list')
        util.mkdirs(expr_dir)
        file_name = os.path.join(expr_dir, '{train_opt.txt')
        with open(file_name, 'wt') as opt_file:
            opt_file.write(message)
            opt_file.write('\n')

    def parse(self):
        """Parse our options, create checkpoints directory suffix, and set up gpu device."""
        opt = self.gather_options()
        opt.isTrain = self.isTrain   # train or test

        self.print_options(opt)

        self.opt = opt
        return self.opt


options/BasicOptions.py
ZhangJD-ong/Iterative-Cycle-consistent-Semi-supervised-Learning-for-fibroglandular-tissue-segmentation
from options.BasicOptions import BaseOptions


class Options_x(BaseOptions):
    """This class includes training options.

    It also includes shared options defined in BaseOptions.
    """

    def initialize(self, parser):
        parser = BaseOptions.initialize(self, parser)
        # visdom and HTML visualization parameters
        parser.add_argument('--name', type=str, default='Tumor_seg', help='name_of_the_project')
        self.isTrain = True
        return parser

options/Options.py
ZhangJD-ong/Iterative-Cycle-consistent-Semi-supervised-Learning-for-fibroglandular-tissue-segmentation
import torch
import gc
import numpy as np
import torch.optim as optim
import SimpleITK as sitk
from options.Options import Options_x
from tqdm import tqdm
from Model.runet import RUNet_seg
from torch.utils.data import DataLoader
from utils import logger,util
from utils.metrics import seg_metric
import torch.nn as nn
import os
from dataset.dataset_lits_test import Test_all_Datasets,Recompone_tool
from collections import OrderedDict

def load(file):
    itkimage = sitk.ReadImage(file)
    image = sitk.GetArrayFromImage(itkimage)
    return image


def test_all(model_name='model_200.pth'):
    opt = Options_x().parse()  # get training options
    device = torch.device('cuda' if torch.cuda.is_available() else "cpu")


    model = RUNet_seg(2,1,16).to(device)
    ckpt = torch.load(opt.checkpoints_dir + '/' + opt.task_name + '/model/' + model_name,map_location=device)
    model.load_state_dict(ckpt['model'])

    save_result_path = os.path.join(opt.checkpoints_dir, opt.task_name, 'results')
    util.mkdir(save_result_path)
    model.eval()
    log_test = logger.Test_Logger(save_result_path,"results")
    cut_param = {'patch_s': opt.patch_size[0], 'patch_h': opt.patch_size[1], 'patch_w': opt.patch_size[2],
                 'stride_s': opt.patch_stride[0], 'stride_h': opt.patch_stride[1], 'stride_w': opt.patch_stride[2]}
    datasets = Test_all_Datasets(opt.datapath, cut_param)

    for img_dataset, original_shape, new_shape, itkimage, file_idx, gt, Whole_breast, breast_erode in datasets:
        save_tool = Recompone_tool(original_shape, new_shape, cut_param)
        dataloader = DataLoader(img_dataset, batch_size=opt.test_batch, num_workers=opt.num_threads, shuffle=False)
        with torch.no_grad():
            for pre,pos in tqdm(dataloader):
                pre,pos = pre.unsqueeze(1).to(device), pos.unsqueeze(1).to(device)
                output = model(pre,pos)

                save_tool.add_result(output.detach().cpu())

        pred = save_tool.recompone_overlap()
        recon = (pred.numpy() > 0.5).astype(np.uint16)*breast_erode.astype(np.int16)

        DSC, PPV, SEN, HD = seg_metric(recon, gt,itkimage)

        index_results = OrderedDict({'DSC': DSC,'PPV': PPV,'SEN': SEN,'HD': HD})
        log_test.update(file_idx,index_results)

        Pred = sitk.GetImageFromArray(np.array(recon))

        sitk.WriteImage(Pred,os.path.join(save_result_path,file_idx+'.nii.gz'))
        del pred,recon,Pred,save_tool,gt
        gc.collect()


if __name__ == '__main__':
    test_all('latest_model.pth')
                

test.py
ZhangJD-ong/Iterative-Cycle-consistent-Semi-supervised-Learning-for-fibroglandular-tissue-segmentation
import torch
import numpy as np
import torch.optim as optim
import SimpleITK as sitk
from options.Options import Options_x
from dataset.dataset_lits_train import Lits_DataSet
from Model.runet import RUNet_seg,RUNet_recon
from torch.utils.data import DataLoader
from utils.common import adjust_learning_rate
from utils import logger,util
import torch.nn as nn
import torch.nn.functional as F
from utils.metrics import LossAverage, DiceLoss, DiceLossWeight, weight_cal
from test import test_all
import os
from collections import OrderedDict


def train(train_dataloader,epoch):
    print("=======Epoch:{}======Learning_rate:{}=========".format(epoch, optimizer.param_groups[0]['lr']))

    Loss = LossAverage()
    SUPERVISED_Loss = LossAverage()
    SELF_Loss = LossAverage()
    RECON_Loss = LossAverage()
    model.train()
    model_recon.eval()

    for i, (pre, pos, gt,label_index) in enumerate(train_dataloader):  # inner loop within one epoch
        ##main model update param
        pre, pos, gt,label_index = pre.to(device), pos.to(device), gt.type(torch.float32).to(device),label_index.type(torch.float32).to(device)
        pred1 = model(pre,pos)
        recon1 = model_recon(pred1,pre)
        pred2 = model(pre, recon1)
        recon2 = model_recon(pred2,pre)
        pred3 = model(pre, recon2)
        recon3 = model_recon(pred3,pre)
        pred4 = model(pre, recon3)

        Recon_loss = nn.L1Loss()(recon1,pos) + nn.L1Loss()(recon2,pos) + nn.L1Loss()(recon3,pos)
        Supervised_loss = dice_loss(pred1 * label_index, gt * label_index)
        self_loss = dice_loss(pred2, pred1) + dice_loss(pred2, pred3) + dice_loss(pred3, pred4)
        loss = Supervised_loss + opt.weight*Recon_loss + self_loss

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        adjust_learning_rate(optimizer, epoch, opt)

        Loss.update(loss.item(), 1)
        SUPERVISED_Loss.update(Supervised_loss.item(), 1)
        RECON_Loss.update(Recon_loss.item(), 1)
        SELF_Loss.update(self_loss.item(), 1)

    return OrderedDict({'Loss': Loss.avg, 'Supervised_Loss': SUPERVISED_Loss.avg, 'Self_Loss': SELF_Loss.avg, 'Recon_Loss': RECON_Loss.avg})


if __name__ == '__main__':
    opt = Options_x().parse()   # get training options
    device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
    print(device)

    model = RUNet_seg(2, 1, 16).to(device)
    ckpt = torch.load(opt.seg_pretrain_path,map_location=device)
    model.load_state_dict(ckpt['model'])
    model_recon = RUNet_recon().to(device)
    ckpt = torch.load(opt.recon_pretrain_path,map_location=device)
    model_recon.load_state_dict(ckpt['model'])


    save_path = opt.checkpoints_dir
    dice_loss = DiceLoss()
    dice_loss_weight = DiceLossWeight()
    bce_loss = torch.nn.BCELoss()
        
    save_result_path = os.path.join(save_path,opt.task_name)
    util.mkdir(save_result_path)

    optimizer = optim.Adam(model.parameters(), lr=opt.lr,weight_decay=1e-5)

    model_save_path = os.path.join(save_result_path,'model')
    util.mkdir(model_save_path)
    logger_save_path = os.path.join(save_result_path,'logger')
    util.mkdir(logger_save_path)
    log_train = logger.Train_Logger(logger_save_path,"train_log")

    train_dataset = Lits_DataSet(opt.datapath)
    train_dataloader = DataLoader(dataset=train_dataset, batch_size=opt.batch_size, \
                                  num_workers=opt.num_threads, shuffle=True)
    for epoch in range(opt.epoch):
        epoch = epoch + 1
        train_log = train(train_dataloader, epoch)
        log_train.update(epoch, train_log)

        state = {'model': model.state_dict(), 'epoch': epoch}
        torch.save(state, os.path.join(model_save_path, 'latest_model.pth'))

        if epoch % opt.model_save_fre == 0:
            torch.save(state, os.path.join(model_save_path, 'model_' + np.str(epoch) + '.pth'))

        torch.cuda.empty_cache()

    test_all('latest_model.pth')




 
        

            
            

            
            
            
            
            
            

train.py
ZhangJD-ong/Iterative-Cycle-consistent-Semi-supervised-Learning-for-fibroglandular-tissue-segmentation
import SimpleITK as sitk
import numpy as np
#import pytorch_ssim
#from skimage.metrics import structural_similarity
import math
from torch.autograd import Variable
from scipy import ndimage
import torch, random



def normalization(img):
    out=(img - np.min(img))/(np.max(img) - np.min(img) + 0.000001 )
    return out


def normalization_test (img):
    out=(img - np.min(img))/(np.max(img) - np.min(img) + 0.000001 )
    return out, np.max(img), np.min(img)

def center_crop_3d(img, label, slice_num=16):
    if img.shape[0] < slice_num:
        return None
    left_x = img.shape[0]//2 - slice_num//2
    right_x = img.shape[0]//2 + slice_num//2

    crop_img = img[left_x:right_x]
    crop_label = label[left_x:right_x]
    return crop_img, crop_label

def load_file_name_list(file_path):
    file_name_list = []
    with open(file_path, 'r') as file_to_read:
        while True:
            lines = file_to_read.readline().strip()  # 整行读取数据
            if not lines:
                break
                pass
            file_name_list.append(lines)
            pass
    return file_name_list


def MaskContour(image, position='xy', line=1):
    itkimage = sitk.GetImageFromArray(image)
    if position == 'xy':
        erode_m = [line, line, 0]
    elif position == 'yz':
        erode_m = [0, line, line]
    elif position == 'zx':
        erode_m = [line, 0, line]
    else:
        erode_m = [line, line, 0]

    mask = sitk.GetArrayFromImage(sitk.BinaryErode(itkimage, erode_m))
    boundary = image - mask
    out = sitk.GetImageFromArray(boundary)
    return out


def print_network(net):
    num_params = 0
    for param in net.parameters():
        num_params += param.numel()
    print(net)
    print('Total number of parameters: %d' % num_params)

def adjust_learning_rate(optimizer, epoch, opt):
    """Sets the learning rate to the initial LR decayed by 10 every 10 epochs"""
    lr = opt.lr * (0.5 ** (epoch // opt.step))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

def adjust_learning_rate_V2(optimizer, lr):
    """Sets the learning rate to a fixed number"""
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
        
        
        
def get_mse(img1, img2):
    mse = np.mean( (img1 - img2) ** 2 )

    return mse


def get_psnr(img1, img2):
    mse = np.mean( (img1 - img2) ** 2 )
    if mse == 0:
        return 100
    PIXEL_MAX = 1.0
    return 10 * math.log10(PIXEL_MAX / math.sqrt(mse))

'''
def get_ssim(img1,img2):
    n=img1.shape[0]
    out = 0
    for i in range(n):
        out+=structural_similarity(img1[i].squeeze(),img2[i].squeeze())
        
    return out/n
'''    
    

def save_result(low_dose, high_dose, output, i, epoch):
    def save_img(img, name):
        # img = SimpleITK.GetImageFromArray(img[0,0].cpu().detach().numpy())
        img = sitk.GetImageFromArray(img)
        sitk.WriteImage(img, 'result/image/'+name+'.nii.gz')
        
    save_img(low_dose, 'low_dose_epoch_'+str(epoch) + "_" + str(i))
    save_img(high_dose, 'high_dose_epoch_'+str(epoch) + "_" + str(i))
    save_img(output, 'output_epoch_'+str(epoch) + "_" + str(i))


def de_normalization(img,max_x,min_x):
    return img*(max_x - min_x) + min_x
utils/common.py
ZhangJD-ong/Iterative-Cycle-consistent-Semi-supervised-Learning-for-fibroglandular-tissue-segmentation
import pandas as pd
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
import torch,random
import numpy as np
from collections import OrderedDict

class Train_Logger():
    def __init__(self,save_path,save_name):
        self.log = None
        self.summary = None
        self.save_path = save_path
        self.save_name = save_name

    def update(self,epoch,train_log):
        item = OrderedDict({'epoch':epoch})
        item.update(train_log)
        # item = dict_round(item,4)
        print("\033[0;33mTrain:\033[0m",train_log)
        self.update_csv(item)
        self.update_tensorboard(item)

    def update_csv(self,item):
        tmp = pd.DataFrame(item,index=[0])
        if self.log is not None:
            self.log = self.log.append(tmp, ignore_index=True)
        else:
            self.log = tmp
        self.log.to_csv('%s/%s.csv' %(self.save_path,self.save_name), index=False)

    def update_tensorboard(self,item):
        if self.summary is None:
            self.summary = SummaryWriter('%s/' % self.save_path)
        epoch = item['epoch']
        for key,value in item.items():
            if key != 'epoch': self.summary.add_scalar(key, value, epoch)

class Test_Logger():
    def __init__(self,save_path,save_name):
        self.log = None
        self.summary = None
        self.save_path = save_path
        self.save_name = save_name

    def update(self,name,log):
        item = OrderedDict({'img_name':name})
        item.update(log)
        print("\033[0;33mTest:\033[0m",log)
        self.update_csv(item)

    def update_csv(self,item):
        tmp = pd.DataFrame(item,index=[0])
        if self.log is not None:
            self.log = self.log.append(tmp, ignore_index=True)
        else:
            self.log = tmp
        self.log.to_csv('%s/%s.csv' %(self.save_path,self.save_name), index=False)

def setpu_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    torch.backends.cudnn.deterministic=True
    random.seed(seed)

def dict_round(dic,num):
    for key,value in dic.items():
        dic[key] = round(value,num)
    return dic
utils/logger.py
ZhangJD-ong/Iterative-Cycle-consistent-Semi-supervised-Learning-for-fibroglandular-tissue-segmentation
import torch.nn as nn
import torch.nn.functional as F
import torch
import SimpleITK as sitk
import numpy as np
import sys
from scipy.ndimage import morphology
sys.dont_write_bytecode = True  # don't generate the binray python file .pyc

hdcomputer = sitk.HausdorffDistanceImageFilter()

def weight_cal(recon,pos):
    a1 = nn.L1Loss(size_average=False, reduce=False, reduction='mean')(recon, pos).detach()
    a2 = torch.mean(a1,dim = [1,2,3,4]).view([recon.shape[0],1,1,1,1])

    return a2


class LossAverage(object):
    """Computes and stores the average and current value for calculate average loss"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = round(self.sum / self.count, 4)
        # print(self.val)



class DiceLoss(nn.Module):
    """
    define the dice loss
    """
    def __init__(self):
        super(DiceLoss, self).__init__()

    def forward(self, input, target):
        smooth = 1.
        iflat = input.contiguous().view(-1)
        tflat = target.contiguous().view(-1)
        intersection = (iflat * tflat).sum()

        A_sum = torch.sum(iflat * iflat)
        B_sum = torch.sum(tflat * tflat)

        return 1-((2. * intersection + smooth) / (A_sum + B_sum + smooth))



class DiceLossWeight(nn.Module):
    """
    define the dice loss
    """
    def __init__(self):
        super(DiceLossWeight, self).__init__()

    def forward(self, input, target,weight_adapt):
        smooth = 1.
        iflat = input.contiguous().view([input.shape[0],-1])
        tflat = target.contiguous().view([input.shape[0],-1])

        intersection = (iflat * tflat).sum(-1)

        A_sum = torch.sum(iflat * iflat,-1)
        B_sum = torch.sum(tflat * tflat,-1)
        loss_batch = 1-((2. * intersection + smooth) / (A_sum + B_sum + smooth))
        loss = loss_batch*weight.squeeze()

        return loss.sum()


"""dice coefficient"""
def dice(pre, gt, tid=1):
    pre=pre==tid   #make it boolean
    gt=gt==tid     #make it boolean
    pre=np.asarray(pre).astype(np.bool)
    gt=np.asarray(gt).astype(np.bool)

    if pre.shape != gt.shape:
        raise ValueError("Shape mismatch: prediction and ground truth must have the same shape.")

    intersection = np.logical_and(pre, gt)
    dsc=(2. * intersection.sum() + 1e-07) / (pre.sum() + gt.sum() + 1e-07)

    return dsc

"""positive predictive value"""
def pospreval(pre,gt,tid=1):
    pre=pre==tid #make it boolean
    gt=gt==tid   #make it boolean
    pre=np.asarray(pre).astype(np.bool)
    gt=np.asarray(gt).astype(np.bool)

    if pre.shape != gt.shape:
        raise ValueError("Shape mismatch: prediction and ground truth must have the same shape.")

    intersection = np.logical_and(pre, gt)
    ppv=(1.0*intersection.sum() + 1e-07) / (pre.sum()+1e-07)

    return ppv

"""sensitivity"""
def sensitivity(pre,gt,tid=1):
    pre=pre==tid #make it boolean
    gt=gt==tid   #make it boolean
    pre=np.asarray(pre).astype(np.bool)
    gt=np.asarray(gt).astype(np.bool)

    if pre.shape != gt.shape:
        raise ValueError("Shape mismatch: prediction and ground truth must have the same shape.")

    intersection = np.logical_and(pre, gt)
    sen=(1.0*intersection.sum()+1e-07) / (gt.sum()+1e-07)

    return sen

def seg_metric(pre,gt,itk_info):

    fake = (pre>0.5).astype(np.float32)
    real = (gt>0.5).astype(np.float32)
    DSC = dice(fake,real)
    PPV = pospreval(fake,real)
    SEN = sensitivity(fake,real)
    real_itk = sitk.GetImageFromArray(real)
    fake_itk = sitk.GetImageFromArray(fake)
    if np.sum(fake) !=0:

        real_itk.SetOrigin(itk_info.GetOrigin())
        real_itk.SetSpacing(itk_info.GetSpacing())
        real_itk.SetDirection(itk_info.GetDirection())

        fake_itk.SetOrigin(itk_info.GetOrigin())
        fake_itk.SetSpacing(itk_info.GetSpacing())
        fake_itk.SetDirection(itk_info.GetDirection())

        hdcomputer.Execute(real_itk>0.5, fake_itk>0.5)
        HD = hdcomputer.GetAverageHausdorffDistance()
    else:
        HD = 100

    return DSC,PPV,SEN,HD





utils/metrics.py
ZhangJD-ong/Iterative-Cycle-consistent-Semi-supervised-Learning-for-fibroglandular-tissue-segmentation
"""This module contains simple helper functions """
from __future__ import print_function
import torch
import numpy as np
from PIL import Image
import os
import SimpleITK as sitk


def connected_domain(img,voxels = 10):
    itk_mask = sitk.GetImageFromArray(img)
    cc_filter = sitk.ConnectedComponentImageFilter()
    cc_filter.SetFullyConnected(True)
    output_mask = cc_filter.Execute(itk_mask)

    lss_filter = sitk.LabelShapeStatisticsImageFilter()
    lss_filter.Execute(output_mask)

    num_connected_label = cc_filter.GetObjectCount()

    np_output_mask = sitk.GetArrayFromImage(output_mask)
    res_mask = np.zeros_like(np_output_mask)+1


    for i in range(1, num_connected_label + 1):
        area = lss_filter.GetNumberOfPixels(i)
        if area < voxels:
            res_mask[np_output_mask == i] = 0

    return res_mask*img


def tensor2im(input_image, imtype=np.uint8):
    """"Converts a Tensor array into a numpy image array.

    Parameters:
        input_image (tensor) --  the input image tensor array
        imtype (type)        --  the desired type of the converted numpy array
    """
    if not isinstance(input_image, np.ndarray):
        if isinstance(input_image, torch.Tensor):  # get the data from a variable
            image_tensor = input_image.data
        else:
            return input_image
        image_numpy = image_tensor[0].cpu().float().numpy()  # convert it into a numpy array
        if image_numpy.shape[0] == 1:  # grayscale to RGB
            image_numpy = np.tile(image_numpy, (3, 1, 1))
        image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0  # post-processing: tranpose and scaling
    else:  # if it is a numpy array, do nothing
        image_numpy = input_image
    return image_numpy.astype(imtype)


def diagnose_network(net, name='network'):
    """Calculate and print the mean of average absolute(gradients)

    Parameters:
        net (torch network) -- Torch network
        name (str) -- the name of the network
    """
    mean = 0.0
    count = 0
    for param in net.parameters():
        if param.grad is not None:
            mean += torch.mean(torch.abs(param.grad.data))
            count += 1
    if count > 0:
        mean = mean / count
    print(name)
    print(mean)


def save_image(image_numpy, image_path, aspect_ratio=1.0):
    """Save a numpy image to the disk

    Parameters:
        image_numpy (numpy array) -- input numpy array
        image_path (str)          -- the path of the image
    """

    image_pil = Image.fromarray(image_numpy)
    h, w, _ = image_numpy.shape

    if aspect_ratio > 1.0:
        image_pil = image_pil.resize((h, int(w * aspect_ratio)), Image.BICUBIC)
    if aspect_ratio < 1.0:
        image_pil = image_pil.resize((int(h / aspect_ratio), w), Image.BICUBIC)
    image_pil.save(image_path)


def print_numpy(x, val=True, shp=False):
    """Print the mean, min, max, median, std, and size of a numpy array

    Parameters:
        val (bool) -- if print the values of the numpy array
        shp (bool) -- if print the shape of the numpy array
    """
    x = x.astype(np.float64)
    if shp:
        print('shape,', x.shape)
    if val:
        x = x.flatten()
        print('mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f' % (
            np.mean(x), np.min(x), np.max(x), np.median(x), np.std(x)))


def mkdirs(paths):
    """create empty directories if they don't exist

    Parameters:
        paths (str list) -- a list of directory paths
    """
    if isinstance(paths, list) and not isinstance(paths, str):
        for path in paths:
            mkdir(path)
    else:
        mkdir(paths)


def mkdir(path):
    """create a single empty directory if it didn't exist

    Parameters:
        path (str) -- a single directory path
    """
    
    if not os.path.exists(path):
        os.makedirs(path)

utils/util.py
ZhangJD-ong/Iterative-Cycle-consistent-Semi-supervised-Learning-for-fibroglandular-tissue-segmentation