3 | 
4 |
8 | 9 | **Paper: Inter-slice Context Residual Learning for 3D Medical Image Segmentation.** 10 | (https://ieeexplore.ieee.org/abstract/document/9245569) 11 | (https://arxiv.org/pdf/2011.14155.pdf) 12 | 13 | ## Requirements 14 | Python 3.6
15 | Torch==1.4.0
16 | Apex==0.1
17 | 18 | ## Usage 19 | 20 | ### 0. Installation 21 | * Clone this repo 22 | ``` 23 | git clone https://github.com/jianpengz/ConResNet.git 24 | cd ConResNet 25 | ``` 26 | ### 1. Data Preparation 27 | * Put the data and image_id_list in `dataset/`. 28 | 29 | ### 2. Training 30 | * Run `run.sh` to start the training. 31 | 32 | ### 3. Evaluation 33 | * Run `python test.py` to start the evaluation. 34 | 35 | ### 7. Citation 36 | If this code is helpful for your study, please cite: 37 | 38 | ``` 39 | @article{zhang2020conresnet, 40 | title={Inter-slice Context Residual Learning for 3D Medical Image Segmentation}, 41 | author={Zhang, Jianpeng and Xie, Yutong and Wang, Yan and Xia, Yong}, 42 | journal={IEEE Transactions on Medical Imaging}, 43 | volume={40}, 44 | number={2}, 45 | pages={661-672}, 46 | year={2021}, 47 | publisher={IEEE} 48 | } 49 | ``` 50 | 51 | ### Contact 52 | Jianpeng Zhang (james.zhang@mail.nwpu.edu.cn) 53 | -------------------------------------------------------------------------------- /dataset/BraTSDataSet.py: -------------------------------------------------------------------------------- 1 | import os.path as osp 2 | import numpy as np 3 | import random 4 | from torch.utils import data 5 | import nibabel as nib 6 | from skimage.transform import resize 7 | 8 | 9 | class BraTSDataSet(data.Dataset): 10 | def __init__(self, root, list_path, max_iters=None, crop_size=(128, 160, 200), scale=True, mirror=True, ignore_label=255): 11 | self.root = root 12 | self.list_path = list_path 13 | self.crop_d, self.crop_h, self.crop_w = crop_size 14 | self.scale = scale 15 | self.ignore_label = ignore_label 16 | self.is_mirror = mirror 17 | self.img_ids = [i_id.strip().split() for i_id in open(self.root + self.list_path)] 18 | 19 | if not max_iters==None: 20 | self.img_ids = self.img_ids * int(np.ceil(float(max_iters) / len(self.img_ids))) 21 | self.files = [] 22 | for item in self.img_ids: 23 | filepath = item[0] +'/'+ osp.splitext(osp.basename(item[0]))[0] 24 | flair_path = filepath + '_flair.nii.gz' 25 | t1_path = filepath + '_t1.nii.gz' 26 | t1ce_path = filepath + '_t1ce.nii.gz' 27 | t2_path = filepath + '_t2.nii.gz' 28 | label_path = filepath + '_seg.nii.gz' 29 | name = osp.splitext(osp.basename(filepath))[0] 30 | flair_file = osp.join(self.root, flair_path) 31 | t1_file = osp.join(self.root, t1_path) 32 | t1ce_file = osp.join(self.root, t1ce_path) 33 | t2_file = osp.join(self.root, t2_path) 34 | label_file = osp.join(self.root, label_path) 35 | self.files.append({ 36 | "flair": flair_file, 37 | "t1": t1_file, 38 | "t1ce": t1ce_file, 39 | "t2": t2_file, 40 | "label": label_file, 41 | "name": name 42 | }) 43 | print('{} images are loaded!'.format(len(self.img_ids))) 44 | 45 | def __len__(self): 46 | return len(self.files) 47 | 48 | def id2trainId(self, label): 49 | shape = label.shape 50 | results_map = np.zeros((3, shape[0], shape[1], shape[2])) 51 | 52 | NCR_NET = (label == 1) 53 | ET = (label == 4) 54 | WT = (label >= 1) 55 | TC = np.logical_or(NCR_NET, ET) 56 | 57 | results_map[0,:,:,:] = np.where(ET, 1, 0) 58 | results_map[1, :, :, :] = np.where(WT, 1, 0) 59 | results_map[2, :, :, :] = np.where(TC, 1, 0) 60 | return results_map 61 | 62 | def truncate(self, MRI): 63 | Hist, _ = np.histogram(MRI, bins=int(MRI.max())) 64 | idexs = np.argwhere(Hist >= 50) 65 | idex_max = np.float32(idexs[-1, 0]) 66 | MRI[np.where(MRI >= idex_max)] = idex_max 67 | sig = MRI[0, 0, 0] 68 | MRI = np.where(MRI != sig, MRI - np.mean(MRI[MRI != sig]), 0 * MRI) 69 | MRI = np.where(MRI != sig, MRI / np.std(MRI[MRI != sig] + 1e-7), 0 * MRI) 70 | return MRI 71 | 72 | def __getitem__(self, index): 73 | datafiles = self.files[index] 74 | flairNII = nib.load(datafiles["flair"]) 75 | t1NII = nib.load(datafiles["t1"]) 76 | t1ceNII = nib.load(datafiles["t1ce"]) 77 | t2NII = nib.load(datafiles["t2"]) 78 | labelNII = nib.load(datafiles["label"]) 79 | flair = self.truncate(flairNII.get_data()) 80 | t1 = self.truncate(t1NII.get_data()) 81 | t1ce = self.truncate(t1ceNII.get_data()) 82 | t2 = self.truncate(t2NII.get_data()) 83 | image = np.array([flair, t1, t1ce, t2]) # 4x240x240x150 84 | label = labelNII.get_data() 85 | image = image.astype(np.float32) 86 | label = label.astype(np.float32) 87 | 88 | if self.scale: 89 | scaler = np.random.uniform(0.9, 1.1) 90 | else: 91 | scaler = 1 92 | scale_d = int(self.crop_d * scaler) 93 | scale_h = int(self.crop_h * scaler) 94 | scale_w = int(self.crop_w * scaler) 95 | 96 | img_h, img_w, img_d = label.shape 97 | d_off = random.randint(0, img_d - scale_d) 98 | h_off = random.randint(15, img_h-15 - scale_h) 99 | w_off = random.randint(10, img_w-10 - scale_w) 100 | 101 | image = image[:, h_off: h_off + scale_h, w_off: w_off + scale_w, d_off: d_off + scale_d] 102 | label = label[h_off: h_off + scale_h, w_off: w_off + scale_w, d_off: d_off + scale_d] 103 | 104 | label = self.id2trainId(label) 105 | 106 | image = image.transpose((0, 3, 1, 2)) # Channel x Depth x H x W 107 | label = label.transpose((0, 3, 1, 2)) # Depth x H x W 108 | 109 | if self.is_mirror: 110 | randi = np.random.rand(1) 111 | if randi <= 0.3: 112 | pass 113 | elif randi <= 0.4: 114 | image = image[:, :, :, ::-1] 115 | label = label[:, :, :, ::-1] 116 | elif randi <= 0.5: 117 | image = image[:, :, ::-1, :] 118 | label = label[:, :, ::-1, :] 119 | elif randi <= 0.6: 120 | image = image[:, ::-1, :, :] 121 | label = label[:, ::-1, :, :] 122 | elif randi <= 0.7: 123 | image = image[:, :, ::-1, ::-1] 124 | label = label[:, :, ::-1, ::-1] 125 | elif randi <= 0.8: 126 | image = image[:, ::-1, :, ::-1] 127 | label = label[:, ::-1, :, ::-1] 128 | elif randi <= 0.9: 129 | image = image[:, ::-1, ::-1, :] 130 | label = label[:, ::-1, ::-1, :] 131 | else: 132 | image = image[:, ::-1, ::-1, ::-1] 133 | label = label[:, ::-1, ::-1, ::-1] 134 | 135 | if self.scale: 136 | image = resize(image, (4, self.crop_d, self.crop_h, self.crop_w), order=1, mode='constant', cval=0, clip=True, preserve_range=True) 137 | label = resize(label, (3, self.crop_d, self.crop_h, self.crop_w), order=0, mode='edge', cval=0, clip=True, preserve_range=True) 138 | image = image.astype(np.float32) 139 | label = label.astype(np.float32) 140 | 141 | # image -> res 142 | image_copy = np.zeros((4, self.crop_d, self.crop_h, self.crop_w)).astype(np.float32) 143 | image_copy[:, 1:, :, :] = image[:, 0:self.crop_d - 1, :, :] 144 | image_res = image - image_copy 145 | image_res[:, 0, :, :] = 0 146 | image_res = np.abs(image_res) 147 | 148 | # label -> res 149 | label_copy = np.zeros((3, self.crop_d, self.crop_h, self.crop_w)).astype(np.float32) 150 | label_copy[:, 1:, :, :] = label[:, 0:self.crop_d - 1, :, :] 151 | label_res = label - label_copy 152 | label_res[np.where(label_res == 0)] = 0 153 | label_res[np.where(label_res != 0)] = 1 154 | 155 | return image.copy(), image_res.copy(), label.copy(), label_res.copy() 156 | 157 | class BraTSValDataSet(data.Dataset): 158 | def __init__(self, root, list_path): 159 | self.root = root 160 | self.list_path = list_path 161 | self.img_ids = [i_id.strip().split() for i_id in open(self.root + self.list_path)] 162 | self.files = [] 163 | for item in self.img_ids: 164 | filepath = item[0] +'/'+ osp.splitext(osp.basename(item[0]))[0] 165 | flair_path = filepath + '_flair.nii.gz' 166 | t1_path = filepath + '_t1.nii.gz' 167 | t1ce_path = filepath + '_t1ce.nii.gz' 168 | t2_path = filepath + '_t2.nii.gz' 169 | label_path = filepath + '_seg.nii.gz' 170 | name = osp.splitext(osp.basename(filepath))[0] 171 | flair_file = osp.join(self.root, flair_path) 172 | t1_file = osp.join(self.root, t1_path) 173 | t1ce_file = osp.join(self.root, t1ce_path) 174 | t2_file = osp.join(self.root, t2_path) 175 | label_file = osp.join(self.root, label_path) 176 | self.files.append({ 177 | "flair": flair_file, 178 | "t1": t1_file, 179 | "t1ce": t1ce_file, 180 | "t2": t2_file, 181 | "label": label_file, 182 | "name": name 183 | }) 184 | print('{} images are loaded!'.format(len(self.img_ids))) 185 | 186 | def __len__(self): 187 | return len(self.files) 188 | 189 | def id2trainId(self, label): 190 | shape = label.shape 191 | results_map = np.zeros((3, shape[0], shape[1], shape[2])) 192 | 193 | NCR_NET = (label == 1) 194 | ET = (label == 4) 195 | WT = (label >= 1) 196 | TC = np.logical_or(NCR_NET, ET) 197 | 198 | results_map[0, :, :, :] = np.where(ET, 1, 0) 199 | results_map[1, :, :, :] = np.where(WT, 1, 0) 200 | results_map[2, :, :, :] = np.where(TC, 1, 0) 201 | return results_map 202 | 203 | 204 | def truncate(self, MRI): 205 | Hist, _ = np.histogram(MRI, bins=int(MRI.max())) 206 | idexs = np.argwhere(Hist >= 50) 207 | idex_max = np.float32(idexs[-1, 0]) 208 | MRI[np.where(MRI >= idex_max)] = idex_max 209 | sig = MRI[0, 0, 0] 210 | MRI = np.where(MRI != sig, MRI - np.mean(MRI[MRI != sig]), 0 * MRI) 211 | MRI = np.where(MRI != sig, MRI / np.std(MRI[MRI != sig] + 1e-7), 0 * MRI) 212 | return MRI 213 | 214 | def __getitem__(self, index): 215 | datafiles = self.files[index] 216 | 217 | flairNII = nib.load(datafiles["flair"]) 218 | t1NII = nib.load(datafiles["t1"]) 219 | t1ceNII = nib.load(datafiles["t1ce"]) 220 | t2NII = nib.load(datafiles["t2"]) 221 | labelNII = nib.load(datafiles["label"]) 222 | 223 | flair = self.truncate(flairNII.get_data()) 224 | t1 = self.truncate(t1NII.get_data()) 225 | t1ce = self.truncate(t1ceNII.get_data()) 226 | t2 = self.truncate(t2NII.get_data()) 227 | image = np.array([flair, t1, t1ce, t2]) # 4x240x240x150 228 | label = labelNII.get_data() 229 | name = datafiles["name"] 230 | 231 | label = self.id2trainId(label) 232 | 233 | image = image.transpose((0, 3, 1, 2)) # Channel x Depth x H x W 234 | label = label.transpose((0, 3, 1, 2)) # Depth x H x W 235 | image = image.astype(np.float32) 236 | label = label.astype(np.float32) 237 | 238 | size = image.shape[1:] 239 | affine = labelNII.affine 240 | 241 | # image -> res 242 | cha, dep, hei, wei = image.shape 243 | image_copy = np.zeros((cha, dep, hei, wei)).astype(np.float32) 244 | image_copy[:, 1:, :, :] = image[:, 0:dep - 1, :, :] 245 | image_res = image - image_copy 246 | image_res[:, 0, :, :] = 0 247 | image_res = np.abs(image_res) 248 | 249 | return image.copy(), image_res.copy(), label.copy(), np.array(size), name, affine 250 | -------------------------------------------------------------------------------- /dataset/PancreasDataSet.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import random 3 | import torch 4 | from torch.utils import data 5 | from skimage.transform import resize 6 | 7 | class PancreasDataSet(data.Dataset): 8 | def __init__(self, list_path, max_iters=None, crop_size=(64, 120, 120), mean=(128, 128, 128), scale=True, mirror=True, ignore_label=255): 9 | self.list_path = list_path 10 | self.crop_d, self.crop_h, self.crop_w = crop_size 11 | self.scale = scale 12 | self.ignore_label = ignore_label 13 | self.mean = mean 14 | self.is_mirror = mirror 15 | self.img_ids = [i_id.strip().split() for i_id in open(list_path)] 16 | if not max_iters==None: 17 | self.img_ids = self.img_ids * int(np.ceil(float(max_iters) / len(self.img_ids))) 18 | self.files = [] 19 | for item in self.img_ids: 20 | filepath = item[0][0:-4] + 'images' + '/' + item[0][-4:] 21 | label_path = item[0][0:-4] + 'labels' + '/' + item[0][-4:] 22 | name = item[0][-4:] 23 | 24 | self.files.append({ 25 | "img": filepath, 26 | "label": label_path, 27 | "name": name 28 | }) 29 | print('{} images are loaded!'.format(len(self.img_ids))) 30 | 31 | def __len__(self): 32 | return len(self.files) 33 | 34 | def id2trainId(self, label): 35 | shape = label.shape 36 | results_map = np.zeros((2, shape[0], shape[1], shape[2])) 37 | 38 | pancreas = (label==1) 39 | background = np.logical_not(pancreas) 40 | 41 | results_map[0,:,:,:] = np.where(background, 1, 0) 42 | results_map[1, :, :, :] = np.where(pancreas, 1, 0) 43 | return results_map 44 | 45 | def pre_precessing(self, image): 46 | image[image <= -100] = -100 47 | image[image >= 240] = 240 48 | image += 100 49 | image = image / 340 50 | return image 51 | 52 | def __getitem__(self, index): 53 | datafiles = self.files[index] 54 | # read nii file 55 | image = np.load(datafiles["img"] + '.npy') 56 | label = np.load(datafiles["label"] + '.npy') 57 | size = image.shape 58 | name = datafiles["name"] 59 | 60 | axes_index = np.argwhere(label == 1) 61 | one, two, three = axes_index[:, 0], axes_index[:, 1], axes_index[:, 2] 62 | min_x = np.min(one) 63 | max_x = np.max(one) 64 | min_x = min_x if min_x < 40 else min_x - 40 65 | max_x = size[0] if max_x >= size[0] - 40 - 1 else max_x + 40 + 1 66 | 67 | min_y = np.min(two) 68 | max_y = np.max(two) 69 | min_y = min_y if min_y < 40 else min_y - 40 70 | max_y = size[1] if max_y >= size[1] - 40 - 1 else max_y + 40 + 1 71 | 72 | min_z = np.min(three) 73 | max_z = np.max(three) 74 | min_z = min_z if min_z < 40 else min_z - 40 75 | max_z = size[2] if max_z >= size[2] - 40 - 1 else max_z + 40 + 1 76 | 77 | image = image[min_x:max_x, min_y:max_y, min_z:max_z] 78 | label = label[min_x:max_x, min_y:max_y, min_z:max_z] 79 | image = self.pre_precessing(image) 80 | 81 | if self.scale: 82 | scaler = np.random.uniform(0.9, 1.1) 83 | else: 84 | scaler = 1 85 | 86 | scale_d = int(self.crop_d * scaler) 87 | scale_h = int(self.crop_h * scaler) 88 | scale_w = int(self.crop_w * scaler) 89 | 90 | img_h, img_w, img_d = label.shape 91 | d_off = random.randint(0, img_d - scale_d) 92 | h_off = random.randint(0, img_h - scale_h) 93 | w_off = random.randint(0, img_w - scale_w) 94 | 95 | image = image[h_off: h_off + scale_h, w_off: w_off + scale_w, d_off: d_off + scale_d] 96 | label = label[h_off: h_off + scale_h, w_off: w_off + scale_w, d_off: d_off + scale_d] 97 | 98 | image = image.transpose((2, 0, 1)) 99 | label = label.transpose((2, 0, 1)) 100 | 101 | 102 | if self.is_mirror: 103 | randi = np.random.rand(1) 104 | if randi <= 0.3: 105 | pass 106 | elif randi <= 0.4: 107 | image = image[:, :, ::-1] 108 | label = label[:, :, ::-1] 109 | elif randi <= 0.5: 110 | image = image[:, ::-1, :] 111 | label = label[:, ::-1, :] 112 | elif randi <= 0.6: 113 | image = image[::-1, :, :] 114 | label = label[::-1, :, :] 115 | elif randi <= 0.7: 116 | image = image[:, ::-1, ::-1] 117 | label = label[:, ::-1, ::-1] 118 | elif randi <= 0.8: 119 | image = image[::-1, :, ::-1] 120 | label = label[::-1, :, ::-1] 121 | elif randi <= 0.9: 122 | image = image[::-1, ::-1, :] 123 | label = label[::-1, ::-1, :] 124 | else: 125 | image = image[::-1, ::-1, ::-1] 126 | label = label[::-1, ::-1, ::-1] 127 | 128 | if self.scale: 129 | image = resize(image, (self.crop_d, self.crop_h, self.crop_w), order=1, mode='constant', cval=0, clip=True, preserve_range=True) 130 | label = resize(label, (self.crop_d, self.crop_h, self.crop_w), order=0, mode='edge', cval=0, clip=True, preserve_range=True) 131 | 132 | image = np.array([image]) 133 | label = np.array([label]) 134 | 135 | image = image.astype(np.float32) 136 | label = label.astype(np.float32) 137 | 138 | # image -> res 139 | image_copy = np.zeros((1, self.crop_d, self.crop_h, self.crop_w)).astype(np.float32) 140 | image_copy[:, 1:, :, :] = image[:, 0:self.crop_d - 1, :, :] 141 | image_res = image - image_copy 142 | image_res[:, 0, :, :] = 0 143 | image_res = np.abs(image_res) 144 | 145 | # label -> res 146 | label_copy = np.zeros((1, self.crop_d, self.crop_h, self.crop_w)).astype(np.float32) 147 | label_copy[:, 1:, :, :] = label[:, 0:self.crop_d - 1, :, :] 148 | label_res = label - label_copy 149 | label_res[np.where(label_res == 0)] = 0 150 | label_res[np.where(label_res != 0)] = 1 151 | 152 | return image.copy(), image_res.copy(), label.copy(), label_res.copy(), np.array(size), name 153 | 154 | class PancreasValDataSet(data.Dataset): 155 | def __init__(self, list_path): 156 | self.list_path = list_path 157 | self.img_ids = [i_id.strip().split() for i_id in open(list_path)] 158 | self.files = [] 159 | for item in self.img_ids: 160 | filepath = item[0][0:-4] +'images' + '/' + item[0][-4:] 161 | label_path = item[0][0:-4] +'labels' + '/' + item[0][-4:] 162 | name = item[0][-4:] 163 | 164 | self.files.append({ 165 | "img": filepath, 166 | "label": label_path, 167 | "name": name 168 | }) 169 | print('{} images are loaded!'.format(len(self.img_ids))) 170 | 171 | def __len__(self): 172 | return len(self.files) 173 | 174 | def id2trainId(self, label): 175 | shape = label.shape 176 | results_map = np.zeros((2, shape[0], shape[1], shape[2])) 177 | 178 | pancreas = (label == 1) 179 | background = np.logical_not(pancreas) 180 | 181 | results_map[0, :, :, :] = np.where(background, 1, 0) 182 | results_map[1, :, :, :] = np.where(pancreas, 1, 0) 183 | return results_map 184 | 185 | def pre_precessing(self, image): 186 | image[image <= -100] = -100 187 | image[image >= 240] = 240 188 | image += 100 189 | image = image / 340 190 | return image 191 | 192 | def __getitem__(self, index): 193 | datafiles = self.files[index] 194 | img = np.load(datafiles["img"] + '.npy') 195 | label = np.load(datafiles["label"] + '.npy') 196 | image = np.array([img]) 197 | size = image.shape 198 | name = datafiles["name"] 199 | 200 | image = image.astype(np.float32) 201 | label = label.astype(np.float32) 202 | image[0, :, :, :] = self.pre_precessing(image[0, :, :, :]) 203 | 204 | label = np.array([label]) 205 | 206 | image = image.transpose((0, 3, 1, 2)) 207 | label = label.transpose((0, 3, 1, 2)) 208 | image = image.astype(np.float32) 209 | label = label.astype(np.float32) 210 | 211 | size = image.shape[1:] 212 | 213 | # image -> res 214 | cha, dep, hei, wei = image.shape 215 | image_copy = np.zeros((cha, dep, hei, wei)).astype(np.float32) 216 | image_copy[:, 1:, :, :] = image[:, 0:dep - 1, :, :] 217 | image_res = image - image_copy 218 | image_res[:, 0, :, :] = 0 219 | image_res = np.abs(image_res) 220 | 221 | return image.copy(), image_res.copy(), label.copy(), np.array(size), name 222 | 223 | -------------------------------------------------------------------------------- /models/ConResNet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from torch.nn import functional as F 3 | import torch 4 | import numpy as np 5 | 6 | 7 | class Conv3d(nn.Conv3d): 8 | 9 | def __init__(self, in_channels, out_channels, kernel_size, stride=(1,1,1), padding=(0,0,0), dilation=(1,1,1), groups=1, bias=False): 10 | super(Conv3d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias) 11 | 12 | def forward(self, x): 13 | weight = self.weight 14 | weight_mean = weight.mean(dim=1, keepdim=True).mean(dim=2, keepdim=True).mean(dim=3, keepdim=True).mean(dim=4, keepdim=True) 15 | weight = weight - weight_mean 16 | std = torch.sqrt(torch.var(weight.view(weight.size(0), -1), dim=1) + 1e-12).view(-1, 1, 1, 1, 1) 17 | weight = weight / std.expand_as(weight) 18 | return F.conv3d(x, weight, self.bias, self.stride, self.padding, self.dilation, self.groups) 19 | 20 | def conv3x3x3(in_planes, out_planes, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1,1,1), dilation=(1,1,1), bias=False, 21 | weight_std=False): 22 | "3x3x3 convolution with padding" 23 | if weight_std: 24 | return Conv3d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, 25 | bias=bias) 26 | else: 27 | return nn.Conv3d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, 28 | dilation=dilation, bias=bias) 29 | 30 | 31 | class ConResAtt(nn.Module): 32 | def __init__(self, in_planes, out_planes, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), 33 | dilation=(1, 1, 1), bias=False, weight_std=False, first_layer=False): 34 | super(ConResAtt, self).__init__() 35 | self.weight_std = weight_std 36 | self.stride = stride 37 | self.in_planes = in_planes 38 | self.out_planes = out_planes 39 | self.first_layer = first_layer 40 | 41 | self.relu = nn.ReLU(inplace=True) 42 | 43 | self.gn_seg = nn.GroupNorm(8, in_planes) 44 | self.conv_seg = conv3x3x3(in_planes, out_planes, kernel_size=(kernel_size[0], kernel_size[1], kernel_size[2]), 45 | stride=(stride[0], stride[1], stride[2]), padding=(padding[0], padding[1], padding[2]), 46 | dilation=(dilation[0], dilation[1], dilation[2]), bias=bias, weight_std=self.weight_std) 47 | 48 | self.gn_res = nn.GroupNorm(8, out_planes) 49 | self.conv_res = conv3x3x3(out_planes, out_planes, kernel_size=(1,1,1), 50 | stride=(1, 1, 1), padding=(0,0,0), 51 | dilation=(dilation[0], dilation[1], dilation[2]), bias=bias, weight_std=self.weight_std) 52 | 53 | self.gn_res1 = nn.GroupNorm(8, out_planes) 54 | self.conv_res1 = conv3x3x3(out_planes, out_planes, kernel_size=(kernel_size[0], kernel_size[1], kernel_size[2]), 55 | stride=(1, 1, 1), padding=(padding[0], padding[1], padding[2]), 56 | dilation=(dilation[0], dilation[1], dilation[2]), bias=bias, weight_std=self.weight_std) 57 | self.gn_res2 = nn.GroupNorm(8, out_planes) 58 | self.conv_res2 = conv3x3x3(out_planes, out_planes, kernel_size=(kernel_size[0], kernel_size[1], kernel_size[2]), 59 | stride=(1, 1, 1), padding=(padding[0], padding[1], padding[2]), 60 | dilation=(dilation[0], dilation[1], dilation[2]), bias=bias, weight_std=self.weight_std) 61 | 62 | self.gn_mp = nn.GroupNorm(8, in_planes) 63 | self.conv_mp_first = conv3x3x3(4, out_planes, kernel_size=(kernel_size[0], kernel_size[1], kernel_size[2]), 64 | stride=(stride[0], stride[1], stride[2]), padding=(padding[0], padding[1], padding[2]), 65 | dilation=(dilation[0], dilation[1], dilation[2]), bias=bias, weight_std=self.weight_std) 66 | self.conv_mp = conv3x3x3(in_planes, out_planes, kernel_size=(kernel_size[0], kernel_size[1], kernel_size[2]), 67 | stride=(stride[0], stride[1], stride[2]), padding=(padding[0], padding[1], padding[2]), 68 | dilation=(dilation[0], dilation[1], dilation[2]), bias=bias, weight_std=self.weight_std) 69 | 70 | def _res(self, x): # bs, channel, D, W, H 71 | 72 | bs, channel, depth, heigt, width = x.shape 73 | x_copy = torch.zeros_like(x).cuda() 74 | x_copy[:, :, 1:, :, :] = x[:, :, 0: depth - 1, :, :] 75 | res = x - x_copy 76 | res[:, :, 0, :, :] = 0 77 | res = torch.abs(res) 78 | return res 79 | 80 | def forward(self, input): 81 | x1, x2 = input 82 | if self.first_layer: 83 | x1 = self.gn_seg(x1) 84 | x1 = self.relu(x1) 85 | x1 = self.conv_seg(x1) 86 | 87 | res = torch.sigmoid(x1) 88 | res = self._res(res) 89 | res = self.conv_res(res) 90 | 91 | x2 = self.conv_mp_first(x2) 92 | x2 = x2 + res 93 | 94 | else: 95 | x1 = self.gn_seg(x1) 96 | x1 = self.relu(x1) 97 | x1 = self.conv_seg(x1) 98 | 99 | res = torch.sigmoid(x1) 100 | res = self._res(res) 101 | res = self.conv_res(res) 102 | 103 | 104 | if self.in_planes != self.out_planes: 105 | x2 = self.gn_mp(x2) 106 | x2 = self.relu(x2) 107 | x2 = self.conv_mp(x2) 108 | 109 | x2 = x2 + res 110 | 111 | x2 = self.gn_res1(x2) 112 | x2 = self.relu(x2) 113 | x2 = self.conv_res1(x2) 114 | 115 | x1 = x1*(1 + torch.sigmoid(x2)) 116 | 117 | return [x1, x2] 118 | 119 | 120 | class NoBottleneck(nn.Module): 121 | def __init__(self, inplanes, planes, stride=(1, 1, 1), dilation=(1, 1, 1), downsample=None, fist_dilation=1, 122 | multi_grid=1, weight_std=False): 123 | super(NoBottleneck, self).__init__() 124 | self.weight_std = weight_std 125 | self.relu = nn.ReLU(inplace=True) 126 | 127 | self.gn1 = nn.GroupNorm(8, inplanes) 128 | self.conv1 = conv3x3x3(inplanes, planes, kernel_size=(3, 3, 3), stride=stride, padding=dilation * multi_grid, 129 | dilation=dilation * multi_grid, bias=False, weight_std=self.weight_std) 130 | 131 | self.gn2 = nn.GroupNorm(8, planes) 132 | self.conv2 = conv3x3x3(planes, planes, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=dilation * multi_grid, 133 | dilation=dilation * multi_grid, bias=False, weight_std=self.weight_std) 134 | 135 | self.downsample = downsample 136 | self.dilation = dilation 137 | self.stride = stride 138 | 139 | def forward(self, x): 140 | skip = x 141 | 142 | seg = self.gn1(x) 143 | seg = self.relu(seg) 144 | seg = self.conv1(seg) 145 | 146 | seg = self.gn2(seg) 147 | seg = self.relu(seg) 148 | seg = self.conv2(seg) 149 | 150 | if self.downsample is not None: 151 | skip = self.downsample(x) 152 | 153 | seg = seg + skip 154 | return seg 155 | 156 | 157 | class conresnet(nn.Module): 158 | def __init__(self, shape, block, layers, num_classes=3, weight_std=False): 159 | self.shape = shape 160 | self.weight_std = weight_std 161 | super(conresnet, self).__init__() 162 | 163 | self.conv_4_32 = nn.Sequential( 164 | conv3x3x3(4, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1), weight_std=self.weight_std)) 165 | 166 | self.conv_32_64 = nn.Sequential( 167 | nn.GroupNorm(8, 32), 168 | nn.ReLU(inplace=True), 169 | conv3x3x3(32, 64, kernel_size=(3, 3, 3), stride=(2, 2, 2), weight_std=self.weight_std)) 170 | 171 | self.conv_64_128 = nn.Sequential( 172 | nn.GroupNorm(8, 64), 173 | nn.ReLU(inplace=True), 174 | conv3x3x3(64, 128, kernel_size=(3, 3, 3), stride=(2, 2, 2), weight_std=self.weight_std)) 175 | 176 | self.conv_128_256 = nn.Sequential( 177 | nn.GroupNorm(8, 128), 178 | nn.ReLU(inplace=True), 179 | conv3x3x3(128, 256, kernel_size=(3, 3, 3), stride=(2, 2, 2), weight_std=self.weight_std)) 180 | 181 | self.layer0 = self._make_layer(block, 32, 32, layers[0], stride=(1, 1, 1)) 182 | self.layer1 = self._make_layer(block, 64, 64, layers[1], stride=(1, 1, 1)) 183 | self.layer2 = self._make_layer(block, 128, 128, layers[2], stride=(1, 1, 1)) 184 | self.layer3 = self._make_layer(block, 256, 256, layers[3], stride=(1, 1, 1)) 185 | self.layer4 = self._make_layer(block, 256, 256, layers[4], stride=(1, 1, 1), dilation=(2,2,2)) 186 | 187 | self.fusionConv = nn.Sequential( 188 | nn.GroupNorm(8, 256), 189 | nn.ReLU(inplace=True), 190 | nn.Dropout3d(0.1), 191 | conv3x3x3(256, 128, kernel_size=(3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1), dilation=(1, 1, 1), weight_std=self.weight_std) 192 | ) 193 | 194 | self.seg_x4 = nn.Sequential( 195 | ConResAtt(128, 64, kernel_size=(3, 3, 3), padding=(1, 1, 1), weight_std=self.weight_std, first_layer=True)) 196 | self.seg_x2 = nn.Sequential( 197 | ConResAtt(64, 32, kernel_size=(3, 3, 3), padding=(1, 1, 1), weight_std=self.weight_std)) 198 | self.seg_x1 = nn.Sequential( 199 | ConResAtt(32, 32, kernel_size=(3, 3, 3), padding=(1, 1, 1), weight_std=self.weight_std)) 200 | 201 | self.seg_cls = nn.Sequential( 202 | nn.Conv3d(32, num_classes, kernel_size=1) 203 | ) 204 | self.res_cls = nn.Sequential( 205 | nn.Conv3d(32, num_classes, kernel_size=1) 206 | ) 207 | self.resx2_cls = nn.Sequential( 208 | nn.Conv3d(32, num_classes, kernel_size=1) 209 | ) 210 | self.resx4_cls = nn.Sequential( 211 | nn.Conv3d(64, num_classes, kernel_size=1) 212 | ) 213 | 214 | def _make_layer(self, block, inplanes, outplanes, blocks, stride=(1, 1, 1), dilation=(1, 1, 1), multi_grid=1): 215 | downsample = None 216 | if stride[0] != 1 or stride[1] != 1 or stride[2] != 1 or inplanes != outplanes: 217 | downsample = nn.Sequential( 218 | nn.GroupNorm(8, inplanes), 219 | nn.ReLU(inplace=True), 220 | conv3x3x3(inplanes, outplanes, kernel_size=(1, 1, 1), stride=stride, padding=(0, 0, 0), 221 | weight_std=self.weight_std) 222 | ) 223 | 224 | layers = [] 225 | generate_multi_grid = lambda index, grids: grids[index % len(grids)] if isinstance(grids, tuple) else 1 226 | layers.append(block(inplanes, outplanes, stride, dilation=dilation, downsample=downsample, 227 | multi_grid=generate_multi_grid(0, multi_grid), weight_std=self.weight_std)) 228 | for i in range(1, blocks): 229 | layers.append( 230 | block(inplanes, outplanes, dilation=dilation, multi_grid=generate_multi_grid(i, multi_grid), 231 | weight_std=self.weight_std)) 232 | return nn.Sequential(*layers) 233 | 234 | 235 | def forward(self, x_list): 236 | x, x_res = x_list 237 | 238 | ## encoder 239 | x = self.conv_4_32(x) 240 | x = self.layer0(x) 241 | skip1 = x 242 | 243 | x = self.conv_32_64(x) 244 | x = self.layer1(x) 245 | skip2 = x 246 | 247 | x = self.conv_64_128(x) 248 | x = self.layer2(x) 249 | skip3 = x 250 | 251 | x = self.conv_128_256(x) 252 | x = self.layer3(x) 253 | 254 | x = self.layer4(x) 255 | 256 | x = self.fusionConv(x) 257 | 258 | ## decoder 259 | res_x4 = F.interpolate(x_res, size=(int(self.shape[0] / 4), int(self.shape[1] / 4), int(self.shape[2] / 4)), mode='trilinear', align_corners=True) 260 | seg_x4 = F.interpolate(x, size=(int(self.shape[0] / 4), int(self.shape[1] / 4), int(self.shape[2] / 4)), mode='trilinear', align_corners=True) 261 | seg_x4 = seg_x4 + skip3 262 | seg_x4, res_x4 = self.seg_x4([seg_x4, res_x4]) 263 | 264 | res_x2 = F.interpolate(res_x4, size=(int(self.shape[0] / 2), int(self.shape[1] / 2), int(self.shape[2] / 2)), mode='trilinear', align_corners=True) 265 | seg_x2 = F.interpolate(seg_x4, size=(int(self.shape[0] / 2), int(self.shape[1] / 2), int(self.shape[2] / 2)), mode='trilinear', align_corners=True) 266 | seg_x2 = seg_x2 + skip2 267 | seg_x2, res_x2 = self.seg_x2([seg_x2, res_x2]) 268 | 269 | res_x1 = F.interpolate(res_x2, size=(int(self.shape[0] / 1), int(self.shape[1] / 1), int(self.shape[2] / 1)), mode='trilinear', align_corners=True) 270 | seg_x1 = F.interpolate(seg_x2, size=(int(self.shape[0] / 1), int(self.shape[1] / 1), int(self.shape[2] / 1)), mode='trilinear', align_corners=True) 271 | seg_x1 = seg_x1 + skip1 272 | seg_x1, res_x1 = self.seg_x1([seg_x1, res_x1]) 273 | 274 | seg = self.seg_cls(seg_x1) 275 | res = self.res_cls(res_x1) 276 | resx2 = self.resx2_cls(res_x2) 277 | resx4 = self.resx4_cls(res_x4) 278 | 279 | resx2 = F.interpolate(resx2, size=(int(self.shape[0] / 1), int(self.shape[1] / 1), int(self.shape[2] / 1)), 280 | mode='trilinear', align_corners=True) 281 | resx4 = F.interpolate(resx4, size=(int(self.shape[0] / 1), int(self.shape[1] / 1), int(self.shape[2] / 1)), 282 | mode='trilinear', align_corners=True) 283 | 284 | return [seg, res, resx2, resx4] 285 | 286 | 287 | def ConResNet(shape, num_classes=3, weight_std=True): 288 | 289 | model = conresnet(shape, NoBottleneck, [1, 2, 2, 2, 2], num_classes, weight_std) 290 | 291 | return model -------------------------------------------------------------------------------- /models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jianpengz/ConResNet/cf32e955401a70340ed16809966711f1e5accdc8/models/__init__.py -------------------------------------------------------------------------------- /models/conresnet.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jianpengz/ConResNet/cf32e955401a70340ed16809966711f1e5accdc8/models/conresnet.png -------------------------------------------------------------------------------- /run.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | CUDA_VISIBLE_DEVICES=2,3 nohup python -m torch.distributed.launch --nproc_per_node=2 --master_port=$RANDOM train_conresnet.py \ 4 | --data_dir='path-to-your-dataset/' \ 5 | --train_list='list/train_list.txt' \ 6 | --val_list='list/val_list.txt' \ 7 | --snapshot_dir='path-to-save-checkpoint/' \ 8 | --input_size='80,160,160' \ 9 | --batch_size=2 \ 10 | --num_gpus=2 \ 11 | --num_steps=40000 \ 12 | --val_pred_every=2000 \ 13 | --learning_rate=1e-4 \ 14 | --num_classes=3 \ 15 | --num_workers=4 \ 16 | --random_mirror=True \ 17 | --random_scale=True \ 18 | > path-to-save-log-file/log.file 2>&1 & 19 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import sys 3 | sys.path.append("..") 4 | import numpy as np 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | from torch.utils import data 10 | from models.ConResNet import ConResNet 11 | from dataset.BraTSDataSet import BraTSValDataSet 12 | import os 13 | from math import ceil 14 | import nibabel as nib 15 | 16 | def get_arguments(): 17 | parser = argparse.ArgumentParser(description="ConResNet for 3D medical image segmentation.") 18 | parser.add_argument("--data-dir", type=str, default='path-to-your-dataset/', 19 | help="Path to the directory containing your dataset.") 20 | parser.add_argument("--data-list", type=str, default='list/val.txt', 21 | help="Path to the file listing the images in the dataset.") 22 | parser.add_argument("--input-size", type=str, default='80,160,160', 23 | help="Comma-separated string with depth, height and width of sub-volumnes.") 24 | parser.add_argument("--num-classes", type=int, default=3, 25 | help="Number of classes to predict (ET, WT, TC).") 26 | parser.add_argument("--restore-from", type=str, default='snapshots/conresnet/your_checkpoint_model.pth', 27 | help="Where restore model parameters from.") 28 | parser.add_argument("--gpu", type=str, default='0', 29 | help="choose gpu device.") 30 | parser.add_argument("--weight-std", type=bool, default=True, 31 | help="whether to use weight standarization in CONV layers.") 32 | return parser.parse_args() 33 | 34 | 35 | def pad_image(img, target_size): 36 | """Pad an image up to the target size.""" 37 | deps_missing = target_size[0] - img.shape[2] 38 | rows_missing = target_size[1] - img.shape[3] 39 | cols_missing = target_size[2] - img.shape[4] 40 | padded_img = np.pad(img, ((0, 0), (0, 0),(0, deps_missing), (0, rows_missing), (0, cols_missing)), 'constant') 41 | return padded_img 42 | 43 | def predict_sliding(net, img_list, tile_size, classes): 44 | image, image_res = img_list 45 | interp = nn.Upsample(size=tile_size, mode='trilinear', align_corners=True) 46 | image_size = image.shape 47 | overlap = 1/3 48 | 49 | strideHW = ceil(tile_size[1] * (1 - overlap)) 50 | strideD = ceil(tile_size[0] * (1 - overlap)) 51 | tile_deps = int(ceil((image_size[2] - tile_size[0]) / strideD) + 1) 52 | tile_rows = int(ceil((image_size[3] - tile_size[1]) / strideHW) + 1) # strided convolution formula 53 | tile_cols = int(ceil((image_size[4] - tile_size[2]) / strideHW) + 1) 54 | full_probs = torch.zeros((classes, image_size[2], image_size[3], image_size[4])) 55 | count_predictions = torch.zeros((classes, image_size[2], image_size[3], image_size[4])) 56 | 57 | for dep in range(tile_deps): 58 | for row in range(tile_rows): 59 | for col in range(tile_cols): 60 | d1 = int(dep * strideD) 61 | y1 = int(row * strideHW) 62 | x1 = int(col * strideHW) 63 | d2 = min(d1 + tile_size[0], image_size[2]) 64 | y2 = min(y1 + tile_size[1], image_size[3]) 65 | x2 = min(x1 + tile_size[2], image_size[4]) 66 | d1 = max(int(d2 - tile_size[0]), 0) 67 | y1 = max(int(y2 - tile_size[1]), 0) 68 | x1 = max(int(x2 - tile_size[2]), 0) 69 | 70 | img = image[:, :, d1:d2, y1:y2, x1:x2] 71 | img_res = image_res[:, :, d1:d2, y1:y2, x1:x2] 72 | padded_img = pad_image(img, tile_size) 73 | padded_img_res = pad_image(img_res, tile_size) 74 | padded_prediction = net([torch.from_numpy(padded_img).cuda(), torch.from_numpy(padded_img_res).cuda()]) 75 | padded_prediction = F.sigmoid(padded_prediction[0]) 76 | 77 | padded_prediction = interp(padded_prediction).cpu().data[0] 78 | prediction = padded_prediction[0:img.shape[2],0:img.shape[3], 0:img.shape[4], :] 79 | count_predictions[:, d1:d2, y1:y2, x1:x2] += 1 80 | full_probs[:, d1:d2, y1:y2, x1:x2] += prediction 81 | 82 | full_probs /= count_predictions 83 | full_probs = full_probs.numpy().transpose(1,2,3,0) 84 | return full_probs 85 | 86 | def dice_score(preds, labels): 87 | assert preds.shape[0] == labels.shape[0], "predict & target batch size don't match" 88 | predict = preds.view().reshape(preds.shape[0], -1) 89 | target = labels.view().reshape(labels.shape[0], -1) 90 | 91 | num = np.sum(np.multiply(predict, target), axis=1) 92 | den = np.sum(predict, axis=1) + np.sum(target, axis=1) +1 93 | 94 | dice = 2*num / den 95 | 96 | return dice.mean() 97 | 98 | 99 | 100 | def main(): 101 | 102 | args = get_arguments() 103 | 104 | os.environ["CUDA_VISIBLE_DEVICES"]=args.gpu 105 | d, h, w = map(int, args.input_size.split(',')) 106 | 107 | input_size = (d, h, w) 108 | 109 | model = ConResNet(input_size, num_classes=args.num_classes, weight_std=args.weight_std) 110 | model = nn.DataParallel(model) 111 | 112 | print('loading from checkpoint: {}'.format(args.restore_from)) 113 | if os.path.exists(args.restore_from): 114 | model.load_state_dict(torch.load(args.restore_from, map_location=torch.device('cpu'))) 115 | else: 116 | print('File not exists in the reload path: {}'.format(args.restore_from)) 117 | 118 | model.eval() 119 | model.cuda() 120 | 121 | testloader = data.DataLoader( 122 | BraTSValDataSet(args.data_dir, args.data_list), 123 | batch_size=1, shuffle=False, pin_memory=True) 124 | 125 | if not os.path.exists('outputs'): 126 | os.makedirs('outputs') 127 | 128 | dice_ET = 0 129 | dice_WT = 0 130 | dice_TC = 0 131 | 132 | for index, batch in enumerate(testloader): 133 | image, image_res, label, size, name, affine = batch 134 | size = size[0].numpy() 135 | affine = affine[0].numpy() 136 | name[0]=name[0].replace("Brats17", "Brats18") 137 | with torch.no_grad(): 138 | output = predict_sliding(model, [image.numpy(),image_res.numpy()], input_size, args.num_classes) 139 | 140 | seg_pred_3class = np.asarray(np.around(output), dtype=np.uint8) 141 | 142 | seg_pred_ET = seg_pred_3class[:, :, :, 0] 143 | seg_pred_WT = seg_pred_3class[:, :, :, 1] 144 | seg_pred_TC = seg_pred_3class[:, :, :, 2] 145 | seg_pred = np.zeros_like(seg_pred_ET) 146 | seg_pred = np.where(seg_pred_WT == 1, 2, seg_pred) 147 | seg_pred = np.where(seg_pred_TC == 1, 1, seg_pred) 148 | seg_pred = np.where(seg_pred_ET == 1, 4, seg_pred) 149 | 150 | seg_gt = np.asarray(label[0].numpy()[:size[0], :size[1], :size[2]], dtype=np.int) 151 | seg_gt_ET = seg_gt[0, :, :, :] 152 | seg_gt_WT = seg_gt[1, :, :, :] 153 | seg_gt_TC = seg_gt[2, :, :, :] 154 | 155 | dice_ET_i = dice_score(seg_pred_ET[None, :, :, :], seg_gt_ET[None, :, :, :]) 156 | dice_WT_i = dice_score(seg_pred_WT[None, :, :, :], seg_gt_WT[None, :, :, :]) 157 | dice_TC_i = dice_score(seg_pred_TC[None, :, :, :], seg_gt_TC[None, :, :, :]) 158 | 159 | print('Processing {}: Dice_ET = {:.4}, Dice_WT = {:.4}, Dice_TC = {:.4}'.format(name, dice_ET_i, dice_WT_i, dice_TC_i)) 160 | 161 | dice_ET += dice_ET_i 162 | dice_WT += dice_WT_i 163 | dice_TC += dice_TC_i 164 | 165 | seg_pred = seg_pred.transpose((1,2,0)) 166 | 167 | seg_pred = seg_pred.astype(np.int16) 168 | 169 | seg_pred = nib.Nifti1Image(seg_pred, affine=affine) 170 | seg_save_p = os.path.join('outputs/%s.nii.gz' % (name[0])) 171 | nib.save(seg_pred, seg_save_p) 172 | 173 | dice_ET_avg = dice_ET / (index + 1) 174 | dice_WT_avg = dice_WT / (index + 1) 175 | dice_TC_avg = dice_TC / (index + 1) 176 | 177 | print('Average score: Dice_ET = {:.4}, Dice_WT = {:.4}, Dice_TC = {:.4}'.format(dice_ET_avg, dice_WT_avg, dice_TC_avg)) 178 | 179 | 180 | if __name__ == '__main__': 181 | main() 182 | -------------------------------------------------------------------------------- /train_conresnet.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import sys 3 | sys.path.append("..") 4 | 5 | import torch 6 | import numpy as np 7 | import torch.optim as optim 8 | import torch.backends.cudnn as cudnn 9 | import torch.nn.functional as F 10 | 11 | import os 12 | import os.path as osp 13 | from models.ConResNet import ConResNet 14 | from dataset.BraTSDataSet import BraTSDataSet, BraTSValDataSet 15 | import timeit 16 | from tensorboardX import SummaryWriter 17 | from utils import loss 18 | from utils.engine import Engine 19 | from math import ceil 20 | 21 | start = timeit.default_timer() 22 | 23 | def str2bool(v): 24 | if v.lower() in ('yes', 'true', 't', 'y', '1'): 25 | return True 26 | elif v.lower() in ('no', 'false', 'f', 'n', '0'): 27 | return False 28 | else: 29 | raise argparse.ArgumentTypeError('Boolean value expected.') 30 | 31 | 32 | def get_arguments(): 33 | """ 34 | A list of parsed arguments. 35 | """ 36 | parser = argparse.ArgumentParser(description="ConResNet for 3D Medical Image Segmentation.") 37 | 38 | parser.add_argument("--data_dir", type=str, default='/media/userdisk0/myproject-Seg/BraTS-pro/dataset/') 39 | parser.add_argument("--train_list", type=str, default='list/BraTS2018_old/train.txt') 40 | parser.add_argument("--val_list", type=str, default='list/BraTS2018_old/val.txt') 41 | parser.add_argument("--snapshot_dir", type=str, default='snapshots/conresnet/') 42 | parser.add_argument("--reload_path", type=str, default='snapshots/conresnet/ConResNet_40000.pth') 43 | parser.add_argument("--reload_from_checkpoint", type=str2bool, default=False) 44 | parser.add_argument("--input_size", type=str, default='80,160,160') 45 | parser.add_argument("--batch_size", type=int, default=1) 46 | parser.add_argument("--num_gpus", type=int, default=1) 47 | parser.add_argument('--local_rank', type=int, default=0) 48 | parser.add_argument("--num_steps", type=int, default=40000) 49 | parser.add_argument("--start_iters", type=int, default=0) 50 | parser.add_argument("--val_pred_every", type=int, default=100) 51 | parser.add_argument("--learning_rate", type=float, default=1e-4) 52 | parser.add_argument("--num_classes", type=int, default=3) 53 | parser.add_argument("--num_workers", type=int, default=1) 54 | parser.add_argument("--weight_std", type=str2bool, default=True) 55 | parser.add_argument("--momentum", type=float, default=0.9) 56 | parser.add_argument("--power", type=float, default=0.9) 57 | parser.add_argument("--weight_decay", type=float, default=0.0005) 58 | parser.add_argument("--ignore_label", type=int, default=255) 59 | parser.add_argument("--is_training", action="store_true") 60 | parser.add_argument("--random_mirror", type=str2bool, default=False) 61 | parser.add_argument("--random_scale", type=str2bool, default=False) 62 | parser.add_argument("--random_seed", type=int, default=1234) 63 | 64 | return parser 65 | 66 | 67 | def lr_poly(base_lr, iter, max_iter, power): 68 | return base_lr * ((1 - float(iter) / max_iter) ** (power)) 69 | 70 | def adjust_learning_rate(optimizer, i_iter, lr, num_steps, power): 71 | lr = lr_poly(lr, i_iter, num_steps, power) 72 | optimizer.param_groups[0]['lr'] = lr 73 | return lr 74 | 75 | 76 | def dice_score(preds, labels): 77 | assert preds.shape[0] == labels.shape[0], "predict & target batch size don't match" 78 | predict = preds.contiguous().view(preds.shape[0], -1) 79 | target = labels.contiguous().view(labels.shape[0], -1) 80 | 81 | num = torch.sum(torch.mul(predict, target), dim=1) 82 | den = torch.sum(predict, dim=1) + torch.sum(target, dim=1) + 1 83 | 84 | dice = 2*num / den 85 | 86 | return dice.mean() 87 | 88 | 89 | def compute_dice_score(preds, labels): 90 | 91 | preds = F.sigmoid(preds) 92 | 93 | pred_ET = preds[:, 0, :, :, :] 94 | pred_WT = preds[:, 1, :, :, :] 95 | pred_TC = preds[:, 2, :, :, :] 96 | label_ET = labels[:, 0, :, :, :] 97 | label_WT = labels[:, 1, :, :, :] 98 | label_TC = labels[:, 2, :, :, :] 99 | dice_ET = dice_score(pred_ET, label_ET).cpu().data.numpy() 100 | dice_WT = dice_score(pred_WT, label_WT).cpu().data.numpy() 101 | dice_TC = dice_score(pred_TC, label_TC).cpu().data.numpy() 102 | return dice_ET, dice_WT, dice_TC 103 | 104 | 105 | def predict_sliding(net, imagelist, tile_size, classes): 106 | image, image_res = imagelist 107 | image_size = image.shape 108 | overlap = 1 / 3 109 | 110 | strideHW = ceil(tile_size[1] * (1 - overlap)) 111 | strideD = ceil(tile_size[0] * (1 - overlap)) 112 | tile_deps = int(ceil((image_size[2] - tile_size[0]) / strideD) + 1) 113 | tile_rows = int(ceil((image_size[3] - tile_size[1]) / strideHW) + 1) 114 | tile_cols = int(ceil((image_size[4] - tile_size[2]) / strideHW) + 1) 115 | full_probs = np.zeros((image_size[0], classes, image_size[2], image_size[3], image_size[4])).astype(np.float32) 116 | count_predictions = np.zeros((image_size[0], classes, image_size[2], image_size[3], image_size[4])).astype(np.float32) 117 | full_probs = torch.from_numpy(full_probs).cuda() 118 | count_predictions = torch.from_numpy(count_predictions).cuda() 119 | 120 | for dep in range(tile_deps): 121 | for row in range(tile_rows): 122 | for col in range(tile_cols): 123 | d1 = int(dep * strideD) 124 | x1 = int(col * strideHW) 125 | y1 = int(row * strideHW) 126 | d2 = min(d1 + tile_size[0], image_size[2]) 127 | x2 = min(x1 + tile_size[2], image_size[4]) 128 | y2 = min(y1 + tile_size[1], image_size[3]) 129 | d1 = max(int(d2 - tile_size[0]), 0) 130 | x1 = max(int(x2 - tile_size[2]), 0) 131 | y1 = max(int(y2 - tile_size[1]), 0) 132 | 133 | img = image[:, :, d1:d2, y1:y2, x1:x2] 134 | img_res = image_res[:, :, d1:d2, y1:y2, x1:x2] 135 | 136 | prediction = net([img, img_res]) 137 | prediction = prediction[0] 138 | 139 | count_predictions[:, :, d1:d2, y1:y2, x1:x2] += 1 140 | full_probs[:, :, d1:d2, y1:y2, x1:x2] += prediction 141 | 142 | full_probs /= count_predictions 143 | return full_probs 144 | 145 | 146 | def validate(input_size, model, ValLoader, num_classes): 147 | # start to validate 148 | val_ET = 0.0 149 | val_WT = 0.0 150 | val_TC = 0.0 151 | 152 | for index, batch in enumerate(ValLoader): 153 | print('%d processd'%(index)) 154 | image, image_res, label, size, name, affine = batch 155 | image = image.cuda() 156 | image_res = image_res.cuda() 157 | label = label.cuda() 158 | with torch.no_grad(): 159 | pred = predict_sliding(model, [image, image_res], input_size, num_classes) 160 | dice_ET, dice_WT, dice_TC = compute_dice_score(pred, label) 161 | val_ET += dice_ET 162 | val_WT += dice_WT 163 | val_TC += dice_TC 164 | 165 | return val_ET/(index+1), val_WT/(index+1), val_TC/(index+1) 166 | 167 | def main(): 168 | """Create the ConResNet model and then start the training.""" 169 | parser = get_arguments() 170 | print(parser) 171 | # os.environ["CUDA_VISIBLE_DEVICES"] = '0' 172 | 173 | with Engine(custom_parser=parser) as engine: 174 | args = parser.parse_args() 175 | if args.num_gpus > 1: 176 | torch.cuda.set_device(args.local_rank) 177 | 178 | writer = SummaryWriter(args.snapshot_dir) 179 | 180 | d, h, w = map(int, args.input_size.split(',')) 181 | input_size = (d, h, w) 182 | 183 | cudnn.benchmark = True 184 | seed = args.random_seed 185 | if engine.distributed: 186 | seed = args.local_rank 187 | torch.manual_seed(seed) 188 | if torch.cuda.is_available(): 189 | torch.cuda.manual_seed(seed) 190 | 191 | model = ConResNet(input_size, num_classes=args.num_classes, weight_std=True) 192 | model.train() 193 | device = torch.device('cuda:{}'.format(args.local_rank)) 194 | model.to(device) 195 | 196 | optimizer = optim.Adam( 197 | [{'params': filter(lambda p: p.requires_grad, model.parameters()), 'lr': args.learning_rate}], 198 | lr=args.learning_rate, weight_decay=args.weight_decay) 199 | 200 | if args.num_gpus > 1: 201 | model = engine.data_parallel(model) 202 | 203 | # load checkpoint... 204 | if args.reload_from_checkpoint: 205 | print('loading from checkpoint: {}'.format(args.reload_path)) 206 | if os.path.exists(args.reload_path): 207 | model.load_state_dict(torch.load(args.reload_path, map_location=torch.device('cpu'))) 208 | else: 209 | print('File not exists in the reload path: {}'.format(args.reload_path)) 210 | 211 | loss_D = loss.DiceLoss4BraTS().to(device) 212 | loss_BCE = loss.BCELoss4BraTS().to(device) 213 | 214 | loss_B = loss.BCELossBoud().to(device) 215 | 216 | if not os.path.exists(args.snapshot_dir): 217 | os.makedirs(args.snapshot_dir) 218 | 219 | trainloader, train_sampler = engine.get_train_loader(BraTSDataSet(args.data_dir, args.train_list, max_iters=args.num_steps * args.batch_size, crop_size=input_size, 220 | scale=args.random_scale, mirror=args.random_mirror)) 221 | valloader, val_sampler = engine.get_test_loader(BraTSValDataSet(args.data_dir, args.val_list)) 222 | 223 | for i_iter, batch in enumerate(trainloader): 224 | i_iter += args.start_iters 225 | images, images_res, labels, labels_res = batch 226 | images = images.cuda() 227 | images_res = images_res.cuda() 228 | labels = labels.cuda() 229 | labels_res = labels_res.cuda() 230 | 231 | optimizer.zero_grad() 232 | lr = adjust_learning_rate(optimizer, i_iter, args.learning_rate, args.num_steps, args.power) 233 | 234 | preds= model([images, images_res]) 235 | preds_seg = preds[0] 236 | preds_res = preds[1] 237 | preds_resx2 = preds[2] 238 | preds_resx4 = preds[3] 239 | 240 | term_seg_Dice = loss_D.forward(preds_seg, labels) 241 | term_seg_BCE = loss_BCE.forward(preds_seg, labels) 242 | 243 | term_res_BCE = loss_B.forward(preds_res, labels_res) 244 | term_resx2_BCE = loss_B.forward(preds_resx2, labels_res) 245 | term_resx4_BCE = loss_B.forward(preds_resx4, labels_res) 246 | 247 | term_all = term_seg_Dice + term_seg_BCE + term_res_BCE + 0.5 * (term_resx2_BCE +term_resx4_BCE) 248 | term_all.backward() 249 | 250 | optimizer.step() 251 | 252 | if i_iter % 100 == 0 and (args.local_rank == 0): 253 | writer.add_scalar('learning_rate', lr, i_iter) 254 | writer.add_scalar('loss', term_all.cpu().data.numpy(), i_iter) 255 | 256 | print('iter = {} of {} completed, lr = {:.4}, seg_loss = {:.4}, res_loss = {:.4}'.format( 257 | i_iter, args.num_steps, lr, (term_seg_Dice+term_seg_BCE).cpu().data.numpy(), (term_res_BCE+term_resx2_BCE+term_resx4_BCE).cpu().data.numpy())) 258 | 259 | 260 | if i_iter >= args.num_steps - 1 and (args.local_rank == 0): 261 | print('save last model ...') 262 | torch.save(model.state_dict(), osp.join(args.snapshot_dir, 'ConResNet_' + str(args.num_steps) + '.pth')) 263 | break 264 | 265 | if i_iter % args.val_pred_every == 0 and i_iter!=0 and (args.local_rank == 0): 266 | print('save model ...') 267 | torch.save(model.state_dict(), osp.join(args.snapshot_dir, 'ConResNet_' + str(i_iter) + '.pth')) 268 | 269 | # val 270 | if i_iter % args.val_pred_every == 0: 271 | print('validate ...') 272 | val_ET, val_WT, val_TC = validate(input_size, model, valloader, args.num_classes) 273 | if (args.local_rank == 0): 274 | writer.add_scalar('Val_ET_Dice', val_ET, i_iter) 275 | writer.add_scalar('Val_WT_Dice', val_WT, i_iter) 276 | writer.add_scalar('Val_TC_Dice', val_TC, i_iter) 277 | print('Validate iter = {}, ET = {:.2}, WT = {:.2}, TC = {:.2}'.format(i_iter, val_ET, val_WT, val_TC)) 278 | 279 | end = timeit.default_timer() 280 | print(end - start, 'seconds') 281 | 282 | 283 | if __name__ == '__main__': 284 | main() 285 | 286 | 287 | 288 | -------------------------------------------------------------------------------- /utils/engine.py: -------------------------------------------------------------------------------- 1 | import os 2 | import os.path as osp 3 | import time 4 | import argparse 5 | 6 | import torch 7 | import torch.distributed as dist 8 | 9 | from utils.logger import get_logger 10 | from utils.pyt_utils import all_reduce_tensor, extant_file 11 | 12 | try: 13 | from apex.parallel import DistributedDataParallel, SyncBatchNorm 14 | except ImportError: 15 | raise ImportError( 16 | "Please install apex from https://www.github.com/nvidia/apex .") 17 | 18 | 19 | logger = get_logger() 20 | 21 | 22 | class Engine(object): 23 | def __init__(self, custom_parser=None): 24 | logger.info( 25 | "PyTorch Version {}".format(torch.__version__)) 26 | self.devices = None 27 | self.distributed = False 28 | 29 | if custom_parser is None: 30 | self.parser = argparse.ArgumentParser() 31 | else: 32 | assert isinstance(custom_parser, argparse.ArgumentParser) 33 | self.parser = custom_parser 34 | 35 | self.inject_default_parser() 36 | self.args = self.parser.parse_args() 37 | 38 | self.continue_state_object = self.args.continue_fpath 39 | 40 | if 'WORLD_SIZE' in os.environ: 41 | self.distributed = int(os.environ['WORLD_SIZE']) > 1 42 | print("WORLD_SIZE is %d" % (int(os.environ['WORLD_SIZE']))) 43 | if self.distributed: 44 | self.local_rank = self.args.local_rank 45 | self.world_size = int(os.environ['WORLD_SIZE']) 46 | torch.cuda.set_device(self.local_rank) 47 | dist.init_process_group(backend="nccl", init_method='env://') 48 | self.devices = [i for i in range(self.world_size)] 49 | else: 50 | gpus = os.environ["CUDA_VISIBLE_DEVICES"] 51 | self.devices = [i for i in range(len(gpus.split(',')))] 52 | 53 | def inject_default_parser(self): 54 | p = self.parser 55 | p.add_argument('-d', '--devices', default='', 56 | help='set data parallel training') 57 | p.add_argument('-c', '--continue', type=extant_file, 58 | metavar="FILE", 59 | dest="continue_fpath", 60 | help='continue from one certain checkpoint') 61 | 62 | def data_parallel(self, model): 63 | if self.distributed: 64 | model = DistributedDataParallel(model) 65 | else: 66 | model = torch.nn.DataParallel(model) 67 | return model 68 | 69 | def get_train_loader(self, train_dataset): 70 | train_sampler = None 71 | is_shuffle = True 72 | batch_size = self.args.batch_size 73 | 74 | if self.distributed: 75 | train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) 76 | batch_size = self.args.batch_size // self.world_size 77 | is_shuffle = False 78 | 79 | train_loader = torch.utils.data.DataLoader(train_dataset, 80 | batch_size=batch_size, 81 | num_workers=self.args.num_workers, 82 | drop_last=False, 83 | shuffle=is_shuffle, 84 | pin_memory=True, 85 | sampler=train_sampler) 86 | 87 | return train_loader, train_sampler 88 | 89 | def get_test_loader(self, test_dataset): 90 | test_sampler = None 91 | is_shuffle = False 92 | batch_size = self.args.batch_size 93 | 94 | if self.distributed: 95 | test_sampler = torch.utils.data.distributed.DistributedSampler( 96 | test_dataset) 97 | batch_size = self.args.batch_size // self.world_size 98 | 99 | test_loader = torch.utils.data.DataLoader(test_dataset, 100 | batch_size=1, 101 | num_workers=self.args.num_workers, 102 | drop_last=False, 103 | shuffle=is_shuffle, 104 | pin_memory=True, 105 | sampler=test_sampler) 106 | 107 | return test_loader, test_sampler 108 | 109 | 110 | def all_reduce_tensor(self, tensor, norm=True): 111 | if self.distributed: 112 | return all_reduce_tensor(tensor, world_size=self.world_size, norm=norm) 113 | else: 114 | return torch.mean(tensor) 115 | 116 | 117 | def __enter__(self): 118 | return self 119 | 120 | def __exit__(self, type, value, tb): 121 | torch.cuda.empty_cache() 122 | if type is not None: 123 | logger.warning( 124 | "A exception occurred during Engine initialization, " 125 | "give up running process") 126 | return False 127 | -------------------------------------------------------------------------------- /utils/logger.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import logging 4 | 5 | _default_level_name = os.getenv('ENGINE_LOGGING_LEVEL', 'INFO') 6 | _default_level = logging.getLevelName(_default_level_name.upper()) 7 | 8 | class LogFormatter(logging.Formatter): 9 | log_fout = None 10 | date_full = '[%(asctime)s %(lineno)d@%(filename)s:%(name)s] ' 11 | date = '%(asctime)s ' 12 | msg = '%(message)s' 13 | 14 | def format(self, record): 15 | if record.levelno == logging.DEBUG: 16 | mcl, mtxt = self._color_dbg, 'DBG' 17 | elif record.levelno == logging.WARNING: 18 | mcl, mtxt = self._color_warn, 'WRN' 19 | elif record.levelno == logging.ERROR: 20 | mcl, mtxt = self._color_err, 'ERR' 21 | else: 22 | mcl, mtxt = self._color_normal, '' 23 | 24 | if mtxt: 25 | mtxt += ' ' 26 | 27 | if self.log_fout: 28 | self.__set_fmt(self.date_full + mtxt + self.msg) 29 | formatted = super(LogFormatter, self).format(record) 30 | # self.log_fout.write(formatted) 31 | # self.log_fout.write('\n') 32 | # self.log_fout.flush() 33 | return formatted 34 | 35 | self.__set_fmt(self._color_date(self.date) + mcl(mtxt + self.msg)) 36 | formatted = super(LogFormatter, self).format(record) 37 | 38 | return formatted 39 | 40 | if sys.version_info.major < 3: 41 | def __set_fmt(self, fmt): 42 | self._fmt = fmt 43 | else: 44 | def __set_fmt(self, fmt): 45 | self._style._fmt = fmt 46 | 47 | @staticmethod 48 | def _color_dbg(msg): 49 | return '\x1b[36m{}\x1b[0m'.format(msg) 50 | 51 | @staticmethod 52 | def _color_warn(msg): 53 | return '\x1b[1;31m{}\x1b[0m'.format(msg) 54 | 55 | @staticmethod 56 | def _color_err(msg): 57 | return '\x1b[1;4;31m{}\x1b[0m'.format(msg) 58 | 59 | @staticmethod 60 | def _color_omitted(msg): 61 | return '\x1b[35m{}\x1b[0m'.format(msg) 62 | 63 | @staticmethod 64 | def _color_normal(msg): 65 | return msg 66 | 67 | @staticmethod 68 | def _color_date(msg): 69 | return '\x1b[32m{}\x1b[0m'.format(msg) 70 | 71 | 72 | def get_logger(log_dir=None, log_file=None, formatter=LogFormatter): 73 | logger = logging.getLogger() 74 | logger.setLevel(_default_level) 75 | del logger.handlers[:] 76 | 77 | if log_dir and log_file: 78 | if not os.path.isdir(log_dir): 79 | os.makedirs(log_dir) 80 | LogFormatter.log_fout = True 81 | file_handler = logging.FileHandler(log_file, mode='a') 82 | file_handler.setLevel(logging.INFO) 83 | file_handler.setFormatter(formatter) 84 | logger.addHandler(file_handler) 85 | 86 | stream_handler = logging.StreamHandler() 87 | stream_handler.setFormatter(formatter(datefmt='%d %H:%M:%S')) 88 | stream_handler.setLevel(0) 89 | logger.addHandler(stream_handler) 90 | return logger 91 | -------------------------------------------------------------------------------- /utils/loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | import torch.nn as nn 4 | import numpy as np 5 | 6 | class BinaryDiceLoss(nn.Module): 7 | def __init__(self, smooth=1, p=2, reduction='mean'): 8 | super(BinaryDiceLoss, self).__init__() 9 | self.smooth = smooth 10 | self.p = p 11 | self.reduction = reduction 12 | 13 | def forward(self, predict, target): 14 | assert predict.shape[0] == target.shape[0], "predict & target batch size don't match" 15 | predict = predict.contiguous().view(predict.shape[0], -1) 16 | target = target.contiguous().view(target.shape[0], -1) 17 | 18 | num = torch.sum(torch.mul(predict, target), dim=1) 19 | den = torch.sum(predict, dim=1) + torch.sum(target, dim=1) + self.smooth 20 | 21 | dice_score = 2*num / den 22 | loss_avg = 1 - dice_score.mean() 23 | 24 | return loss_avg 25 | 26 | class DiceLoss4BraTS(nn.Module): 27 | def __init__(self, weight=None, ignore_index=None, **kwargs): 28 | super(DiceLoss4BraTS, self).__init__() 29 | self.kwargs = kwargs 30 | self.weight = weight 31 | self.ignore_index = ignore_index 32 | 33 | def forward(self, predict, target): 34 | assert predict.shape == target.shape, 'predict %s & target %s shape do not match' % (predict.shape, target.shape) 35 | dice = BinaryDiceLoss(**self.kwargs) 36 | total_loss = 0 37 | predict = F.sigmoid(predict) 38 | 39 | for i in range(target.shape[1]): 40 | if i != self.ignore_index: 41 | dice_loss = dice(predict[:, i], target[:, i]) 42 | if self.weight is not None: 43 | assert self.weight.shape[0] == target.shape[1], \ 44 | 'Expect weight shape [{}], get[{}]'.format(target.shape[1], self.weight.shape[0]) 45 | dice_loss *= self.weights[i] 46 | total_loss += dice_loss 47 | 48 | return total_loss/(target.shape[1]-1 if self.ignore_index!=None else target.shape[1]) 49 | 50 | 51 | class BCELoss4BraTS(nn.Module): 52 | def __init__(self, ignore_index=None, **kwargs): 53 | super(BCELoss4BraTS, self).__init__() 54 | self.kwargs = kwargs 55 | self.ignore_index = ignore_index 56 | self.criterion = nn.BCEWithLogitsLoss() 57 | 58 | def weighted_BCE_cross_entropy(self, output, target, weights = None): 59 | if weights is not None: 60 | assert len(weights) == 2 61 | output = torch.clamp(output, min=1e-7, max=1-1e-7) 62 | bce = weights[1] * (target * torch.log(output)) + \ 63 | weights[0] * ((1-target) * torch.log((1-output))) 64 | else: 65 | output = torch.clamp(output, min=1e-3, max=1 - 1e-3) 66 | bce = target * torch.log(output) + (1-target) * torch.log((1-output)) 67 | return torch.neg(torch.mean(bce)) 68 | 69 | def forward(self, predict, target): 70 | assert predict.shape == target.shape, 'predict & target shape do not match' 71 | total_loss = 0 72 | for i in range(target.shape[1]): 73 | if i != self.ignore_index: 74 | bce_loss = self.criterion(predict[:, i], target[:, i]) 75 | total_loss += bce_loss 76 | 77 | return total_loss.mean() 78 | 79 | 80 | class BCELossBoud(nn.Module): 81 | def __init__(self, weight=None, ignore_index=None, **kwargs): 82 | super(BCELossBoud, self).__init__() 83 | self.kwargs = kwargs 84 | self.weight = weight 85 | self.ignore_index = ignore_index 86 | self.criterion = nn.BCEWithLogitsLoss() 87 | 88 | def weighted_BCE_cross_entropy(self, output, target, weights = None): 89 | if weights is not None: 90 | assert len(weights) == 2 91 | output = torch.clamp(output, min=1e-3, max=1-1e-3) 92 | bce = weights[1] * (target * torch.log(output)) + \ 93 | weights[0] * ((1-target) * torch.log((1-output))) 94 | else: 95 | output = torch.clamp(output, min=1e-3, max=1 - 1e-3) 96 | bce = target * torch.log(output) + (1-target) * torch.log((1-output)) 97 | return torch.neg(torch.mean(bce)) 98 | 99 | def forward(self, predict, target): 100 | 101 | bs, category, depth, width, heigt = target.shape 102 | bce_loss = [] 103 | for i in range(predict.shape[1]): 104 | pred_i = predict[:,i] 105 | targ_i = target[:,i] 106 | tt = np.log(depth * width * heigt / (target[:, i].cpu().data.numpy().sum()+1)) 107 | bce_i = self.weighted_BCE_cross_entropy(pred_i, targ_i, weights=[1, tt]) 108 | bce_loss.append(bce_i) 109 | 110 | bce_loss = torch.stack(bce_loss) 111 | total_loss = bce_loss.mean() 112 | return total_loss -------------------------------------------------------------------------------- /utils/pyt_utils.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | import os 3 | import sys 4 | import time 5 | import argparse 6 | from collections import OrderedDict, defaultdict 7 | 8 | import torch 9 | import torch.utils.model_zoo as model_zoo 10 | import torch.distributed as dist 11 | 12 | from .logger import get_logger 13 | 14 | logger = get_logger() 15 | 16 | 17 | def reduce_tensor(tensor, dst=0, op=dist.ReduceOp.SUM, world_size=1): 18 | tensor = tensor.clone() 19 | dist.reduce(tensor, dst, op) 20 | if dist.get_rank() == dst: 21 | tensor.div_(world_size) 22 | 23 | return tensor 24 | 25 | 26 | def all_reduce_tensor(tensor, op=dist.ReduceOp.SUM, world_size=1, norm=True): 27 | tensor = tensor.clone() 28 | dist.all_reduce(tensor, op) 29 | if norm: 30 | tensor.div_(world_size) 31 | 32 | return tensor 33 | 34 | 35 | def extant_file(x): 36 | """ 37 | 'Type' for argparse - checks that file exists but does not open. 38 | """ 39 | if not os.path.exists(x): 40 | # Argparse uses the ArgumentTypeError to give a rejection message like: 41 | # error: argument input: x does not exist 42 | raise argparse.ArgumentTypeError("{0} does not exist".format(x)) 43 | return x 44 | 45 | 46 | 47 | --------------------------------------------------------------------------------