From 1fd50af92367fe7e1ac09cde00e5eb1756234541 Mon Sep 17 00:00:00 2001 From: luoxd Date: Sun, 6 Mar 2022 15:20:05 +0800 Subject: [PATCH 1/7] update 5 fold cross validation code --- code/README.md | 46 - code/config.py | 229 ----- .../swin_tiny_patch4_window7_224_lite.yaml | 12 - code/dataloaders/acdc_data_processing.py | 37 - code/dataloaders/brats2019.py | 237 ----- code/dataloaders/brats_proprecessing.py | 110 --- code/dataloaders/dataset.py | 155 ++- .../prostate_dataset_preprocessing.py | 100 ++ code/dataloaders/utils.py | 214 ----- code/networks/VoxResNet.py | 0 code/networks/attention_unet.py | 0 code/networks/config.py | 229 ----- code/networks/discriminator.py | 0 code/networks/enet.py | 614 ------------ code/networks/grid_attention_layer.py | 0 code/networks/net_factory.py | 87 +- code/networks/net_factory_3d.py | 3 - code/networks/networks_other.py | 0 code/networks/neural_network.py | 903 ------------------ code/networks/nnunet.py | 535 ----------- ...ransformer_unet_skip_expand_decoder_sys.py | 804 ---------------- code/networks/unet.py | 78 +- code/networks/unet_3D.py | 0 code/networks/unet_3D_dv_semi.py | 112 --- code/networks/unet_multitask.py | 349 +++++++ code/networks/utils.py | 0 code/networks/vision_transformer.py | 90 -- code/networks/vnet.py | 0 code/pretrained_ckpt/readme.txt | 1 - code/test.py | 149 +++ code/test_2D_fully.py | 117 --- code/test_3D.py | 41 - code/test_3D_util.py | 152 --- code/test_acdc_unet_semi_seg.sh | 8 - code/test_brats2019_semi_seg.sh | 10 - code/test_urpc.py | 55 -- code/test_urpc_util.py | 161 ---- code/train_acdc_unet_semi_seg.sh | 8 - code/train_adversarial_network_2D.py | 283 ------ code/train_adversarial_network_3D.py | 271 ------ code/train_brats2019_semi_seg.sh | 10 - code/train_cross_consistency_training_2D.py | 277 ------ code/train_cross_pseudo_supervision_2D.py | 356 ------- code/train_cross_pseudo_supervision_3D.py | 321 ------- ...oss_teaching_between_cnn_transformer_2D.py | 413 -------- code/train_deep_co_training_2D.py | 267 ------ code/train_entropy_minimization_2D.py | 258 ----- code/train_entropy_minimization_3D.py | 246 ----- code/train_fully_supervised_2D.py | 218 ----- code/train_fully_supervised_3D.py | 203 ---- ...ain_interpolation_consistency_training.py} | 591 ++++++------ ...n_interpolation_consistency_training_3D.py | 284 ------ ...an_teacher_2D.py => train_mean_teacher.py} | 541 +++++------ code/train_mean_teacher_3D.py | 265 ----- ...> train_uncertainty_aware_mean_teacher.py} | 581 ++++++----- ...train_uncertainty_aware_mean_teacher_3D.py | 288 ------ ...tainty_rectified_pyramid_consistency_2D.py | 311 ------ ...tainty_rectified_pyramid_consistency_3D.py | 313 ------ code/utils/MIPloss.py | 125 +++ code/utils/gate_crf_loss.py | 205 ++++ code/utils/losses.py | 29 + code/utils/metrics.py | 0 code/utils/ramps.py | 0 code/utils/util.py | 0 code/val_2D.py | 117 ++- code/val_3D.py | 107 --- code/val_urpc_util.py | 107 --- 67 files changed, 1969 insertions(+), 10664 deletions(-) delete mode 100644 code/README.md delete mode 100644 code/config.py delete mode 100644 code/configs/swin_tiny_patch4_window7_224_lite.yaml delete mode 100644 code/dataloaders/acdc_data_processing.py delete mode 100644 code/dataloaders/brats2019.py delete mode 100644 code/dataloaders/brats_proprecessing.py create mode 100644 code/dataloaders/prostate_dataset_preprocessing.py delete mode 100644 code/dataloaders/utils.py mode change 100644 => 100755 code/networks/VoxResNet.py mode change 100644 => 100755 code/networks/attention_unet.py delete mode 100644 code/networks/config.py mode change 100644 => 100755 code/networks/discriminator.py delete mode 100644 code/networks/enet.py mode change 100644 => 100755 code/networks/grid_attention_layer.py mode change 100644 => 100755 code/networks/net_factory.py mode change 100644 => 100755 code/networks/net_factory_3d.py mode change 100644 => 100755 code/networks/networks_other.py delete mode 100644 code/networks/neural_network.py delete mode 100644 code/networks/nnunet.py delete mode 100644 code/networks/swin_transformer_unet_skip_expand_decoder_sys.py mode change 100644 => 100755 code/networks/unet.py mode change 100644 => 100755 code/networks/unet_3D.py delete mode 100644 code/networks/unet_3D_dv_semi.py create mode 100755 code/networks/unet_multitask.py mode change 100644 => 100755 code/networks/utils.py delete mode 100644 code/networks/vision_transformer.py mode change 100644 => 100755 code/networks/vnet.py delete mode 100644 code/pretrained_ckpt/readme.txt create mode 100644 code/test.py delete mode 100644 code/test_2D_fully.py delete mode 100644 code/test_3D.py delete mode 100644 code/test_3D_util.py delete mode 100644 code/test_acdc_unet_semi_seg.sh delete mode 100644 code/test_brats2019_semi_seg.sh delete mode 100644 code/test_urpc.py delete mode 100644 code/test_urpc_util.py delete mode 100644 code/train_acdc_unet_semi_seg.sh delete mode 100644 code/train_adversarial_network_2D.py delete mode 100644 code/train_adversarial_network_3D.py delete mode 100644 code/train_brats2019_semi_seg.sh delete mode 100644 code/train_cross_consistency_training_2D.py delete mode 100644 code/train_cross_pseudo_supervision_2D.py delete mode 100644 code/train_cross_pseudo_supervision_3D.py delete mode 100644 code/train_cross_teaching_between_cnn_transformer_2D.py delete mode 100644 code/train_deep_co_training_2D.py delete mode 100644 code/train_entropy_minimization_2D.py delete mode 100644 code/train_entropy_minimization_3D.py delete mode 100644 code/train_fully_supervised_2D.py delete mode 100644 code/train_fully_supervised_3D.py rename code/{train_interpolation_consistency_training_2D.py => train_interpolation_consistency_training.py} (60%) delete mode 100644 code/train_interpolation_consistency_training_3D.py rename code/{train_mean_teacher_2D.py => train_mean_teacher.py} (68%) delete mode 100644 code/train_mean_teacher_3D.py rename code/{train_uncertainty_aware_mean_teacher_2D.py => train_uncertainty_aware_mean_teacher.py} (68%) delete mode 100644 code/train_uncertainty_aware_mean_teacher_3D.py delete mode 100644 code/train_uncertainty_rectified_pyramid_consistency_2D.py delete mode 100644 code/train_uncertainty_rectified_pyramid_consistency_3D.py create mode 100755 code/utils/MIPloss.py create mode 100755 code/utils/gate_crf_loss.py mode change 100644 => 100755 code/utils/losses.py mode change 100644 => 100755 code/utils/metrics.py mode change 100644 => 100755 code/utils/ramps.py mode change 100644 => 100755 code/utils/util.py mode change 100644 => 100755 code/val_2D.py delete mode 100644 code/val_3D.py delete mode 100644 code/val_urpc_util.py diff --git a/code/README.md b/code/README.md deleted file mode 100644 index 40c6631..0000000 --- a/code/README.md +++ /dev/null @@ -1,46 +0,0 @@ -## Semi-supervised Learning for Medical Image Segmentation (**SSL4MIS**) - -## Requirements -Some important required packages include: -* [Pytorch][torch_link] version >=0.4.1. -* TensorBoardX -* Python == 3.6 -* Efficientnet-Pytorch `pip install efficientnet_pytorch` -* Some basic python packages such as Numpy, Scikit-image, SimpleITK, Scipy ...... - -Follow official guidance to install [Pytorch][torch_link]. - -[torch_link]:https://pytorch.org/ - -# Usage - -1. Clone the repo: -``` -git clone https://https://github.com/HiLab-git/SSL4MIS.git -cd SSL4MIS -``` -2. Download the processed data and put the data in `../data/BraTS2019` or `../data/ACDC`, please read and follow the [README](https://github.com/Luoxd1996/SSL4MIS/tree/master/data/). - -3. Train the model -``` -cd code -python train_XXXXX_3D.py or python train_XXXXX_2D.py or bash train_acdc_XXXXX.sh -``` - -4. Test the model -``` -python test_XXXXX.py -``` -# Reimplemented methods -* [Mean Teacher](https://papers.nips.cc/paper/6719-mean-teachers-are-better-role-models-weight-averaged-consistency-targets-improve-semi-supervised-deep-learning-results.pdf)[[2D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_mean_teacher_2D.py)/[3D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_mean_teacher_3D.py)] -* [Entropy Minimization](https://openaccess.thecvf.com/content_CVPR_2019/papers/Vu_ADVENT_Adversarial_Entropy_Minimization_for_Domain_Adaptation_in_Semantic_Segmentation_CVPR_2019_paper.pdf)[[2D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_entropy_minimization_2D.py)/[3D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_entropy_minimization_3D.py)] -* [Deep Adversarial Networks](https://link.springer.com/chapter/10.1007/978-3-319-66179-7_47)[[2D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_adversarial_network_2D.py)/[3D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_adversarial_network_3D.py)] -* [Uncertainty Aware Mean Teacher](https://arxiv.org/pdf/1907.07034.pdf)[[2D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_uncertainty_aware_mean_teacher_2D.py)/[3D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_uncertainty_aware_mean_teacher_3D.py)] -* [Interpolation Consistency Training](https://arxiv.org/pdf/1903.03825.pdf)[[2D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_interpolation_consistency_training_2D.py)/[3D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_interpolation_consistency_training_3D.py)] -* [Uncertainty Rectified Pyramid Consistency](https://arxiv.org/pdf/2012.07042.pdf)[[2D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_uncertainty_rectified_pyramid_consistency_2D.py)/[3D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_uncertainty_rectified_pyramid_consistency_3D.py)] -* [Cross Pseudo Supervision](https://arxiv.org/abs/2106.01226)[[2D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_cross_pseudo_supervision_2D.py)/[3D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_cross_pseudo_supervision_3D.py)] -* [Cross Consistency Training](https://openaccess.thecvf.com/content_CVPR_2020/papers/Ouali_Semi-Supervised_Semantic_Segmentation_With_Cross-Consistency_Training_CVPR_2020_paper.pdf)[[2D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_cross_consistency_training_2D.py)] -* [Deep Co-Training](https://openaccess.thecvf.com/content_ECCV_2018/papers/Siyuan_Qiao_Deep_Co-Training_for_ECCV_2018_paper.pdf)[[2D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_deep_co_training_2D.py)] -* [Cross Teaching between CNN and Transformer](https://arxiv.org/pdf/2112.04894.pdf)[[2D](https://github.com/HiLab-git/SSL4MIS/blob/master/code/train_cross_teaching_between_cnn_transformer_2D.py)] -## Acknowledgement -* Part of the code is adapted from open-source codebase and original implementations of algorithms, we thank these author for their fantastic and efficient codebase, such as, [UA-MT](https://github.com/yulequan/UA-MT), [Attention-Gated-Networks](https://github.com/ozan-oktay/Attention-Gated-Networks) and [segmentatic_segmentation.pytorch](https://github.com/qubvel/segmentation_models.pytorch) . diff --git a/code/config.py b/code/config.py deleted file mode 100644 index 35bf199..0000000 --- a/code/config.py +++ /dev/null @@ -1,229 +0,0 @@ -# -------------------------------------------------------- -# Swin Transformer -# Copyright (c) 2021 Microsoft -# Licensed under The MIT License [see LICENSE for details] -# Written by Ze Liu -# --------------------------------------------------------' - -import os -import yaml -from yacs.config import CfgNode as CN - -_C = CN() - -# Base config files -_C.BASE = [''] - -# ----------------------------------------------------------------------------- -# Data settings -# ----------------------------------------------------------------------------- -_C.DATA = CN() -# Batch size for a single GPU, could be overwritten by command line argument -_C.DATA.BATCH_SIZE = 128 -# Path to dataset, could be overwritten by command line argument -_C.DATA.DATA_PATH = '' -# Dataset name -_C.DATA.DATASET = 'imagenet' -# Input image size -_C.DATA.IMG_SIZE = 224 -# Interpolation to resize image (random, bilinear, bicubic) -_C.DATA.INTERPOLATION = 'bicubic' -# Use zipped dataset instead of folder dataset -# could be overwritten by command line argument -_C.DATA.ZIP_MODE = False -# Cache Data in Memory, could be overwritten by command line argument -_C.DATA.CACHE_MODE = 'part' -# Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU. -_C.DATA.PIN_MEMORY = True -# Number of data loading threads -_C.DATA.NUM_WORKERS = 8 - -# ----------------------------------------------------------------------------- -# Model settings -# ----------------------------------------------------------------------------- -_C.MODEL = CN() -# Model type -_C.MODEL.TYPE = 'swin' -# Model name -_C.MODEL.NAME = 'swin_tiny_patch4_window7_224' -# Checkpoint to resume, could be overwritten by command line argument -_C.MODEL.PRETRAIN_CKPT = './pretrained_ckpt/swin_tiny_patch4_window7_224.pth' -_C.MODEL.RESUME = '' -# Number of classes, overwritten in data preparation -_C.MODEL.NUM_CLASSES = 1000 -# Dropout rate -_C.MODEL.DROP_RATE = 0.0 -# Drop path rate -_C.MODEL.DROP_PATH_RATE = 0.1 -# Label Smoothing -_C.MODEL.LABEL_SMOOTHING = 0.1 - -# Swin Transformer parameters -_C.MODEL.SWIN = CN() -_C.MODEL.SWIN.PATCH_SIZE = 4 -_C.MODEL.SWIN.IN_CHANS = 3 -_C.MODEL.SWIN.EMBED_DIM = 96 -_C.MODEL.SWIN.DEPTHS = [2, 2, 6, 2] -_C.MODEL.SWIN.DECODER_DEPTHS = [2, 2, 6, 2] -_C.MODEL.SWIN.NUM_HEADS = [3, 6, 12, 24] -_C.MODEL.SWIN.WINDOW_SIZE = 7 -_C.MODEL.SWIN.MLP_RATIO = 4. -_C.MODEL.SWIN.QKV_BIAS = True -_C.MODEL.SWIN.QK_SCALE = None -_C.MODEL.SWIN.APE = False -_C.MODEL.SWIN.PATCH_NORM = True -_C.MODEL.SWIN.FINAL_UPSAMPLE= "expand_first" - -# ----------------------------------------------------------------------------- -# Training settings -# ----------------------------------------------------------------------------- -_C.TRAIN = CN() -_C.TRAIN.START_EPOCH = 0 -_C.TRAIN.EPOCHS = 300 -_C.TRAIN.WARMUP_EPOCHS = 20 -_C.TRAIN.WEIGHT_DECAY = 0.05 -_C.TRAIN.BASE_LR = 5e-4 -_C.TRAIN.WARMUP_LR = 5e-7 -_C.TRAIN.MIN_LR = 5e-6 -# Clip gradient norm -_C.TRAIN.CLIP_GRAD = 5.0 -# Auto resume from latest checkpoint -_C.TRAIN.AUTO_RESUME = True -# Gradient accumulation steps -# could be overwritten by command line argument -_C.TRAIN.ACCUMULATION_STEPS = 0 -# Whether to use gradient checkpointing to save memory -# could be overwritten by command line argument -_C.TRAIN.USE_CHECKPOINT = False - -# LR scheduler -_C.TRAIN.LR_SCHEDULER = CN() -_C.TRAIN.LR_SCHEDULER.NAME = 'cosine' -# Epoch interval to decay LR, used in StepLRScheduler -_C.TRAIN.LR_SCHEDULER.DECAY_EPOCHS = 30 -# LR decay rate, used in StepLRScheduler -_C.TRAIN.LR_SCHEDULER.DECAY_RATE = 0.1 - -# Optimizer -_C.TRAIN.OPTIMIZER = CN() -_C.TRAIN.OPTIMIZER.NAME = 'adamw' -# Optimizer Epsilon -_C.TRAIN.OPTIMIZER.EPS = 1e-8 -# Optimizer Betas -_C.TRAIN.OPTIMIZER.BETAS = (0.9, 0.999) -# SGD momentum -_C.TRAIN.OPTIMIZER.MOMENTUM = 0.9 - -# ----------------------------------------------------------------------------- -# Augmentation settings -# ----------------------------------------------------------------------------- -_C.AUG = CN() -# Color jitter factor -_C.AUG.COLOR_JITTER = 0.4 -# Use AutoAugment policy. "v0" or "original" -_C.AUG.AUTO_AUGMENT = 'rand-m9-mstd0.5-inc1' -# Random erase prob -_C.AUG.REPROB = 0.25 -# Random erase mode -_C.AUG.REMODE = 'pixel' -# Random erase count -_C.AUG.RECOUNT = 1 -# Mixup alpha, mixup enabled if > 0 -_C.AUG.MIXUP = 0.8 -# Cutmix alpha, cutmix enabled if > 0 -_C.AUG.CUTMIX = 1.0 -# Cutmix min/max ratio, overrides alpha and enables cutmix if set -_C.AUG.CUTMIX_MINMAX = None -# Probability of performing mixup or cutmix when either/both is enabled -_C.AUG.MIXUP_PROB = 1.0 -# Probability of switching to cutmix when both mixup and cutmix enabled -_C.AUG.MIXUP_SWITCH_PROB = 0.5 -# How to apply mixup/cutmix params. Per "batch", "pair", or "elem" -_C.AUG.MIXUP_MODE = 'batch' - -# ----------------------------------------------------------------------------- -# Testing settings -# ----------------------------------------------------------------------------- -_C.TEST = CN() -# Whether to use center crop when testing -_C.TEST.CROP = True - -# ----------------------------------------------------------------------------- -# Misc -# ----------------------------------------------------------------------------- -# Mixed precision opt level, if O0, no amp is used ('O0', 'O1', 'O2') -# overwritten by command line argument -_C.AMP_OPT_LEVEL = '' -# Path to output folder, overwritten by command line argument -_C.OUTPUT = '' -# Tag of experiment, overwritten by command line argument -_C.TAG = 'default' -# Frequency to save checkpoint -_C.SAVE_FREQ = 1 -# Frequency to logging info -_C.PRINT_FREQ = 10 -# Fixed random seed -_C.SEED = 0 -# Perform evaluation only, overwritten by command line argument -_C.EVAL_MODE = False -# Test throughput only, overwritten by command line argument -_C.THROUGHPUT_MODE = False -# local rank for DistributedDataParallel, given by command line argument -_C.LOCAL_RANK = 0 - - -def _update_config_from_file(config, cfg_file): - config.defrost() - with open(cfg_file, 'r') as f: - yaml_cfg = yaml.load(f, Loader=yaml.FullLoader) - - for cfg in yaml_cfg.setdefault('BASE', ['']): - if cfg: - _update_config_from_file( - config, os.path.join(os.path.dirname(cfg_file), cfg) - ) - print('=> merge config from {}'.format(cfg_file)) - config.merge_from_file(cfg_file) - config.freeze() - - -def update_config(config, args): - _update_config_from_file(config, args.cfg) - - config.defrost() - if args.opts: - config.merge_from_list(args.opts) - - # merge from specific arguments - if args.batch_size: - config.DATA.BATCH_SIZE = args.batch_size - if args.zip: - config.DATA.ZIP_MODE = True - if args.cache_mode: - config.DATA.CACHE_MODE = args.cache_mode - if args.resume: - config.MODEL.RESUME = args.resume - if args.accumulation_steps: - config.TRAIN.ACCUMULATION_STEPS = args.accumulation_steps - if args.use_checkpoint: - config.TRAIN.USE_CHECKPOINT = True - if args.amp_opt_level: - config.AMP_OPT_LEVEL = args.amp_opt_level - if args.tag: - config.TAG = args.tag - if args.eval: - config.EVAL_MODE = True - if args.throughput: - config.THROUGHPUT_MODE = True - - config.freeze() - - -def get_config(args): - """Get a yacs CfgNode object with default values.""" - # Return a clone so that the defaults will not be altered - # This is for the "local variable" use pattern - config = _C.clone() - update_config(config, args) - - return config diff --git a/code/configs/swin_tiny_patch4_window7_224_lite.yaml b/code/configs/swin_tiny_patch4_window7_224_lite.yaml deleted file mode 100644 index 599b4f0..0000000 --- a/code/configs/swin_tiny_patch4_window7_224_lite.yaml +++ /dev/null @@ -1,12 +0,0 @@ -MODEL: - TYPE: swin - NAME: swin_tiny_patch4_window7_224 - DROP_PATH_RATE: 0.2 - PRETRAIN_CKPT: "../code/pretrained_ckpt/swin_tiny_patch4_window7_224.pth" - SWIN: - FINAL_UPSAMPLE: "expand_first" - EMBED_DIM: 96 - DEPTHS: [ 2, 2, 2, 2 ] - DECODER_DEPTHS: [ 2, 2, 2, 1] - NUM_HEADS: [ 3, 6, 12, 24 ] - WINDOW_SIZE: 7 \ No newline at end of file diff --git a/code/dataloaders/acdc_data_processing.py b/code/dataloaders/acdc_data_processing.py deleted file mode 100644 index 0471a65..0000000 --- a/code/dataloaders/acdc_data_processing.py +++ /dev/null @@ -1,37 +0,0 @@ -import glob -import os - -import h5py -import numpy as np -import SimpleITK as sitk - -slice_num = 0 -mask_path = sorted(glob.glob("/home/xdluo/data/ACDC/image/*.nii.gz")) -for case in mask_path: - img_itk = sitk.ReadImage(case) - origin = img_itk.GetOrigin() - spacing = img_itk.GetSpacing() - direction = img_itk.GetDirection() - image = sitk.GetArrayFromImage(img_itk) - msk_path = case.replace("image", "label").replace(".nii.gz", "_gt.nii.gz") - if os.path.exists(msk_path): - print(msk_path) - msk_itk = sitk.ReadImage(msk_path) - mask = sitk.GetArrayFromImage(msk_itk) - image = (image - image.min()) / (image.max() - image.min()) - print(image.shape) - image = image.astype(np.float32) - item = case.split("/")[-1].split(".")[0] - if image.shape != mask.shape: - print("Error") - print(item) - for slice_ind in range(image.shape[0]): - f = h5py.File( - '/home/xdluo/data/ACDC/data/{}_slice_{}.h5'.format(item, slice_ind), 'w') - f.create_dataset( - 'image', data=image[slice_ind], compression="gzip") - f.create_dataset('label', data=mask[slice_ind], compression="gzip") - f.close() - slice_num += 1 -print("Converted all ACDC volumes to 2D slices") -print("Total {} slices".format(slice_num)) diff --git a/code/dataloaders/brats2019.py b/code/dataloaders/brats2019.py deleted file mode 100644 index 585efce..0000000 --- a/code/dataloaders/brats2019.py +++ /dev/null @@ -1,237 +0,0 @@ -import os -import torch -import numpy as np -from glob import glob -from torch.utils.data import Dataset -import h5py -import itertools -from torch.utils.data.sampler import Sampler - - -class BraTS2019(Dataset): - """ BraTS2019 Dataset """ - - def __init__(self, base_dir=None, split='train', num=None, transform=None): - self._base_dir = base_dir - self.transform = transform - self.sample_list = [] - - train_path = self._base_dir+'/train.txt' - test_path = self._base_dir+'/val.txt' - - if split == 'train': - with open(train_path, 'r') as f: - self.image_list = f.readlines() - elif split == 'test': - with open(test_path, 'r') as f: - self.image_list = f.readlines() - - self.image_list = [item.replace('\n', '').split(",")[0] for item in self.image_list] - if num is not None: - self.image_list = self.image_list[:num] - print("total {} samples".format(len(self.image_list))) - - def __len__(self): - return len(self.image_list) - - def __getitem__(self, idx): - image_name = self.image_list[idx] - h5f = h5py.File(self._base_dir + "/data/{}.h5".format(image_name), 'r') - image = h5f['image'][:] - label = h5f['label'][:] - sample = {'image': image, 'label': label.astype(np.uint8)} - if self.transform: - sample = self.transform(sample) - return sample - - -class CenterCrop(object): - def __init__(self, output_size): - self.output_size = output_size - - def __call__(self, sample): - image, label = sample['image'], sample['label'] - - # pad the sample if necessary - if label.shape[0] <= self.output_size[0] or label.shape[1] <= self.output_size[1] or label.shape[2] <= \ - self.output_size[2]: - pw = max((self.output_size[0] - label.shape[0]) // 2 + 3, 0) - ph = max((self.output_size[1] - label.shape[1]) // 2 + 3, 0) - pd = max((self.output_size[2] - label.shape[2]) // 2 + 3, 0) - image = np.pad(image, [(pw, pw), (ph, ph), (pd, pd)], - mode='constant', constant_values=0) - label = np.pad(label, [(pw, pw), (ph, ph), (pd, pd)], - mode='constant', constant_values=0) - - (w, h, d) = image.shape - - w1 = int(round((w - self.output_size[0]) / 2.)) - h1 = int(round((h - self.output_size[1]) / 2.)) - d1 = int(round((d - self.output_size[2]) / 2.)) - - label = label[w1:w1 + self.output_size[0], h1:h1 + - self.output_size[1], d1:d1 + self.output_size[2]] - image = image[w1:w1 + self.output_size[0], h1:h1 + - self.output_size[1], d1:d1 + self.output_size[2]] - - return {'image': image, 'label': label} - - -class RandomCrop(object): - """ - Crop randomly the image in a sample - Args: - output_size (int): Desired output size - """ - - def __init__(self, output_size, with_sdf=False): - self.output_size = output_size - self.with_sdf = with_sdf - - def __call__(self, sample): - image, label = sample['image'], sample['label'] - if self.with_sdf: - sdf = sample['sdf'] - - # pad the sample if necessary - if label.shape[0] <= self.output_size[0] or label.shape[1] <= self.output_size[1] or label.shape[2] <= \ - self.output_size[2]: - pw = max((self.output_size[0] - label.shape[0]) // 2 + 3, 0) - ph = max((self.output_size[1] - label.shape[1]) // 2 + 3, 0) - pd = max((self.output_size[2] - label.shape[2]) // 2 + 3, 0) - image = np.pad(image, [(pw, pw), (ph, ph), (pd, pd)], - mode='constant', constant_values=0) - label = np.pad(label, [(pw, pw), (ph, ph), (pd, pd)], - mode='constant', constant_values=0) - if self.with_sdf: - sdf = np.pad(sdf, [(pw, pw), (ph, ph), (pd, pd)], - mode='constant', constant_values=0) - - (w, h, d) = image.shape - # if np.random.uniform() > 0.33: - # w1 = np.random.randint((w - self.output_size[0])//4, 3*(w - self.output_size[0])//4) - # h1 = np.random.randint((h - self.output_size[1])//4, 3*(h - self.output_size[1])//4) - # else: - w1 = np.random.randint(0, w - self.output_size[0]) - h1 = np.random.randint(0, h - self.output_size[1]) - d1 = np.random.randint(0, d - self.output_size[2]) - - label = label[w1:w1 + self.output_size[0], h1:h1 + - self.output_size[1], d1:d1 + self.output_size[2]] - image = image[w1:w1 + self.output_size[0], h1:h1 + - self.output_size[1], d1:d1 + self.output_size[2]] - if self.with_sdf: - sdf = sdf[w1:w1 + self.output_size[0], h1:h1 + - self.output_size[1], d1:d1 + self.output_size[2]] - return {'image': image, 'label': label, 'sdf': sdf} - else: - return {'image': image, 'label': label} - - -class RandomRotFlip(object): - """ - Crop randomly flip the dataset in a sample - Args: - output_size (int): Desired output size - """ - - def __call__(self, sample): - image, label = sample['image'], sample['label'] - k = np.random.randint(0, 4) - image = np.rot90(image, k) - label = np.rot90(label, k) - axis = np.random.randint(0, 2) - image = np.flip(image, axis=axis).copy() - label = np.flip(label, axis=axis).copy() - - return {'image': image, 'label': label} - - -class RandomNoise(object): - def __init__(self, mu=0, sigma=0.1): - self.mu = mu - self.sigma = sigma - - def __call__(self, sample): - image, label = sample['image'], sample['label'] - noise = np.clip(self.sigma * np.random.randn( - image.shape[0], image.shape[1], image.shape[2]), -2*self.sigma, 2*self.sigma) - noise = noise + self.mu - image = image + noise - return {'image': image, 'label': label} - - -class CreateOnehotLabel(object): - def __init__(self, num_classes): - self.num_classes = num_classes - - def __call__(self, sample): - image, label = sample['image'], sample['label'] - onehot_label = np.zeros( - (self.num_classes, label.shape[0], label.shape[1], label.shape[2]), dtype=np.float32) - for i in range(self.num_classes): - onehot_label[i, :, :, :] = (label == i).astype(np.float32) - return {'image': image, 'label': label, 'onehot_label': onehot_label} - - -class ToTensor(object): - """Convert ndarrays in sample to Tensors.""" - - def __call__(self, sample): - image = sample['image'] - image = image.reshape( - 1, image.shape[0], image.shape[1], image.shape[2]).astype(np.float32) - if 'onehot_label' in sample: - return {'image': torch.from_numpy(image), 'label': torch.from_numpy(sample['label']).long(), - 'onehot_label': torch.from_numpy(sample['onehot_label']).long()} - else: - return {'image': torch.from_numpy(image), 'label': torch.from_numpy(sample['label']).long()} - - -class TwoStreamBatchSampler(Sampler): - """Iterate two sets of indices - - An 'epoch' is one iteration through the primary indices. - During the epoch, the secondary indices are iterated through - as many times as needed. - """ - - def __init__(self, primary_indices, secondary_indices, batch_size, secondary_batch_size): - self.primary_indices = primary_indices - self.secondary_indices = secondary_indices - self.secondary_batch_size = secondary_batch_size - self.primary_batch_size = batch_size - secondary_batch_size - - assert len(self.primary_indices) >= self.primary_batch_size > 0 - assert len(self.secondary_indices) >= self.secondary_batch_size > 0 - - def __iter__(self): - primary_iter = iterate_once(self.primary_indices) - secondary_iter = iterate_eternally(self.secondary_indices) - return ( - primary_batch + secondary_batch - for (primary_batch, secondary_batch) - in zip(grouper(primary_iter, self.primary_batch_size), - grouper(secondary_iter, self.secondary_batch_size)) - ) - - def __len__(self): - return len(self.primary_indices) // self.primary_batch_size - - -def iterate_once(iterable): - return np.random.permutation(iterable) - - -def iterate_eternally(indices): - def infinite_shuffles(): - while True: - yield np.random.permutation(indices) - return itertools.chain.from_iterable(infinite_shuffles()) - - -def grouper(iterable, n): - "Collect data into fixed-length chunks or blocks" - # grouper('ABCDEFG', 3) --> ABC DEF" - args = [iter(iterable)] * n - return zip(*args) \ No newline at end of file diff --git a/code/dataloaders/brats_proprecessing.py b/code/dataloaders/brats_proprecessing.py deleted file mode 100644 index a634cfc..0000000 --- a/code/dataloaders/brats_proprecessing.py +++ /dev/null @@ -1,110 +0,0 @@ -import numpy as np -from PIL import Image -import matplotlib.pyplot as plt -from skimage import measure -import nibabel as nib -import SimpleITK as sitk -import glob - - -def brain_bbox(data, gt): - mask = (data != 0) - brain_voxels = np.where(mask != 0) - minZidx = int(np.min(brain_voxels[0])) - maxZidx = int(np.max(brain_voxels[0])) - minXidx = int(np.min(brain_voxels[1])) - maxXidx = int(np.max(brain_voxels[1])) - minYidx = int(np.min(brain_voxels[2])) - maxYidx = int(np.max(brain_voxels[2])) - data_bboxed = data[minZidx:maxZidx, minXidx:maxXidx, minYidx:maxYidx] - gt_bboxed = gt[minZidx:maxZidx, minXidx:maxXidx, minYidx:maxYidx] - return data_bboxed, gt_bboxed - - -def volume_bounding_box(data, gt, expend=0, status="train"): - data, gt = brain_bbox(data, gt) - print(data.shape) - mask = (gt != 0) - brain_voxels = np.where(mask != 0) - z, x, y = data.shape - minZidx = int(np.min(brain_voxels[0])) - maxZidx = int(np.max(brain_voxels[0])) - minXidx = int(np.min(brain_voxels[1])) - maxXidx = int(np.max(brain_voxels[1])) - minYidx = int(np.min(brain_voxels[2])) - maxYidx = int(np.max(brain_voxels[2])) - - minZidx_jitterd = max(minZidx - expend, 0) - maxZidx_jitterd = min(maxZidx + expend, z) - minXidx_jitterd = max(minXidx - expend, 0) - maxXidx_jitterd = min(maxXidx + expend, x) - minYidx_jitterd = max(minYidx - expend, 0) - maxYidx_jitterd = min(maxYidx + expend, y) - - data_bboxed = data[minZidx_jitterd:maxZidx_jitterd, - minXidx_jitterd:maxXidx_jitterd, minYidx_jitterd:maxYidx_jitterd] - print([minZidx, maxZidx, minXidx, maxXidx, minYidx, maxYidx]) - print([minZidx_jitterd, maxZidx_jitterd, - minXidx_jitterd, maxXidx_jitterd, minYidx_jitterd, maxYidx_jitterd]) - - if status == "train": - gt_bboxed = np.zeros_like(data_bboxed, dtype=np.uint8) - gt_bboxed[expend:maxZidx_jitterd-expend, expend:maxXidx_jitterd - - expend, expend:maxYidx_jitterd - expend] = 1 - return data_bboxed, gt_bboxed - - if status == "test": - gt_bboxed = gt[minZidx_jitterd:maxZidx_jitterd, - minXidx_jitterd:maxXidx_jitterd, minYidx_jitterd:maxYidx_jitterd] - return data_bboxed, gt_bboxed - - -def itensity_normalize_one_volume(volume): - """ - normalize the itensity of an nd volume based on the mean and std of nonzeor region - inputs: - volume: the input nd volume - outputs: - out: the normalized nd volume - """ - - pixels = volume[volume > 0] - mean = pixels.mean() - std = pixels.std() - out = (volume - mean)/std - out_random = np.random.normal(0, 1, size=volume.shape) -# out[volume == 0] = out_random[volume == 0] - out = out.astype(np.float32) - return out - - -class MedicalImageDeal(object): - def __init__(self, img, percent=1): - self.img = img - self.percent = percent - - @property - def valid_img(self): - from skimage import exposure - cdf = exposure.cumulative_distribution(self.img) - watershed = cdf[1][cdf[0] >= self.percent][0] - return np.clip(self.img, self.img.min(), watershed) - - @property - def norm_img(self): - return (self.img - self.img.min()) / (self.img.max() - self.img.min()) - - -all_flair = glob.glob("flair/*_flair.nii.gz") -for p in all_flair: - data = sitk.GetArrayFromImage(sitk.ReadImage(p)) - lab = sitk.GetArrayFromImage(sitk.ReadImage(p.replace("flair", "seg"))) - img, lab = brain_bbox(data, lab) - img = MedicalImageDeal(img, percent=0.999).valid_img - img = itensity_normalize_one_volume(img) - lab[lab > 0] = 1 - uid = p.split("/")[-1] - sitk.WriteImage(sitk.GetImageFromArray( - img), "/media/xdluo/Data/brats19/data/flair/{}".format(uid)) - sitk.WriteImage(sitk.GetImageFromArray( - lab), "/media/xdluo/Data/brats19/data/label/{}".format(uid)) diff --git a/code/dataloaders/dataset.py b/code/dataloaders/dataset.py index 91dc573..1724e00 100644 --- a/code/dataloaders/dataset.py +++ b/code/dataloaders/dataset.py @@ -1,37 +1,66 @@ +import itertools import os -import cv2 -import torch import random -import numpy as np +import re from glob import glob -from torch.utils.data import Dataset + +import cv2 import h5py -from scipy.ndimage.interpolation import zoom -import itertools +import numpy as np +import torch from scipy import ndimage -from torch.utils.data.sampler import Sampler +from scipy.ndimage.interpolation import zoom +from torch.utils.data import Dataset +from sklearn.model_selection import KFold class BaseDataSets(Dataset): - def __init__(self, base_dir=None, split='train', num=None, transform=None): + def __init__(self, base_dir=None, labeled_type="labeled", labeled_ratio=10, split='train', transform=None, fold=1): self._base_dir = base_dir self.sample_list = [] self.split = split self.transform = transform + self.labeled_type = labeled_type + self.all_volumes = sorted(os.listdir(self._base_dir + "/all_volumes")) + train_ids, test_ids = self._get_fold_ids(fold) + all_labeled_ids = train_ids[::labeled_ratio] if self.split == 'train': - with open(self._base_dir + '/train_slices.list', 'r') as f1: - self.sample_list = f1.readlines() - self.sample_list = [item.replace('\n', '') - for item in self.sample_list] + self.all_slices = os.listdir(self._base_dir + "/all_slices") + self.sample_list = [] + labeled_ids = [i for i in all_labeled_ids if i in train_ids] + unlabeled_ids = [i for i in train_ids if i not in labeled_ids] + if self.labeled_type == "labeled": + print("Labeled patients IDs", labeled_ids) + for ids in labeled_ids: + new_data_list = list(filter(lambda x: re.match( + '{}.*'.format(ids.replace(".h5", "")), x) != None, self.all_slices)) + self.sample_list.extend(new_data_list) + print("total labeled {} samples".format(len(self.sample_list))) + else: + print("Unlabeled patients IDs", unlabeled_ids) + for ids in unlabeled_ids: + new_data_list = list(filter(lambda x: re.match( + '{}.*'.format(ids.replace(".h5", "")), x) != None, self.all_slices)) + self.sample_list.extend(new_data_list) + print("total unlabeled {} samples".format(len(self.sample_list))) elif self.split == 'val': - with open(self._base_dir + '/val.list', 'r') as f: - self.sample_list = f.readlines() - self.sample_list = [item.replace('\n', '') - for item in self.sample_list] - if num is not None and self.split == "train": - self.sample_list = self.sample_list[:num] - print("total {} samples".format(len(self.sample_list))) + print("test_ids", test_ids) + self.all_volumes = os.listdir( + self._base_dir + "/all_volumes") + self.sample_list = [] + for ids in test_ids: + new_data_list = list(filter(lambda x: re.match( + '{}.*'.format(ids.replace(".h5", "")), x) != None, self.all_volumes)) + self.sample_list.extend(new_data_list) + + def _get_fold_ids(self, fold): + folds = KFold(n_splits=5, shuffle=False) + all_cases = np.array(self.all_volumes) + k_fold_data = [] + for trn_idx, val_idx in folds.split(all_cases): + k_fold_data.append([all_cases[trn_idx], all_cases[val_idx]]) + return k_fold_data[fold][0], k_fold_data[fold][1] def __len__(self): return len(self.sample_list) @@ -40,15 +69,20 @@ def __getitem__(self, idx): case = self.sample_list[idx] if self.split == "train": h5f = h5py.File(self._base_dir + - "/data/slices/{}.h5".format(case), 'r') + "/all_slices/{}".format(case), 'r') else: - h5f = h5py.File(self._base_dir + "/data/{}.h5".format(case), 'r') - image = h5f['image'][:] - label = h5f['label'][:] - sample = {'image': image, 'label': label} + h5f = h5py.File(self._base_dir + + "/all_volumes/{}".format(case), 'r') if self.split == "train": + image = h5f['image'][:] + label = h5f["label"][:] + sample = {'image': image, 'label': label} sample = self.transform(sample) - sample["idx"] = idx + else: + image = h5f['image'][:] + label = h5f['label'][:].astype(np.int16) + sample = {'image': image, 'label': label} + sample["idx"] = case.split("_")[0] return sample @@ -62,10 +96,19 @@ def random_rot_flip(image, label): return image, label -def random_rotate(image, label): +def random_rotate(image, label, cval): angle = np.random.randint(-20, 20) image = ndimage.rotate(image, angle, order=0, reshape=False) - label = ndimage.rotate(label, angle, order=0, reshape=False) + label = ndimage.rotate(label, angle, order=0, + reshape=False, mode="constant", cval=cval) + return image, label + + +def random_noise(image, label, mu=0, sigma=0.1): + noise = np.clip(sigma * np.random.randn(image.shape[0], image.shape[1]), + -2 * sigma, 2 * sigma) + noise = noise + mu + image = image + noise return image, label @@ -75,13 +118,12 @@ def __init__(self, output_size): def __call__(self, sample): image, label = sample['image'], sample['label'] - # ind = random.randrange(0, img.shape[0]) - # image = img[ind, ...] - # label = lab[ind, ...] if random.random() > 0.5: image, label = random_rot_flip(image, label) - elif random.random() > 0.5: - image, label = random_rotate(image, label) + if random.random() > 0.5: + image, label = random_rotate(image, label, cval=0) + if random.random() > 0.5: + image, label = random_noise(image, label) x, y = image.shape image = zoom( image, (self.output_size[0] / x, self.output_size[1] / y), order=0) @@ -92,52 +134,3 @@ def __call__(self, sample): label = torch.from_numpy(label.astype(np.uint8)) sample = {'image': image, 'label': label} return sample - - -class TwoStreamBatchSampler(Sampler): - """Iterate two sets of indices - - An 'epoch' is one iteration through the primary indices. - During the epoch, the secondary indices are iterated through - as many times as needed. - """ - - def __init__(self, primary_indices, secondary_indices, batch_size, secondary_batch_size): - self.primary_indices = primary_indices - self.secondary_indices = secondary_indices - self.secondary_batch_size = secondary_batch_size - self.primary_batch_size = batch_size - secondary_batch_size - - assert len(self.primary_indices) >= self.primary_batch_size > 0 - assert len(self.secondary_indices) >= self.secondary_batch_size > 0 - - def __iter__(self): - primary_iter = iterate_once(self.primary_indices) - secondary_iter = iterate_eternally(self.secondary_indices) - return ( - primary_batch + secondary_batch - for (primary_batch, secondary_batch) - in zip(grouper(primary_iter, self.primary_batch_size), - grouper(secondary_iter, self.secondary_batch_size)) - ) - - def __len__(self): - return len(self.primary_indices) // self.primary_batch_size - - -def iterate_once(iterable): - return np.random.permutation(iterable) - - -def iterate_eternally(indices): - def infinite_shuffles(): - while True: - yield np.random.permutation(indices) - return itertools.chain.from_iterable(infinite_shuffles()) - - -def grouper(iterable, n): - "Collect data into fixed-length chunks or blocks" - # grouper('ABCDEFG', 3) --> ABC DEF" - args = [iter(iterable)] * n - return zip(*args) diff --git a/code/dataloaders/prostate_dataset_preprocessing.py b/code/dataloaders/prostate_dataset_preprocessing.py new file mode 100644 index 0000000..a1af761 --- /dev/null +++ b/code/dataloaders/prostate_dataset_preprocessing.py @@ -0,0 +1,100 @@ +# save images in slice level +import glob +import os + +import h5py +import numpy as np +import SimpleITK as sitk + + +class MedicalImageDeal(object): + def __init__(self, img, percent=1): + self.img = img + self.percent = percent + + @property + def valid_img(self): + from skimage import exposure + cdf = exposure.cumulative_distribution(self.img) + watershed = cdf[1][cdf[0] >= self.percent][0] + return np.clip(self.img, self.img.min(), watershed) + + @property + def norm_img(self): + return (self.img - self.img.min()) / (self.img.max() - self.img.min()) + +# slice_num = 0 +# mask_path = sorted( +# glob.glob("/home/SENSETIME/luoxiangde.vendor/Desktop/SSL4MIS_5Fold/data/prostate_zonal_nii/*_lab.nii.gz")) +# for image_path in mask_path: +# image_itk = sitk.ReadImage(image_path.replace("_lab", "")) +# image = sitk.GetArrayFromImage(image_itk) + +# image = MedicalImageDeal(image, percent=0.99).valid_img +# image = (image - image.min()) / (image.max() - image.min()) +# norm_img_itk = sitk.GetImageFromArray(image) +# norm_img_itk.CopyInformation(image_itk) +# sitk.WriteImage(norm_img_itk, image_path.replace("_lab", "")) + + +# saving images in slice level + +slice_num = 0 +mask_path = sorted( + glob.glob("/home/SENSETIME/luoxiangde.vendor/Desktop/SSL4MIS_5Fold/data/CHAOS_NII/label/*.nii.gz")) +for case in mask_path: + label_itk = sitk.ReadImage(case) + label = sitk.GetArrayFromImage(label_itk) + + image_path = case.replace("/label/", "/image/") + image_itk = sitk.ReadImage(image_path) + image = sitk.GetArrayFromImage(image_itk) + spacing = image_itk.GetSpacing() + + image = MedicalImageDeal(image, percent=0.99).valid_img + image = (image - image.min()) / (image.max() - image.min()) + print(image.shape) + image = image.astype(np.float32) + item = case.split("/")[-1].split(".")[0].replace("_gt", "") + if image.shape != label.shape: + print("Error") + print(item) + + f = h5py.File( + '/home/SENSETIME/luoxiangde.vendor/Desktop/SSL4MIS_5Fold/data/CHAOS/all_volumes/{}.h5'.format(item), 'w') + f.create_dataset( + 'image', data=image, compression="gzip") + f.create_dataset('label', data=label, compression="gzip") + f.create_dataset('spacing', data=np.array(spacing), compression="gzip") + f.close() +print("Converted all ACDC volumes to 2D slices") +print("Total {} slices".format(slice_num)) +# # saving images in volume level + +# slice_num = 0 +# mask_path = sorted( +# glob.glob("/home/SENSETIME/luoxiangde.vendor/Desktop/SSL4MIS_5Fold/data/prostate_zonal_nii/*_lab.nii.gz")) +# for case in mask_path: +# label_itk = sitk.ReadImage(case) +# label = sitk.GetArrayFromImage(label_itk) + +# image_path = case.replace("_lab", "") +# image_itk = sitk.ReadImage(image_path) +# image = sitk.GetArrayFromImage(image_itk) +# spacing = image_itk.GetSpacing() + +# image = image.astype(np.float32) +# item = case.split("/")[-1].split(".")[0].replace("_lab", "") +# if image.shape != label.shape: +# print("Error") +# print(item) + +# f = h5py.File( +# '/home/SENSETIME/luoxiangde.vendor/Desktop/SSL4MIS_5Fold/data/ProstateX/all_volumes/{}.h5'.format(item), 'w') +# f.create_dataset( +# 'image', data=image, compression="gzip") +# f.create_dataset('label', data=label, compression="gzip") +# f.create_dataset('spacing', data=np.array(spacing), compression="gzip") +# f.close() +# print("Converted all Prostate volumes to 2D slices") +# print("Total {} slices".format(slice_num)) \ No newline at end of file diff --git a/code/dataloaders/utils.py b/code/dataloaders/utils.py deleted file mode 100644 index 9117078..0000000 --- a/code/dataloaders/utils.py +++ /dev/null @@ -1,214 +0,0 @@ -import os -import torch -import numpy as np -import torch.nn as nn -# import matplotlib.pyplot as plt -from skimage import measure -import scipy.ndimage as nd - - -def recursive_glob(rootdir='.', suffix=''): - """Performs recursive glob with given suffix and rootdir - :param rootdir is the root directory - :param suffix is the suffix to be searched - """ - return [os.path.join(looproot, filename) - for looproot, _, filenames in os.walk(rootdir) - for filename in filenames if filename.endswith(suffix)] - -def get_cityscapes_labels(): - return np.array([ - # [ 0, 0, 0], - [128, 64, 128], - [244, 35, 232], - [70, 70, 70], - [102, 102, 156], - [190, 153, 153], - [153, 153, 153], - [250, 170, 30], - [220, 220, 0], - [107, 142, 35], - [152, 251, 152], - [0, 130, 180], - [220, 20, 60], - [255, 0, 0], - [0, 0, 142], - [0, 0, 70], - [0, 60, 100], - [0, 80, 100], - [0, 0, 230], - [119, 11, 32]]) - -def get_pascal_labels(): - """Load the mapping that associates pascal classes with label colors - Returns: - np.ndarray with dimensions (21, 3) - """ - return np.asarray([[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0], - [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128], - [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0], - [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128], - [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0], - [0, 64, 128]]) - - -def encode_segmap(mask): - """Encode segmentation label images as pascal classes - Args: - mask (np.ndarray): raw segmentation label image of dimension - (M, N, 3), in which the Pascal classes are encoded as colours. - Returns: - (np.ndarray): class map with dimensions (M,N), where the value at - a given location is the integer denoting the class index. - """ - mask = mask.astype(int) - label_mask = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.int16) - for ii, label in enumerate(get_pascal_labels()): - label_mask[np.where(np.all(mask == label, axis=-1))[:2]] = ii - label_mask = label_mask.astype(int) - return label_mask - - -def decode_seg_map_sequence(label_masks, dataset='pascal'): - rgb_masks = [] - for label_mask in label_masks: - rgb_mask = decode_segmap(label_mask, dataset) - rgb_masks.append(rgb_mask) - rgb_masks = torch.from_numpy(np.array(rgb_masks).transpose([0, 3, 1, 2])) - return rgb_masks - -def decode_segmap(label_mask, dataset, plot=False): - """Decode segmentation class labels into a color image - Args: - label_mask (np.ndarray): an (M,N) array of integer values denoting - the class label at each spatial location. - plot (bool, optional): whether to show the resulting color image - in a figure. - Returns: - (np.ndarray, optional): the resulting decoded color image. - """ - if dataset == 'pascal': - n_classes = 21 - label_colours = get_pascal_labels() - elif dataset == 'cityscapes': - n_classes = 19 - label_colours = get_cityscapes_labels() - else: - raise NotImplementedError - - r = label_mask.copy() - g = label_mask.copy() - b = label_mask.copy() - for ll in range(0, n_classes): - r[label_mask == ll] = label_colours[ll, 0] - g[label_mask == ll] = label_colours[ll, 1] - b[label_mask == ll] = label_colours[ll, 2] - rgb = np.zeros((label_mask.shape[0], label_mask.shape[1], 3)) - rgb[:, :, 0] = r / 255.0 - rgb[:, :, 1] = g / 255.0 - rgb[:, :, 2] = b / 255.0 - if plot: - plt.imshow(rgb) - plt.show() - else: - return rgb - -def generate_param_report(logfile, param): - log_file = open(logfile, 'w') - # for key, val in param.items(): - # log_file.write(key + ':' + str(val) + '\n') - log_file.write(str(param)) - log_file.close() - -def cross_entropy2d(logit, target, ignore_index=255, weight=None, size_average=True, batch_average=True): - n, c, h, w = logit.size() - # logit = logit.permute(0, 2, 3, 1) - target = target.squeeze(1) - if weight is None: - criterion = nn.CrossEntropyLoss(weight=weight, ignore_index=ignore_index, size_average=False) - else: - criterion = nn.CrossEntropyLoss(weight=torch.from_numpy(np.array(weight)).float().cuda(), ignore_index=ignore_index, size_average=False) - loss = criterion(logit, target.long()) - - if size_average: - loss /= (h * w) - - if batch_average: - loss /= n - - return loss - -def lr_poly(base_lr, iter_, max_iter=100, power=0.9): - return base_lr * ((1 - float(iter_) / max_iter) ** power) - - -def get_iou(pred, gt, n_classes=21): - total_iou = 0.0 - for i in range(len(pred)): - pred_tmp = pred[i] - gt_tmp = gt[i] - - intersect = [0] * n_classes - union = [0] * n_classes - for j in range(n_classes): - match = (pred_tmp == j) + (gt_tmp == j) - - it = torch.sum(match == 2).item() - un = torch.sum(match > 0).item() - - intersect[j] += it - union[j] += un - - iou = [] - for k in range(n_classes): - if union[k] == 0: - continue - iou.append(intersect[k] / union[k]) - - img_iou = (sum(iou) / len(iou)) - total_iou += img_iou - - return total_iou - -def get_dice(pred, gt): - total_dice = 0.0 - pred = pred.long() - gt = gt.long() - for i in range(len(pred)): - pred_tmp = pred[i] - gt_tmp = gt[i] - dice = 2.0*torch.sum(pred_tmp*gt_tmp).item()/(1.0+torch.sum(pred_tmp**2)+torch.sum(gt_tmp**2)).item() - print(dice) - total_dice += dice - - return total_dice - -def get_mc_dice(pred, gt, num=2): - # num is the total number of classes, include the background - total_dice = np.zeros(num-1) - pred = pred.long() - gt = gt.long() - for i in range(len(pred)): - for j in range(1, num): - pred_tmp = (pred[i]==j) - gt_tmp = (gt[i]==j) - dice = 2.0*torch.sum(pred_tmp*gt_tmp).item()/(1.0+torch.sum(pred_tmp**2)+torch.sum(gt_tmp**2)).item() - total_dice[j-1] +=dice - return total_dice - -def post_processing(prediction): - prediction = nd.binary_fill_holes(prediction) - label_cc, num_cc = measure.label(prediction,return_num=True) - total_cc = np.sum(prediction) - measure.regionprops(label_cc) - for cc in range(1,num_cc+1): - single_cc = (label_cc==cc) - single_vol = np.sum(single_cc) - if single_vol/total_cc<0.2: - prediction[single_cc]=0 - - return prediction - - - - diff --git a/code/networks/VoxResNet.py b/code/networks/VoxResNet.py old mode 100644 new mode 100755 diff --git a/code/networks/attention_unet.py b/code/networks/attention_unet.py old mode 100644 new mode 100755 diff --git a/code/networks/config.py b/code/networks/config.py deleted file mode 100644 index 35bf199..0000000 --- a/code/networks/config.py +++ /dev/null @@ -1,229 +0,0 @@ -# -------------------------------------------------------- -# Swin Transformer -# Copyright (c) 2021 Microsoft -# Licensed under The MIT License [see LICENSE for details] -# Written by Ze Liu -# --------------------------------------------------------' - -import os -import yaml -from yacs.config import CfgNode as CN - -_C = CN() - -# Base config files -_C.BASE = [''] - -# ----------------------------------------------------------------------------- -# Data settings -# ----------------------------------------------------------------------------- -_C.DATA = CN() -# Batch size for a single GPU, could be overwritten by command line argument -_C.DATA.BATCH_SIZE = 128 -# Path to dataset, could be overwritten by command line argument -_C.DATA.DATA_PATH = '' -# Dataset name -_C.DATA.DATASET = 'imagenet' -# Input image size -_C.DATA.IMG_SIZE = 224 -# Interpolation to resize image (random, bilinear, bicubic) -_C.DATA.INTERPOLATION = 'bicubic' -# Use zipped dataset instead of folder dataset -# could be overwritten by command line argument -_C.DATA.ZIP_MODE = False -# Cache Data in Memory, could be overwritten by command line argument -_C.DATA.CACHE_MODE = 'part' -# Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU. -_C.DATA.PIN_MEMORY = True -# Number of data loading threads -_C.DATA.NUM_WORKERS = 8 - -# ----------------------------------------------------------------------------- -# Model settings -# ----------------------------------------------------------------------------- -_C.MODEL = CN() -# Model type -_C.MODEL.TYPE = 'swin' -# Model name -_C.MODEL.NAME = 'swin_tiny_patch4_window7_224' -# Checkpoint to resume, could be overwritten by command line argument -_C.MODEL.PRETRAIN_CKPT = './pretrained_ckpt/swin_tiny_patch4_window7_224.pth' -_C.MODEL.RESUME = '' -# Number of classes, overwritten in data preparation -_C.MODEL.NUM_CLASSES = 1000 -# Dropout rate -_C.MODEL.DROP_RATE = 0.0 -# Drop path rate -_C.MODEL.DROP_PATH_RATE = 0.1 -# Label Smoothing -_C.MODEL.LABEL_SMOOTHING = 0.1 - -# Swin Transformer parameters -_C.MODEL.SWIN = CN() -_C.MODEL.SWIN.PATCH_SIZE = 4 -_C.MODEL.SWIN.IN_CHANS = 3 -_C.MODEL.SWIN.EMBED_DIM = 96 -_C.MODEL.SWIN.DEPTHS = [2, 2, 6, 2] -_C.MODEL.SWIN.DECODER_DEPTHS = [2, 2, 6, 2] -_C.MODEL.SWIN.NUM_HEADS = [3, 6, 12, 24] -_C.MODEL.SWIN.WINDOW_SIZE = 7 -_C.MODEL.SWIN.MLP_RATIO = 4. -_C.MODEL.SWIN.QKV_BIAS = True -_C.MODEL.SWIN.QK_SCALE = None -_C.MODEL.SWIN.APE = False -_C.MODEL.SWIN.PATCH_NORM = True -_C.MODEL.SWIN.FINAL_UPSAMPLE= "expand_first" - -# ----------------------------------------------------------------------------- -# Training settings -# ----------------------------------------------------------------------------- -_C.TRAIN = CN() -_C.TRAIN.START_EPOCH = 0 -_C.TRAIN.EPOCHS = 300 -_C.TRAIN.WARMUP_EPOCHS = 20 -_C.TRAIN.WEIGHT_DECAY = 0.05 -_C.TRAIN.BASE_LR = 5e-4 -_C.TRAIN.WARMUP_LR = 5e-7 -_C.TRAIN.MIN_LR = 5e-6 -# Clip gradient norm -_C.TRAIN.CLIP_GRAD = 5.0 -# Auto resume from latest checkpoint -_C.TRAIN.AUTO_RESUME = True -# Gradient accumulation steps -# could be overwritten by command line argument -_C.TRAIN.ACCUMULATION_STEPS = 0 -# Whether to use gradient checkpointing to save memory -# could be overwritten by command line argument -_C.TRAIN.USE_CHECKPOINT = False - -# LR scheduler -_C.TRAIN.LR_SCHEDULER = CN() -_C.TRAIN.LR_SCHEDULER.NAME = 'cosine' -# Epoch interval to decay LR, used in StepLRScheduler -_C.TRAIN.LR_SCHEDULER.DECAY_EPOCHS = 30 -# LR decay rate, used in StepLRScheduler -_C.TRAIN.LR_SCHEDULER.DECAY_RATE = 0.1 - -# Optimizer -_C.TRAIN.OPTIMIZER = CN() -_C.TRAIN.OPTIMIZER.NAME = 'adamw' -# Optimizer Epsilon -_C.TRAIN.OPTIMIZER.EPS = 1e-8 -# Optimizer Betas -_C.TRAIN.OPTIMIZER.BETAS = (0.9, 0.999) -# SGD momentum -_C.TRAIN.OPTIMIZER.MOMENTUM = 0.9 - -# ----------------------------------------------------------------------------- -# Augmentation settings -# ----------------------------------------------------------------------------- -_C.AUG = CN() -# Color jitter factor -_C.AUG.COLOR_JITTER = 0.4 -# Use AutoAugment policy. "v0" or "original" -_C.AUG.AUTO_AUGMENT = 'rand-m9-mstd0.5-inc1' -# Random erase prob -_C.AUG.REPROB = 0.25 -# Random erase mode -_C.AUG.REMODE = 'pixel' -# Random erase count -_C.AUG.RECOUNT = 1 -# Mixup alpha, mixup enabled if > 0 -_C.AUG.MIXUP = 0.8 -# Cutmix alpha, cutmix enabled if > 0 -_C.AUG.CUTMIX = 1.0 -# Cutmix min/max ratio, overrides alpha and enables cutmix if set -_C.AUG.CUTMIX_MINMAX = None -# Probability of performing mixup or cutmix when either/both is enabled -_C.AUG.MIXUP_PROB = 1.0 -# Probability of switching to cutmix when both mixup and cutmix enabled -_C.AUG.MIXUP_SWITCH_PROB = 0.5 -# How to apply mixup/cutmix params. Per "batch", "pair", or "elem" -_C.AUG.MIXUP_MODE = 'batch' - -# ----------------------------------------------------------------------------- -# Testing settings -# ----------------------------------------------------------------------------- -_C.TEST = CN() -# Whether to use center crop when testing -_C.TEST.CROP = True - -# ----------------------------------------------------------------------------- -# Misc -# ----------------------------------------------------------------------------- -# Mixed precision opt level, if O0, no amp is used ('O0', 'O1', 'O2') -# overwritten by command line argument -_C.AMP_OPT_LEVEL = '' -# Path to output folder, overwritten by command line argument -_C.OUTPUT = '' -# Tag of experiment, overwritten by command line argument -_C.TAG = 'default' -# Frequency to save checkpoint -_C.SAVE_FREQ = 1 -# Frequency to logging info -_C.PRINT_FREQ = 10 -# Fixed random seed -_C.SEED = 0 -# Perform evaluation only, overwritten by command line argument -_C.EVAL_MODE = False -# Test throughput only, overwritten by command line argument -_C.THROUGHPUT_MODE = False -# local rank for DistributedDataParallel, given by command line argument -_C.LOCAL_RANK = 0 - - -def _update_config_from_file(config, cfg_file): - config.defrost() - with open(cfg_file, 'r') as f: - yaml_cfg = yaml.load(f, Loader=yaml.FullLoader) - - for cfg in yaml_cfg.setdefault('BASE', ['']): - if cfg: - _update_config_from_file( - config, os.path.join(os.path.dirname(cfg_file), cfg) - ) - print('=> merge config from {}'.format(cfg_file)) - config.merge_from_file(cfg_file) - config.freeze() - - -def update_config(config, args): - _update_config_from_file(config, args.cfg) - - config.defrost() - if args.opts: - config.merge_from_list(args.opts) - - # merge from specific arguments - if args.batch_size: - config.DATA.BATCH_SIZE = args.batch_size - if args.zip: - config.DATA.ZIP_MODE = True - if args.cache_mode: - config.DATA.CACHE_MODE = args.cache_mode - if args.resume: - config.MODEL.RESUME = args.resume - if args.accumulation_steps: - config.TRAIN.ACCUMULATION_STEPS = args.accumulation_steps - if args.use_checkpoint: - config.TRAIN.USE_CHECKPOINT = True - if args.amp_opt_level: - config.AMP_OPT_LEVEL = args.amp_opt_level - if args.tag: - config.TAG = args.tag - if args.eval: - config.EVAL_MODE = True - if args.throughput: - config.THROUGHPUT_MODE = True - - config.freeze() - - -def get_config(args): - """Get a yacs CfgNode object with default values.""" - # Return a clone so that the defaults will not be altered - # This is for the "local variable" use pattern - config = _C.clone() - update_config(config, args) - - return config diff --git a/code/networks/discriminator.py b/code/networks/discriminator.py old mode 100644 new mode 100755 diff --git a/code/networks/enet.py b/code/networks/enet.py deleted file mode 100644 index f47016c..0000000 --- a/code/networks/enet.py +++ /dev/null @@ -1,614 +0,0 @@ -import torch.nn as nn -import torch - - -class InitialBlock(nn.Module): - """The initial block is composed of two branches: - 1. a main branch which performs a regular convolution with stride 2; - 2. an extension branch which performs max-pooling. - Doing both operations in parallel and concatenating their results - allows for efficient downsampling and expansion. The main branch - outputs 13 feature maps while the extension branch outputs 3, for a - total of 16 feature maps after concatenation. - Keyword arguments: - - in_channels (int): the number of input channels. - - out_channels (int): the number output channels. - - kernel_size (int, optional): the kernel size of the filters used in - the convolution layer. Default: 3. - - padding (int, optional): zero-padding added to both sides of the - input. Default: 0. - - bias (bool, optional): Adds a learnable bias to the output if - ``True``. Default: False. - - relu (bool, optional): When ``True`` ReLU is used as the activation - function; otherwise, PReLU is used. Default: True. - """ - - def __init__(self, - in_channels, - out_channels, - bias=False, - relu=True): - super().__init__() - - if relu: - activation = nn.ReLU - else: - activation = nn.PReLU - - # Main branch - As stated above the number of output channels for this - # branch is the total minus 3, since the remaining channels come from - # the extension branch - self.main_branch = nn.Conv2d( - in_channels, - out_channels - in_channels, - kernel_size=3, - stride=2, - padding=1, - bias=bias) - - # Extension branch - self.ext_branch = nn.MaxPool2d(3, stride=2, padding=1) - - # Initialize batch normalization to be used after concatenation - self.batch_norm = nn.BatchNorm2d(out_channels) - - # PReLU layer to apply after concatenating the branches - self.out_activation = activation() - - def forward(self, x): - main = self.main_branch(x) - ext = self.ext_branch(x) - - # Concatenate branches - out = torch.cat((main, ext), 1) - - # Apply batch normalization - out = self.batch_norm(out) - - return self.out_activation(out) - - -class RegularBottleneck(nn.Module): - """Regular bottlenecks are the main building block of ENet. - Main branch: - 1. Shortcut connection. - Extension branch: - 1. 1x1 convolution which decreases the number of channels by - ``internal_ratio``, also called a projection; - 2. regular, dilated or asymmetric convolution; - 3. 1x1 convolution which increases the number of channels back to - ``channels``, also called an expansion; - 4. dropout as a regularizer. - Keyword arguments: - - channels (int): the number of input and output channels. - - internal_ratio (int, optional): a scale factor applied to - ``channels`` used to compute the number of - channels after the projection. eg. given ``channels`` equal to 128 and - internal_ratio equal to 2 the number of channels after the projection - is 64. Default: 4. - - kernel_size (int, optional): the kernel size of the filters used in - the convolution layer described above in item 2 of the extension - branch. Default: 3. - - padding (int, optional): zero-padding added to both sides of the - input. Default: 0. - - dilation (int, optional): spacing between kernel elements for the - convolution described in item 2 of the extension branch. Default: 1. - asymmetric (bool, optional): flags if the convolution described in - item 2 of the extension branch is asymmetric or not. Default: False. - - dropout_prob (float, optional): probability of an element to be - zeroed. Default: 0 (no dropout). - - bias (bool, optional): Adds a learnable bias to the output if - ``True``. Default: False. - - relu (bool, optional): When ``True`` ReLU is used as the activation - function; otherwise, PReLU is used. Default: True. - """ - - def __init__(self, - channels, - internal_ratio=4, - kernel_size=3, - padding=0, - dilation=1, - asymmetric=False, - dropout_prob=0, - bias=False, - relu=True): - super().__init__() - - # Check in the internal_scale parameter is within the expected range - # [1, channels] - if internal_ratio <= 1 or internal_ratio > channels: - raise RuntimeError("Value out of range. Expected value in the " - "interval [1, {0}], got internal_scale={1}." - .format(channels, internal_ratio)) - - internal_channels = channels // internal_ratio - - if relu: - activation = nn.ReLU - else: - activation = nn.PReLU - - # Main branch - shortcut connection - - # Extension branch - 1x1 convolution, followed by a regular, dilated or - # asymmetric convolution, followed by another 1x1 convolution, and, - # finally, a regularizer (spatial dropout). Number of channels is constant. - - # 1x1 projection convolution - self.ext_conv1 = nn.Sequential( - nn.Conv2d( - channels, - internal_channels, - kernel_size=1, - stride=1, - bias=bias), nn.BatchNorm2d(internal_channels), activation()) - - # If the convolution is asymmetric we split the main convolution in - # two. Eg. for a 5x5 asymmetric convolution we have two convolution: - # the first is 5x1 and the second is 1x5. - if asymmetric: - self.ext_conv2 = nn.Sequential( - nn.Conv2d( - internal_channels, - internal_channels, - kernel_size=(kernel_size, 1), - stride=1, - padding=(padding, 0), - dilation=dilation, - bias=bias), nn.BatchNorm2d(internal_channels), activation(), - nn.Conv2d( - internal_channels, - internal_channels, - kernel_size=(1, kernel_size), - stride=1, - padding=(0, padding), - dilation=dilation, - bias=bias), nn.BatchNorm2d(internal_channels), activation()) - else: - self.ext_conv2 = nn.Sequential( - nn.Conv2d( - internal_channels, - internal_channels, - kernel_size=kernel_size, - stride=1, - padding=padding, - dilation=dilation, - bias=bias), nn.BatchNorm2d(internal_channels), activation()) - - # 1x1 expansion convolution - self.ext_conv3 = nn.Sequential( - nn.Conv2d( - internal_channels, - channels, - kernel_size=1, - stride=1, - bias=bias), nn.BatchNorm2d(channels), activation()) - - self.ext_regul = nn.Dropout2d(p=dropout_prob) - - # PReLU layer to apply after adding the branches - self.out_activation = activation() - - def forward(self, x): - # Main branch shortcut - main = x - - # Extension branch - ext = self.ext_conv1(x) - ext = self.ext_conv2(ext) - ext = self.ext_conv3(ext) - ext = self.ext_regul(ext) - - # Add main and extension branches - out = main + ext - - return self.out_activation(out) - - -class DownsamplingBottleneck(nn.Module): - """Downsampling bottlenecks further downsample the feature map size. - Main branch: - 1. max pooling with stride 2; indices are saved to be used for - unpooling later. - Extension branch: - 1. 2x2 convolution with stride 2 that decreases the number of channels - by ``internal_ratio``, also called a projection; - 2. regular convolution (by default, 3x3); - 3. 1x1 convolution which increases the number of channels to - ``out_channels``, also called an expansion; - 4. dropout as a regularizer. - Keyword arguments: - - in_channels (int): the number of input channels. - - out_channels (int): the number of output channels. - - internal_ratio (int, optional): a scale factor applied to ``channels`` - used to compute the number of channels after the projection. eg. given - ``channels`` equal to 128 and internal_ratio equal to 2 the number of - channels after the projection is 64. Default: 4. - - return_indices (bool, optional): if ``True``, will return the max - indices along with the outputs. Useful when unpooling later. - - dropout_prob (float, optional): probability of an element to be - zeroed. Default: 0 (no dropout). - - bias (bool, optional): Adds a learnable bias to the output if - ``True``. Default: False. - - relu (bool, optional): When ``True`` ReLU is used as the activation - function; otherwise, PReLU is used. Default: True. - """ - - def __init__(self, - in_channels, - out_channels, - internal_ratio=4, - return_indices=False, - dropout_prob=0, - bias=False, - relu=True): - super().__init__() - - # Store parameters that are needed later - self.return_indices = return_indices - - # Check in the internal_scale parameter is within the expected range - # [1, channels] - if internal_ratio <= 1 or internal_ratio > in_channels: - raise RuntimeError("Value out of range. Expected value in the " - "interval [1, {0}], got internal_scale={1}. " - .format(in_channels, internal_ratio)) - - internal_channels = in_channels // internal_ratio - - if relu: - activation = nn.ReLU - else: - activation = nn.PReLU - - # Main branch - max pooling followed by feature map (channels) padding - self.main_max1 = nn.MaxPool2d( - 2, - stride=2, - return_indices=return_indices) - - # Extension branch - 2x2 convolution, followed by a regular, dilated or - # asymmetric convolution, followed by another 1x1 convolution. Number - # of channels is doubled. - - # 2x2 projection convolution with stride 2 - self.ext_conv1 = nn.Sequential( - nn.Conv2d( - in_channels, - internal_channels, - kernel_size=2, - stride=2, - bias=bias), nn.BatchNorm2d(internal_channels), activation()) - - # Convolution - self.ext_conv2 = nn.Sequential( - nn.Conv2d( - internal_channels, - internal_channels, - kernel_size=3, - stride=1, - padding=1, - bias=bias), nn.BatchNorm2d(internal_channels), activation()) - - # 1x1 expansion convolution - self.ext_conv3 = nn.Sequential( - nn.Conv2d( - internal_channels, - out_channels, - kernel_size=1, - stride=1, - bias=bias), nn.BatchNorm2d(out_channels), activation()) - - self.ext_regul = nn.Dropout2d(p=dropout_prob) - - # PReLU layer to apply after concatenating the branches - self.out_activation = activation() - - def forward(self, x): - # Main branch shortcut - if self.return_indices: - main, max_indices = self.main_max1(x) - else: - main = self.main_max1(x) - - # Extension branch - ext = self.ext_conv1(x) - ext = self.ext_conv2(ext) - ext = self.ext_conv3(ext) - ext = self.ext_regul(ext) - - # Main branch channel padding - n, ch_ext, h, w = ext.size() - ch_main = main.size()[1] - padding = torch.zeros(n, ch_ext - ch_main, h, w) - - # Before concatenating, check if main is on the CPU or GPU and - # convert padding accordingly - if main.is_cuda: - padding = padding.cuda() - - # Concatenate - main = torch.cat((main, padding), 1) - - # Add main and extension branches - out = main + ext - - return self.out_activation(out), max_indices - - -class UpsamplingBottleneck(nn.Module): - """The upsampling bottlenecks upsample the feature map resolution using max - pooling indices stored from the corresponding downsampling bottleneck. - Main branch: - 1. 1x1 convolution with stride 1 that decreases the number of channels by - ``internal_ratio``, also called a projection; - 2. max unpool layer using the max pool indices from the corresponding - downsampling max pool layer. - Extension branch: - 1. 1x1 convolution with stride 1 that decreases the number of channels by - ``internal_ratio``, also called a projection; - 2. transposed convolution (by default, 3x3); - 3. 1x1 convolution which increases the number of channels to - ``out_channels``, also called an expansion; - 4. dropout as a regularizer. - Keyword arguments: - - in_channels (int): the number of input channels. - - out_channels (int): the number of output channels. - - internal_ratio (int, optional): a scale factor applied to ``in_channels`` - used to compute the number of channels after the projection. eg. given - ``in_channels`` equal to 128 and ``internal_ratio`` equal to 2 the number - of channels after the projection is 64. Default: 4. - - dropout_prob (float, optional): probability of an element to be zeroed. - Default: 0 (no dropout). - - bias (bool, optional): Adds a learnable bias to the output if ``True``. - Default: False. - - relu (bool, optional): When ``True`` ReLU is used as the activation - function; otherwise, PReLU is used. Default: True. - """ - - def __init__(self, - in_channels, - out_channels, - internal_ratio=4, - dropout_prob=0, - bias=False, - relu=True): - super().__init__() - - # Check in the internal_scale parameter is within the expected range - # [1, channels] - if internal_ratio <= 1 or internal_ratio > in_channels: - raise RuntimeError("Value out of range. Expected value in the " - "interval [1, {0}], got internal_scale={1}. " - .format(in_channels, internal_ratio)) - - internal_channels = in_channels // internal_ratio - - if relu: - activation = nn.ReLU - else: - activation = nn.PReLU - - # Main branch - max pooling followed by feature map (channels) padding - self.main_conv1 = nn.Sequential( - nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias), - nn.BatchNorm2d(out_channels)) - - # Remember that the stride is the same as the kernel_size, just like - # the max pooling layers - self.main_unpool1 = nn.MaxUnpool2d(kernel_size=2) - - # Extension branch - 1x1 convolution, followed by a regular, dilated or - # asymmetric convolution, followed by another 1x1 convolution. Number - # of channels is doubled. - - # 1x1 projection convolution with stride 1 - self.ext_conv1 = nn.Sequential( - nn.Conv2d( - in_channels, internal_channels, kernel_size=1, bias=bias), - nn.BatchNorm2d(internal_channels), activation()) - - # Transposed convolution - self.ext_tconv1 = nn.ConvTranspose2d( - internal_channels, - internal_channels, - kernel_size=2, - stride=2, - bias=bias) - self.ext_tconv1_bnorm = nn.BatchNorm2d(internal_channels) - self.ext_tconv1_activation = activation() - - # 1x1 expansion convolution - self.ext_conv2 = nn.Sequential( - nn.Conv2d( - internal_channels, out_channels, kernel_size=1, bias=bias), - nn.BatchNorm2d(out_channels), activation()) - - self.ext_regul = nn.Dropout2d(p=dropout_prob) - - # PReLU layer to apply after concatenating the branches - self.out_activation = activation() - - def forward(self, x, max_indices, output_size): - # Main branch shortcut - main = self.main_conv1(x) - main = self.main_unpool1( - main, max_indices, output_size=output_size) - - # Extension branch - ext = self.ext_conv1(x) - ext = self.ext_tconv1(ext, output_size=output_size) - ext = self.ext_tconv1_bnorm(ext) - ext = self.ext_tconv1_activation(ext) - ext = self.ext_conv2(ext) - ext = self.ext_regul(ext) - - # Add main and extension branches - out = main + ext - - return self.out_activation(out) - - -class ENet(nn.Module): - """Generate the ENet model. - Keyword arguments: - - num_classes (int): the number of classes to segment. - - encoder_relu (bool, optional): When ``True`` ReLU is used as the - activation function in the encoder blocks/layers; otherwise, PReLU - is used. Default: False. - - decoder_relu (bool, optional): When ``True`` ReLU is used as the - activation function in the decoder blocks/layers; otherwise, PReLU - is used. Default: True. - """ - - def __init__(self, in_channels, num_classes, encoder_relu=False, decoder_relu=True): - super().__init__() - - self.initial_block = InitialBlock(in_channels, 16, relu=encoder_relu) - - # Stage 1 - Encoder - self.downsample1_0 = DownsamplingBottleneck( - 16, - 64, - return_indices=True, - dropout_prob=0.01, - relu=encoder_relu) - self.regular1_1 = RegularBottleneck( - 64, padding=1, dropout_prob=0.01, relu=encoder_relu) - self.regular1_2 = RegularBottleneck( - 64, padding=1, dropout_prob=0.01, relu=encoder_relu) - self.regular1_3 = RegularBottleneck( - 64, padding=1, dropout_prob=0.01, relu=encoder_relu) - self.regular1_4 = RegularBottleneck( - 64, padding=1, dropout_prob=0.01, relu=encoder_relu) - - # Stage 2 - Encoder - self.downsample2_0 = DownsamplingBottleneck( - 64, - 128, - return_indices=True, - dropout_prob=0.1, - relu=encoder_relu) - self.regular2_1 = RegularBottleneck( - 128, padding=1, dropout_prob=0.1, relu=encoder_relu) - self.dilated2_2 = RegularBottleneck( - 128, dilation=2, padding=2, dropout_prob=0.1, relu=encoder_relu) - self.asymmetric2_3 = RegularBottleneck( - 128, - kernel_size=5, - padding=2, - asymmetric=True, - dropout_prob=0.1, - relu=encoder_relu) - self.dilated2_4 = RegularBottleneck( - 128, dilation=4, padding=4, dropout_prob=0.1, relu=encoder_relu) - self.regular2_5 = RegularBottleneck( - 128, padding=1, dropout_prob=0.1, relu=encoder_relu) - self.dilated2_6 = RegularBottleneck( - 128, dilation=8, padding=8, dropout_prob=0.1, relu=encoder_relu) - self.asymmetric2_7 = RegularBottleneck( - 128, - kernel_size=5, - asymmetric=True, - padding=2, - dropout_prob=0.1, - relu=encoder_relu) - self.dilated2_8 = RegularBottleneck( - 128, dilation=16, padding=16, dropout_prob=0.1, relu=encoder_relu) - - # Stage 3 - Encoder - self.regular3_0 = RegularBottleneck( - 128, padding=1, dropout_prob=0.1, relu=encoder_relu) - self.dilated3_1 = RegularBottleneck( - 128, dilation=2, padding=2, dropout_prob=0.1, relu=encoder_relu) - self.asymmetric3_2 = RegularBottleneck( - 128, - kernel_size=5, - padding=2, - asymmetric=True, - dropout_prob=0.1, - relu=encoder_relu) - self.dilated3_3 = RegularBottleneck( - 128, dilation=4, padding=4, dropout_prob=0.1, relu=encoder_relu) - self.regular3_4 = RegularBottleneck( - 128, padding=1, dropout_prob=0.1, relu=encoder_relu) - self.dilated3_5 = RegularBottleneck( - 128, dilation=8, padding=8, dropout_prob=0.1, relu=encoder_relu) - self.asymmetric3_6 = RegularBottleneck( - 128, - kernel_size=5, - asymmetric=True, - padding=2, - dropout_prob=0.1, - relu=encoder_relu) - self.dilated3_7 = RegularBottleneck( - 128, dilation=16, padding=16, dropout_prob=0.1, relu=encoder_relu) - - # Stage 4 - Decoder - self.upsample4_0 = UpsamplingBottleneck( - 128, 64, dropout_prob=0.1, relu=decoder_relu) - self.regular4_1 = RegularBottleneck( - 64, padding=1, dropout_prob=0.1, relu=decoder_relu) - self.regular4_2 = RegularBottleneck( - 64, padding=1, dropout_prob=0.1, relu=decoder_relu) - - # Stage 5 - Decoder - self.upsample5_0 = UpsamplingBottleneck( - 64, 16, dropout_prob=0.1, relu=decoder_relu) - self.regular5_1 = RegularBottleneck( - 16, padding=1, dropout_prob=0.1, relu=decoder_relu) - self.transposed_conv = nn.ConvTranspose2d( - 16, - num_classes, - kernel_size=3, - stride=2, - padding=1, - bias=False) - - def forward(self, x): - # Initial block - input_size = x.size() - x = self.initial_block(x) - - # Stage 1 - Encoder - stage1_input_size = x.size() - x, max_indices1_0 = self.downsample1_0(x) - x = self.regular1_1(x) - x = self.regular1_2(x) - x = self.regular1_3(x) - x = self.regular1_4(x) - - # Stage 2 - Encoder - stage2_input_size = x.size() - x, max_indices2_0 = self.downsample2_0(x) - x = self.regular2_1(x) - x = self.dilated2_2(x) - x = self.asymmetric2_3(x) - x = self.dilated2_4(x) - x = self.regular2_5(x) - x = self.dilated2_6(x) - x = self.asymmetric2_7(x) - x = self.dilated2_8(x) - - # Stage 3 - Encoder - x = self.regular3_0(x) - x = self.dilated3_1(x) - x = self.asymmetric3_2(x) - x = self.dilated3_3(x) - x = self.regular3_4(x) - x = self.dilated3_5(x) - x = self.asymmetric3_6(x) - x = self.dilated3_7(x) - - # Stage 4 - Decoder - x = self.upsample4_0(x, max_indices2_0, output_size=stage2_input_size) - x = self.regular4_1(x) - x = self.regular4_2(x) - - # Stage 5 - Decoder - x = self.upsample5_0(x, max_indices1_0, output_size=stage1_input_size) - x = self.regular5_1(x) - x = self.transposed_conv(x, output_size=input_size) - - return x diff --git a/code/networks/grid_attention_layer.py b/code/networks/grid_attention_layer.py old mode 100644 new mode 100755 diff --git a/code/networks/net_factory.py b/code/networks/net_factory.py old mode 100644 new mode 100755 index 7eb0733..4743ab4 --- a/code/networks/net_factory.py +++ b/code/networks/net_factory.py @@ -1,98 +1,23 @@ from networks.efficientunet import Effi_UNet -from networks.enet import ENet from networks.pnet import PNet2D -from networks.unet import UNet, UNet_DS, UNet_URPC, UNet_CCT -import argparse -from networks.vision_transformer import SwinUnet as ViT_seg -from networks.config import get_config -from networks.nnunet import initialize_network - - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/ACDC', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='ACDC/Cross_Supervision_CNN_Trans2D', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=8, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[224, 224], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') -parser.add_argument('--num_classes', type=int, default=4, - help='output channel of network') -parser.add_argument( - '--cfg', type=str, default="../code/configs/swin_tiny_patch4_window7_224_lite.yaml", help='path to config file', ) -parser.add_argument( - "--opts", - help="Modify config options by adding 'KEY VALUE' pairs. ", - default=None, - nargs='+', -) -parser.add_argument('--zip', action="iframe.php?url=https%3A%2F%2Fgithub.com%2Fstore_true", - help='use zipped dataset instead of folder dataset') -parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'], - help='no: no cache, ' - 'full: cache all data, ' - 'part: sharding the dataset into nonoverlapping pieces and only cache one piece') -parser.add_argument('--resume', help='resume from checkpoint') -parser.add_argument('--accumulation-steps', type=int, - help="gradient accumulation steps") -parser.add_argument('--use-checkpoint', action="iframe.php?url=https%3A%2F%2Fgithub.com%2Fstore_true", - help="whether to use gradient checkpointing to save memory") -parser.add_argument('--amp-opt-level', type=str, default='O1', choices=['O0', 'O1', 'O2'], - help='mixed precision opt level, if O0, no amp is used') -parser.add_argument('--tag', help='tag of experiment') -parser.add_argument('--eval', action="iframe.php?url=https%3A%2F%2Fgithub.com%2Fstore_true", - help='Perform evaluation only') -parser.add_argument('--throughput', action="iframe.php?url=https%3A%2F%2Fgithub.com%2Fstore_true", - help='Test throughput only') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=4, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=7, - help='labeled data') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() -config = get_config(args) +from networks.unet import UNet, UNet_DS, UNet_URDS +from networks.unet_multitask import UNet_Multitask def net_factory(net_type="unet", in_chns=1, class_num=3): if net_type == "unet": net = UNet(in_chns=in_chns, class_num=class_num).cuda() - elif net_type == "enet": - net = ENet(in_channels=in_chns, num_classes=class_num).cuda() + elif net_type == "unet_multitask": + net = UNet_Multitask(in_chns=in_chns, class_num=class_num).cuda() elif net_type == "unet_ds": net = UNet_DS(in_chns=in_chns, class_num=class_num).cuda() - elif net_type == "unet_cct": - net = UNet_CCT(in_chns=in_chns, class_num=class_num).cuda() - elif net_type == "unet_urpc": - net = UNet_URPC(in_chns=in_chns, class_num=class_num).cuda() + elif net_type == "unet_urds": + net = UNet_URDS(in_chns=in_chns, class_num=class_num).cuda() elif net_type == "efficient_unet": net = Effi_UNet('efficientnet-b3', encoder_weights='imagenet', in_channels=in_chns, classes=class_num).cuda() - elif net_type == "ViT_Seg": - net = ViT_seg(config, img_size=args.patch_size, - num_classes=args.num_classes).cuda() elif net_type == "pnet": net = PNet2D(in_chns, class_num, 64, [1, 2, 4, 8, 16]).cuda() - elif net_type == "nnUNet": - net = initialize_network(num_classes=class_num).cuda() else: net = None return net diff --git a/code/networks/net_factory_3d.py b/code/networks/net_factory_3d.py old mode 100644 new mode 100755 index 1ad4ace..3eee904 --- a/code/networks/net_factory_3d.py +++ b/code/networks/net_factory_3d.py @@ -2,7 +2,6 @@ from networks.vnet import VNet from networks.VoxResNet import VoxResNet from networks.attention_unet import Attention_UNet -from networks.nnunet import initialize_network def net_factory_3d(net_type="unet_3D", in_chns=1, class_num=2): @@ -16,8 +15,6 @@ def net_factory_3d(net_type="unet_3D", in_chns=1, class_num=2): elif net_type == "vnet": net = VNet(n_channels=in_chns, n_classes=class_num, normalization='batchnorm', has_dropout=True).cuda() - elif net_type == "nnUNet": - net = initialize_network(num_classes=class_num).cuda() else: net = None return net diff --git a/code/networks/networks_other.py b/code/networks/networks_other.py old mode 100644 new mode 100755 diff --git a/code/networks/neural_network.py b/code/networks/neural_network.py deleted file mode 100644 index ffda64b..0000000 --- a/code/networks/neural_network.py +++ /dev/null @@ -1,903 +0,0 @@ -# Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -# from torch.cuda.amp import autocast -import numpy as np -from batchgenerators.augmentations.utils import pad_nd_image -from torch import nn -import torch -from scipy.ndimage.filters import gaussian_filter -from typing import Union, Tuple, List - - -class no_op(object): - def __enter__(self): - pass - - def __exit__(self, *args): - pass - - -def maybe_to_torch(d): - if isinstance(d, list): - d = [maybe_to_torch(i) if not isinstance( - i, torch.Tensor) else i for i in d] - elif not isinstance(d, torch.Tensor): - d = torch.from_numpy(d).float() - return d - - -def to_cuda(data, non_blocking=True, gpu_id=0): - if isinstance(data, list): - data = [i.cuda(gpu_id, non_blocking=non_blocking) for i in data] - else: - data = data.cuda(gpu_id, non_blocking=non_blocking) - return data - - -class NeuralNetwork(nn.Module): - def __init__(self): - super(NeuralNetwork, self).__init__() - - def get_device(self): - if next(self.parameters()).device == "cpu": - return "cpu" - else: - return next(self.parameters()).device.index - - def set_device(self, device): - if device == "cpu": - self.cpu() - else: - self.cuda(device) - - def forward(self, x): - raise NotImplementedError - - -class SegmentationNetwork(NeuralNetwork): - def __init__(self): - super(NeuralNetwork, self).__init__() - - # if we have 5 pooling then our patch size must be divisible by 2**5 - # for example in a 2d network that does 5 pool in x and 6 pool - self.input_shape_must_be_divisible_by = None - # in y this would be (32, 64) - - # we need to know this because we need to know if we are a 2d or a 3d netowrk - self.conv_op = None # nn.Conv2d or nn.Conv3d - - # this tells us how many channely we have in the output. Important for preallocation in inference - self.num_classes = None # number of channels in the output - - # depending on the loss, we do not hard code a nonlinearity into the architecture. To aggregate predictions - # during inference, we need to apply the nonlinearity, however. So it is important to let the newtork know what - # to apply in inference. For the most part this will be softmax - self.inference_apply_nonlin = lambda x: x # softmax_helper - - # This is for saving a gaussian importance map for inference. It weights voxels higher that are closer to the - # center. Prediction at the borders are often less accurate and are thus downweighted. Creating these Gaussians - # can be expensive, so it makes sense to save and reuse them. - self._gaussian_3d = self._patch_size_for_gaussian_3d = None - self._gaussian_2d = self._patch_size_for_gaussian_2d = None - - def predict_3D(self, x: np.ndarray, do_mirroring: bool, mirror_axes: Tuple[int, ...] = (0, 1, 2), - use_sliding_window: bool = False, - step_size: float = 0.5, patch_size: Tuple[int, ...] = None, regions_class_order: Tuple[int, ...] = None, - use_gaussian: bool = False, pad_border_mode: str = "constant", - pad_kwargs: dict = None, all_in_gpu: bool = False, - verbose: bool = True, mixed_precision: bool = True) -> Tuple[np.ndarray, np.ndarray]: - """ - Use this function to predict a 3D image. It does not matter whether the network is a 2D or 3D U-Net, it will - detect that automatically and run the appropriate code. - When running predictions, you need to specify whether you want to run fully convolutional of sliding window - based inference. We very strongly recommend you use sliding window with the default settings. - It is the responsibility of the user to make sure the network is in the proper mode (eval for inference!). If - the network is not in eval mode it will print a warning. - :param x: Your input data. Must be a nd.ndarray of shape (c, x, y, z). - :param do_mirroring: If True, use test time data augmentation in the form of mirroring - :param mirror_axes: Determines which axes to use for mirroing. Per default, mirroring is done along all three - axes - :param use_sliding_window: if True, run sliding window prediction. Heavily recommended! This is also the default - :param step_size: When running sliding window prediction, the step size determines the distance between adjacent - predictions. The smaller the step size, the denser the predictions (and the longer it takes!). Step size is given - as a fraction of the patch_size. 0.5 is the default and means that wen advance by patch_size * 0.5 between - predictions. step_size cannot be larger than 1! - :param patch_size: The patch size that was used for training the network. Do not use different patch sizes here, - this will either crash or give potentially less accurate segmentations - :param regions_class_order: Fabian only - :param use_gaussian: (Only applies to sliding window prediction) If True, uses a Gaussian importance weighting - to weigh predictions closer to the center of the current patch higher than those at the borders. The reason - behind this is that the segmentation accuracy decreases towards the borders. Default (and recommended): True - :param pad_border_mode: leave this alone - :param pad_kwargs: leave this alone - :param all_in_gpu: experimental. You probably want to leave this as is it - :param verbose: Do you want a wall of text? If yes then set this to True - :param mixed_precision: if True, will run inference in mixed precision with autocast() - :return: - """ - torch.cuda.empty_cache() - - assert step_size <= 1, 'step_size must be smaller than 1. Otherwise there will be a gap between consecutive ' \ - 'predictions' - - if verbose: - print("debug: mirroring", do_mirroring, "mirror_axes", mirror_axes) - - assert self.get_device() != "cpu", "CPU not implemented" - - if pad_kwargs is None: - pad_kwargs = {'constant_values': 0} - - # A very long time ago the mirror axes were (2, 3, 4) for a 3d network. This is just to intercept any old - # code that uses this convention - if len(mirror_axes): - if self.conv_op == nn.Conv2d: - if max(mirror_axes) > 1: - raise ValueError("mirror axes. duh") - if self.conv_op == nn.Conv3d: - if max(mirror_axes) > 2: - raise ValueError("mirror axes. duh") - - if self.training: - print( - 'WARNING! Network is in train mode during inference. This may be intended, or not...') - - assert len(x.shape) == 4, "data must have shape (c,x,y,z)" - - if mixed_precision: - context = autocast - else: - context = no_op - - with context(): - with torch.no_grad(): - if self.conv_op == nn.Conv3d: - if use_sliding_window: - res = self._internal_predict_3D_3Dconv_tiled(x, step_size, do_mirroring, mirror_axes, patch_size, - regions_class_order, use_gaussian, pad_border_mode, - pad_kwargs=pad_kwargs, all_in_gpu=all_in_gpu, - verbose=verbose) - else: - res = self._internal_predict_3D_3Dconv(x, patch_size, do_mirroring, mirror_axes, regions_class_order, - pad_border_mode, pad_kwargs=pad_kwargs, verbose=verbose) - elif self.conv_op == nn.Conv2d: - if use_sliding_window: - res = self._internal_predict_3D_2Dconv_tiled(x, patch_size, do_mirroring, mirror_axes, step_size, - regions_class_order, use_gaussian, pad_border_mode, - pad_kwargs, all_in_gpu, False) - else: - res = self._internal_predict_3D_2Dconv(x, patch_size, do_mirroring, mirror_axes, regions_class_order, - pad_border_mode, pad_kwargs, all_in_gpu, False) - else: - raise RuntimeError( - "Invalid conv op, cannot determine what dimensionality (2d/3d) the network is") - - return res - - def predict_2D(self, x, do_mirroring: bool, mirror_axes: tuple = (0, 1, 2), use_sliding_window: bool = False, - step_size: float = 0.5, patch_size: tuple = None, regions_class_order: tuple = None, - use_gaussian: bool = False, pad_border_mode: str = "constant", - pad_kwargs: dict = None, all_in_gpu: bool = False, - verbose: bool = True, mixed_precision: bool = True) -> Tuple[np.ndarray, np.ndarray]: - """ - Use this function to predict a 2D image. If this is a 3D U-Net it will crash because you cannot predict a 2D - image with that (you dummy). - When running predictions, you need to specify whether you want to run fully convolutional of sliding window - based inference. We very strongly recommend you use sliding window with the default settings. - It is the responsibility of the user to make sure the network is in the proper mode (eval for inference!). If - the network is not in eval mode it will print a warning. - :param x: Your input data. Must be a nd.ndarray of shape (c, x, y). - :param do_mirroring: If True, use test time data augmentation in the form of mirroring - :param mirror_axes: Determines which axes to use for mirroing. Per default, mirroring is done along all three - axes - :param use_sliding_window: if True, run sliding window prediction. Heavily recommended! This is also the default - :param step_size: When running sliding window prediction, the step size determines the distance between adjacent - predictions. The smaller the step size, the denser the predictions (and the longer it takes!). Step size is given - as a fraction of the patch_size. 0.5 is the default and means that wen advance by patch_size * 0.5 between - predictions. step_size cannot be larger than 1! - :param patch_size: The patch size that was used for training the network. Do not use different patch sizes here, - this will either crash or give potentially less accurate segmentations - :param regions_class_order: Fabian only - :param use_gaussian: (Only applies to sliding window prediction) If True, uses a Gaussian importance weighting - to weigh predictions closer to the center of the current patch higher than those at the borders. The reason - behind this is that the segmentation accuracy decreases towards the borders. Default (and recommended): True - :param pad_border_mode: leave this alone - :param pad_kwargs: leave this alone - :param all_in_gpu: experimental. You probably want to leave this as is it - :param verbose: Do you want a wall of text? If yes then set this to True - :return: - """ - torch.cuda.empty_cache() - - assert step_size <= 1, 'step_size must be smaler than 1. Otherwise there will be a gap between consecutive ' \ - 'predictions' - - if self.conv_op == nn.Conv3d: - raise RuntimeError( - "Cannot predict 2d if the network is 3d. Dummy.") - - if verbose: - print("debug: mirroring", do_mirroring, "mirror_axes", mirror_axes) - - assert self.get_device() != "cpu", "CPU not implemented" - - if pad_kwargs is None: - pad_kwargs = {'constant_values': 0} - - # A very long time ago the mirror axes were (2, 3) for a 2d network. This is just to intercept any old - # code that uses this convention - if len(mirror_axes): - if max(mirror_axes) > 1: - raise ValueError("mirror axes. duh") - - if self.training: - print( - 'WARNING! Network is in train mode during inference. This may be intended, or not...') - - assert len(x.shape) == 3, "data must have shape (c,x,y)" - - if mixed_precision: - context = autocast - else: - context = no_op - - with context(): - with torch.no_grad(): - if self.conv_op == nn.Conv2d: - if use_sliding_window: - res = self._internal_predict_2D_2Dconv_tiled(x, step_size, do_mirroring, mirror_axes, patch_size, - regions_class_order, use_gaussian, pad_border_mode, - pad_kwargs, all_in_gpu, verbose) - else: - res = self._internal_predict_2D_2Dconv(x, patch_size, do_mirroring, mirror_axes, regions_class_order, - pad_border_mode, pad_kwargs, verbose) - else: - raise RuntimeError( - "Invalid conv op, cannot determine what dimensionality (2d/3d) the network is") - - return res - - @staticmethod - def _get_gaussian(patch_size, sigma_scale=1. / 8) -> np.ndarray: - tmp = np.zeros(patch_size) - center_coords = [i // 2 for i in patch_size] - sigmas = [i * sigma_scale for i in patch_size] - tmp[tuple(center_coords)] = 1 - gaussian_importance_map = gaussian_filter( - tmp, sigmas, 0, mode='constant', cval=0) - gaussian_importance_map = gaussian_importance_map / \ - np.max(gaussian_importance_map) * 1 - gaussian_importance_map = gaussian_importance_map.astype(np.float32) - - # gaussian_importance_map cannot be 0, otherwise we may end up with nans! - gaussian_importance_map[gaussian_importance_map == 0] = np.min( - gaussian_importance_map[gaussian_importance_map != 0]) - - return gaussian_importance_map - - @staticmethod - def _compute_steps_for_sliding_window(patch_size: Tuple[int, ...], image_size: Tuple[int, ...], step_size: float) -> List[List[int]]: - assert [i >= j for i, j in zip( - image_size, patch_size)], "image size must be as large or larger than patch_size" - assert 0 < step_size <= 1, 'step_size must be larger than 0 and smaller or equal to 1' - - # our step width is patch_size*step_size at most, but can be narrower. For example if we have image size of - # 110, patch size of 64 and step_size of 0.5, then we want to make 3 steps starting at coordinate 0, 23, 46 - target_step_sizes_in_voxels = [i * step_size for i in patch_size] - - num_steps = [int(np.ceil((i - k) / j)) + 1 for i, j, - k in zip(image_size, target_step_sizes_in_voxels, patch_size)] - - steps = [] - for dim in range(len(patch_size)): - # the highest step value for this dimension is - max_step_value = image_size[dim] - patch_size[dim] - if num_steps[dim] > 1: - actual_step_size = max_step_value / (num_steps[dim] - 1) - else: - # does not matter because there is only one step at 0 - actual_step_size = 99999999999 - - steps_here = [int(np.round(actual_step_size * i)) - for i in range(num_steps[dim])] - - steps.append(steps_here) - - return steps - - def _internal_predict_3D_3Dconv_tiled(self, x: np.ndarray, step_size: float, do_mirroring: bool, mirror_axes: tuple, - patch_size: tuple, regions_class_order: tuple, use_gaussian: bool, - pad_border_mode: str, pad_kwargs: dict, all_in_gpu: bool, - verbose: bool) -> Tuple[np.ndarray, np.ndarray]: - # better safe than sorry - assert len(x.shape) == 4, "x must be (c, x, y, z)" - assert self.get_device() != "cpu" - if verbose: - print("step_size:", step_size) - if verbose: - print("do mirror:", do_mirroring) - - assert patch_size is not None, "patch_size cannot be None for tiled prediction" - - # for sliding window inference the image must at least be as large as the patch size. It does not matter - # whether the shape is divisible by 2**num_pool as long as the patch size is - data, slicer = pad_nd_image( - x, patch_size, pad_border_mode, pad_kwargs, True, None) - data_shape = data.shape # still c, x, y, z - - # compute the steps for sliding window - steps = self._compute_steps_for_sliding_window( - patch_size, data_shape[1:], step_size) - num_tiles = len(steps[0]) * len(steps[1]) * len(steps[2]) - - if verbose: - print("data shape:", data_shape) - print("patch size:", patch_size) - print("steps (x, y, and z):", steps) - print("number of tiles:", num_tiles) - - # we only need to compute that once. It can take a while to compute this due to the large sigma in - # gaussian_filter - if use_gaussian and num_tiles > 1: - if self._gaussian_3d is None or not all( - [i == j for i, j in zip(patch_size, self._patch_size_for_gaussian_3d)]): - if verbose: - print('computing Gaussian') - gaussian_importance_map = self._get_gaussian( - patch_size, sigma_scale=1. / 8) - - self._gaussian_3d = gaussian_importance_map - self._patch_size_for_gaussian_3d = patch_size - else: - if verbose: - print("using precomputed Gaussian") - gaussian_importance_map = self._gaussian_3d - - gaussian_importance_map = torch.from_numpy(gaussian_importance_map).cuda(self.get_device(), - non_blocking=True) - - else: - gaussian_importance_map = None - - if all_in_gpu: - # If we run the inference in GPU only (meaning all tensors are allocated on the GPU, this reduces - # CPU-GPU communication but required more GPU memory) we need to preallocate a few things on GPU - - if use_gaussian and num_tiles > 1: - # half precision for the outputs should be good enough. If the outputs here are half, the - # gaussian_importance_map should be as well - gaussian_importance_map = gaussian_importance_map.half() - - # make sure we did not round anything to 0 - gaussian_importance_map[gaussian_importance_map == 0] = gaussian_importance_map[ - gaussian_importance_map != 0].min() - - add_for_nb_of_preds = gaussian_importance_map - else: - add_for_nb_of_preds = torch.ones( - data.shape[1:], device=self.get_device()) - - if verbose: - print("initializing result array (on GPU)") - aggregated_results = torch.zeros([self.num_classes] + list(data.shape[1:]), dtype=torch.half, - device=self.get_device()) - - if verbose: - print("moving data to GPU") - data = torch.from_numpy(data).cuda( - self.get_device(), non_blocking=True) - - if verbose: - print("initializing result_numsamples (on GPU)") - aggregated_nb_of_predictions = torch.zeros([self.num_classes] + list(data.shape[1:]), dtype=torch.half, - device=self.get_device()) - else: - if use_gaussian and num_tiles > 1: - add_for_nb_of_preds = self._gaussian_3d - else: - add_for_nb_of_preds = np.ones(data.shape[1:], dtype=np.float32) - aggregated_results = np.zeros( - [self.num_classes] + list(data.shape[1:]), dtype=np.float32) - aggregated_nb_of_predictions = np.zeros( - [self.num_classes] + list(data.shape[1:]), dtype=np.float32) - - for x in steps[0]: - lb_x = x - ub_x = x + patch_size[0] - for y in steps[1]: - lb_y = y - ub_y = y + patch_size[1] - for z in steps[2]: - lb_z = z - ub_z = z + patch_size[2] - - predicted_patch = self._internal_maybe_mirror_and_pred_3D( - data[None, :, lb_x:ub_x, lb_y:ub_y, - lb_z:ub_z], mirror_axes, do_mirroring, - gaussian_importance_map)[0] - - if all_in_gpu: - predicted_patch = predicted_patch.half() - else: - predicted_patch = predicted_patch.cpu().numpy() - - aggregated_results[:, lb_x:ub_x, - lb_y:ub_y, lb_z:ub_z] += predicted_patch - aggregated_nb_of_predictions[:, lb_x:ub_x, - lb_y:ub_y, lb_z:ub_z] += add_for_nb_of_preds - - # we reverse the padding here (remeber that we padded the input to be at least as large as the patch size - slicer = tuple( - [slice(0, aggregated_results.shape[i]) for i in - range(len(aggregated_results.shape) - (len(slicer) - 1))] + slicer[1:]) - aggregated_results = aggregated_results[slicer] - aggregated_nb_of_predictions = aggregated_nb_of_predictions[slicer] - - # computing the class_probabilities by dividing the aggregated result with result_numsamples - class_probabilities = aggregated_results / aggregated_nb_of_predictions - - if regions_class_order is None: - predicted_segmentation = class_probabilities.argmax(0) - else: - if all_in_gpu: - class_probabilities_here = class_probabilities.detach().cpu().numpy() - else: - class_probabilities_here = class_probabilities - predicted_segmentation = np.zeros( - class_probabilities_here.shape[1:], dtype=np.float32) - for i, c in enumerate(regions_class_order): - predicted_segmentation[class_probabilities_here[i] > 0.5] = c - - if all_in_gpu: - if verbose: - print("copying results to CPU") - - if regions_class_order is None: - predicted_segmentation = predicted_segmentation.detach().cpu().numpy() - - class_probabilities = class_probabilities.detach().cpu().numpy() - - if verbose: - print("prediction done") - return predicted_segmentation, class_probabilities - - def _internal_predict_2D_2Dconv(self, x: np.ndarray, min_size: Tuple[int, int], do_mirroring: bool, - mirror_axes: tuple = (0, 1, 2), regions_class_order: tuple = None, - pad_border_mode: str = "constant", pad_kwargs: dict = None, - verbose: bool = True) -> Tuple[np.ndarray, np.ndarray]: - """ - This one does fully convolutional inference. No sliding window - """ - assert len(x.shape) == 3, "x must be (c, x, y)" - assert self.get_device() != "cpu" - assert self.input_shape_must_be_divisible_by is not None, 'input_shape_must_be_divisible_by must be set to ' \ - 'run _internal_predict_2D_2Dconv' - if verbose: - print("do mirror:", do_mirroring) - - data, slicer = pad_nd_image(x, min_size, pad_border_mode, pad_kwargs, True, - self.input_shape_must_be_divisible_by) - - predicted_probabilities = self._internal_maybe_mirror_and_pred_2D(data[None], mirror_axes, do_mirroring, - None)[0] - - slicer = tuple( - [slice(0, predicted_probabilities.shape[i]) for i in range(len(predicted_probabilities.shape) - - (len(slicer) - 1))] + slicer[1:]) - predicted_probabilities = predicted_probabilities[slicer] - - if regions_class_order is None: - predicted_segmentation = predicted_probabilities.argmax(0) - predicted_segmentation = predicted_segmentation.detach().cpu().numpy() - predicted_probabilities = predicted_probabilities.detach().cpu().numpy() - else: - predicted_probabilities = predicted_probabilities.detach().cpu().numpy() - predicted_segmentation = np.zeros( - predicted_probabilities.shape[1:], dtype=np.float32) - for i, c in enumerate(regions_class_order): - predicted_segmentation[predicted_probabilities[i] > 0.5] = c - - return predicted_segmentation, predicted_probabilities - - def _internal_predict_3D_3Dconv(self, x: np.ndarray, min_size: Tuple[int, ...], do_mirroring: bool, - mirror_axes: tuple = (0, 1, 2), regions_class_order: tuple = None, - pad_border_mode: str = "constant", pad_kwargs: dict = None, - verbose: bool = True) -> Tuple[np.ndarray, np.ndarray]: - """ - This one does fully convolutional inference. No sliding window - """ - assert len(x.shape) == 4, "x must be (c, x, y, z)" - assert self.get_device() != "cpu" - assert self.input_shape_must_be_divisible_by is not None, 'input_shape_must_be_divisible_by must be set to ' \ - 'run _internal_predict_3D_3Dconv' - if verbose: - print("do mirror:", do_mirroring) - - data, slicer = pad_nd_image(x, min_size, pad_border_mode, pad_kwargs, True, - self.input_shape_must_be_divisible_by) - - predicted_probabilities = self._internal_maybe_mirror_and_pred_3D(data[None], mirror_axes, do_mirroring, - None)[0] - - slicer = tuple( - [slice(0, predicted_probabilities.shape[i]) for i in range(len(predicted_probabilities.shape) - - (len(slicer) - 1))] + slicer[1:]) - predicted_probabilities = predicted_probabilities[slicer] - - if regions_class_order is None: - predicted_segmentation = predicted_probabilities.argmax(0) - predicted_segmentation = predicted_segmentation.detach().cpu().numpy() - predicted_probabilities = predicted_probabilities.detach().cpu().numpy() - else: - predicted_probabilities = predicted_probabilities.detach().cpu().numpy() - predicted_segmentation = np.zeros( - predicted_probabilities.shape[1:], dtype=np.float32) - for i, c in enumerate(regions_class_order): - predicted_segmentation[predicted_probabilities[i] > 0.5] = c - - return predicted_segmentation, predicted_probabilities - - def _internal_maybe_mirror_and_pred_3D(self, x: Union[np.ndarray, torch.tensor], mirror_axes: tuple, - do_mirroring: bool = True, - mult: np.ndarray or torch.tensor = None) -> torch.tensor: - assert len(x.shape) == 5, 'x must be (b, c, x, y, z)' - # everything in here takes place on the GPU. If x and mult are not yet on GPU this will be taken care of here - # we now return a cuda tensor! Not numpy array! - - x = to_cuda(maybe_to_torch(x), gpu_id=self.get_device()) - result_torch = torch.zeros([1, self.num_classes] + list(x.shape[2:]), - dtype=torch.float).cuda(self.get_device(), non_blocking=True) - - if mult is not None: - mult = to_cuda(maybe_to_torch(mult), gpu_id=self.get_device()) - - if do_mirroring: - mirror_idx = 8 - num_results = 2 ** len(mirror_axes) - else: - mirror_idx = 1 - num_results = 1 - - for m in range(mirror_idx): - if m == 0: - pred = self.inference_apply_nonlin(self(x)) - result_torch += 1 / num_results * pred - - if m == 1 and (2 in mirror_axes): - pred = self.inference_apply_nonlin(self(torch.flip(x, (4, )))) - result_torch += 1 / num_results * torch.flip(pred, (4,)) - - if m == 2 and (1 in mirror_axes): - pred = self.inference_apply_nonlin(self(torch.flip(x, (3, )))) - result_torch += 1 / num_results * torch.flip(pred, (3,)) - - if m == 3 and (2 in mirror_axes) and (1 in mirror_axes): - pred = self.inference_apply_nonlin(self(torch.flip(x, (4, 3)))) - result_torch += 1 / num_results * torch.flip(pred, (4, 3)) - - if m == 4 and (0 in mirror_axes): - pred = self.inference_apply_nonlin(self(torch.flip(x, (2, )))) - result_torch += 1 / num_results * torch.flip(pred, (2,)) - - if m == 5 and (0 in mirror_axes) and (2 in mirror_axes): - pred = self.inference_apply_nonlin(self(torch.flip(x, (4, 2)))) - result_torch += 1 / num_results * torch.flip(pred, (4, 2)) - - if m == 6 and (0 in mirror_axes) and (1 in mirror_axes): - pred = self.inference_apply_nonlin(self(torch.flip(x, (3, 2)))) - result_torch += 1 / num_results * torch.flip(pred, (3, 2)) - - if m == 7 and (0 in mirror_axes) and (1 in mirror_axes) and (2 in mirror_axes): - pred = self.inference_apply_nonlin( - self(torch.flip(x, (4, 3, 2)))) - result_torch += 1 / num_results * torch.flip(pred, (4, 3, 2)) - - if mult is not None: - result_torch[:, :] *= mult - - return result_torch - - def _internal_maybe_mirror_and_pred_2D(self, x: Union[np.ndarray, torch.tensor], mirror_axes: tuple, - do_mirroring: bool = True, - mult: np.ndarray or torch.tensor = None) -> torch.tensor: - # everything in here takes place on the GPU. If x and mult are not yet on GPU this will be taken care of here - # we now return a cuda tensor! Not numpy array! - assert len(x.shape) == 4, 'x must be (b, c, x, y)' - - x = to_cuda(maybe_to_torch(x), gpu_id=self.get_device()) - result_torch = torch.zeros([x.shape[0], self.num_classes] + list(x.shape[2:]), - dtype=torch.float).cuda(self.get_device(), non_blocking=True) - - if mult is not None: - mult = to_cuda(maybe_to_torch(mult), gpu_id=self.get_device()) - - if do_mirroring: - mirror_idx = 4 - num_results = 2 ** len(mirror_axes) - else: - mirror_idx = 1 - num_results = 1 - - for m in range(mirror_idx): - if m == 0: - pred = self.inference_apply_nonlin(self(x)) - result_torch += 1 / num_results * pred - - if m == 1 and (1 in mirror_axes): - pred = self.inference_apply_nonlin(self(torch.flip(x, (3, )))) - result_torch += 1 / num_results * torch.flip(pred, (3, )) - - if m == 2 and (0 in mirror_axes): - pred = self.inference_apply_nonlin(self(torch.flip(x, (2, )))) - result_torch += 1 / num_results * torch.flip(pred, (2, )) - - if m == 3 and (0 in mirror_axes) and (1 in mirror_axes): - pred = self.inference_apply_nonlin(self(torch.flip(x, (3, 2)))) - result_torch += 1 / num_results * torch.flip(pred, (3, 2)) - - if mult is not None: - result_torch[:, :] *= mult - - return result_torch - - def _internal_predict_2D_2Dconv_tiled(self, x: np.ndarray, step_size: float, do_mirroring: bool, mirror_axes: tuple, - patch_size: tuple, regions_class_order: tuple, use_gaussian: bool, - pad_border_mode: str, pad_kwargs: dict, all_in_gpu: bool, - verbose: bool) -> Tuple[np.ndarray, np.ndarray]: - # better safe than sorry - assert len(x.shape) == 3, "x must be (c, x, y)" - assert self.get_device() != "cpu" - if verbose: - print("step_size:", step_size) - if verbose: - print("do mirror:", do_mirroring) - - assert patch_size is not None, "patch_size cannot be None for tiled prediction" - - # for sliding window inference the image must at least be as large as the patch size. It does not matter - # whether the shape is divisible by 2**num_pool as long as the patch size is - data, slicer = pad_nd_image( - x, patch_size, pad_border_mode, pad_kwargs, True, None) - data_shape = data.shape # still c, x, y - - # compute the steps for sliding window - steps = self._compute_steps_for_sliding_window( - patch_size, data_shape[1:], step_size) - num_tiles = len(steps[0]) * len(steps[1]) - - if verbose: - print("data shape:", data_shape) - print("patch size:", patch_size) - print("steps (x, y, and z):", steps) - print("number of tiles:", num_tiles) - - # we only need to compute that once. It can take a while to compute this due to the large sigma in - # gaussian_filter - if use_gaussian and num_tiles > 1: - if self._gaussian_2d is None or not all( - [i == j for i, j in zip(patch_size, self._patch_size_for_gaussian_2d)]): - if verbose: - print('computing Gaussian') - gaussian_importance_map = self._get_gaussian( - patch_size, sigma_scale=1. / 8) - - self._gaussian_2d = gaussian_importance_map - self._patch_size_for_gaussian_2d = patch_size - else: - if verbose: - print("using precomputed Gaussian") - gaussian_importance_map = self._gaussian_2d - - gaussian_importance_map = torch.from_numpy(gaussian_importance_map).cuda(self.get_device(), - non_blocking=True) - else: - gaussian_importance_map = None - - if all_in_gpu: - # If we run the inference in GPU only (meaning all tensors are allocated on the GPU, this reduces - # CPU-GPU communication but required more GPU memory) we need to preallocate a few things on GPU - - if use_gaussian and num_tiles > 1: - # half precision for the outputs should be good enough. If the outputs here are half, the - # gaussian_importance_map should be as well - gaussian_importance_map = gaussian_importance_map.half() - - # make sure we did not round anything to 0 - gaussian_importance_map[gaussian_importance_map == 0] = gaussian_importance_map[ - gaussian_importance_map != 0].min() - - add_for_nb_of_preds = gaussian_importance_map - else: - add_for_nb_of_preds = torch.ones( - data.shape[1:], device=self.get_device()) - - if verbose: - print("initializing result array (on GPU)") - aggregated_results = torch.zeros([self.num_classes] + list(data.shape[1:]), dtype=torch.half, - device=self.get_device()) - - if verbose: - print("moving data to GPU") - data = torch.from_numpy(data).cuda( - self.get_device(), non_blocking=True) - - if verbose: - print("initializing result_numsamples (on GPU)") - aggregated_nb_of_predictions = torch.zeros([self.num_classes] + list(data.shape[1:]), dtype=torch.half, - device=self.get_device()) - else: - if use_gaussian and num_tiles > 1: - add_for_nb_of_preds = self._gaussian_2d - else: - add_for_nb_of_preds = np.ones(data.shape[1:], dtype=np.float32) - aggregated_results = np.zeros( - [self.num_classes] + list(data.shape[1:]), dtype=np.float32) - aggregated_nb_of_predictions = np.zeros( - [self.num_classes] + list(data.shape[1:]), dtype=np.float32) - - for x in steps[0]: - lb_x = x - ub_x = x + patch_size[0] - for y in steps[1]: - lb_y = y - ub_y = y + patch_size[1] - - predicted_patch = self._internal_maybe_mirror_and_pred_2D( - data[None, :, lb_x:ub_x, lb_y:ub_y], mirror_axes, do_mirroring, - gaussian_importance_map)[0] - - if all_in_gpu: - predicted_patch = predicted_patch.half() - else: - predicted_patch = predicted_patch.cpu().numpy() - - aggregated_results[:, lb_x:ub_x, lb_y:ub_y] += predicted_patch - aggregated_nb_of_predictions[:, lb_x:ub_x, - lb_y:ub_y] += add_for_nb_of_preds - - # we reverse the padding here (remeber that we padded the input to be at least as large as the patch size - slicer = tuple( - [slice(0, aggregated_results.shape[i]) for i in - range(len(aggregated_results.shape) - (len(slicer) - 1))] + slicer[1:]) - aggregated_results = aggregated_results[slicer] - aggregated_nb_of_predictions = aggregated_nb_of_predictions[slicer] - - # computing the class_probabilities by dividing the aggregated result with result_numsamples - class_probabilities = aggregated_results / aggregated_nb_of_predictions - - if regions_class_order is None: - predicted_segmentation = class_probabilities.argmax(0) - else: - if all_in_gpu: - class_probabilities_here = class_probabilities.detach().cpu().numpy() - else: - class_probabilities_here = class_probabilities - predicted_segmentation = np.zeros( - class_probabilities_here.shape[1:], dtype=np.float32) - for i, c in enumerate(regions_class_order): - predicted_segmentation[class_probabilities_here[i] > 0.5] = c - - if all_in_gpu: - if verbose: - print("copying results to CPU") - - if regions_class_order is None: - predicted_segmentation = predicted_segmentation.detach().cpu().numpy() - - class_probabilities = class_probabilities.detach().cpu().numpy() - - if verbose: - print("prediction done") - return predicted_segmentation, class_probabilities - - def _internal_predict_3D_2Dconv(self, x: np.ndarray, min_size: Tuple[int, int], do_mirroring: bool, - mirror_axes: tuple = (0, 1), regions_class_order: tuple = None, - pad_border_mode: str = "constant", pad_kwargs: dict = None, - all_in_gpu: bool = False, verbose: bool = True) -> Tuple[np.ndarray, np.ndarray]: - if all_in_gpu: - raise NotImplementedError - assert len(x.shape) == 4, "data must be c, x, y, z" - predicted_segmentation = [] - softmax_pred = [] - for s in range(x.shape[1]): - pred_seg, softmax_pres = self._internal_predict_2D_2Dconv( - x[:, s], min_size, do_mirroring, mirror_axes, regions_class_order, pad_border_mode, pad_kwargs, verbose) - predicted_segmentation.append(pred_seg[None]) - softmax_pred.append(softmax_pres[None]) - predicted_segmentation = np.vstack(predicted_segmentation) - softmax_pred = np.vstack(softmax_pred).transpose((1, 0, 2, 3)) - return predicted_segmentation, softmax_pred - - def predict_3D_pseudo3D_2Dconv(self, x: np.ndarray, min_size: Tuple[int, int], do_mirroring: bool, - mirror_axes: tuple = (0, 1), regions_class_order: tuple = None, - pseudo3D_slices: int = 5, all_in_gpu: bool = False, - pad_border_mode: str = "constant", pad_kwargs: dict = None, - verbose: bool = True) -> Tuple[np.ndarray, np.ndarray]: - if all_in_gpu: - raise NotImplementedError - assert len(x.shape) == 4, "data must be c, x, y, z" - assert pseudo3D_slices % 2 == 1, "pseudo3D_slices must be odd" - extra_slices = (pseudo3D_slices - 1) // 2 - - shp_for_pad = np.array(x.shape) - shp_for_pad[1] = extra_slices - - pad = np.zeros(shp_for_pad, dtype=np.float32) - data = np.concatenate((pad, x, pad), 1) - - predicted_segmentation = [] - softmax_pred = [] - for s in range(extra_slices, data.shape[1] - extra_slices): - d = data[:, (s - extra_slices):(s + extra_slices + 1)] - d = d.reshape((-1, d.shape[-2], d.shape[-1])) - pred_seg, softmax_pres = \ - self._internal_predict_2D_2Dconv(d, min_size, do_mirroring, mirror_axes, - regions_class_order, pad_border_mode, pad_kwargs, verbose) - predicted_segmentation.append(pred_seg[None]) - softmax_pred.append(softmax_pres[None]) - predicted_segmentation = np.vstack(predicted_segmentation) - softmax_pred = np.vstack(softmax_pred).transpose((1, 0, 2, 3)) - - return predicted_segmentation, softmax_pred - - def _internal_predict_3D_2Dconv_tiled(self, x: np.ndarray, patch_size: Tuple[int, int], do_mirroring: bool, - mirror_axes: tuple = (0, 1), step_size: float = 0.5, - regions_class_order: tuple = None, use_gaussian: bool = False, - pad_border_mode: str = "edge", pad_kwargs: dict = None, - all_in_gpu: bool = False, - verbose: bool = True) -> Tuple[np.ndarray, np.ndarray]: - if all_in_gpu: - raise NotImplementedError - - assert len(x.shape) == 4, "data must be c, x, y, z" - - predicted_segmentation = [] - softmax_pred = [] - - for s in range(x.shape[1]): - pred_seg, softmax_pres = self._internal_predict_2D_2Dconv_tiled( - x[:, s], step_size, do_mirroring, mirror_axes, patch_size, regions_class_order, use_gaussian, - pad_border_mode, pad_kwargs, all_in_gpu, verbose) - - predicted_segmentation.append(pred_seg[None]) - softmax_pred.append(softmax_pres[None]) - - predicted_segmentation = np.vstack(predicted_segmentation) - softmax_pred = np.vstack(softmax_pred).transpose((1, 0, 2, 3)) - - return predicted_segmentation, softmax_pred - - -if __name__ == '__main__': - print(SegmentationNetwork._compute_steps_for_sliding_window( - (30, 224, 224), (162, 529, 529), 0.5)) - print(SegmentationNetwork._compute_steps_for_sliding_window( - (30, 224, 224), (162, 529, 529), 1)) - print(SegmentationNetwork._compute_steps_for_sliding_window( - (30, 224, 224), (162, 529, 529), 0.1)) - - print(SegmentationNetwork._compute_steps_for_sliding_window( - (30, 224, 224), (60, 448, 224), 1)) - print(SegmentationNetwork._compute_steps_for_sliding_window( - (30, 224, 224), (60, 448, 224), 0.5)) - - print(SegmentationNetwork._compute_steps_for_sliding_window( - (30, 224, 224), (30, 224, 224), 1)) - print(SegmentationNetwork._compute_steps_for_sliding_window( - (30, 224, 224), (30, 224, 224), 0.125)) - - print(SegmentationNetwork._compute_steps_for_sliding_window( - (123, 54, 123), (246, 162, 369), 0.25)) diff --git a/code/networks/nnunet.py b/code/networks/nnunet.py deleted file mode 100644 index 906bd32..0000000 --- a/code/networks/nnunet.py +++ /dev/null @@ -1,535 +0,0 @@ -# Copyright 2020 Division of Medical Image Computing, German Cancer Research Center (DKFZ), Heidelberg, Germany -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -from copy import deepcopy - -import torch.nn.functional as F -from torch import nn -import torch -import numpy as np - -from networks.neural_network import SegmentationNetwork -import torch.nn.functional - - -def softmax_helper(x): return F.softmax(x, 1) - - -class InitWeights_He(object): - def __init__(self, neg_slope=1e-2): - self.neg_slope = neg_slope - - def __call__(self, module): - if isinstance(module, nn.Conv3d) or isinstance(module, nn.Conv2d) or isinstance(module, nn.ConvTranspose2d) or isinstance(module, nn.ConvTranspose3d): - module.weight = nn.init.kaiming_normal_( - module.weight, a=self.neg_slope) - if module.bias is not None: - module.bias = nn.init.constant_(module.bias, 0) - - -class ConvDropoutNormNonlin(nn.Module): - """ - fixes a bug in ConvDropoutNormNonlin where lrelu was used regardless of nonlin. Bad. - """ - - def __init__(self, input_channels, output_channels, - conv_op=nn.Conv2d, conv_kwargs=None, - norm_op=nn.BatchNorm2d, norm_op_kwargs=None, - dropout_op=nn.Dropout2d, dropout_op_kwargs=None, - nonlin=nn.LeakyReLU, nonlin_kwargs=None): - super(ConvDropoutNormNonlin, self).__init__() - if nonlin_kwargs is None: - nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True} - if dropout_op_kwargs is None: - dropout_op_kwargs = {'p': 0.5, 'inplace': True} - if norm_op_kwargs is None: - norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1} - if conv_kwargs is None: - conv_kwargs = {'kernel_size': 3, 'stride': 1, - 'padding': 1, 'dilation': 1, 'bias': True} - - self.nonlin_kwargs = nonlin_kwargs - self.nonlin = nonlin - self.dropout_op = dropout_op - self.dropout_op_kwargs = dropout_op_kwargs - self.norm_op_kwargs = norm_op_kwargs - self.conv_kwargs = conv_kwargs - self.conv_op = conv_op - self.norm_op = norm_op - - self.conv = self.conv_op( - input_channels, output_channels, **self.conv_kwargs) - if self.dropout_op is not None and self.dropout_op_kwargs['p'] is not None and self.dropout_op_kwargs[ - 'p'] > 0: - self.dropout = self.dropout_op(**self.dropout_op_kwargs) - else: - self.dropout = None - self.instnorm = self.norm_op(output_channels, **self.norm_op_kwargs) - self.lrelu = self.nonlin(**self.nonlin_kwargs) - - def forward(self, x): - x = self.conv(x) - if self.dropout is not None: - x = self.dropout(x) - return self.lrelu(self.instnorm(x)) - - -class ConvDropoutNonlinNorm(ConvDropoutNormNonlin): - def forward(self, x): - x = self.conv(x) - if self.dropout is not None: - x = self.dropout(x) - return self.instnorm(self.lrelu(x)) - - -class StackedConvLayers(nn.Module): - def __init__(self, input_feature_channels, output_feature_channels, num_convs, - conv_op=nn.Conv2d, conv_kwargs=None, - norm_op=nn.BatchNorm2d, norm_op_kwargs=None, - dropout_op=nn.Dropout2d, dropout_op_kwargs=None, - nonlin=nn.LeakyReLU, nonlin_kwargs=None, first_stride=None, basic_block=ConvDropoutNormNonlin): - ''' - stacks ConvDropoutNormLReLU layers. initial_stride will only be applied to first layer in the stack. The other parameters affect all layers - :param input_feature_channels: - :param output_feature_channels: - :param num_convs: - :param dilation: - :param kernel_size: - :param padding: - :param dropout: - :param initial_stride: - :param conv_op: - :param norm_op: - :param dropout_op: - :param inplace: - :param neg_slope: - :param norm_affine: - :param conv_bias: - ''' - self.input_channels = input_feature_channels - self.output_channels = output_feature_channels - - if nonlin_kwargs is None: - nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True} - if dropout_op_kwargs is None: - dropout_op_kwargs = {'p': 0.5, 'inplace': True} - if norm_op_kwargs is None: - norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1} - if conv_kwargs is None: - conv_kwargs = {'kernel_size': 3, 'stride': 1, - 'padding': 1, 'dilation': 1, 'bias': True} - - self.nonlin_kwargs = nonlin_kwargs - self.nonlin = nonlin - self.dropout_op = dropout_op - self.dropout_op_kwargs = dropout_op_kwargs - self.norm_op_kwargs = norm_op_kwargs - self.conv_kwargs = conv_kwargs - self.conv_op = conv_op - self.norm_op = norm_op - - if first_stride is not None: - self.conv_kwargs_first_conv = deepcopy(conv_kwargs) - self.conv_kwargs_first_conv['stride'] = first_stride - else: - self.conv_kwargs_first_conv = conv_kwargs - - super(StackedConvLayers, self).__init__() - self.blocks = nn.Sequential( - *([basic_block(input_feature_channels, output_feature_channels, self.conv_op, - self.conv_kwargs_first_conv, - self.norm_op, self.norm_op_kwargs, self.dropout_op, self.dropout_op_kwargs, - self.nonlin, self.nonlin_kwargs)] + - [basic_block(output_feature_channels, output_feature_channels, self.conv_op, - self.conv_kwargs, - self.norm_op, self.norm_op_kwargs, self.dropout_op, self.dropout_op_kwargs, - self.nonlin, self.nonlin_kwargs) for _ in range(num_convs - 1)])) - - def forward(self, x): - return self.blocks(x) - - -def print_module_training_status(module): - if isinstance(module, nn.Conv2d) or isinstance(module, nn.Conv3d) or isinstance(module, nn.Dropout3d) or \ - isinstance(module, nn.Dropout2d) or isinstance(module, nn.Dropout) or isinstance(module, nn.InstanceNorm3d) \ - or isinstance(module, nn.InstanceNorm2d) or isinstance(module, nn.InstanceNorm1d) \ - or isinstance(module, nn.BatchNorm2d) or isinstance(module, nn.BatchNorm3d) or isinstance(module, - nn.BatchNorm1d): - print(str(module), module.training) - - -class Upsample(nn.Module): - def __init__(self, size=None, scale_factor=None, mode='nearest', align_corners=False): - super(Upsample, self).__init__() - self.align_corners = align_corners - self.mode = mode - self.scale_factor = scale_factor - self.size = size - - def forward(self, x): - return nn.functional.interpolate(x, size=self.size, scale_factor=self.scale_factor, mode=self.mode, - align_corners=self.align_corners) - - -class Generic_UNet(SegmentationNetwork): - DEFAULT_BATCH_SIZE_3D = 2 - DEFAULT_PATCH_SIZE_3D = (64, 192, 160) - SPACING_FACTOR_BETWEEN_STAGES = 2 - BASE_NUM_FEATURES_3D = 30 - MAX_NUMPOOL_3D = 999 - MAX_NUM_FILTERS_3D = 320 - - DEFAULT_PATCH_SIZE_2D = (256, 256) - BASE_NUM_FEATURES_2D = 30 - DEFAULT_BATCH_SIZE_2D = 50 - MAX_NUMPOOL_2D = 999 - MAX_FILTERS_2D = 480 - - use_this_for_batch_size_computation_2D = 19739648 - use_this_for_batch_size_computation_3D = 520000000 # 505789440 - - def __init__(self, input_channels, base_num_features, num_classes, num_pool, num_conv_per_stage=2, - feat_map_mul_on_downscale=2, conv_op=nn.Conv2d, - norm_op=nn.BatchNorm2d, norm_op_kwargs=None, - dropout_op=nn.Dropout2d, dropout_op_kwargs=None, - nonlin=nn.LeakyReLU, nonlin_kwargs=None, deep_supervision=True, dropout_in_localization=False, - final_nonlin=softmax_helper, weightInitializer=InitWeights_He(1e-2), pool_op_kernel_sizes=None, - conv_kernel_sizes=None, - upscale_logits=False, convolutional_pooling=False, convolutional_upsampling=False, - max_num_features=None, basic_block=ConvDropoutNormNonlin, - seg_output_use_bias=False): - """ - basically more flexible than v1, architecture is the same - - Does this look complicated? Nah bro. Functionality > usability - - This does everything you need, including world peace. - - Questions? -> f.isensee@dkfz.de - """ - super(Generic_UNet, self).__init__() - self.convolutional_upsampling = convolutional_upsampling - self.convolutional_pooling = convolutional_pooling - self.upscale_logits = upscale_logits - if nonlin_kwargs is None: - nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True} - if dropout_op_kwargs is None: - dropout_op_kwargs = {'p': 0.5, 'inplace': True} - if norm_op_kwargs is None: - norm_op_kwargs = {'eps': 1e-5, 'affine': True, 'momentum': 0.1} - - self.conv_kwargs = {'stride': 1, 'dilation': 1, 'bias': True} - - self.nonlin = nonlin - self.nonlin_kwargs = nonlin_kwargs - self.dropout_op_kwargs = dropout_op_kwargs - self.norm_op_kwargs = norm_op_kwargs - self.weightInitializer = weightInitializer - self.conv_op = conv_op - self.norm_op = norm_op - self.dropout_op = dropout_op - self.num_classes = num_classes - self.final_nonlin = final_nonlin - self._deep_supervision = deep_supervision - self.do_ds = deep_supervision - - if conv_op == nn.Conv2d: - upsample_mode = 'bilinear' - pool_op = nn.MaxPool2d - transpconv = nn.ConvTranspose2d - if pool_op_kernel_sizes is None: - pool_op_kernel_sizes = [(2, 2)] * num_pool - if conv_kernel_sizes is None: - conv_kernel_sizes = [(3, 3)] * (num_pool + 1) - elif conv_op == nn.Conv3d: - upsample_mode = 'trilinear' - pool_op = nn.MaxPool3d - transpconv = nn.ConvTranspose3d - if pool_op_kernel_sizes is None: - pool_op_kernel_sizes = [(2, 2, 2)] * num_pool - if conv_kernel_sizes is None: - conv_kernel_sizes = [(3, 3, 3)] * (num_pool + 1) - else: - raise ValueError( - "unknown convolution dimensionality, conv op: %s" % str(conv_op)) - - self.input_shape_must_be_divisible_by = np.prod( - pool_op_kernel_sizes, 0, dtype=np.int64) - self.pool_op_kernel_sizes = pool_op_kernel_sizes - self.conv_kernel_sizes = conv_kernel_sizes - - self.conv_pad_sizes = [] - for krnl in self.conv_kernel_sizes: - self.conv_pad_sizes.append([1 if i == 3 else 0 for i in krnl]) - - if max_num_features is None: - if self.conv_op == nn.Conv3d: - self.max_num_features = self.MAX_NUM_FILTERS_3D - else: - self.max_num_features = self.MAX_FILTERS_2D - else: - self.max_num_features = max_num_features - - self.conv_blocks_context = [] - self.conv_blocks_localization = [] - self.td = [] - self.tu = [] - self.seg_outputs = [] - - output_features = base_num_features - input_features = input_channels - - for d in range(num_pool): - # determine the first stride - if d != 0 and self.convolutional_pooling: - first_stride = pool_op_kernel_sizes[d - 1] - else: - first_stride = None - - self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[d] - self.conv_kwargs['padding'] = self.conv_pad_sizes[d] - # add convolutions - self.conv_blocks_context.append(StackedConvLayers(input_features, output_features, num_conv_per_stage, - self.conv_op, self.conv_kwargs, self.norm_op, - self.norm_op_kwargs, self.dropout_op, - self.dropout_op_kwargs, self.nonlin, self.nonlin_kwargs, - first_stride, basic_block=basic_block)) - if not self.convolutional_pooling: - self.td.append(pool_op(pool_op_kernel_sizes[d])) - input_features = output_features - output_features = int( - np.round(output_features * feat_map_mul_on_downscale)) - - output_features = min(output_features, self.max_num_features) - - # now the bottleneck. - # determine the first stride - if self.convolutional_pooling: - first_stride = pool_op_kernel_sizes[-1] - else: - first_stride = None - - # the output of the last conv must match the number of features from the skip connection if we are not using - # convolutional upsampling. If we use convolutional upsampling then the reduction in feature maps will be - # done by the transposed conv - if self.convolutional_upsampling: - final_num_features = output_features - else: - final_num_features = self.conv_blocks_context[-1].output_channels - - self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool] - self.conv_kwargs['padding'] = self.conv_pad_sizes[num_pool] - self.conv_blocks_context.append(nn.Sequential( - StackedConvLayers(input_features, output_features, num_conv_per_stage - 1, self.conv_op, self.conv_kwargs, - self.norm_op, self.norm_op_kwargs, self.dropout_op, self.dropout_op_kwargs, self.nonlin, - self.nonlin_kwargs, first_stride, basic_block=basic_block), - StackedConvLayers(output_features, final_num_features, 1, self.conv_op, self.conv_kwargs, - self.norm_op, self.norm_op_kwargs, self.dropout_op, self.dropout_op_kwargs, self.nonlin, - self.nonlin_kwargs, basic_block=basic_block))) - - # if we don't want to do dropout in the localization pathway then we set the dropout prob to zero here - if not dropout_in_localization: - old_dropout_p = self.dropout_op_kwargs['p'] - self.dropout_op_kwargs['p'] = 0.0 - - # now lets build the localization pathway - for u in range(num_pool): - nfeatures_from_down = final_num_features - nfeatures_from_skip = self.conv_blocks_context[ - -(2 + u)].output_channels # self.conv_blocks_context[-1] is bottleneck, so start with -2 - n_features_after_tu_and_concat = nfeatures_from_skip * 2 - - # the first conv reduces the number of features to match those of skip - # the following convs work on that number of features - # if not convolutional upsampling then the final conv reduces the num of features again - if u != num_pool - 1 and not self.convolutional_upsampling: - final_num_features = self.conv_blocks_context[-( - 3 + u)].output_channels - else: - final_num_features = nfeatures_from_skip - - if not self.convolutional_upsampling: - self.tu.append( - Upsample(scale_factor=pool_op_kernel_sizes[-(u + 1)], mode=upsample_mode)) - else: - self.tu.append(transpconv(nfeatures_from_down, nfeatures_from_skip, pool_op_kernel_sizes[-(u + 1)], - pool_op_kernel_sizes[-(u + 1)], bias=False)) - - self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[- (u + 1)] - self.conv_kwargs['padding'] = self.conv_pad_sizes[- (u + 1)] - self.conv_blocks_localization.append(nn.Sequential( - StackedConvLayers(n_features_after_tu_and_concat, nfeatures_from_skip, num_conv_per_stage - 1, - self.conv_op, self.conv_kwargs, self.norm_op, self.norm_op_kwargs, self.dropout_op, - self.dropout_op_kwargs, self.nonlin, self.nonlin_kwargs, basic_block=basic_block), - StackedConvLayers(nfeatures_from_skip, final_num_features, 1, self.conv_op, self.conv_kwargs, - self.norm_op, self.norm_op_kwargs, self.dropout_op, self.dropout_op_kwargs, - self.nonlin, self.nonlin_kwargs, basic_block=basic_block) - )) - - for ds in range(len(self.conv_blocks_localization)): - self.seg_outputs.append(conv_op(self.conv_blocks_localization[ds][-1].output_channels, num_classes, - 1, 1, 0, 1, 1, seg_output_use_bias)) - - self.upscale_logits_ops = [] - cum_upsample = np.cumprod( - np.vstack(pool_op_kernel_sizes), axis=0)[::-1] - for usl in range(num_pool - 1): - if self.upscale_logits: - self.upscale_logits_ops.append(Upsample(scale_factor=tuple([int(i) for i in cum_upsample[usl + 1]]), - mode=upsample_mode)) - else: - self.upscale_logits_ops.append(lambda x: x) - - if not dropout_in_localization: - self.dropout_op_kwargs['p'] = old_dropout_p - - # register all modules properly - self.conv_blocks_localization = nn.ModuleList( - self.conv_blocks_localization) - self.conv_blocks_context = nn.ModuleList(self.conv_blocks_context) - self.td = nn.ModuleList(self.td) - self.tu = nn.ModuleList(self.tu) - self.seg_outputs = nn.ModuleList(self.seg_outputs) - if self.upscale_logits: - self.upscale_logits_ops = nn.ModuleList( - self.upscale_logits_ops) # lambda x:x is not a Module so we need to distinguish here - - if self.weightInitializer is not None: - self.apply(self.weightInitializer) - # self.apply(print_module_training_status) - - def forward(self, x): - skips = [] - seg_outputs = [] - for d in range(len(self.conv_blocks_context) - 1): - x = self.conv_blocks_context[d](x) - skips.append(x) - if not self.convolutional_pooling: - x = self.td[d](x) - - x = self.conv_blocks_context[-1](x) - - for u in range(len(self.tu)): - x = self.tu[u](x) - x = torch.cat((x, skips[-(u + 1)]), dim=1) - x = self.conv_blocks_localization[u](x) - seg_outputs.append(self.final_nonlin(self.seg_outputs[u](x))) - - if self._deep_supervision and self.do_ds: - return tuple([seg_outputs[-1]] + [i(j) for i, j in - zip(list(self.upscale_logits_ops)[::-1], seg_outputs[:-1][::-1])]) - else: - return seg_outputs[-1] - - @staticmethod - def compute_approx_vram_consumption(patch_size, num_pool_per_axis, base_num_features, max_num_features, - num_modalities, num_classes, pool_op_kernel_sizes, deep_supervision=False, - conv_per_stage=2): - """ - This only applies for num_conv_per_stage and convolutional_upsampling=True - not real vram consumption. just a constant term to which the vram consumption will be approx proportional - (+ offset for parameter storage) - :param deep_supervision: - :param patch_size: - :param num_pool_per_axis: - :param base_num_features: - :param max_num_features: - :param num_modalities: - :param num_classes: - :param pool_op_kernel_sizes: - :return: - """ - if not isinstance(num_pool_per_axis, np.ndarray): - num_pool_per_axis = np.array(num_pool_per_axis) - - npool = len(pool_op_kernel_sizes) - - map_size = np.array(patch_size) - tmp = np.int64((conv_per_stage * 2 + 1) * np.prod(map_size, dtype=np.int64) * base_num_features + - num_modalities * np.prod(map_size, dtype=np.int64) + - num_classes * np.prod(map_size, dtype=np.int64)) - - num_feat = base_num_features - - for p in range(npool): - for pi in range(len(num_pool_per_axis)): - map_size[pi] /= pool_op_kernel_sizes[p][pi] - num_feat = min(num_feat * 2, max_num_features) - # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv - num_blocks = (conv_per_stage * 2 + - 1) if p < (npool - 1) else conv_per_stage - tmp += num_blocks * np.prod(map_size, dtype=np.int64) * num_feat - if deep_supervision and p < (npool - 2): - tmp += np.prod(map_size, dtype=np.int64) * num_classes - # print(p, map_size, num_feat, tmp) - return tmp - - -default_dict = { - "base_num_features": 16, - "conv_per_stage": 2, - "initial_lr": 0.01, - "lr_scheduler": None, - "lr_scheduler_eps": 0.001, - "lr_scheduler_patience": 30, - "lr_threshold": 1e-06, - "max_num_epochs": 1000, - "net_conv_kernel_sizes": [[1, 3, 3], [1, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3], [3, 3, 3]], - "net_num_pool_op_kernel_sizes": [[1, 2, 2], [1, 2, 2], [2, 2, 2], [2, 2, 2], [1, 2, 2], [1, 2, 2]], - "net_pool_per_axis": [2, 6, 6], - "num_batches_per_epoch": 250, - "num_classes": 3, - "num_input_channels": 1, - "transpose_backward": [0, 1, 2], - "transpose_forward": [0, 1, 2], -} - - -def initialize_network(threeD=True, num_classes=2): - """ - This is specific to the U-Net and must be adapted for other network architectures - :return: - """ - # self.print_to_log_file(self.net_num_pool_op_kernel_sizes) - # self.print_to_log_file(self.net_conv_kernel_sizes) - - if threeD: - conv_op = nn.Conv3d - dropout_op = nn.Dropout3d - norm_op = nn.InstanceNorm3d - else: - conv_op = nn.Conv2d - dropout_op = nn.Dropout2d - norm_op = nn.InstanceNorm2d - default_dict["num_classes"] = num_classes - norm_op_kwargs = {'eps': 1e-5, 'affine': True} - dropout_op_kwargs = {'p': 0, 'inplace': True} - net_nonlin = nn.LeakyReLU - net_nonlin_kwargs = {'negative_slope': 1e-2, 'inplace': True} - network = Generic_UNet(default_dict["num_input_channels"], default_dict["base_num_features"], default_dict["num_classes"], len(default_dict["net_num_pool_op_kernel_sizes"]), - default_dict["conv_per_stage"], 2, conv_op, norm_op, norm_op_kwargs, dropout_op, - dropout_op_kwargs, - net_nonlin, net_nonlin_kwargs, False, False, lambda x: x, InitWeights_He( - 1e-2), - default_dict["net_num_pool_op_kernel_sizes"], default_dict["net_conv_kernel_sizes"], False, True, True) - print("nnUNet have {} paramerters in total".format( - sum(x.numel() for x in network.parameters()))) - return network.cuda() - -# input = torch.FloatTensor(1, 1, 32, 192, 192) -# input_var = input.cuda() -# model = initialize_network(threeD=True) -# out = model(input_var) -# print(out.size()) \ No newline at end of file diff --git a/code/networks/swin_transformer_unet_skip_expand_decoder_sys.py b/code/networks/swin_transformer_unet_skip_expand_decoder_sys.py deleted file mode 100644 index a85885e..0000000 --- a/code/networks/swin_transformer_unet_skip_expand_decoder_sys.py +++ /dev/null @@ -1,804 +0,0 @@ -# This file borrowed from Swin-UNet: https://github.com/HuCaoFighting/Swin-Unet -import torch -import torch.nn as nn -import torch.utils.checkpoint as checkpoint -from einops import rearrange -from timm.models.layers import DropPath, to_2tuple, trunc_normal_ - - -class Mlp(nn.Module): - def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features - self.fc1 = nn.Linear(in_features, hidden_features) - self.act = act_layer() - self.fc2 = nn.Linear(hidden_features, out_features) - self.drop = nn.Dropout(drop) - - def forward(self, x): - x = self.fc1(x) - x = self.act(x) - x = self.drop(x) - x = self.fc2(x) - x = self.drop(x) - return x - - -def window_partition(x, window_size): - """ - Args: - x: (B, H, W, C) - window_size (int): window size - - Returns: - windows: (num_windows*B, window_size, window_size, C) - """ - B, H, W, C = x.shape - x = x.view(B, H // window_size, window_size, - W // window_size, window_size, C) - windows = x.permute(0, 1, 3, 2, 4, 5).contiguous( - ).view(-1, window_size, window_size, C) - return windows - - -def window_reverse(windows, window_size, H, W): - """ - Args: - windows: (num_windows*B, window_size, window_size, C) - window_size (int): Window size - H (int): Height of image - W (int): Width of image - - Returns: - x: (B, H, W, C) - """ - B = int(windows.shape[0] / (H * W / window_size / window_size)) - x = windows.view(B, H // window_size, W // window_size, - window_size, window_size, -1) - x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) - return x - - -class WindowAttention(nn.Module): - r""" Window based multi-head self attention (W-MSA) module with relative position bias. - It supports both of shifted and non-shifted window. - - Args: - dim (int): Number of input channels. - window_size (tuple[int]): The height and width of the window. - num_heads (int): Number of attention heads. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set - attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 - proj_drop (float, optional): Dropout ratio of output. Default: 0.0 - """ - - def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): - - super().__init__() - self.dim = dim - self.window_size = window_size # Wh, Ww - self.num_heads = num_heads - head_dim = dim // num_heads - self.scale = qk_scale or head_dim ** -0.5 - - # define a parameter table of relative position bias - self.relative_position_bias_table = nn.Parameter( - torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH - - # get pair-wise relative position index for each token inside the window - coords_h = torch.arange(self.window_size[0]) - coords_w = torch.arange(self.window_size[1]) - coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww - coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww - relative_coords = coords_flatten[:, :, None] - \ - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww - relative_coords = relative_coords.permute( - 1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 - relative_coords[:, :, 0] += self.window_size[0] - \ - 1 # shift to start from 0 - relative_coords[:, :, 1] += self.window_size[1] - 1 - relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 - relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww - self.register_buffer("relative_position_index", - relative_position_index) - - self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) - self.attn_drop = nn.Dropout(attn_drop) - self.proj = nn.Linear(dim, dim) - self.proj_drop = nn.Dropout(proj_drop) - - trunc_normal_(self.relative_position_bias_table, std=.02) - self.softmax = nn.Softmax(dim=-1) - - def forward(self, x, mask=None): - """ - Args: - x: input features with shape of (num_windows*B, N, C) - mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None - """ - B_, N, C = x.shape - qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // - self.num_heads).permute(2, 0, 3, 1, 4) - # make torchscript happy (cannot use tensor as tuple) - q, k, v = qkv[0], qkv[1], qkv[2] - - q = q * self.scale - attn = (q @ k.transpose(-2, -1)) - - relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( - self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH - relative_position_bias = relative_position_bias.permute( - 2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww - attn = attn + relative_position_bias.unsqueeze(0) - - if mask is not None: - nW = mask.shape[0] - attn = attn.view(B_ // nW, nW, self.num_heads, N, - N) + mask.unsqueeze(1).unsqueeze(0) - attn = attn.view(-1, self.num_heads, N, N) - attn = self.softmax(attn) - else: - attn = self.softmax(attn) - - attn = self.attn_drop(attn) - - x = (attn @ v).transpose(1, 2).reshape(B_, N, C) - x = self.proj(x) - x = self.proj_drop(x) - return x - - def extra_repr(self) -> str: - return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}' - - def flops(self, N): - # calculate flops for 1 window with token length of N - flops = 0 - # qkv = self.qkv(x) - flops += N * self.dim * 3 * self.dim - # attn = (q @ k.transpose(-2, -1)) - flops += self.num_heads * N * (self.dim // self.num_heads) * N - # x = (attn @ v) - flops += self.num_heads * N * N * (self.dim // self.num_heads) - # x = self.proj(x) - flops += N * self.dim * self.dim - return flops - - -class SwinTransformerBlock(nn.Module): - r""" Swin Transformer Block. - - Args: - dim (int): Number of input channels. - input_resolution (tuple[int]): Input resulotion. - num_heads (int): Number of attention heads. - window_size (int): Window size. - shift_size (int): Shift size for SW-MSA. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. - drop (float, optional): Dropout rate. Default: 0.0 - attn_drop (float, optional): Attention dropout rate. Default: 0.0 - drop_path (float, optional): Stochastic depth rate. Default: 0.0 - act_layer (nn.Module, optional): Activation layer. Default: nn.GELU - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - """ - - def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0, - mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., - act_layer=nn.GELU, norm_layer=nn.LayerNorm): - super().__init__() - self.dim = dim - self.input_resolution = input_resolution - self.num_heads = num_heads - self.window_size = window_size - self.shift_size = shift_size - self.mlp_ratio = mlp_ratio - if min(self.input_resolution) <= self.window_size: - # if window size is larger than input resolution, we don't partition windows - self.shift_size = 0 - self.window_size = min(self.input_resolution) - assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" - - self.norm1 = norm_layer(dim) - self.attn = WindowAttention( - dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, - qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) - - self.drop_path = DropPath( - drop_path) if drop_path > 0. else nn.Identity() - self.norm2 = norm_layer(dim) - mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, - act_layer=act_layer, drop=drop) - - if self.shift_size > 0: - # calculate attention mask for SW-MSA - H, W = self.input_resolution - img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1 - h_slices = (slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None)) - w_slices = (slice(0, -self.window_size), - slice(-self.window_size, -self.shift_size), - slice(-self.shift_size, None)) - cnt = 0 - for h in h_slices: - for w in w_slices: - img_mask[:, h, w, :] = cnt - cnt += 1 - - # nW, window_size, window_size, 1 - mask_windows = window_partition(img_mask, self.window_size) - mask_windows = mask_windows.view(-1, - self.window_size * self.window_size) - attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) - attn_mask = attn_mask.masked_fill( - attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) - else: - attn_mask = None - - self.register_buffer("attn_mask", attn_mask) - - def forward(self, x): - H, W = self.input_resolution - B, L, C = x.shape - assert L == H * W, "input feature has wrong size" - - shortcut = x - x = self.norm1(x) - x = x.view(B, H, W, C) - - # cyclic shift - if self.shift_size > 0: - shifted_x = torch.roll( - x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) - else: - shifted_x = x - - # partition windows - # nW*B, window_size, window_size, C - x_windows = window_partition(shifted_x, self.window_size) - # nW*B, window_size*window_size, C - x_windows = x_windows.view(-1, self.window_size * self.window_size, C) - - # W-MSA/SW-MSA - # nW*B, window_size*window_size, C - attn_windows = self.attn(x_windows, mask=self.attn_mask) - - # merge windows - attn_windows = attn_windows.view(-1, - self.window_size, self.window_size, C) - shifted_x = window_reverse( - attn_windows, self.window_size, H, W) # B H' W' C - - # reverse cyclic shift - if self.shift_size > 0: - x = torch.roll(shifted_x, shifts=( - self.shift_size, self.shift_size), dims=(1, 2)) - else: - x = shifted_x - x = x.view(B, H * W, C) - - # FFN - x = shortcut + self.drop_path(x) - x = x + self.drop_path(self.mlp(self.norm2(x))) - - return x - - def extra_repr(self) -> str: - return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \ - f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" - - def flops(self): - flops = 0 - H, W = self.input_resolution - # norm1 - flops += self.dim * H * W - # W-MSA/SW-MSA - nW = H * W / self.window_size / self.window_size - flops += nW * self.attn.flops(self.window_size * self.window_size) - # mlp - flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio - # norm2 - flops += self.dim * H * W - return flops - - -class PatchMerging(nn.Module): - r""" Patch Merging Layer. - - Args: - input_resolution (tuple[int]): Resolution of input feature. - dim (int): Number of input channels. - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - """ - - def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm): - super().__init__() - self.input_resolution = input_resolution - self.dim = dim - self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) - self.norm = norm_layer(4 * dim) - - def forward(self, x): - """ - x: B, H*W, C - """ - H, W = self.input_resolution - B, L, C = x.shape - assert L == H * W, "input feature has wrong size" - assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even." - - x = x.view(B, H, W, C) - - x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C - x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C - x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C - x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C - x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C - x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C - - x = self.norm(x) - x = self.reduction(x) - - return x - - def extra_repr(self) -> str: - return f"input_resolution={self.input_resolution}, dim={self.dim}" - - def flops(self): - H, W = self.input_resolution - flops = H * W * self.dim - flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim - return flops - - -class PatchExpand(nn.Module): - def __init__(self, input_resolution, dim, dim_scale=2, norm_layer=nn.LayerNorm): - super().__init__() - self.input_resolution = input_resolution - self.dim = dim - self.expand = nn.Linear( - dim, 2*dim, bias=False) if dim_scale == 2 else nn.Identity() - self.norm = norm_layer(dim // dim_scale) - - def forward(self, x): - """ - x: B, H*W, C - """ - H, W = self.input_resolution - x = self.expand(x) - B, L, C = x.shape - assert L == H * W, "input feature has wrong size" - - x = x.view(B, H, W, C) - x = rearrange(x, 'b h w (p1 p2 c)-> b (h p1) (w p2) c', - p1=2, p2=2, c=C//4) - x = x.view(B, -1, C//4) - x = self.norm(x) - - return x - - -class FinalPatchExpand_X4(nn.Module): - def __init__(self, input_resolution, dim, dim_scale=4, norm_layer=nn.LayerNorm): - super().__init__() - self.input_resolution = input_resolution - self.dim = dim - self.dim_scale = dim_scale - self.expand = nn.Linear(dim, 16*dim, bias=False) - self.output_dim = dim - self.norm = norm_layer(self.output_dim) - - def forward(self, x): - """ - x: B, H*W, C - """ - H, W = self.input_resolution - x = self.expand(x) - B, L, C = x.shape - assert L == H * W, "input feature has wrong size" - - x = x.view(B, H, W, C) - x = rearrange(x, 'b h w (p1 p2 c)-> b (h p1) (w p2) c', - p1=self.dim_scale, p2=self.dim_scale, c=C//(self.dim_scale**2)) - x = x.view(B, -1, self.output_dim) - x = self.norm(x) - - return x - - -class BasicLayer(nn.Module): - """ A basic Swin Transformer layer for one stage. - - Args: - dim (int): Number of input channels. - input_resolution (tuple[int]): Input resolution. - depth (int): Number of blocks. - num_heads (int): Number of attention heads. - window_size (int): Local window size. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. - drop (float, optional): Dropout rate. Default: 0.0 - attn_drop (float, optional): Attention dropout rate. Default: 0.0 - drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None - use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. - """ - - def __init__(self, dim, input_resolution, depth, num_heads, window_size, - mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False): - - super().__init__() - self.dim = dim - self.input_resolution = input_resolution - self.depth = depth - self.use_checkpoint = use_checkpoint - - # build blocks - self.blocks = nn.ModuleList([ - SwinTransformerBlock(dim=dim, input_resolution=input_resolution, - num_heads=num_heads, window_size=window_size, - shift_size=0 if ( - i % 2 == 0) else window_size // 2, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop, attn_drop=attn_drop, - drop_path=drop_path[i] if isinstance( - drop_path, list) else drop_path, - norm_layer=norm_layer) - for i in range(depth)]) - - # patch merging layer - if downsample is not None: - self.downsample = downsample( - input_resolution, dim=dim, norm_layer=norm_layer) - else: - self.downsample = None - - def forward(self, x): - for blk in self.blocks: - if self.use_checkpoint: - x = checkpoint.checkpoint(blk, x) - else: - x = blk(x) - if self.downsample is not None: - x = self.downsample(x) - return x - - def extra_repr(self) -> str: - return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}" - - def flops(self): - flops = 0 - for blk in self.blocks: - flops += blk.flops() - if self.downsample is not None: - flops += self.downsample.flops() - return flops - - -class BasicLayer_up(nn.Module): - """ A basic Swin Transformer layer for one stage. - - Args: - dim (int): Number of input channels. - input_resolution (tuple[int]): Input resolution. - depth (int): Number of blocks. - num_heads (int): Number of attention heads. - window_size (int): Local window size. - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. - qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. - drop (float, optional): Dropout rate. Default: 0.0 - attn_drop (float, optional): Attention dropout rate. Default: 0.0 - drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 - norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm - downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None - use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. - """ - - def __init__(self, dim, input_resolution, depth, num_heads, window_size, - mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., - drop_path=0., norm_layer=nn.LayerNorm, upsample=None, use_checkpoint=False): - - super().__init__() - self.dim = dim - self.input_resolution = input_resolution - self.depth = depth - self.use_checkpoint = use_checkpoint - - # build blocks - self.blocks = nn.ModuleList([ - SwinTransformerBlock(dim=dim, input_resolution=input_resolution, - num_heads=num_heads, window_size=window_size, - shift_size=0 if ( - i % 2 == 0) else window_size // 2, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop, attn_drop=attn_drop, - drop_path=drop_path[i] if isinstance( - drop_path, list) else drop_path, - norm_layer=norm_layer) - for i in range(depth)]) - - # patch merging layer - if upsample is not None: - self.upsample = PatchExpand( - input_resolution, dim=dim, dim_scale=2, norm_layer=norm_layer) - else: - self.upsample = None - - def forward(self, x): - for blk in self.blocks: - if self.use_checkpoint: - x = checkpoint.checkpoint(blk, x) - else: - x = blk(x) - if self.upsample is not None: - x = self.upsample(x) - return x - - -class PatchEmbed(nn.Module): - r""" Image to Patch Embedding - - Args: - img_size (int): Image size. Default: 224. - patch_size (int): Patch token size. Default: 4. - in_chans (int): Number of input image channels. Default: 3. - embed_dim (int): Number of linear projection output channels. Default: 96. - norm_layer (nn.Module, optional): Normalization layer. Default: None - """ - - def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): - super().__init__() - img_size = to_2tuple(img_size) - patch_size = to_2tuple(patch_size) - patches_resolution = [img_size[0] // - patch_size[0], img_size[1] // patch_size[1]] - self.img_size = img_size - self.patch_size = patch_size - self.patches_resolution = patches_resolution - self.num_patches = patches_resolution[0] * patches_resolution[1] - - self.in_chans = in_chans - self.embed_dim = embed_dim - - self.proj = nn.Conv2d(in_chans, embed_dim, - kernel_size=patch_size, stride=patch_size) - if norm_layer is not None: - self.norm = norm_layer(embed_dim) - else: - self.norm = None - - def forward(self, x): - B, C, H, W = x.shape - # FIXME look at relaxing size constraints - assert H == self.img_size[0] and W == self.img_size[1], \ - f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." - x = self.proj(x).flatten(2).transpose(1, 2) # B Ph*Pw C - if self.norm is not None: - x = self.norm(x) - return x - - def flops(self): - Ho, Wo = self.patches_resolution - flops = Ho * Wo * self.embed_dim * self.in_chans * \ - (self.patch_size[0] * self.patch_size[1]) - if self.norm is not None: - flops += Ho * Wo * self.embed_dim - return flops - - -class SwinTransformerSys(nn.Module): - r""" Swin Transformer - A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - - https://arxiv.org/pdf/2103.14030 - - Args: - img_size (int | tuple(int)): Input image size. Default 224 - patch_size (int | tuple(int)): Patch size. Default: 4 - in_chans (int): Number of input image channels. Default: 3 - num_classes (int): Number of classes for classification head. Default: 1000 - embed_dim (int): Patch embedding dimension. Default: 96 - depths (tuple(int)): Depth of each Swin Transformer layer. - num_heads (tuple(int)): Number of attention heads in different layers. - window_size (int): Window size. Default: 7 - mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4 - qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True - qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None - drop_rate (float): Dropout rate. Default: 0 - attn_drop_rate (float): Attention dropout rate. Default: 0 - drop_path_rate (float): Stochastic depth rate. Default: 0.1 - norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. - ape (bool): If True, add absolute position embedding to the patch embedding. Default: False - patch_norm (bool): If True, add normalization after patch embedding. Default: True - use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False - """ - - def __init__(self, img_size=224, patch_size=4, in_chans=3, num_classes=1000, - embed_dim=96, depths=[2, 2, 2, 2], depths_decoder=[1, 2, 2, 2], num_heads=[3, 6, 12, 24], - window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None, - drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1, - norm_layer=nn.LayerNorm, ape=False, patch_norm=True, - use_checkpoint=False, final_upsample="expand_first", **kwargs): - super().__init__() - - print("SwinTransformerSys expand initial----depths:{};depths_decoder:{};drop_path_rate:{};num_classes:{}".format(depths, - depths_decoder, drop_path_rate, num_classes)) - - self.num_classes = num_classes - self.num_layers = len(depths) - self.embed_dim = embed_dim - self.ape = ape - self.patch_norm = patch_norm - self.num_features = int(embed_dim * 2 ** (self.num_layers - 1)) - self.num_features_up = int(embed_dim * 2) - self.mlp_ratio = mlp_ratio - self.final_upsample = final_upsample - - # split image into non-overlapping patches - self.patch_embed = PatchEmbed( - img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, - norm_layer=norm_layer if self.patch_norm else None) - num_patches = self.patch_embed.num_patches - patches_resolution = self.patch_embed.patches_resolution - self.patches_resolution = patches_resolution - - # absolute position embedding - if self.ape: - self.absolute_pos_embed = nn.Parameter( - torch.zeros(1, num_patches, embed_dim)) - trunc_normal_(self.absolute_pos_embed, std=.02) - - self.pos_drop = nn.Dropout(p=drop_rate) - - # stochastic depth - dpr = [x.item() for x in torch.linspace(0, drop_path_rate, - sum(depths))] # stochastic depth decay rule - - # build encoder and bottleneck layers - self.layers = nn.ModuleList() - for i_layer in range(self.num_layers): - layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer), - input_resolution=(patches_resolution[0] // (2 ** i_layer), - patches_resolution[1] // (2 ** i_layer)), - depth=depths[i_layer], - num_heads=num_heads[i_layer], - window_size=window_size, - mlp_ratio=self.mlp_ratio, - qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, - drop_path=dpr[sum(depths[:i_layer]):sum( - depths[:i_layer + 1])], - norm_layer=norm_layer, - downsample=PatchMerging if ( - i_layer < self.num_layers - 1) else None, - use_checkpoint=use_checkpoint) - self.layers.append(layer) - - # build decoder layers - self.layers_up = nn.ModuleList() - self.concat_back_dim = nn.ModuleList() - for i_layer in range(self.num_layers): - concat_linear = nn.Linear(2*int(embed_dim*2**(self.num_layers-1-i_layer)), - int(embed_dim*2**(self.num_layers-1-i_layer))) if i_layer > 0 else nn.Identity() - if i_layer == 0: - layer_up = PatchExpand(input_resolution=(patches_resolution[0] // (2 ** (self.num_layers-1-i_layer)), - patches_resolution[1] // (2 ** (self.num_layers-1-i_layer))), dim=int(embed_dim * 2 ** (self.num_layers-1-i_layer)), dim_scale=2, norm_layer=norm_layer) - else: - layer_up = BasicLayer_up(dim=int(embed_dim * 2 ** (self.num_layers-1-i_layer)), - input_resolution=(patches_resolution[0] // (2 ** (self.num_layers-1-i_layer)), - patches_resolution[1] // (2 ** (self.num_layers-1-i_layer))), - depth=depths[( - self.num_layers-1-i_layer)], - num_heads=num_heads[( - self.num_layers-1-i_layer)], - window_size=window_size, - mlp_ratio=self.mlp_ratio, - qkv_bias=qkv_bias, qk_scale=qk_scale, - drop=drop_rate, attn_drop=attn_drop_rate, - drop_path=dpr[sum(depths[:( - self.num_layers-1-i_layer)]):sum(depths[:(self.num_layers-1-i_layer) + 1])], - norm_layer=norm_layer, - upsample=PatchExpand if ( - i_layer < self.num_layers - 1) else None, - use_checkpoint=use_checkpoint) - self.layers_up.append(layer_up) - self.concat_back_dim.append(concat_linear) - - self.norm = norm_layer(self.num_features) - self.norm_up = norm_layer(self.embed_dim) - - if self.final_upsample == "expand_first": - print("---final upsample expand_first---") - self.up = FinalPatchExpand_X4(input_resolution=( - img_size//patch_size, img_size//patch_size), dim_scale=4, dim=embed_dim) - self.output = nn.Conv2d( - in_channels=embed_dim, out_channels=self.num_classes, kernel_size=1, bias=False) - - self.apply(self._init_weights) - - def _init_weights(self, m): - if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=.02) - if isinstance(m, nn.Linear) and m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.LayerNorm): - nn.init.constant_(m.bias, 0) - nn.init.constant_(m.weight, 1.0) - - @torch.jit.ignore - def no_weight_decay(self): - return {'absolute_pos_embed'} - - @torch.jit.ignore - def no_weight_decay_keywords(self): - return {'relative_position_bias_table'} - - #Encoder and Bottleneck - def forward_features(self, x): - x = self.patch_embed(x) - if self.ape: - x = x + self.absolute_pos_embed - x = self.pos_drop(x) - x_downsample = [] - - for layer in self.layers: - x_downsample.append(x) - x = layer(x) - - x = self.norm(x) # B L C - - return x, x_downsample - - # Dencoder and Skip connection - def forward_up_features(self, x, x_downsample): - for inx, layer_up in enumerate(self.layers_up): - if inx == 0: - x = layer_up(x) - else: - x = torch.cat([x, x_downsample[3-inx]], -1) - x = self.concat_back_dim[inx](x) - x = layer_up(x) - - x = self.norm_up(x) # B L C - - return x - - def up_x4(self, x): - H, W = self.patches_resolution - B, L, C = x.shape - assert L == H*W, "input features has wrong size" - - if self.final_upsample == "expand_first": - x = self.up(x) - x = x.view(B, 4*H, 4*W, -1) - x = x.permute(0, 3, 1, 2) # B,C,H,W - x = self.output(x) - - return x - - def forward(self, x): - x, x_downsample = self.forward_features(x) - x = self.forward_up_features(x, x_downsample) - x = self.up_x4(x) - - return x - - def flops(self): - flops = 0 - flops += self.patch_embed.flops() - for i, layer in enumerate(self.layers): - flops += layer.flops() - flops += self.num_features * \ - self.patches_resolution[0] * \ - self.patches_resolution[1] // (2 ** self.num_layers) - flops += self.num_features * self.num_classes - return flops diff --git a/code/networks/unet.py b/code/networks/unet.py old mode 100644 new mode 100755 index 64c28c8..4ab4b9d --- a/code/networks/unet.py +++ b/code/networks/unet.py @@ -9,25 +9,7 @@ import torch.nn as nn from torch.distributions.uniform import Uniform -def kaiming_normal_init_weight(model): - for m in model.modules(): - if isinstance(m, nn.Conv3d): - torch.nn.init.kaiming_normal_(m.weight) - elif isinstance(m, nn.BatchNorm3d): - m.weight.data.fill_(1) - m.bias.data.zero_() - return model - -def sparse_init_weight(model): - for m in model.modules(): - if isinstance(m, nn.Conv3d): - torch.nn.init.sparse_(m.weight, sparsity=0.1) - elif isinstance(m, nn.BatchNorm3d): - m.weight.data.fill_(1) - m.bias.data.zero_() - return model - - + class ConvBlock(nn.Module): """two convolution layers with batch norm and leaky relu""" @@ -206,9 +188,9 @@ def forward(self, feature, shape): return dp0_out_seg, dp1_out_seg, dp2_out_seg, dp3_out_seg -class Decoder_URPC(nn.Module): +class Decoder_URDS(nn.Module): def __init__(self, params): - super(Decoder_URPC, self).__init__() + super(Decoder_URDS, self).__init__() self.params = params self.in_chns = self.params['in_chns'] self.ft_chns = self.params['feature_chns'] @@ -269,8 +251,8 @@ def forward(self, feature, shape): return dp0_out_seg, dp1_out_seg, dp2_out_seg, dp3_out_seg -def Dropout(x, p=0.3): - x = torch.nn.functional.dropout(x, p) +def Dropout(x, p=0.5): + x = torch.nn.functional.dropout2d(x, p) return x @@ -321,37 +303,9 @@ def forward(self, x): return output -class UNet_CCT(nn.Module): - def __init__(self, in_chns, class_num): - super(UNet_CCT, self).__init__() - - params = {'in_chns': in_chns, - 'feature_chns': [16, 32, 64, 128, 256], - 'dropout': [0.05, 0.1, 0.2, 0.3, 0.5], - 'class_num': class_num, - 'bilinear': False, - 'acti_func': 'relu'} - self.encoder = Encoder(params) - self.main_decoder = Decoder(params) - self.aux_decoder1 = Decoder(params) - self.aux_decoder2 = Decoder(params) - self.aux_decoder3 = Decoder(params) - - def forward(self, x): - feature = self.encoder(x) - main_seg = self.main_decoder(feature) - aux1_feature = [FeatureNoise()(i) for i in feature] - aux_seg1 = self.aux_decoder1(aux1_feature) - aux2_feature = [Dropout(i) for i in feature] - aux_seg2 = self.aux_decoder2(aux2_feature) - aux3_feature = [FeatureDropout(i) for i in feature] - aux_seg3 = self.aux_decoder3(aux3_feature) - return main_seg, aux_seg1, aux_seg2, aux_seg3 - - -class UNet_URPC(nn.Module): +class UNet_DS(nn.Module): def __init__(self, in_chns, class_num): - super(UNet_URPC, self).__init__() + super(UNet_DS, self).__init__() params = {'in_chns': in_chns, 'feature_chns': [16, 32, 64, 128, 256], @@ -360,19 +314,19 @@ def __init__(self, in_chns, class_num): 'bilinear': False, 'acti_func': 'relu'} self.encoder = Encoder(params) - self.decoder = Decoder_URPC(params) + self.decoder = Decoder_DS(params) def forward(self, x): shape = x.shape[2:] feature = self.encoder(x) - dp1_out_seg, dp2_out_seg, dp3_out_seg, dp4_out_seg = self.decoder( + dp0_out_seg, dp1_out_seg, dp2_out_seg, dp3_out_seg = self.decoder( feature, shape) - return dp1_out_seg, dp2_out_seg, dp3_out_seg, dp4_out_seg + return dp0_out_seg, dp1_out_seg, dp2_out_seg, dp3_out_seg -class UNet_DS(nn.Module): +class UNet_URDS(nn.Module): def __init__(self, in_chns, class_num): - super(UNet_DS, self).__init__() + super(UNet_URDS, self).__init__() params = {'in_chns': in_chns, 'feature_chns': [16, 32, 64, 128, 256], @@ -381,13 +335,11 @@ def __init__(self, in_chns, class_num): 'bilinear': False, 'acti_func': 'relu'} self.encoder = Encoder(params) - self.decoder = Decoder_DS(params) + self.decoder = Decoder_URDS(params) def forward(self, x): shape = x.shape[2:] feature = self.encoder(x) - dp0_out_seg, dp1_out_seg, dp2_out_seg, dp3_out_seg = self.decoder( + dp1_out_seg, dp2_out_seg, dp3_out_seg, dp4_out_seg = self.decoder( feature, shape) - return dp0_out_seg, dp1_out_seg, dp2_out_seg, dp3_out_seg - - + return dp1_out_seg, dp2_out_seg, dp3_out_seg, dp4_out_seg diff --git a/code/networks/unet_3D.py b/code/networks/unet_3D.py old mode 100644 new mode 100755 diff --git a/code/networks/unet_3D_dv_semi.py b/code/networks/unet_3D_dv_semi.py deleted file mode 100644 index c8aa058..0000000 --- a/code/networks/unet_3D_dv_semi.py +++ /dev/null @@ -1,112 +0,0 @@ -""" -This file is adapted from https://github.com/ozan-oktay/Attention-Gated-Networks -""" - -import math -import torch -import torch.nn as nn -from networks.utils import UnetConv3, UnetUp3, UnetUp3_CT, UnetDsv3 -import torch.nn.functional as F -from networks.networks_other import init_weights - - -class unet_3D_dv_semi(nn.Module): - - def __init__(self, feature_scale=4, n_classes=21, is_deconv=True, in_channels=3, is_batchnorm=True): - super(unet_3D_dv_semi, self).__init__() - self.is_deconv = is_deconv - self.in_channels = in_channels - self.is_batchnorm = is_batchnorm - self.feature_scale = feature_scale - - filters = [64, 128, 256, 512, 1024] - filters = [int(x / self.feature_scale) for x in filters] - - # downsampling - self.conv1 = UnetConv3(self.in_channels, filters[0], self.is_batchnorm, kernel_size=( - 3, 3, 3), padding_size=(1, 1, 1)) - self.maxpool1 = nn.MaxPool3d(kernel_size=(2, 2, 2)) - - self.conv2 = UnetConv3(filters[0], filters[1], self.is_batchnorm, kernel_size=( - 3, 3, 3), padding_size=(1, 1, 1)) - self.maxpool2 = nn.MaxPool3d(kernel_size=(2, 2, 2)) - - self.conv3 = UnetConv3(filters[1], filters[2], self.is_batchnorm, kernel_size=( - 3, 3, 3), padding_size=(1, 1, 1)) - self.maxpool3 = nn.MaxPool3d(kernel_size=(2, 2, 2)) - - self.conv4 = UnetConv3(filters[2], filters[3], self.is_batchnorm, kernel_size=( - 3, 3, 3), padding_size=(1, 1, 1)) - self.maxpool4 = nn.MaxPool3d(kernel_size=(2, 2, 2)) - - self.center = UnetConv3(filters[3], filters[4], self.is_batchnorm, kernel_size=( - 3, 3, 3), padding_size=(1, 1, 1)) - - # upsampling - self.up_concat4 = UnetUp3_CT(filters[4], filters[3], is_batchnorm) - self.up_concat3 = UnetUp3_CT(filters[3], filters[2], is_batchnorm) - self.up_concat2 = UnetUp3_CT(filters[2], filters[1], is_batchnorm) - self.up_concat1 = UnetUp3_CT(filters[1], filters[0], is_batchnorm) - - # deep supervision - self.dsv4 = UnetDsv3( - in_size=filters[3], out_size=n_classes, scale_factor=8) - self.dsv3 = UnetDsv3( - in_size=filters[2], out_size=n_classes, scale_factor=4) - self.dsv2 = UnetDsv3( - in_size=filters[1], out_size=n_classes, scale_factor=2) - self.dsv1 = nn.Conv3d( - in_channels=filters[0], out_channels=n_classes, kernel_size=1) - - self.dropout1 = nn.Dropout3d(p=0.5) - self.dropout2 = nn.Dropout3d(p=0.3) - self.dropout3 = nn.Dropout3d(p=0.2) - self.dropout4 = nn.Dropout3d(p=0.1) - - # initialise weights - for m in self.modules(): - if isinstance(m, nn.Conv3d): - init_weights(m, init_type='kaiming') - elif isinstance(m, nn.BatchNorm3d): - init_weights(m, init_type='kaiming') - - def forward(self, inputs): - conv1 = self.conv1(inputs) - maxpool1 = self.maxpool1(conv1) - - conv2 = self.conv2(maxpool1) - maxpool2 = self.maxpool2(conv2) - - conv3 = self.conv3(maxpool2) - maxpool3 = self.maxpool3(conv3) - - conv4 = self.conv4(maxpool3) - maxpool4 = self.maxpool4(conv4) - - center = self.center(maxpool4) - - up4 = self.up_concat4(conv4, center) - up4 = self.dropout1(up4) - - up3 = self.up_concat3(conv3, up4) - up3 = self.dropout2(up3) - - up2 = self.up_concat2(conv2, up3) - up2 = self.dropout3(up2) - - up1 = self.up_concat1(conv1, up2) - up1 = self.dropout4(up1) - - # Deep Supervision - dsv4 = self.dsv4(up4) - dsv3 = self.dsv3(up3) - dsv2 = self.dsv2(up2) - dsv1 = self.dsv1(up1) - - return dsv1, dsv2, dsv3, dsv4 - - @staticmethod - def apply_argmax_softmax(pred): - log_p = F.softmax(pred, dim=1) - - return log_p diff --git a/code/networks/unet_multitask.py b/code/networks/unet_multitask.py new file mode 100755 index 0000000..11961f1 --- /dev/null +++ b/code/networks/unet_multitask.py @@ -0,0 +1,349 @@ +# -*- coding: utf-8 -*- +""" +The implementation is borrowed from: https://github.com/HiLab-git/PyMIC +""" +from __future__ import division, print_function + +import numpy as np +import torch +import torch.nn as nn +from torch.distributions.uniform import Uniform + + +class ConvBlock(nn.Module): + """two convolution layers with batch norm and leaky relu""" + + def __init__(self, in_channels, out_channels, dropout_p): + super(ConvBlock, self).__init__() + self.conv_conv = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), + nn.BatchNorm2d(out_channels), + nn.LeakyReLU(), + nn.Dropout(dropout_p), + nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), + nn.BatchNorm2d(out_channels), + nn.LeakyReLU() + ) + + def forward(self, x): + return self.conv_conv(x) + + +class DownBlock(nn.Module): + """Downsampling followed by ConvBlock""" + + def __init__(self, in_channels, out_channels, dropout_p): + super(DownBlock, self).__init__() + self.maxpool_conv = nn.Sequential( + nn.MaxPool2d(2), + ConvBlock(in_channels, out_channels, dropout_p) + + ) + + def forward(self, x): + return self.maxpool_conv(x) + + +class UpBlock(nn.Module): + """Upssampling followed by ConvBlock""" + + def __init__(self, in_channels1, in_channels2, out_channels, dropout_p, + bilinear=True): + super(UpBlock, self).__init__() + self.bilinear = bilinear + if bilinear: + self.conv1x1 = nn.Conv2d(in_channels1, in_channels2, kernel_size=1) + self.up = nn.Upsample( + scale_factor=2, mode='bilinear', align_corners=True) + else: + self.up = nn.ConvTranspose2d( + in_channels1, in_channels2, kernel_size=2, stride=2) + self.conv = ConvBlock(in_channels2 * 2, out_channels, dropout_p) + + def forward(self, x1, x2): + if self.bilinear: + x1 = self.conv1x1(x1) + x1 = self.up(x1) + x = torch.cat([x2, x1], dim=1) + return self.conv(x) + + +class Encoder(nn.Module): + def __init__(self, params): + super(Encoder, self).__init__() + self.params = params + self.in_chns = self.params['in_chns'] + self.ft_chns = self.params['feature_chns'] + self.n_class = self.params['class_num'] + self.bilinear = self.params['bilinear'] + self.dropout = self.params['dropout'] + assert (len(self.ft_chns) == 5) + self.in_conv = ConvBlock( + self.in_chns, self.ft_chns[0], self.dropout[0]) + self.down1 = DownBlock( + self.ft_chns[0], self.ft_chns[1], self.dropout[1]) + self.down2 = DownBlock( + self.ft_chns[1], self.ft_chns[2], self.dropout[2]) + self.down3 = DownBlock( + self.ft_chns[2], self.ft_chns[3], self.dropout[3]) + self.down4 = DownBlock( + self.ft_chns[3], self.ft_chns[4], self.dropout[4]) + + def forward(self, x): + x0 = self.in_conv(x) + x1 = self.down1(x0) + x2 = self.down2(x1) + x3 = self.down3(x2) + x4 = self.down4(x3) + return [x0, x1, x2, x3, x4] + + +class Decoder(nn.Module): + def __init__(self, params): + super(Decoder, self).__init__() + self.params = params + self.in_chns = self.params['in_chns'] + self.ft_chns = self.params['feature_chns'] + self.n_class = self.params['class_num'] + self.bilinear = self.params['bilinear'] + assert (len(self.ft_chns) == 5) + + self.up1 = UpBlock( + self.ft_chns[4], self.ft_chns[3], self.ft_chns[3], dropout_p=0.0) + self.up2 = UpBlock( + self.ft_chns[3], self.ft_chns[2], self.ft_chns[2], dropout_p=0.0) + self.up3 = UpBlock( + self.ft_chns[2], self.ft_chns[1], self.ft_chns[1], dropout_p=0.0) + self.up4 = UpBlock( + self.ft_chns[1], self.ft_chns[0], self.ft_chns[0], dropout_p=0.0) + + self.out_conv = nn.Conv2d(self.ft_chns[0], self.n_class, + kernel_size=3, padding=1) + self.out_conv_regression = nn.Conv2d(self.ft_chns[0], 1, kernel_size=3, padding=1) + self.sigmoid = nn.Sigmoid() + + def forward(self, feature): + x0 = feature[0] + x1 = feature[1] + x2 = feature[2] + x3 = feature[3] + x4 = feature[4] + + x = self.up1(x4, x3) + x = self.up2(x, x2) + x = self.up3(x, x1) + x = self.up4(x, x0) + output_seg = self.out_conv(x) + output_reg = self.sigmoid(self.out_conv_regression(x)) + return output_seg, output_reg + + +class Decoder_DS(nn.Module): + def __init__(self, params): + super(Decoder_DS, self).__init__() + self.params = params + self.in_chns = self.params['in_chns'] + self.ft_chns = self.params['feature_chns'] + self.n_class = self.params['class_num'] + self.bilinear = self.params['bilinear'] + assert (len(self.ft_chns) == 5) + + self.up1 = UpBlock( + self.ft_chns[4], self.ft_chns[3], self.ft_chns[3], dropout_p=0.0) + self.up2 = UpBlock( + self.ft_chns[3], self.ft_chns[2], self.ft_chns[2], dropout_p=0.0) + self.up3 = UpBlock( + self.ft_chns[2], self.ft_chns[1], self.ft_chns[1], dropout_p=0.0) + self.up4 = UpBlock( + self.ft_chns[1], self.ft_chns[0], self.ft_chns[0], dropout_p=0.0) + + self.out_conv = nn.Conv2d(self.ft_chns[0], self.n_class, + kernel_size=3, padding=1) + self.out_conv_dp4 = nn.Conv2d(self.ft_chns[4], self.n_class, + kernel_size=3, padding=1) + self.out_conv_dp3 = nn.Conv2d(self.ft_chns[3], self.n_class, + kernel_size=3, padding=1) + self.out_conv_dp2 = nn.Conv2d(self.ft_chns[2], self.n_class, + kernel_size=3, padding=1) + self.out_conv_dp1 = nn.Conv2d(self.ft_chns[1], self.n_class, + kernel_size=3, padding=1) + + def forward(self, feature, shape): + x0 = feature[0] + x1 = feature[1] + x2 = feature[2] + x3 = feature[3] + x4 = feature[4] + x = self.up1(x4, x3) + dp3_out_seg = self.out_conv_dp3(x) + dp3_out_seg = torch.nn.functional.interpolate(dp3_out_seg, shape) + + x = self.up2(x, x2) + dp2_out_seg = self.out_conv_dp2(x) + dp2_out_seg = torch.nn.functional.interpolate(dp2_out_seg, shape) + + x = self.up3(x, x1) + dp1_out_seg = self.out_conv_dp1(x) + dp1_out_seg = torch.nn.functional.interpolate(dp1_out_seg, shape) + + x = self.up4(x, x0) + dp0_out_seg = self.out_conv(x) + return dp0_out_seg, dp1_out_seg, dp2_out_seg, dp3_out_seg + + +class Decoder_URDS(nn.Module): + def __init__(self, params): + super(Decoder_URDS, self).__init__() + self.params = params + self.in_chns = self.params['in_chns'] + self.ft_chns = self.params['feature_chns'] + self.n_class = self.params['class_num'] + self.bilinear = self.params['bilinear'] + assert (len(self.ft_chns) == 5) + + self.up1 = UpBlock( + self.ft_chns[4], self.ft_chns[3], self.ft_chns[3], dropout_p=0.0) + self.up2 = UpBlock( + self.ft_chns[3], self.ft_chns[2], self.ft_chns[2], dropout_p=0.0) + self.up3 = UpBlock( + self.ft_chns[2], self.ft_chns[1], self.ft_chns[1], dropout_p=0.0) + self.up4 = UpBlock( + self.ft_chns[1], self.ft_chns[0], self.ft_chns[0], dropout_p=0.0) + + self.out_conv = nn.Conv2d(self.ft_chns[0], self.n_class, + kernel_size=3, padding=1) + self.out_conv_dp4 = nn.Conv2d(self.ft_chns[4], self.n_class, + kernel_size=3, padding=1) + self.out_conv_dp3 = nn.Conv2d(self.ft_chns[3], self.n_class, + kernel_size=3, padding=1) + self.out_conv_dp2 = nn.Conv2d(self.ft_chns[2], self.n_class, + kernel_size=3, padding=1) + self.out_conv_dp1 = nn.Conv2d(self.ft_chns[1], self.n_class, + kernel_size=3, padding=1) + self.feature_noise = FeatureNoise() + + def forward(self, feature, shape): + x0 = feature[0] + x1 = feature[1] + x2 = feature[2] + x3 = feature[3] + x4 = feature[4] + x = self.up1(x4, x3) + if self.training: + dp3_out_seg = self.out_conv_dp3(Dropout(x, p=0.5)) + else: + dp3_out_seg = self.out_conv_dp3(x) + dp3_out_seg = torch.nn.functional.interpolate(dp3_out_seg, shape) + + x = self.up2(x, x2) + if self.training: + dp2_out_seg = self.out_conv_dp2(FeatureDropout(x)) + else: + dp2_out_seg = self.out_conv_dp2(x) + dp2_out_seg = torch.nn.functional.interpolate(dp2_out_seg, shape) + + x = self.up3(x, x1) + if self.training: + dp1_out_seg = self.out_conv_dp1(self.feature_noise(x)) + else: + dp1_out_seg = self.out_conv_dp1(x) + dp1_out_seg = torch.nn.functional.interpolate(dp1_out_seg, shape) + + x = self.up4(x, x0) + dp0_out_seg = self.out_conv(x) + return dp0_out_seg, dp1_out_seg, dp2_out_seg, dp3_out_seg + + +def Dropout(x, p=0.5): + x = torch.nn.functional.dropout2d(x, p) + return x + + +def FeatureDropout(x): + attention = torch.mean(x, dim=1, keepdim=True) + max_val, _ = torch.max(attention.view( + x.size(0), -1), dim=1, keepdim=True) + threshold = max_val * np.random.uniform(0.7, 0.9) + threshold = threshold.view(x.size(0), 1, 1, 1).expand_as(attention) + drop_mask = (attention < threshold).float() + x = x.mul(drop_mask) + return x + + +class FeatureNoise(nn.Module): + def __init__(self, uniform_range=0.3): + super(FeatureNoise, self).__init__() + self.uni_dist = Uniform(-uniform_range, uniform_range) + + def feature_based_noise(self, x): + noise_vector = self.uni_dist.sample( + x.shape[1:]).to(x.device).unsqueeze(0) + x_noise = x.mul(noise_vector) + x + return x_noise + + def forward(self, x): + x = self.feature_based_noise(x) + return x + + +class UNet_Multitask(nn.Module): + def __init__(self, in_chns, class_num): + super(UNet_Multitask, self).__init__() + + params = {'in_chns': in_chns, + 'feature_chns': [16, 32, 64, 128, 256], + 'dropout': [0.05, 0.1, 0.2, 0.3, 0.5], + 'class_num': class_num, + 'bilinear': False, + 'acti_func': 'relu'} + + self.encoder = Encoder(params) + self.decoder = Decoder(params) + + def forward(self, x): + feature = self.encoder(x) + output_seg, output_reg = self.decoder(feature) + return output_seg, output_reg + + +class UNet_DS(nn.Module): + def __init__(self, in_chns, class_num): + super(UNet_DS, self).__init__() + + params = {'in_chns': in_chns, + 'feature_chns': [16, 32, 64, 128, 256], + 'dropout': [0.05, 0.1, 0.2, 0.3, 0.5], + 'class_num': class_num, + 'bilinear': False, + 'acti_func': 'relu'} + self.encoder = Encoder(params) + self.decoder = Decoder_DS(params) + + def forward(self, x): + shape = x.shape[2:] + feature = self.encoder(x) + dp0_out_seg, dp1_out_seg, dp2_out_seg, dp3_out_seg = self.decoder( + feature, shape) + return dp0_out_seg, dp1_out_seg, dp2_out_seg, dp3_out_seg + + +class UNet_URDS(nn.Module): + def __init__(self, in_chns, class_num): + super(UNet_URDS, self).__init__() + + params = {'in_chns': in_chns, + 'feature_chns': [16, 32, 64, 128, 256], + 'dropout': [0.05, 0.1, 0.2, 0.3, 0.5], + 'class_num': class_num, + 'bilinear': False, + 'acti_func': 'relu'} + self.encoder = Encoder(params) + self.decoder = Decoder_URDS(params) + + def forward(self, x): + shape = x.shape[2:] + feature = self.encoder(x) + dp1_out_seg, dp2_out_seg, dp3_out_seg, dp4_out_seg = self.decoder( + feature, shape) + return dp1_out_seg, dp2_out_seg, dp3_out_seg, dp4_out_seg + diff --git a/code/networks/utils.py b/code/networks/utils.py old mode 100644 new mode 100755 diff --git a/code/networks/vision_transformer.py b/code/networks/vision_transformer.py deleted file mode 100644 index 927ee2a..0000000 --- a/code/networks/vision_transformer.py +++ /dev/null @@ -1,90 +0,0 @@ -# coding=utf-8 -# This file borrowed from Swin-UNet: https://github.com/HuCaoFighting/Swin-Unet -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import copy -import logging -import math - -from os.path import join as pjoin - -import torch -import torch.nn as nn -import numpy as np - -from torch.nn import CrossEntropyLoss, Dropout, Softmax, Linear, Conv2d, LayerNorm -from torch.nn.modules.utils import _pair -from scipy import ndimage -from networks.swin_transformer_unet_skip_expand_decoder_sys import SwinTransformerSys - -logger = logging.getLogger(__name__) - -class SwinUnet(nn.Module): - def __init__(self, config, img_size=224, num_classes=21843, zero_head=False, vis=False): - super(SwinUnet, self).__init__() - self.num_classes = num_classes - self.zero_head = zero_head - self.config = config - - self.swin_unet = SwinTransformerSys(img_size=config.DATA.IMG_SIZE, - patch_size=config.MODEL.SWIN.PATCH_SIZE, - in_chans=config.MODEL.SWIN.IN_CHANS, - num_classes=self.num_classes, - embed_dim=config.MODEL.SWIN.EMBED_DIM, - depths=config.MODEL.SWIN.DEPTHS, - num_heads=config.MODEL.SWIN.NUM_HEADS, - window_size=config.MODEL.SWIN.WINDOW_SIZE, - mlp_ratio=config.MODEL.SWIN.MLP_RATIO, - qkv_bias=config.MODEL.SWIN.QKV_BIAS, - qk_scale=config.MODEL.SWIN.QK_SCALE, - drop_rate=config.MODEL.DROP_RATE, - drop_path_rate=config.MODEL.DROP_PATH_RATE, - ape=config.MODEL.SWIN.APE, - patch_norm=config.MODEL.SWIN.PATCH_NORM, - use_checkpoint=config.TRAIN.USE_CHECKPOINT) - - def forward(self, x): - if x.size()[1] == 1: - x = x.repeat(1,3,1,1) - logits = self.swin_unet(x) - return logits - - def load_from(self, config): - pretrained_path = config.MODEL.PRETRAIN_CKPT - if pretrained_path is not None: - print("pretrained_path:{}".format(pretrained_path)) - device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - pretrained_dict = torch.load(pretrained_path, map_location=device) - if "model" not in pretrained_dict: - print("---start load pretrained modle by splitting---") - pretrained_dict = {k[17:]:v for k,v in pretrained_dict.items()} - for k in list(pretrained_dict.keys()): - if "output" in k: - print("delete key:{}".format(k)) - del pretrained_dict[k] - msg = self.swin_unet.load_state_dict(pretrained_dict,strict=False) - # print(msg) - return - pretrained_dict = pretrained_dict['model'] - print("---start load pretrained modle of swin encoder---") - - model_dict = self.swin_unet.state_dict() - full_dict = copy.deepcopy(pretrained_dict) - for k, v in pretrained_dict.items(): - if "layers." in k: - current_layer_num = 3-int(k[7:8]) - current_k = "layers_up." + str(current_layer_num) + k[8:] - full_dict.update({current_k:v}) - for k in list(full_dict.keys()): - if k in model_dict: - if full_dict[k].shape != model_dict[k].shape: - print("delete:{};shape pretrain:{};shape model:{}".format(k,v.shape,model_dict[k].shape)) - del full_dict[k] - - msg = self.swin_unet.load_state_dict(full_dict, strict=False) - # print(msg) - else: - print("none pretrain") - \ No newline at end of file diff --git a/code/networks/vnet.py b/code/networks/vnet.py old mode 100644 new mode 100755 diff --git a/code/pretrained_ckpt/readme.txt b/code/pretrained_ckpt/readme.txt deleted file mode 100644 index a1691a6..0000000 --- a/code/pretrained_ckpt/readme.txt +++ /dev/null @@ -1 +0,0 @@ -download pre-trained model to this folder, link:https://drive.google.com/drive/folders/1UC3XOoezeum0uck4KBVGa8osahs6rKUY diff --git a/code/test.py b/code/test.py new file mode 100644 index 0000000..e3566a4 --- /dev/null +++ b/code/test.py @@ -0,0 +1,149 @@ +import argparse +import os +import re +import shutil + +import h5py +from matplotlib.pyplot import axis +import nibabel as nib +import numpy as np +import SimpleITK as sitk +import torch +from medpy import metric +from scipy.ndimage import zoom +from scipy.ndimage.interpolation import zoom +from sklearn.model_selection import KFold +from tqdm import tqdm +from networks.net_factory import net_factory + +parser = argparse.ArgumentParser() +parser.add_argument('--root_path', type=str, + default='../data/ProstateX', help='Name of Experiment') +parser.add_argument('--exp', type=str, + default='ProstateX/Mean_Teacher', help='experiment_name') +parser.add_argument('--model', type=str, + default='unet', help='model_name') +parser.add_argument('--labeled_ratio', type=int, default=8, + help='1/labeled_ratio data is provided mask') +parser.add_argument('--fold', type=int, + default=1, help='fold') +parser.add_argument('--patch_size', type=list, default=[256, 256], + help='patch size of network input') +parser.add_argument('--num_classes', type=int, default=3, + help='output channel of network') +parser.add_argument('--sup_type', type=str, default="label", + help='label') + + +def get_fold_ids(FLAGS): + all_volumes = sorted(os.listdir(FLAGS.root_path + "/all_volumes")) + folds = KFold(n_splits=5, shuffle=False) + all_cases = np.array(all_volumes) + k_fold_data = [] + for trn_idx, val_idx in folds.split(all_cases): + k_fold_data.append([all_cases[trn_idx], all_cases[val_idx]]) + return k_fold_data[FLAGS.fold][0], k_fold_data[FLAGS.fold][1] + + +def calculate_metric_percase(pred, gt, spacing): + if pred.sum() > 0 and gt.sum() > 0: + pred[pred > 0] = 1 + gt[gt > 0] = 1 + dice = metric.binary.dc(pred, gt) + asd = metric.binary.asd(pred, gt, voxelspacing=spacing) + hd95 = metric.binary.hd95(pred, gt, voxelspacing=spacing) + else: + dice = 0.0 + hd95 = 100.0 + asd = 20.0 + return dice, hd95, asd + + +def test_single_volume(case, net, test_save_path, FLAGS): + h5f = h5py.File(FLAGS.root_path + + "/all_volumes/{}".format(case), 'r') + image = h5f['image'][:] + label = h5f['label'][:] + spacing = h5f['spacing'][:] + prediction = np.zeros_like(label) + for ind in range(image.shape[0]): + slice = image[ind, :, :] + x, y = slice.shape[0], slice.shape[1] + slice = zoom(slice, (FLAGS.patch_size / x, FLAGS.patch_size / y), order=0) + input = torch.from_numpy(slice).unsqueeze( + 0).unsqueeze(0).float().cuda() + net.eval() + with torch.no_grad(): + out_main = net(input) + out = torch.argmax(torch.softmax( + out_main, dim=1), dim=1).squeeze(0) + out = out.cpu().detach().numpy() + pred = zoom(out, (x / FLAGS.patch_size, y / FLAGS.patch_size), order=0) + prediction[ind] = pred + case = case.replace(".h5", "") + + metric_list = [] + for i in range(1, FLAGS.num_classes): + metric_list.append(calculate_metric_percase( + prediction == i, label == i, spacing=(spacing[2], spacing[0], spacing[1]))) + img_itk = sitk.GetImageFromArray(image.astype(np.float32)) + img_itk.SetSpacing(spacing) + prd_itk = sitk.GetImageFromArray(prediction.astype(np.float32)) + prd_itk.SetSpacing(spacing) + lab_itk = sitk.GetImageFromArray(label.astype(np.float32)) + lab_itk.SetSpacing(spacing) + sitk.WriteImage(prd_itk, test_save_path + case + "_pred.nii.gz") + sitk.WriteImage(img_itk, test_save_path + case + "_img.nii.gz") + sitk.WriteImage(lab_itk, test_save_path + case + "_gt.nii.gz") + return np.array(metric_list) + + +def Inference(FLAGS): + train_ids, test_ids = get_fold_ids(FLAGS) + all_volumes = os.listdir( + FLAGS.root_path + "/all_volumes") + image_list = [] + for ids in test_ids: + new_data_list = list(filter(lambda x: re.match( + '{}.*'.format(ids), x) != None, all_volumes)) + image_list.extend(new_data_list) + snapshot_path = snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( + FLAGS.exp, FLAGS.labeled_ratio, FLAGS.fold) + test_save_path = "../model/{}/1_of_{}_labeled/fold{}/prediction/".format( + FLAGS.exp, FLAGS.labeled_ratio, FLAGS.fold) + if os.path.exists(test_save_path): + shutil.rmtree(test_save_path) + os.makedirs(test_save_path) + net = net_factory(net_type=FLAGS.model, in_chns=1, + class_num=FLAGS.num_classes) + save_mode_path = os.path.join( + snapshot_path, '{}_best_model.pth'.format(FLAGS.model)) + # save_mode_path = os.path.join( + # snapshot_path, 'iter_60000.pth') + net.load_state_dict(torch.load(save_mode_path)) + print("init weight from {}".format(save_mode_path)) + net.eval() + + metric_array = np.zeros((len(image_list), FLAGS.num_classes-1, 3)) + for ind, case in enumerate(tqdm(image_list)): + print(case) + cases_metric = test_single_volume( + case, net, test_save_path, FLAGS) + print(cases_metric) + metric_array[ind, ...] = cases_metric + np.save("../model/{}/1_of_{}_labeled/fold{}/prediction/Results.npy".format( + FLAGS.exp, FLAGS.labeled_ratio, FLAGS.fold), metric_array) + return metric_array + + +if __name__ == '__main__': + FLAGS = parser.parse_args() + total = 0.0 + for i in [3]: + FLAGS.fold = i + print("Inference fold{}".format(i)) + metric_array = Inference(FLAGS) + print("mean class results:", np.mean(metric_array, axis=0)) + print("mean case results:", np.mean(metric_array, axis=0).mean(axis=0)) + print(total/1) + \ No newline at end of file diff --git a/code/test_2D_fully.py b/code/test_2D_fully.py deleted file mode 100644 index 0d66984..0000000 --- a/code/test_2D_fully.py +++ /dev/null @@ -1,117 +0,0 @@ -import argparse -import os -import shutil - -import h5py -import nibabel as nib -import numpy as np -import SimpleITK as sitk -import torch -from medpy import metric -from scipy.ndimage import zoom -from scipy.ndimage.interpolation import zoom -from tqdm import tqdm - -# from networks.efficientunet import UNet -from networks.net_factory import net_factory - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/ACDC', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='ACDC/Fully_Supervised', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet', help='model_name') -parser.add_argument('--num_classes', type=int, default=4, - help='output channel of network') -parser.add_argument('--labeled_num', type=int, default=3, - help='labeled data') - - -def calculate_metric_percase(pred, gt): - pred[pred > 0] = 1 - gt[gt > 0] = 1 - dice = metric.binary.dc(pred, gt) - asd = metric.binary.asd(pred, gt) - hd95 = metric.binary.hd95(pred, gt) - return dice, hd95, asd - - -def test_single_volume(case, net, test_save_path, FLAGS): - h5f = h5py.File(FLAGS.root_path + "/data/{}.h5".format(case), 'r') - image = h5f['image'][:] - label = h5f['label'][:] - prediction = np.zeros_like(label) - for ind in range(image.shape[0]): - slice = image[ind, :, :] - x, y = slice.shape[0], slice.shape[1] - slice = zoom(slice, (256 / x, 256 / y), order=0) - input = torch.from_numpy(slice).unsqueeze( - 0).unsqueeze(0).float().cuda() - net.eval() - with torch.no_grad(): - if FLAGS.model == "unet_urds": - out_main, _, _, _ = net(input) - else: - out_main = net(input) - out = torch.argmax(torch.softmax( - out_main, dim=1), dim=1).squeeze(0) - out = out.cpu().detach().numpy() - pred = zoom(out, (x / 256, y / 256), order=0) - prediction[ind] = pred - - first_metric = calculate_metric_percase(prediction == 1, label == 1) - second_metric = calculate_metric_percase(prediction == 2, label == 2) - third_metric = calculate_metric_percase(prediction == 3, label == 3) - - img_itk = sitk.GetImageFromArray(image.astype(np.float32)) - img_itk.SetSpacing((1, 1, 10)) - prd_itk = sitk.GetImageFromArray(prediction.astype(np.float32)) - prd_itk.SetSpacing((1, 1, 10)) - lab_itk = sitk.GetImageFromArray(label.astype(np.float32)) - lab_itk.SetSpacing((1, 1, 10)) - sitk.WriteImage(prd_itk, test_save_path + case + "_pred.nii.gz") - sitk.WriteImage(img_itk, test_save_path + case + "_img.nii.gz") - sitk.WriteImage(lab_itk, test_save_path + case + "_gt.nii.gz") - return first_metric, second_metric, third_metric - - -def Inference(FLAGS): - with open(FLAGS.root_path + '/test.list', 'r') as f: - image_list = f.readlines() - image_list = sorted([item.replace('\n', '').split(".")[0] - for item in image_list]) - snapshot_path = "../model/{}_{}_labeled/{}".format( - FLAGS.exp, FLAGS.labeled_num, FLAGS.model) - test_save_path = "../model/{}_{}_labeled/{}_predictions/".format( - FLAGS.exp, FLAGS.labeled_num, FLAGS.model) - if os.path.exists(test_save_path): - shutil.rmtree(test_save_path) - os.makedirs(test_save_path) - net = net_factory(net_type=FLAGS.model, in_chns=1, - class_num=FLAGS.num_classes) - save_mode_path = os.path.join( - snapshot_path, '{}_best_model.pth'.format(FLAGS.model)) - net.load_state_dict(torch.load(save_mode_path)) - print("init weight from {}".format(save_mode_path)) - net.eval() - - first_total = 0.0 - second_total = 0.0 - third_total = 0.0 - for case in tqdm(image_list): - first_metric, second_metric, third_metric = test_single_volume( - case, net, test_save_path, FLAGS) - first_total += np.asarray(first_metric) - second_total += np.asarray(second_metric) - third_total += np.asarray(third_metric) - avg_metric = [first_total / len(image_list), second_total / - len(image_list), third_total / len(image_list)] - return avg_metric - - -if __name__ == '__main__': - FLAGS = parser.parse_args() - metric = Inference(FLAGS) - print(metric) - print((metric[0]+metric[1]+metric[2])/3) diff --git a/code/test_3D.py b/code/test_3D.py deleted file mode 100644 index 82e7ef9..0000000 --- a/code/test_3D.py +++ /dev/null @@ -1,41 +0,0 @@ -import argparse -import os -import shutil -from glob import glob - -import torch - -from networks.unet_3D import unet_3D -from test_3D_util import test_all_case - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/BraTS2019', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='BraTS2019/Interpolation_Consistency_Training_25', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet_3D', help='model_name') - - -def Inference(FLAGS): - snapshot_path = "../model/{}/{}".format(FLAGS.exp, FLAGS.model) - num_classes = 2 - test_save_path = "../model/{}/Prediction".format(FLAGS.exp) - if os.path.exists(test_save_path): - shutil.rmtree(test_save_path) - os.makedirs(test_save_path) - net = unet_3D(n_classes=num_classes, in_channels=1).cuda() - save_mode_path = os.path.join( - snapshot_path, '{}_best_model.pth'.format(FLAGS.model)) - net.load_state_dict(torch.load(save_mode_path)) - print("init weight from {}".format(save_mode_path)) - net.eval() - avg_metric = test_all_case(net, base_dir=FLAGS.root_path, method=FLAGS.model, test_list="test.txt", num_classes=num_classes, - patch_size=(96, 96, 96), stride_xy=64, stride_z=64, test_save_path=test_save_path) - return avg_metric - - -if __name__ == '__main__': - FLAGS = parser.parse_args() - metric = Inference(FLAGS) - print(metric) diff --git a/code/test_3D_util.py b/code/test_3D_util.py deleted file mode 100644 index a2f18a3..0000000 --- a/code/test_3D_util.py +++ /dev/null @@ -1,152 +0,0 @@ -import math - -import h5py -import nibabel as nib -import numpy as np -import SimpleITK as sitk -import torch -import torch.nn.functional as F -from medpy import metric -from skimage.measure import label -from tqdm import tqdm - - -def test_single_case(net, image, stride_xy, stride_z, patch_size, num_classes=1): - w, h, d = image.shape - - # if the size of image is less than patch_size, then padding it - add_pad = False - if w < patch_size[0]: - w_pad = patch_size[0]-w - add_pad = True - else: - w_pad = 0 - if h < patch_size[1]: - h_pad = patch_size[1]-h - add_pad = True - else: - h_pad = 0 - if d < patch_size[2]: - d_pad = patch_size[2]-d - add_pad = True - else: - d_pad = 0 - wl_pad, wr_pad = w_pad//2, w_pad-w_pad//2 - hl_pad, hr_pad = h_pad//2, h_pad-h_pad//2 - dl_pad, dr_pad = d_pad//2, d_pad-d_pad//2 - if add_pad: - image = np.pad(image, [(wl_pad, wr_pad), (hl_pad, hr_pad), - (dl_pad, dr_pad)], mode='constant', constant_values=0) - ww, hh, dd = image.shape - - sx = math.ceil((ww - patch_size[0]) / stride_xy) + 1 - sy = math.ceil((hh - patch_size[1]) / stride_xy) + 1 - sz = math.ceil((dd - patch_size[2]) / stride_z) + 1 - # print("{}, {}, {}".format(sx, sy, sz)) - score_map = np.zeros((num_classes, ) + image.shape).astype(np.float32) - cnt = np.zeros(image.shape).astype(np.float32) - - for x in range(0, sx): - xs = min(stride_xy*x, ww-patch_size[0]) - for y in range(0, sy): - ys = min(stride_xy * y, hh-patch_size[1]) - for z in range(0, sz): - zs = min(stride_z * z, dd-patch_size[2]) - test_patch = image[xs:xs+patch_size[0], - ys:ys+patch_size[1], zs:zs+patch_size[2]] - test_patch = np.expand_dims(np.expand_dims( - test_patch, axis=0), axis=0).astype(np.float32) - test_patch = torch.from_numpy(test_patch).cuda() - - with torch.no_grad(): - y1 = net(test_patch) - # ensemble - y = torch.softmax(y1, dim=1) - y = y.cpu().data.numpy() - y = y[0, :, :, :, :] - score_map[:, xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] \ - = score_map[:, xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] + y - cnt[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] \ - = cnt[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] + 1 - score_map = score_map/np.expand_dims(cnt, axis=0) - label_map = np.argmax(score_map, axis=0) - - if add_pad: - label_map = label_map[wl_pad:wl_pad+w, - hl_pad:hl_pad+h, dl_pad:dl_pad+d] - score_map = score_map[:, wl_pad:wl_pad + - w, hl_pad:hl_pad+h, dl_pad:dl_pad+d] - return label_map - - -def cal_metric(gt, pred): - if pred.sum() > 0 and gt.sum() > 0: - dice = metric.binary.dc(pred, gt) - hd95 = metric.binary.hd95(pred, gt) - return np.array([dice, hd95]) - else: - return np.zeros(2) - - -def test_all_case(net, base_dir, method="unet_3D", test_list="full_test.list", num_classes=4, patch_size=(48, 160, 160), stride_xy=32, stride_z=24, test_save_path=None): - with open(base_dir + '/{}'.format(test_list), 'r') as f: - image_list = f.readlines() - image_list = [base_dir + "/data/{}.h5".format( - item.replace('\n', '').split(",")[0]) for item in image_list] - total_metric = np.zeros((num_classes-1, 4)) - print("Testing begin") - with open(test_save_path + "/{}.txt".format(method), "a") as f: - for image_path in tqdm(image_list): - ids = image_path.split("/")[-1].replace(".h5", "") - h5f = h5py.File(image_path, 'r') - image = h5f['image'][:] - label = h5f['label'][:] - prediction = test_single_case( - net, image, stride_xy, stride_z, patch_size, num_classes=num_classes) - metric = calculate_metric_percase(label == 1, prediction == 1) - total_metric[0, :] += metric - f.writelines("{},{},{},{},{}\n".format( - ids, metric[0], metric[1], metric[2], metric[3])) - - pred_itk = sitk.GetImageFromArray(prediction.astype(np.uint8)) - pred_itk.SetSpacing((1.0, 1.0, 1.0)) - sitk.WriteImage(pred_itk, test_save_path + - "/{}_pred.nii.gz".format(ids)) - - img_itk = sitk.GetImageFromArray(image) - img_itk.SetSpacing((1.0, 1.0, 1.0)) - sitk.WriteImage(img_itk, test_save_path + - "/{}_img.nii.gz".format(ids)) - - lab_itk = sitk.GetImageFromArray(label.astype(np.uint8)) - lab_itk.SetSpacing((1.0, 1.0, 1.0)) - sitk.WriteImage(lab_itk, test_save_path + - "/{}_lab.nii.gz".format(ids)) - f.writelines("Mean metrics,{},{},{},{}".format(total_metric[0, 0] / len(image_list), total_metric[0, 1] / len( - image_list), total_metric[0, 2] / len(image_list), total_metric[0, 3] / len(image_list))) - f.close() - print("Testing end") - return total_metric / len(image_list) - - -def cal_dice(prediction, label, num=2): - total_dice = np.zeros(num-1) - for i in range(1, num): - prediction_tmp = (prediction == i) - label_tmp = (label == i) - prediction_tmp = prediction_tmp.astype(np.float) - label_tmp = label_tmp.astype(np.float) - - dice = 2 * np.sum(prediction_tmp * label_tmp) / \ - (np.sum(prediction_tmp) + np.sum(label_tmp)) - total_dice[i - 1] += dice - - return total_dice - - -def calculate_metric_percase(pred, gt): - dice = metric.binary.dc(pred, gt) - ravd = abs(metric.binary.ravd(pred, gt)) - hd = metric.binary.hd95(pred, gt) - asd = metric.binary.asd(pred, gt) - return np.array([dice, ravd, hd, asd]) diff --git a/code/test_acdc_unet_semi_seg.sh b/code/test_acdc_unet_semi_seg.sh deleted file mode 100644 index 6b687bb..0000000 --- a/code/test_acdc_unet_semi_seg.sh +++ /dev/null @@ -1,8 +0,0 @@ -CUDA_VISIBLE_DEVICES=0 python test_2D_fully.py --root_path ../data/ACDC --exp ACDC/Fully_Supervised --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python test_2D_fully.py --root_path ../data/ACDC --exp ACDC/Entropy_Minimization --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python test_2D_fully.py --root_path ../data/ACDC --exp ACDC/Interpolation_Consistency_Training --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python test_2D_fully.py --root_path ../data/ACDC --exp ACDC/Mean_Teacher --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python test_2D_fully.py --root_path ../data/ACDC --exp ACDC/Uncertainty_Aware_Mean_Teacher --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python test_2D_fully.py --root_path ../data/ACDC --exp ACDC/Adversarial_Network --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python test_2D_fully.py --root_path ../data/ACDC --exp ACDC/Uncertainty_Rectified_Pyramid_Consistency --model unet_urpc --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python test_2D_fully.py --root_path ../data/ACDC --exp ACDC/Fully_Supervised --num_classes 4 --labeled_num 140 \ No newline at end of file diff --git a/code/test_brats2019_semi_seg.sh b/code/test_brats2019_semi_seg.sh deleted file mode 100644 index 0a944bb..0000000 --- a/code/test_brats2019_semi_seg.sh +++ /dev/null @@ -1,10 +0,0 @@ -# & means run these methods at the same time, and && means run these methods one by one -python -u test_3D.py --root_path ../data/BraTS2019 --exp BraTS2019/Fully_supervised_25 --model unet_3D && -python -u test_3D.py --root_path ../data/BraTS2019 --exp BraTS2019/Fully_supervised_250 --model unet_3D && -python -u test_3D.py --root_path ../data/BraTS2019 --exp BraTS2019/Mean_Teacher_25 --model unet_3D && -python -u test_3D.py --root_path ../data/BraTS2019 --exp BraTS2019/Uncertainty_Aware_Mean_Teacher_25 --model unet_3D && -python -u test_3D.py --root_path ../data/BraTS2019 --exp BraTS2019/Interpolation_Consistency_Training_25 --model unet_3D && -python -u test_3D.py --root_path ../data/BraTS2019 --exp BraTS2019/Entropy_Minimization_25 --model unet_3D && -python -u test_3D.py --root_path ../data/BraTS2019 --exp BraTS2019/Cross_Pseudo_Supervision_25 --model unet_3D && -python -u test_3D.py --root_path ../data/BraTS2019 --exp BraTS2019/Adversarial_Network_25 --model unet_3D && -python -u test_3D.py --root_path ../data/BraTS2019 --exp BraTS2019/Uncertainty_Rectified_Pyramid_Consistency_25 --model unet_3D_dv_semi \ No newline at end of file diff --git a/code/test_urpc.py b/code/test_urpc.py deleted file mode 100644 index 44bc0c6..0000000 --- a/code/test_urpc.py +++ /dev/null @@ -1,55 +0,0 @@ -import argparse -import os -import shutil -from glob import glob -import numpy - -import torch - -from networks.unet_3D_dv_semi import unet_3D_dv_semi -from networks.unet_3D import unet_3D -from test_urpc_util import test_all_case - - -def net_factory(net_type="unet_3D", num_classes=3, in_channels=1): - if net_type == "unet_3D": - net = unet_3D(n_classes=num_classes, in_channels=in_channels).cuda() - elif net_type == "unet_3D_dv_semi": - net = unet_3D_dv_semi(n_classes=num_classes, - in_channels=in_channels).cuda() - else: - net = None - return net - - -def Inference(FLAGS): - snapshot_path = "../model/{}/{}".format(FLAGS.exp, FLAGS.model) - num_classes = 2 - test_save_path = "../model/{}/Prediction".format(FLAGS.exp) - if os.path.exists(test_save_path): - shutil.rmtree(test_save_path) - os.makedirs(test_save_path) - net = net_factory(FLAGS.model, num_classes, in_channels=1) - save_mode_path = os.path.join( - snapshot_path, '{}_best_model.pth'.format(FLAGS.model)) - net.load_state_dict(torch.load(save_mode_path)) - print("init weight from {}".format(save_mode_path)) - net.eval() - avg_metric = test_all_case(net, base_dir=FLAGS.root_path, method=FLAGS.model, test_list="test.txt", num_classes=num_classes, - patch_size=(96, 96, 96), stride_xy=64, stride_z=64, test_save_path=test_save_path) - return avg_metric - - -if __name__ == '__main__': - - parser = argparse.ArgumentParser() - parser.add_argument('--root_path', type=str, - default='../data/BraTS2019', help='Name of Experiment') - parser.add_argument('--exp', type=str, - default="BraTS2019/Uncertainty_Rectified_Pyramid_Consistency_25_labeled", help='experiment_name') - parser.add_argument('--model', type=str, - default="unet_3D_dv_semi", help='model_name') - FLAGS = parser.parse_args() - - metric = Inference(FLAGS) - print(metric) diff --git a/code/test_urpc_util.py b/code/test_urpc_util.py deleted file mode 100644 index 155be25..0000000 --- a/code/test_urpc_util.py +++ /dev/null @@ -1,161 +0,0 @@ -import math - -import h5py -import nibabel as nib -import numpy as np -import SimpleITK as sitk -import torch -import torch.nn.functional as F -from medpy import metric -from skimage.measure import label -from tqdm import tqdm - - -def test_single_case(net, image, stride_xy, stride_z, patch_size, num_classes=1): - w, h, d = image.shape - - # if the size of image is less than patch_size, then padding it - add_pad = False - if w < patch_size[0]: - w_pad = patch_size[0]-w - add_pad = True - else: - w_pad = 0 - if h < patch_size[1]: - h_pad = patch_size[1]-h - add_pad = True - else: - h_pad = 0 - if d < patch_size[2]: - d_pad = patch_size[2]-d - add_pad = True - else: - d_pad = 0 - wl_pad, wr_pad = w_pad//2, w_pad-w_pad//2 - hl_pad, hr_pad = h_pad//2, h_pad-h_pad//2 - dl_pad, dr_pad = d_pad//2, d_pad-d_pad//2 - if add_pad: - image = np.pad(image, [(wl_pad, wr_pad), (hl_pad, hr_pad), - (dl_pad, dr_pad)], mode='constant', constant_values=0) - ww, hh, dd = image.shape - - sx = math.ceil((ww - patch_size[0]) / stride_xy) + 1 - sy = math.ceil((hh - patch_size[1]) / stride_xy) + 1 - sz = math.ceil((dd - patch_size[2]) / stride_z) + 1 - # print("{}, {}, {}".format(sx, sy, sz)) - score_map = np.zeros((num_classes, ) + image.shape).astype(np.float32) - cnt = np.zeros(image.shape).astype(np.float32) - - for x in range(0, sx): - xs = min(stride_xy*x, ww-patch_size[0]) - for y in range(0, sy): - ys = min(stride_xy * y, hh-patch_size[1]) - for z in range(0, sz): - zs = min(stride_z * z, dd-patch_size[2]) - test_patch = image[xs:xs+patch_size[0], - ys:ys+patch_size[1], zs:zs+patch_size[2]] - test_patch = np.expand_dims(np.expand_dims( - test_patch, axis=0), axis=0).astype(np.float32) - test_patch = torch.from_numpy(test_patch).cuda() - - with torch.no_grad(): - y_main, y_aux1, y_aux2, y_aux3 = net(test_patch) - # ensemble - y_main = torch.softmax(y_main, dim=1) - y_aux1 = torch.softmax(y_aux1, dim=1) - y_aux2 = torch.softmax(y_aux2, dim=1) - y_aux3 = torch.softmax(y_aux3, dim=1) - y = y_main - # y = (y_main+y_aux1+y_aux2+y_aux3) - y = y.cpu().data.numpy() - y = y[0, :, :, :, :] - score_map[:, xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] \ - = score_map[:, xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] + y - cnt[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] \ - = cnt[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] + 1 - score_map = score_map/np.expand_dims(cnt, axis=0) - label_map = np.argmax(score_map, axis=0) - - if add_pad: - label_map = label_map[wl_pad:wl_pad+w, - hl_pad:hl_pad+h, dl_pad:dl_pad+d] - score_map = score_map[:, wl_pad:wl_pad + - w, hl_pad:hl_pad+h, dl_pad:dl_pad+d] - return label_map - - -def cal_metric(gt, pred): - if pred.sum() > 0 and gt.sum() > 0: - dice = metric.binary.dc(pred, gt) - hd95 = metric.binary.hd95(pred, gt) - return np.array([dice, hd95]) - else: - return np.zeros(2) - - -def test_all_case(net, base_dir, method="unet_3D", test_list="full_test.list", num_classes=4, patch_size=(48, 160, 160), stride_xy=32, stride_z=24, test_save_path=None): - with open(base_dir + '/{}'.format(test_list), 'r') as f: - image_list = f.readlines() - image_list = [base_dir + "/data/{}.h5".format( - item.replace('\n', '').split(",")[0]) for item in image_list] - total_metric = np.zeros((num_classes - 1, 4)) - print("Testing begin") - with open(test_save_path + "/{}.txt".format(method), "a") as f: - for image_path in tqdm(image_list): - ids = image_path.split("/")[-1].replace(".h5", "") - h5f = h5py.File(image_path, 'r') - image = h5f['image'][:] - label = h5f['label'][:] - prediction = test_single_case( - net, image, stride_xy, stride_z, patch_size, num_classes=num_classes) - - metric = calculate_metric_percase(label == 1, prediction == 1) - total_metric[0, :] += metric - f.writelines("{},{},{},{},{}\n".format( - ids, metric[0], metric[1], metric[2], metric[3])) - - pred_itk = sitk.GetImageFromArray(prediction.astype(np.uint8)) - pred_itk.SetSpacing((1.0, 1.0, 1.0)) - sitk.WriteImage(pred_itk, test_save_path + - "/{}_pred.nii.gz".format(ids)) - - img_itk = sitk.GetImageFromArray(image) - img_itk.SetSpacing((1.0, 1.0, 1.0)) - sitk.WriteImage(img_itk, test_save_path + - "/{}_img.nii.gz".format(ids)) - - lab_itk = sitk.GetImageFromArray(label.astype(np.uint8)) - lab_itk.SetSpacing((1.0, 1.0, 1.0)) - sitk.WriteImage(lab_itk, test_save_path + - "/{}_lab.nii.gz".format(ids)) - f.writelines("Mean metrics,{},{},{},{}".format(total_metric[0, 0] / len(image_list), total_metric[0, 1] / len( - image_list), total_metric[0, 2] / len(image_list), total_metric[0, 3] / len(image_list))) - f.close() - print("Testing end") - return total_metric / len(image_list) - - -def cal_dice(prediction, label, num=2): - total_dice = np.zeros(num-1) - for i in range(1, num): - prediction_tmp = (prediction == i) - label_tmp = (label == i) - prediction_tmp = prediction_tmp.astype(np.float) - label_tmp = label_tmp.astype(np.float) - - dice = 2 * np.sum(prediction_tmp * label_tmp) / \ - (np.sum(prediction_tmp) + np.sum(label_tmp)) - total_dice[i - 1] += dice - - return total_dice - - -def calculate_metric_percase(pred, gt): - if pred.sum() > 0 and gt.sum() > 0: - dice = metric.binary.dc(pred, gt) - ravd = abs(metric.binary.ravd(pred, gt)) - hd = metric.binary.hd95(pred, gt) - asd = metric.binary.asd(pred, gt) - return np.array([dice, ravd, hd, asd]) - else: - return np.zeros(4) diff --git a/code/train_acdc_unet_semi_seg.sh b/code/train_acdc_unet_semi_seg.sh deleted file mode 100644 index 05a94ae..0000000 --- a/code/train_acdc_unet_semi_seg.sh +++ /dev/null @@ -1,8 +0,0 @@ -CUDA_VISIBLE_DEVICES=0 python train_fully_supervised_2D.py --root_path ../data/ACDC --exp ACDC/Fully_Supervised --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python train_entropy_minimization_2D.py --root_path ../data/ACDC --exp ACDC/Entropy_Minimization --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python train_interpolation_consistency_training_2D.py --root_path ../data/ACDC --exp ACDC/Interpolation_Consistency_Training --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python train_mean_teacher_2D.py --root_path ../data/ACDC --exp ACDC/Mean_Teacher --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python train_uncertainty_aware_mean_teacher_2D.py --root_path ../data/ACDC --exp ACDC/Uncertainty_Aware_Mean_Teacher --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python train_adversarial_network_2D.py --root_path ../data/ACDC --exp ACDC/Adversarial_Network --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python train_uncertainty_rectified_pyramid_consistency_2D.py --root_path ../data/ACDC --exp ACDC/Uncertainty_Rectified_Pyramid_Consistency --num_classes 4 --labeled_num 7 && \ -CUDA_VISIBLE_DEVICES=0 python train_fully_supervised_2D.py --root_path ../data/ACDC --exp ACDC/Fully_Supervised --num_classes 4 --labeled_num 140 \ No newline at end of file diff --git a/code/train_adversarial_network_2D.py b/code/train_adversarial_network_2D.py deleted file mode 100644 index 125566e..0000000 --- a/code/train_adversarial_network_2D.py +++ /dev/null @@ -1,283 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.dataset import (BaseDataSets, RandomGenerator, - TwoStreamBatchSampler) -from networks.discriminator import FCDiscriminator -from networks.net_factory import net_factory -from utils import losses, metrics, ramps -from val_2D import test_single_volume - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/ACDC', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='ACDC/Adversarial_Network', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=24, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--DAN_lr', type=float, default=0.0001, - help='DAN learning rate') -parser.add_argument('--patch_size', type=list, default=[256, 256], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') -parser.add_argument('--num_classes', type=int, default=4, - help='output channel of network') -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=12, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=3, - help='labeled data') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() - - -def patients_to_slices(dataset, patiens_num): - ref_dict = None - if "ACDC" in dataset: - ref_dict = {"3": 68, "7": 136, - "14": 256, "21": 396, "28": 512, "35": 664, "140": 1312} - elif "Prostate": - ref_dict = {"2": 27, "4": 53, "8": 120, - "12": 179, "16": 256, "21": 312, "42": 623} - else: - print("Error") - return ref_dict[str(patiens_num)] - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def train(args, snapshot_path): - base_lr = args.base_lr - num_classes = args.num_classes - batch_size = args.batch_size - max_iterations = args.max_iterations - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - model = net_factory(net_type=args.model, in_chns=1, class_num=num_classes) - - DAN = FCDiscriminator(num_classes=num_classes) - DAN = DAN.cuda() - - db_train = BaseDataSets(base_dir=args.root_path, split="train", num=None, transform=transforms.Compose([ - RandomGenerator(args.patch_size) - ])) - - total_slices = len(db_train) - labeled_slice = patients_to_slices(args.root_path, args.labeled_num) - print("Total silices is: {}, labeled slices is: {}".format( - total_slices, labeled_slice)) - labeled_idxs = list(range(0, labeled_slice)) - unlabeled_idxs = list(range(labeled_slice, total_slices)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=16, pin_memory=True, worker_init_fn=worker_init_fn) - - db_val = BaseDataSets(base_dir=args.root_path, split="val") - valloader = DataLoader(db_val, batch_size=1, shuffle=False, - num_workers=1) - - model.train() - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - DAN_optimizer = optim.Adam( - DAN.parameters(), lr=args.DAN_lr, betas=(0.9, 0.99)) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - - DAN_target = torch.tensor([0] * args.batch_size).cuda() - DAN_target[:args.labeled_bs] = 1 - model.train() - DAN.eval() - - outputs = model(volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - - loss_ce = ce_loss(outputs[:args.labeled_bs], - label_batch[:][:args.labeled_bs].long()) - loss_dice = dice_loss( - outputs_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - supervised_loss = 0.5 * (loss_dice + loss_ce) - - consistency_weight = get_current_consistency_weight(iter_num//150) - DAN_outputs = DAN( - outputs_soft[args.labeled_bs:], volume_batch[args.labeled_bs:]) - - consistency_loss = F.cross_entropy( - DAN_outputs, (DAN_target[:args.labeled_bs]).long()) - loss = supervised_loss + consistency_weight * consistency_loss - optimizer.zero_grad() - loss.backward() - optimizer.step() - - model.eval() - DAN.train() - with torch.no_grad(): - outputs = model(volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - - DAN_outputs = DAN(outputs_soft, volume_batch) - DAN_loss = F.cross_entropy(DAN_outputs, DAN_target.long()) - DAN_optimizer.zero_grad() - DAN_loss.backward() - DAN_optimizer.step() - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - - if iter_num % 20 == 0: - image = volume_batch[1, 0:1, :, :] - writer.add_image('train/Image', image, iter_num) - outputs = torch.argmax(torch.softmax( - outputs, dim=1), dim=1, keepdim=True) - writer.add_image('train/Prediction', - outputs[1, ...] * 50, iter_num) - labs = label_batch[1, ...].unsqueeze(0) * 50 - writer.add_image('train/GroundTruth', labs, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - metric_list = 0.0 - for i_batch, sampled_batch in enumerate(valloader): - metric_i = test_single_volume( - sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) - metric_list += np.array(metric_i) - metric_list = metric_list / len(db_val) - for class_i in range(num_classes-1): - writer.add_scalar('info/val_{}_dice'.format(class_i+1), - metric_list[class_i, 0], iter_num) - writer.add_scalar('info/val_{}_hd95'.format(class_i+1), - metric_list[class_i, 1], iter_num) - - performance = np.mean(metric_list, axis=0)[0] - - mean_hd95 = np.mean(metric_list, axis=0)[1] - writer.add_scalar('info/val_mean_dice', performance, iter_num) - writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) - - if performance > best_performance: - best_performance = performance - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - logging.info( - 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}_labeled/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_adversarial_network_3D.py b/code/train_adversarial_network_3D.py deleted file mode 100644 index 8cea3e2..0000000 --- a/code/train_adversarial_network_3D.py +++ /dev/null @@ -1,271 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.brats2019 import (BraTS2019, CenterCrop, RandomCrop, - RandomRotFlip, ToTensor, - TwoStreamBatchSampler) -from networks.discriminator import FC3DDiscriminator -from networks.net_factory_3d import net_factory_3d -from utils import losses, metrics, ramps -from val_3D import test_all_case - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/BraTS2019', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='BraTs2019_Adversarial_Network', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet_3D', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=4, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--DAN_lr', type=float, default=0.0001, - help='DAN learning rate') -parser.add_argument('--patch_size', type=list, default=[96, 96, 96], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=2, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=25, - help='labeled data') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def update_ema_variables(model, ema_model, alpha, global_step): - # Use the true average until the exponential average is more correct - alpha = min(1 - 1 / (global_step + 1), alpha) - for ema_param, param in zip(ema_model.parameters(), model.parameters()): - ema_param.data.mul_(alpha).add_(1 - alpha, param.data) - - -def train(args, snapshot_path): - num_classes = 2 - base_lr = args.base_lr - train_data_path = args.root_path - batch_size = args.batch_size - max_iterations = args.max_iterations - - net = net_factory_3d(net_type=args.model, in_chns=1, class_num=num_classes) - model = net.cuda() - DAN = FC3DDiscriminator(num_classes=num_classes) - DAN = DAN.cuda() - - db_train = BraTS2019(base_dir=train_data_path, - split='train', - num=None, - transform=transforms.Compose([ - RandomRotFlip(), - RandomCrop(args.patch_size), - ToTensor(), - ])) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - labeled_idxs = list(range(0, args.labeled_num)) - unlabeled_idxs = list(range(args.labeled_num, 250)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - model.train() - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - DAN_optimizer = optim.Adam( - DAN.parameters(), lr=args.DAN_lr, betas=(0.9, 0.99)) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(2) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - - DAN_target = torch.tensor([1, 1, 0, 0]).cuda() - model.train() - DAN.eval() - - outputs = model(volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - - loss_ce = ce_loss(outputs[:args.labeled_bs], - label_batch[:args.labeled_bs]) - loss_dice = dice_loss( - outputs_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - supervised_loss = 0.5 * (loss_dice + loss_ce) - - consistency_weight = get_current_consistency_weight(iter_num//150) - DAN_outputs = DAN( - outputs_soft[args.labeled_bs:], volume_batch[args.labeled_bs:]) - - consistency_loss = F.cross_entropy( - DAN_outputs, (DAN_target[:args.labeled_bs]).long()) - loss = supervised_loss + consistency_weight * consistency_loss - optimizer.zero_grad() - loss.backward() - optimizer.step() - - model.eval() - DAN.train() - with torch.no_grad(): - outputs = model(volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - - DAN_outputs = DAN(outputs_soft, volume_batch) - DAN_loss = F.cross_entropy(DAN_outputs, DAN_target.long()) - DAN_optimizer.zero_grad() - DAN_loss.backward() - DAN_optimizer.step() - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - - if iter_num % 20 == 0: - image = volume_batch[0, 0:1, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=True) - writer.add_image('train/Image', grid_image, iter_num) - - image = outputs_soft[0, 1:2, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Predicted_label', - grid_image, iter_num) - - image = label_batch[0, :, :, 20:61:10].unsqueeze( - 0).permute(3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Groundtruth_label', - grid_image, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - avg_metric = test_all_case( - model, args.root_path, test_list="val.txt", num_classes=2, patch_size=args.patch_size, - stride_xy=64, stride_z=64) - if avg_metric[:, 0].mean() > best_performance: - best_performance = avg_metric[:, 0].mean() - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - writer.add_scalar('info/val_dice_score', - avg_metric[0, 0], iter_num) - writer.add_scalar('info/val_hd95', - avg_metric[0, 1], iter_num) - logging.info( - 'iteration %d : dice_score : %f hd95 : %f' % (iter_num, avg_metric[0, 0].mean(), avg_metric[0, 1].mean())) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_brats2019_semi_seg.sh b/code/train_brats2019_semi_seg.sh deleted file mode 100644 index 38cc581..0000000 --- a/code/train_brats2019_semi_seg.sh +++ /dev/null @@ -1,10 +0,0 @@ -# & means run these methods at the same time, and && means run these methods one by one -python -u train_fully_supervised_3D.py --labeled_num 25 --root_path ../data/BraTS2019 --max_iterations 30000 --exp BraTS2019/Fully_supervised --base_lr 0.1 && -python -u train_fully_supervised_3D.py --labeled_num 250 --root_path ../data/BraTS2019 --max_iterations 30000 --exp BraTS2019/Fully_supervised --base_lr 0.1 && -python -u train_adversarial_network_3D.py --labeled_num 25 --total_num 250 --root_path ../data/BraTS2019 --max_iterations 30000 --exp BraTS2019/Adversarial_Network --base_lr 0.1 && -python -u train_entropy_minimization_3D.py --labeled_num 25 --total_num 250 --root_path ../data/BraTS2019 --max_iterations 30000 --exp BraTS2019/Entropy_Minimization --base_lr 0.1 && -python -u train_interpolation_consistency_training_3D.py --labeled_num 25 --total_num 250 --root_path ../data/BraTS2019 --max_iterations 30000 --base_lr 0.1 --exp BraTS2019/Interpolation_Consistency_Training && -python -u train_mean_teacher_3D.py --labeled_num 25 --total_num 250 --root_path ../data/BraTS2019 --max_iterations 30000 --exp BraTS2019/Mean_Teacher --base_lr 0.1 && -python -u train_uncertainty_aware_mean_teacher_3D.py --labeled_num 25 --total_num 250 --root_path ../data/BraTS2019 --max_iterations 30000 --base_lr 0.1 --exp BraTS2019/Uncertainty_Aware_Mean_Teacher && -python -u train_uncertainty_rectified_pyramid_consistency_3D.py --labeled_num 25 --total_num 250 --root_path ../data/BraTS2019 --max_iterations 30000 --base_lr 0.1 --exp BraTS2019/Uncertainty_Rectified_Pyramid_Consistency && -python -u train_cross_pseudo_supervision_3D.py --labeled_num 25 --total_num 250 --root_path ../data/BraTS2019 --max_iterations 30000 --base_lr 0.1 --exp BraTS2019/Cross_Pseudo_Supervision \ No newline at end of file diff --git a/code/train_cross_consistency_training_2D.py b/code/train_cross_consistency_training_2D.py deleted file mode 100644 index e89be74..0000000 --- a/code/train_cross_consistency_training_2D.py +++ /dev/null @@ -1,277 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.dataset import BaseDataSets, RandomGenerator, TwoStreamBatchSampler -from utils import losses, metrics, ramps -from val_2D import test_single_volume_ds -from networks.net_factory import net_factory - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/ACDC', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='ACDC/Cross_Consistency_Training', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet_cct', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=24, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[256, 256], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') -parser.add_argument('--num_classes', type=int, default=4, - help='output channel of network') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=12, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=7, - help='labeled data') -# costs -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() - - -def patients_to_slices(dataset, patiens_num): - ref_dict = None - if "ACDC" in dataset: - ref_dict = {"3": 68, "7": 136, - "14": 256, "21": 396, "28": 512, "35": 664, "140": 1312} - elif "Prostate": - ref_dict = {"2": 27, "4": 53, "8": 120, - "12": 179, "16": 256, "21": 312, "42": 623} - else: - print("Error") - return ref_dict[str(patiens_num)] - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def train(args, snapshot_path): - base_lr = args.base_lr - num_classes = args.num_classes - batch_size = args.batch_size - max_iterations = args.max_iterations - - model = net_factory(net_type=args.model, in_chns=1, - class_num=num_classes) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - db_train = BaseDataSets(base_dir=args.root_path, split="train", num=None, transform=transforms.Compose([ - RandomGenerator(args.patch_size) - ])) - db_val = BaseDataSets(base_dir=args.root_path, split="val") - total_slices = len(db_train) - labeled_slice = patients_to_slices(args.root_path, args.labeled_num) - print("Total silices is: {}, labeled slices is: {}".format( - total_slices, labeled_slice)) - labeled_idxs = list(range(0, labeled_slice)) - unlabeled_idxs = list(range(labeled_slice, total_slices)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size - args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - model.train() - - valloader = DataLoader(db_val, batch_size=1, shuffle=False, - num_workers=1) - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - kl_distance = nn.KLDivLoss(reduction='none') - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - - outputs, outputs_aux1, outputs_aux2, outputs_aux3 = model( - volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - outputs_aux1_soft = torch.softmax(outputs_aux1, dim=1) - outputs_aux2_soft = torch.softmax(outputs_aux2, dim=1) - outputs_aux3_soft = torch.softmax(outputs_aux3, dim=1) - - loss_ce = ce_loss(outputs[:args.labeled_bs], - label_batch[:args.labeled_bs][:].long()) - loss_ce_aux1 = ce_loss(outputs_aux1[:args.labeled_bs], - label_batch[:args.labeled_bs][:].long()) - loss_ce_aux2 = ce_loss(outputs_aux2[:args.labeled_bs], - label_batch[:args.labeled_bs][:].long()) - loss_ce_aux3 = ce_loss(outputs_aux3[:args.labeled_bs], - label_batch[:args.labeled_bs][:].long()) - - loss_dice = dice_loss( - outputs_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - loss_dice_aux1 = dice_loss( - outputs_aux1_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - loss_dice_aux2 = dice_loss( - outputs_aux2_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - loss_dice_aux3 = dice_loss( - outputs_aux3_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - - supervised_loss = (loss_ce + loss_ce_aux1 + loss_ce_aux2 + loss_ce_aux3 + - loss_dice + loss_dice_aux1 + loss_dice_aux2 + loss_dice_aux3) / 8 - - consistency_weight = get_current_consistency_weight(iter_num // 150) - consistency_loss_aux1 = torch.mean( - (outputs_soft[args.labeled_bs:] - outputs_aux1_soft[args.labeled_bs:]) ** 2) - consistency_loss_aux2 = torch.mean( - (outputs_soft[args.labeled_bs:] - outputs_aux2_soft[args.labeled_bs:]) ** 2) - consistency_loss_aux3 = torch.mean( - (outputs_soft[args.labeled_bs:] - outputs_aux3_soft[args.labeled_bs:]) ** 2) - - consistency_loss = (consistency_loss_aux1 + consistency_loss_aux2 + consistency_loss_aux3) / 3 - loss = supervised_loss + consistency_weight * consistency_loss - optimizer.zero_grad() - loss.backward() - optimizer.step() - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - - if iter_num % 20 == 0: - image = volume_batch[1, 0:1, :, :] - writer.add_image('train/Image', image, iter_num) - outputs = torch.argmax(torch.softmax( - outputs, dim=1), dim=1, keepdim=True) - writer.add_image('train/Prediction', - outputs[1, ...] * 50, iter_num) - labs = label_batch[1, ...].unsqueeze(0) * 50 - writer.add_image('train/GroundTruth', labs, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - metric_list = 0.0 - for i_batch, sampled_batch in enumerate(valloader): - metric_i = test_single_volume_ds( - sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) - metric_list += np.array(metric_i) - metric_list = metric_list / len(db_val) - for class_i in range(num_classes - 1): - writer.add_scalar('info/val_{}_dice'.format(class_i + 1), - metric_list[class_i, 0], iter_num) - writer.add_scalar('info/val_{}_hd95'.format(class_i + 1), - metric_list[class_i, 1], iter_num) - - performance = np.mean(metric_list, axis=0)[0] - - mean_hd95 = np.mean(metric_list, axis=0)[1] - writer.add_scalar('info/val_mean_dice', performance, iter_num) - writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) - - if performance > best_performance: - best_performance = performance - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - logging.info( - 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}_labeled/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_cross_pseudo_supervision_2D.py b/code/train_cross_pseudo_supervision_2D.py deleted file mode 100644 index 8ec818a..0000000 --- a/code/train_cross_pseudo_supervision_2D.py +++ /dev/null @@ -1,356 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.dataset import (BaseDataSets, RandomGenerator, - TwoStreamBatchSampler) -from networks.net_factory import net_factory -from utils import losses, metrics, ramps -from val_2D import test_single_volume - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/ACDC', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='ACDC/Cross_Pseudo_Supervision', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=24, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[256, 256], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') -parser.add_argument('--num_classes', type=int, default=4, - help='output channel of network') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=12, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=136, - help='labeled data') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() - -def kaiming_normal_init_weight(model): - for m in model.modules(): - if isinstance(m, nn.Conv2d): - torch.nn.init.kaiming_normal_(m.weight) - elif isinstance(m, nn.BatchNorm2d): - m.weight.data.fill_(1) - m.bias.data.zero_() - return model - -def xavier_normal_init_weight(model): - for m in model.modules(): - if isinstance(m, nn.Conv2d): - torch.nn.init.xavier_normal_(m.weight) - elif isinstance(m, nn.BatchNorm2d): - m.weight.data.fill_(1) - m.bias.data.zero_() - return model - -def patients_to_slices(dataset, patiens_num): - ref_dict = None - if "ACDC" in dataset: - ref_dict = {"3": 68, "7": 136, - "14": 256, "21": 396, "28": 512, "35": 664, "140": 1312} - elif "Prostate": - ref_dict = {"2": 27, "4": 53, "8": 120, - "12": 179, "16": 256, "21": 312, "42": 623} - else: - print("Error") - return ref_dict[str(patiens_num)] - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def update_ema_variables(model, ema_model, alpha, global_step): - # Use the true average until the exponential average is more correct - alpha = min(1 - 1 / (global_step + 1), alpha) - for ema_param, param in zip(ema_model.parameters(), model.parameters()): - ema_param.data.mul_(alpha).add_(1 - alpha, param.data) - - -def train(args, snapshot_path): - base_lr = args.base_lr - num_classes = args.num_classes - batch_size = args.batch_size - max_iterations = args.max_iterations - - def create_model(ema=False): - # Network definition - model = net_factory(net_type=args.model, in_chns=1, - class_num=num_classes) - if ema: - for param in model.parameters(): - param.detach_() - return model - - model1 = create_model() - model2 = create_model() - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - db_train = BaseDataSets(base_dir=args.root_path, split="train", num=None, transform=transforms.Compose([ - RandomGenerator(args.patch_size) - ])) - db_val = BaseDataSets(base_dir=args.root_path, split="val") - - total_slices = len(db_train) - labeled_slice = patients_to_slices(args.root_path, args.labeled_num) - print("Total silices is: {}, labeled slices is: {}".format( - total_slices, labeled_slice)) - labeled_idxs = list(range(0, labeled_slice)) - unlabeled_idxs = list(range(labeled_slice, total_slices)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - model1.train() - model2.train() - - valloader = DataLoader(db_val, batch_size=1, shuffle=False, - num_workers=1) - - optimizer1 = optim.SGD(model1.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - optimizer2 = optim.SGD(model2.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance1 = 0.0 - best_performance2 = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - - outputs1 = model1(volume_batch) - outputs_soft1 = torch.softmax(outputs1, dim=1) - - outputs2 = model2(volume_batch) - outputs_soft2 = torch.softmax(outputs2, dim=1) - consistency_weight = get_current_consistency_weight(iter_num // 150) - - loss1 = 0.5 * (ce_loss(outputs1[:args.labeled_bs], label_batch[:][:args.labeled_bs].long()) + dice_loss( - outputs_soft1[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1))) - loss2 = 0.5 * (ce_loss(outputs2[:args.labeled_bs], label_batch[:][:args.labeled_bs].long()) + dice_loss( - outputs_soft2[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1))) - - pseudo_outputs1 = torch.argmax(outputs_soft1[args.labeled_bs:].detach(), dim=1, keepdim=False) - pseudo_outputs2 = torch.argmax(outputs_soft2[args.labeled_bs:].detach(), dim=1, keepdim=False) - - pseudo_supervision1 = ce_loss(outputs1[args.labeled_bs:], pseudo_outputs2) - pseudo_supervision2 = ce_loss(outputs2[args.labeled_bs:], pseudo_outputs1) - - model1_loss = loss1 + consistency_weight * pseudo_supervision1 - model2_loss = loss2 + consistency_weight * pseudo_supervision2 - - loss = model1_loss + model2_loss - - - optimizer1.zero_grad() - optimizer2.zero_grad() - - loss.backward() - - optimizer1.step() - optimizer2.step() - - iter_num = iter_num + 1 - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer1.param_groups: - param_group['lr'] = lr_ - for param_group in optimizer2.param_groups: - param_group['lr'] = lr_ - - writer.add_scalar('lr', lr_, iter_num) - writer.add_scalar( - 'consistency_weight/consistency_weight', consistency_weight, iter_num) - writer.add_scalar('loss/model1_loss', - model1_loss, iter_num) - writer.add_scalar('loss/model2_loss', - model2_loss, iter_num) - logging.info('iteration %d : model1 loss : %f model2 loss : %f' % (iter_num, model1_loss.item(), model2_loss.item())) - if iter_num % 50 == 0: - image = volume_batch[1, 0:1, :, :] - writer.add_image('train/Image', image, iter_num) - outputs = torch.argmax(torch.softmax( - outputs1, dim=1), dim=1, keepdim=True) - writer.add_image('train/model1_Prediction', - outputs[1, ...] * 50, iter_num) - outputs = torch.argmax(torch.softmax( - outputs2, dim=1), dim=1, keepdim=True) - writer.add_image('train/model2_Prediction', - outputs[1, ...] * 50, iter_num) - labs = label_batch[1, ...].unsqueeze(0) * 50 - writer.add_image('train/GroundTruth', labs, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model1.eval() - metric_list = 0.0 - for i_batch, sampled_batch in enumerate(valloader): - metric_i = test_single_volume( - sampled_batch["image"], sampled_batch["label"], model1, classes=num_classes) - metric_list += np.array(metric_i) - metric_list = metric_list / len(db_val) - for class_i in range(num_classes-1): - writer.add_scalar('info/model1_val_{}_dice'.format(class_i+1), - metric_list[class_i, 0], iter_num) - writer.add_scalar('info/model1_val_{}_hd95'.format(class_i+1), - metric_list[class_i, 1], iter_num) - - performance1 = np.mean(metric_list, axis=0)[0] - - mean_hd951 = np.mean(metric_list, axis=0)[1] - writer.add_scalar('info/model1_val_mean_dice', performance1, iter_num) - writer.add_scalar('info/model1_val_mean_hd95', mean_hd951, iter_num) - - if performance1 > best_performance1: - best_performance1 = performance1 - save_mode_path = os.path.join(snapshot_path, - 'model1_iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance1, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model1.pth'.format(args.model)) - torch.save(model1.state_dict(), save_mode_path) - torch.save(model1.state_dict(), save_best) - - logging.info( - 'iteration %d : model1_mean_dice : %f model1_mean_hd95 : %f' % (iter_num, performance1, mean_hd951)) - model1.train() - - model2.eval() - metric_list = 0.0 - for i_batch, sampled_batch in enumerate(valloader): - metric_i = test_single_volume( - sampled_batch["image"], sampled_batch["label"], model2, classes=num_classes) - metric_list += np.array(metric_i) - metric_list = metric_list / len(db_val) - for class_i in range(num_classes-1): - writer.add_scalar('info/model2_val_{}_dice'.format(class_i+1), - metric_list[class_i, 0], iter_num) - writer.add_scalar('info/model2_val_{}_hd95'.format(class_i+1), - metric_list[class_i, 1], iter_num) - - performance2 = np.mean(metric_list, axis=0)[0] - - mean_hd952 = np.mean(metric_list, axis=0)[1] - writer.add_scalar('info/model2_val_mean_dice', performance2, iter_num) - writer.add_scalar('info/model2_val_mean_hd95', mean_hd952, iter_num) - - if performance2 > best_performance2: - best_performance2 = performance2 - save_mode_path = os.path.join(snapshot_path, - 'model2_iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance2))) - save_best = os.path.join(snapshot_path, - '{}_best_model2.pth'.format(args.model)) - torch.save(model2.state_dict(), save_mode_path) - torch.save(model2.state_dict(), save_best) - - logging.info( - 'iteration %d : model2_mean_dice : %f model2_mean_hd95 : %f' % (iter_num, performance2, mean_hd952)) - model2.train() - - # change lr - if iter_num % 2500 == 0: - lr_ = base_lr * 0.1 ** (iter_num // 2500) - for param_group in optimizer1.param_groups: - param_group['lr'] = lr_ - for param_group in optimizer2.param_groups: - param_group['lr'] = lr_ - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'model1_iter_' + str(iter_num) + '.pth') - torch.save(model1.state_dict(), save_mode_path) - logging.info("save model1 to {}".format(save_mode_path)) - - save_mode_path = os.path.join( - snapshot_path, 'model2_iter_' + str(iter_num) + '.pth') - torch.save(model2.state_dict(), save_mode_path) - logging.info("save model2 to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - time1 = time.time() - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_cross_pseudo_supervision_3D.py b/code/train_cross_pseudo_supervision_3D.py deleted file mode 100644 index 47ec42d..0000000 --- a/code/train_cross_pseudo_supervision_3D.py +++ /dev/null @@ -1,321 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.brats2019 import (BraTS2019, CenterCrop, RandomCrop, - RandomRotFlip, ToTensor, - TwoStreamBatchSampler) -from networks.net_factory_3d import net_factory_3d -from utils import losses, metrics, ramps -from val_3D import test_all_case - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/BraTS2019', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='BraTs2019_Cross_Pseudo_Supervision', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet_3D', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=4, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[96, 96, 96], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=2, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=25, - help='labeled data') - -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def update_ema_variables(model, ema_model, alpha, global_step): - # Use the true average until the exponential average is more correct - alpha = min(1 - 1 / (global_step + 1), alpha) - for ema_param, param in zip(ema_model.parameters(), model.parameters()): - ema_param.data.mul_(alpha).add_(1 - alpha, param.data) - - -def kaiming_normal_init_weight(model): - for m in model.modules(): - if isinstance(m, nn.Conv3d): - torch.nn.init.kaiming_normal_(m.weight) - elif isinstance(m, nn.BatchNorm3d): - m.weight.data.fill_(1) - m.bias.data.zero_() - return model - - -def xavier_normal_init_weight(model): - for m in model.modules(): - if isinstance(m, nn.Conv3d): - torch.nn.init.xavier_normal_(m.weight) - elif isinstance(m, nn.BatchNorm3d): - m.weight.data.fill_(1) - m.bias.data.zero_() - return model - - -def train(args, snapshot_path): - base_lr = args.base_lr - train_data_path = args.root_path - batch_size = args.batch_size - max_iterations = args.max_iterations - num_classes = 2 - - net1 = net_factory_3d(net_type=args.model, in_chns=1, class_num=num_classes).cuda() - net2 = net_factory_3d(net_type=args.model, in_chns=1, class_num=num_classes).cuda() - model1 = kaiming_normal_init_weight(net1) - model2 = xavier_normal_init_weight(net2) - model1.train() - model2.train() - - db_train = BraTS2019(base_dir=train_data_path, - split='train', - num=None, - transform=transforms.Compose([ - RandomRotFlip(), - RandomCrop(args.patch_size), - ToTensor(), - ])) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - labeled_idxs = list(range(0, args.labeled_num)) - unlabeled_idxs = list(range(args.labeled_num, 250)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size - args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - optimizer1 = optim.SGD(model1.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - optimizer2 = optim.SGD(model2.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - best_performance1 = 0.0 - best_performance2 = 0.0 - iter_num = 0 - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - max_epoch = max_iterations // len(trainloader) + 1 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - - outputs1 = model1(volume_batch) - outputs_soft1 = torch.softmax(outputs1, dim=1) - - outputs2 = model2(volume_batch) - outputs_soft2 = torch.softmax(outputs2, dim=1) - consistency_weight = get_current_consistency_weight(iter_num // 150) - - loss1 = 0.5 * (ce_loss(outputs1[:args.labeled_bs], - label_batch[:][:args.labeled_bs].long()) + dice_loss( - outputs_soft1[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1))) - loss2 = 0.5 * (ce_loss(outputs2[:args.labeled_bs], - label_batch[:][:args.labeled_bs].long()) + dice_loss( - outputs_soft2[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1))) - - pseudo_outputs1 = torch.argmax(outputs_soft1[args.labeled_bs:].detach(), dim=1, keepdim=False) - pseudo_outputs2 = torch.argmax(outputs_soft2[args.labeled_bs:].detach(), dim=1, keepdim=False) - - pseudo_supervision1 = ce_loss(outputs1[args.labeled_bs:], pseudo_outputs2) - pseudo_supervision2 = ce_loss(outputs2[args.labeled_bs:], pseudo_outputs1) - - model1_loss = loss1 + consistency_weight * pseudo_supervision1 - model2_loss = loss2 + consistency_weight * pseudo_supervision2 - - loss = model1_loss + model2_loss - - optimizer1.zero_grad() - optimizer2.zero_grad() - - loss.backward() - - optimizer1.step() - optimizer2.step() - - iter_num = iter_num + 1 - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group1 in optimizer1.param_groups: - param_group1['lr'] = lr_ - for param_group2 in optimizer2.param_groups: - param_group2['lr'] = lr_ - - writer.add_scalar('lr', lr_, iter_num) - writer.add_scalar( - 'consistency_weight/consistency_weight', consistency_weight, iter_num) - writer.add_scalar('loss/model1_loss', - model1_loss, iter_num) - writer.add_scalar('loss/model2_loss', - model2_loss, iter_num) - logging.info( - 'iteration %d : model1 loss : %f model2 loss : %f' % (iter_num, model1_loss.item(), model2_loss.item())) - if iter_num % 50 == 0: - image = volume_batch[0, 0:1, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=True) - writer.add_image('train/Image', grid_image, iter_num) - - image = outputs_soft1[0, 0:1, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Model1_Predicted_label', - grid_image, iter_num) - - image = outputs_soft2[0, 0:1, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Model2_Predicted_label', - grid_image, iter_num) - - image = label_batch[0, :, :, 20:61:10].unsqueeze( - 0).permute(3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Groundtruth_label', - grid_image, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model1.eval() - avg_metric1 = test_all_case( - model1, args.root_path, test_list="val.txt", num_classes=2, patch_size=args.patch_size, - stride_xy=64, stride_z=64) - if avg_metric1[:, 0].mean() > best_performance1: - best_performance1 = avg_metric1[:, 0].mean() - save_mode_path = os.path.join(snapshot_path, - 'model1_iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance1, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model1.pth'.format(args.model)) - torch.save(model1.state_dict(), save_mode_path) - torch.save(model1.state_dict(), save_best) - - writer.add_scalar('info/model1_val_dice_score', - avg_metric1[0, 0], iter_num) - writer.add_scalar('info/model1_val_hd95', - avg_metric1[0, 1], iter_num) - logging.info( - 'iteration %d : model1_dice_score : %f model1_hd95 : %f' % ( - iter_num, avg_metric1[0, 0].mean(), avg_metric1[0, 1].mean())) - model1.train() - - model2.eval() - avg_metric2 = test_all_case( - model2, args.root_path, test_list="val.txt", num_classes=2, patch_size=args.patch_size, - stride_xy=64, stride_z=64) - if avg_metric2[:, 0].mean() > best_performance2: - best_performance2 = avg_metric2[:, 0].mean() - save_mode_path = os.path.join(snapshot_path, - 'model2_iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance2, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model2.pth'.format(args.model)) - torch.save(model2.state_dict(), save_mode_path) - torch.save(model2.state_dict(), save_best) - - writer.add_scalar('info/model2_val_dice_score', - avg_metric2[0, 0], iter_num) - writer.add_scalar('info/model2_val_hd95', - avg_metric2[0, 1], iter_num) - logging.info( - 'iteration %d : model2_dice_score : %f model2_hd95 : %f' % ( - iter_num, avg_metric2[0, 0].mean(), avg_metric2[0, 1].mean())) - model2.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'model1_iter_' + str(iter_num) + '.pth') - torch.save(model1.state_dict(), save_mode_path) - logging.info("save model1 to {}".format(save_mode_path)) - - save_mode_path = os.path.join( - snapshot_path, 'model2_iter_' + str(iter_num) + '.pth') - torch.save(model2.state_dict(), save_mode_path) - logging.info("save model2 to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - time1 = time.time() - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_cross_teaching_between_cnn_transformer_2D.py b/code/train_cross_teaching_between_cnn_transformer_2D.py deleted file mode 100644 index 375f7a4..0000000 --- a/code/train_cross_teaching_between_cnn_transformer_2D.py +++ /dev/null @@ -1,413 +0,0 @@ -# -*- coding: utf-8 -*- -# Author: Xiangde Luo -# Date: 16 Dec. 2021 -# Implementation for Semi-Supervised Medical Image Segmentation via Cross Teaching between CNN and Transformer. -# # Reference: -# @article{luo2021ctbct, -# title={Semi-Supervised Medical Image Segmentation via Cross Teaching between CNN and Transformer}, -# author={Luo, Xiangde and Hu, Minhao and Song, Tao and Wang, Guotai and Zhang, Shaoting}, -# journal={arXiv preprint arXiv:2112.04894}, -# year={2021}} -# In the original paper, we don't use the validation set to select checkpoints and use the last iteration to inference for all methods. -# In addition, we combine the validation set and test set to report the results. -# We found that the random data split has some bias (the validation set is very tough and the test set is very easy). -# Actually, this setting is also a fair comparison. -# download pre-trained model to "code/pretrained_ckpt" folder, link:https://drive.google.com/drive/folders/1UC3XOoezeum0uck4KBVGa8osahs6rKUY - -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from tqdm import tqdm - -from config import get_config -from dataloaders import utils -from dataloaders.dataset import (BaseDataSets, RandomGenerator, - TwoStreamBatchSampler) -from networks.net_factory import net_factory -from networks.vision_transformer import SwinUnet as ViT_seg -from utils import losses, metrics, ramps -from val_2D import test_single_volume - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/ACDC', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='ACDC/Cross_Teaching_Between_CNN_Transformer', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=16, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[224, 224], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') -parser.add_argument('--num_classes', type=int, default=4, - help='output channel of network') -parser.add_argument( - '--cfg', type=str, default="../code/configs/swin_tiny_patch4_window7_224_lite.yaml", help='path to config file', ) -parser.add_argument( - "--opts", - help="Modify config options by adding 'KEY VALUE' pairs. ", - default=None, - nargs='+', -) -parser.add_argument('--zip', action="iframe.php?url=https%3A%2F%2Fgithub.com%2Fstore_true", - help='use zipped dataset instead of folder dataset') -parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'], - help='no: no cache, ' - 'full: cache all data, ' - 'part: sharding the dataset into nonoverlapping pieces and only cache one piece') -parser.add_argument('--resume', help='resume from checkpoint') -parser.add_argument('--accumulation-steps', type=int, - help="gradient accumulation steps") -parser.add_argument('--use-checkpoint', action="iframe.php?url=https%3A%2F%2Fgithub.com%2Fstore_true", - help="whether to use gradient checkpointing to save memory") -parser.add_argument('--amp-opt-level', type=str, default='O1', choices=['O0', 'O1', 'O2'], - help='mixed precision opt level, if O0, no amp is used') -parser.add_argument('--tag', help='tag of experiment') -parser.add_argument('--eval', action="iframe.php?url=https%3A%2F%2Fgithub.com%2Fstore_true", - help='Perform evaluation only') -parser.add_argument('--throughput', action="iframe.php?url=https%3A%2F%2Fgithub.com%2Fstore_true", - help='Test throughput only') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=8, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=7, - help='labeled data') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() -config = get_config(args) - - -def kaiming_normal_init_weight(model): - for m in model.modules(): - if isinstance(m, nn.Conv2d): - torch.nn.init.kaiming_normal_(m.weight) - elif isinstance(m, nn.BatchNorm2d): - m.weight.data.fill_(1) - m.bias.data.zero_() - return model - - -def xavier_normal_init_weight(model): - for m in model.modules(): - if isinstance(m, nn.Conv2d): - torch.nn.init.xavier_normal_(m.weight) - elif isinstance(m, nn.BatchNorm2d): - m.weight.data.fill_(1) - m.bias.data.zero_() - return model - - -def patients_to_slices(dataset, patiens_num): - ref_dict = None - if "ACDC" in dataset: - ref_dict = {"3": 68, "7": 136, - "14": 256, "21": 396, "28": 512, "35": 664, "140": 1312} - elif "Prostate": - ref_dict = {"2": 27, "4": 53, "8": 120, - "12": 179, "16": 256, "21": 312, "42": 623} - else: - print("Error") - return ref_dict[str(patiens_num)] - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def update_ema_variables(model, ema_model, alpha, global_step): - # Use the true average until the exponential average is more correct - alpha = min(1 - 1 / (global_step + 1), alpha) - for ema_param, param in zip(ema_model.parameters(), model.parameters()): - ema_param.data.mul_(alpha).add_(1 - alpha, param.data) - - -def train(args, snapshot_path): - base_lr = args.base_lr - num_classes = args.num_classes - batch_size = args.batch_size - max_iterations = args.max_iterations - - def create_model(ema=False): - # Network definition - model = net_factory(net_type=args.model, in_chns=1, - class_num=num_classes) - if ema: - for param in model.parameters(): - param.detach_() - return model - - model1 = create_model() - model2 = ViT_seg(config, img_size=args.patch_size, - num_classes=args.num_classes).cuda() - model2.load_from(config) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - db_train = BaseDataSets(base_dir=args.root_path, split="train", num=None, transform=transforms.Compose([ - RandomGenerator(args.patch_size) - ])) - db_val = BaseDataSets(base_dir=args.root_path, split="val") - - total_slices = len(db_train) - labeled_slice = patients_to_slices(args.root_path, args.labeled_num) - print("Total silices is: {}, labeled slices is: {}".format( - total_slices, labeled_slice)) - labeled_idxs = list(range(0, labeled_slice)) - unlabeled_idxs = list(range(labeled_slice, total_slices)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - model1.train() - model2.train() - - valloader = DataLoader(db_val, batch_size=1, shuffle=False, - num_workers=1) - - optimizer1 = optim.SGD(model1.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - optimizer2 = optim.SGD(model2.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance1 = 0.0 - best_performance2 = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - - outputs1 = model1(volume_batch) - outputs_soft1 = torch.softmax(outputs1, dim=1) - - outputs2 = model2(volume_batch) - outputs_soft2 = torch.softmax(outputs2, dim=1) - consistency_weight = get_current_consistency_weight( - iter_num // 150) - - loss1 = 0.5 * (ce_loss(outputs1[:args.labeled_bs], label_batch[:args.labeled_bs].long()) + dice_loss( - outputs_soft1[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1))) - loss2 = 0.5 * (ce_loss(outputs2[:args.labeled_bs], label_batch[:args.labeled_bs].long()) + dice_loss( - outputs_soft2[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1))) - - pseudo_outputs1 = torch.argmax( - outputs_soft1[args.labeled_bs:].detach(), dim=1, keepdim=False) - pseudo_outputs2 = torch.argmax( - outputs_soft2[args.labeled_bs:].detach(), dim=1, keepdim=False) - - pseudo_supervision1 = dice_loss( - outputs_soft1[args.labeled_bs:], pseudo_outputs2.unsqueeze(1)) - pseudo_supervision2 = dice_loss( - outputs_soft2[args.labeled_bs:], pseudo_outputs1.unsqueeze(1)) - - model1_loss = loss1 + consistency_weight * pseudo_supervision1 - model2_loss = loss2 + consistency_weight * pseudo_supervision2 - - loss = model1_loss + model2_loss - - optimizer1.zero_grad() - optimizer2.zero_grad() - - loss.backward() - - optimizer1.step() - optimizer2.step() - - iter_num = iter_num + 1 - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer1.param_groups: - param_group['lr'] = lr_ - for param_group in optimizer2.param_groups: - param_group['lr'] = lr_ - - writer.add_scalar('lr', lr_, iter_num) - writer.add_scalar( - 'consistency_weight/consistency_weight', consistency_weight, iter_num) - writer.add_scalar('loss/model1_loss', - model1_loss, iter_num) - writer.add_scalar('loss/model2_loss', - model2_loss, iter_num) - logging.info('iteration %d : model1 loss : %f model2 loss : %f' % ( - iter_num, model1_loss.item(), model2_loss.item())) - if iter_num % 50 == 0: - image = volume_batch[1, 0:1, :, :] - writer.add_image('train/Image', image, iter_num) - outputs = torch.argmax(torch.softmax( - outputs1, dim=1), dim=1, keepdim=True) - writer.add_image('train/model1_Prediction', - outputs[1, ...] * 50, iter_num) - outputs = torch.argmax(torch.softmax( - outputs2, dim=1), dim=1, keepdim=True) - writer.add_image('train/model2_Prediction', - outputs[1, ...] * 50, iter_num) - labs = label_batch[1, ...].unsqueeze(0) * 50 - writer.add_image('train/GroundTruth', labs, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model1.eval() - metric_list = 0.0 - for i_batch, sampled_batch in enumerate(valloader): - metric_i = test_single_volume( - sampled_batch["image"], sampled_batch["label"], model1, classes=num_classes, patch_size=args.patch_size) - metric_list += np.array(metric_i) - metric_list = metric_list / len(db_val) - for class_i in range(num_classes-1): - writer.add_scalar('info/model1_val_{}_dice'.format(class_i+1), - metric_list[class_i, 0], iter_num) - writer.add_scalar('info/model1_val_{}_hd95'.format(class_i+1), - metric_list[class_i, 1], iter_num) - - performance1 = np.mean(metric_list, axis=0)[0] - - mean_hd951 = np.mean(metric_list, axis=0)[1] - writer.add_scalar('info/model1_val_mean_dice', - performance1, iter_num) - writer.add_scalar('info/model1_val_mean_hd95', - mean_hd951, iter_num) - - if performance1 > best_performance1: - best_performance1 = performance1 - save_mode_path = os.path.join(snapshot_path, - 'model1_iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance1, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model1.pth'.format(args.model)) - torch.save(model1.state_dict(), save_mode_path) - torch.save(model1.state_dict(), save_best) - - logging.info( - 'iteration %d : model1_mean_dice : %f model1_mean_hd95 : %f' % (iter_num, performance1, mean_hd951)) - model1.train() - - model2.eval() - metric_list = 0.0 - for i_batch, sampled_batch in enumerate(valloader): - metric_i = test_single_volume( - sampled_batch["image"], sampled_batch["label"], model2, classes=num_classes, patch_size=args.patch_size) - metric_list += np.array(metric_i) - metric_list = metric_list / len(db_val) - for class_i in range(num_classes-1): - writer.add_scalar('info/model2_val_{}_dice'.format(class_i+1), - metric_list[class_i, 0], iter_num) - writer.add_scalar('info/model2_val_{}_hd95'.format(class_i+1), - metric_list[class_i, 1], iter_num) - - performance2 = np.mean(metric_list, axis=0)[0] - - mean_hd952 = np.mean(metric_list, axis=0)[1] - writer.add_scalar('info/model2_val_mean_dice', - performance2, iter_num) - writer.add_scalar('info/model2_val_mean_hd95', - mean_hd952, iter_num) - - if performance2 > best_performance2: - best_performance2 = performance2 - save_mode_path = os.path.join(snapshot_path, - 'model2_iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance2, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model2.pth'.format(args.model)) - torch.save(model2.state_dict(), save_mode_path) - torch.save(model2.state_dict(), save_best) - - logging.info( - 'iteration %d : model2_mean_dice : %f model2_mean_hd95 : %f' % (iter_num, performance2, mean_hd952)) - model2.train() - - # change lr - if iter_num % 2500 == 0: - lr_ = base_lr * 0.1 ** (iter_num // 2500) - for param_group in optimizer1.param_groups: - param_group['lr'] = lr_ - for param_group in optimizer2.param_groups: - param_group['lr'] = lr_ - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'model1_iter_' + str(iter_num) + '.pth') - torch.save(model1.state_dict(), save_mode_path) - logging.info("save model1 to {}".format(save_mode_path)) - - save_mode_path = os.path.join( - snapshot_path, 'model2_iter_' + str(iter_num) + '.pth') - torch.save(model2.state_dict(), save_mode_path) - logging.info("save model2 to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - time1 = time.time() - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_deep_co_training_2D.py b/code/train_deep_co_training_2D.py deleted file mode 100644 index 636f3d4..0000000 --- a/code/train_deep_co_training_2D.py +++ /dev/null @@ -1,267 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.dataset import (BaseDataSets, RandomGenerator, - TwoStreamBatchSampler) -from networks.discriminator import FCDiscriminator -from networks.net_factory import net_factory -from utils import losses, metrics, ramps -from val_2D import test_single_volume - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/ACDC', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='ACDC/Deep_Co_Training', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=24, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[256, 256], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') -parser.add_argument('--num_classes', type=int, default=4, - help='output channel of network') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=12, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=3, - help='labeled data') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() - - -def patients_to_slices(dataset, patiens_num): - ref_dict = None - if "ACDC" in dataset: - ref_dict = {"3": 68, "7": 136, - "14": 256, "21": 396, "28": 512, "35": 664, "140": 1312} - elif "Prostate": - ref_dict = {"2": 27, "4": 53, "8": 120, - "12": 179, "16": 256, "21": 312, "42": 623} - else: - print("Error") - return ref_dict[str(patiens_num)] - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def train(args, snapshot_path): - base_lr = args.base_lr - num_classes = args.num_classes - batch_size = args.batch_size - max_iterations = args.max_iterations - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - model = net_factory(net_type=args.model, in_chns=1, class_num=num_classes) - - db_train = BaseDataSets(base_dir=args.root_path, split="train", num=None, transform=transforms.Compose([ - RandomGenerator(args.patch_size) - ])) - - total_slices = len(db_train) - labeled_slice = patients_to_slices(args.root_path, args.labeled_num) - print("Total silices is: {}, labeled slices is: {}".format( - total_slices, labeled_slice)) - labeled_idxs = list(range(0, labeled_slice)) - unlabeled_idxs = list(range(labeled_slice, total_slices)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=16, pin_memory=True, worker_init_fn=worker_init_fn) - - db_val = BaseDataSets(base_dir=args.root_path, split="val") - valloader = DataLoader(db_val, batch_size=1, shuffle=False, - num_workers=1) - - model.train() - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - unlabeled_volume_batch = volume_batch[args.labeled_bs:] - - outputs = model(volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - - rot_times = random.randrange(0,4) - - rotated_unlabeled_volume_batch = torch.rot90(unlabeled_volume_batch, rot_times, [2,3]) - - unlabeled_rot_outputs = model(rotated_unlabeled_volume_batch) - unlabeled_rot_outputs_soft = torch.softmax(unlabeled_rot_outputs, dim=1) - - loss_ce = ce_loss(outputs[:args.labeled_bs], - label_batch[:][:args.labeled_bs].long()) - loss_dice = dice_loss( - outputs_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - supervised_loss = 0.5 * (loss_dice + loss_ce) - - consistency_weight = get_current_consistency_weight(iter_num//150) - - consistency_loss = 0.5 * (torch.mean((unlabeled_rot_outputs_soft.detach() - torch.rot90(outputs_soft[args.labeled_bs:], rot_times, [2,3]))**2) + torch.mean((unlabeled_rot_outputs_soft - torch.rot90(outputs_soft[args.labeled_bs:].detach(), rot_times, [2,3]))**2)) - - loss = supervised_loss + consistency_weight * consistency_loss - optimizer.zero_grad() - loss.backward() - optimizer.step() - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - - if iter_num % 20 == 0: - image = volume_batch[1, 0:1, :, :] - writer.add_image('train/Image', image, iter_num) - outputs = torch.argmax(torch.softmax( - outputs, dim=1), dim=1, keepdim=True) - writer.add_image('train/Prediction', - outputs[1, ...] * 50, iter_num) - labs = label_batch[1, ...].unsqueeze(0) * 50 - writer.add_image('train/GroundTruth', labs, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - metric_list = 0.0 - for i_batch, sampled_batch in enumerate(valloader): - metric_i = test_single_volume( - sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) - metric_list += np.array(metric_i) - metric_list = metric_list / len(db_val) - for class_i in range(num_classes-1): - writer.add_scalar('info/val_{}_dice'.format(class_i+1), - metric_list[class_i, 0], iter_num) - writer.add_scalar('info/val_{}_hd95'.format(class_i+1), - metric_list[class_i, 1], iter_num) - - performance = np.mean(metric_list, axis=0)[0] - - mean_hd95 = np.mean(metric_list, axis=0)[1] - writer.add_scalar('info/val_mean_dice', performance, iter_num) - writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) - - if performance > best_performance: - best_performance = performance - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - logging.info( - 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}_labeled/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_entropy_minimization_2D.py b/code/train_entropy_minimization_2D.py deleted file mode 100644 index a111703..0000000 --- a/code/train_entropy_minimization_2D.py +++ /dev/null @@ -1,258 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.dataset import (BaseDataSets, RandomGenerator, - TwoStreamBatchSampler) -from networks.discriminator import FCDiscriminator -from networks.net_factory import net_factory -from utils import losses, metrics, ramps -from val_2D import test_single_volume - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/ACDC', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='ACDC/Entropy_Minimization', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=24, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[256, 256], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') -parser.add_argument('--num_classes', type=int, default=4, - help='output channel of network') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=12, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=3, - help='labeled data') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() - - -def patients_to_slices(dataset, patiens_num): - ref_dict = None - if "ACDC" in dataset: - ref_dict = {"3": 68, "7": 136, - "14": 256, "21": 396, "28": 512, "35": 664, "140": 1312} - elif "Prostate": - ref_dict = {"2": 27, "4": 53, "8": 120, - "12": 179, "16": 256, "21": 312, "42": 623} - else: - print("Error") - return ref_dict[str(patiens_num)] - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def train(args, snapshot_path): - base_lr = args.base_lr - num_classes = args.num_classes - batch_size = args.batch_size - max_iterations = args.max_iterations - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - model = net_factory(net_type=args.model, in_chns=1, class_num=num_classes) - - db_train = BaseDataSets(base_dir=args.root_path, split="train", num=None, transform=transforms.Compose([ - RandomGenerator(args.patch_size) - ])) - - total_slices = len(db_train) - labeled_slice = patients_to_slices(args.root_path, args.labeled_num) - print("Total silices is: {}, labeled slices is: {}".format( - total_slices, labeled_slice)) - labeled_idxs = list(range(0, labeled_slice)) - unlabeled_idxs = list(range(labeled_slice, total_slices)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=16, pin_memory=True, worker_init_fn=worker_init_fn) - - db_val = BaseDataSets(base_dir=args.root_path, split="val") - valloader = DataLoader(db_val, batch_size=1, shuffle=False, - num_workers=1) - - model.train() - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - unlabeled_volume_batch = volume_batch[args.labeled_bs:] - - outputs = model(volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - - loss_ce = ce_loss(outputs[:args.labeled_bs], - label_batch[:][:args.labeled_bs].long()) - loss_dice = dice_loss( - outputs_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - supervised_loss = 0.5 * (loss_dice + loss_ce) - - consistency_weight = get_current_consistency_weight(iter_num//150) - consistency_loss = losses.entropy_loss(outputs_soft, C=4) - loss = supervised_loss + consistency_weight * consistency_loss - optimizer.zero_grad() - loss.backward() - optimizer.step() - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - - if iter_num % 20 == 0: - image = volume_batch[1, 0:1, :, :] - writer.add_image('train/Image', image, iter_num) - outputs = torch.argmax(torch.softmax( - outputs, dim=1), dim=1, keepdim=True) - writer.add_image('train/Prediction', - outputs[1, ...] * 50, iter_num) - labs = label_batch[1, ...].unsqueeze(0) * 50 - writer.add_image('train/GroundTruth', labs, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - metric_list = 0.0 - for i_batch, sampled_batch in enumerate(valloader): - metric_i = test_single_volume( - sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) - metric_list += np.array(metric_i) - metric_list = metric_list / len(db_val) - for class_i in range(num_classes-1): - writer.add_scalar('info/val_{}_dice'.format(class_i+1), - metric_list[class_i, 0], iter_num) - writer.add_scalar('info/val_{}_hd95'.format(class_i+1), - metric_list[class_i, 1], iter_num) - - performance = np.mean(metric_list, axis=0)[0] - - mean_hd95 = np.mean(metric_list, axis=0)[1] - writer.add_scalar('info/val_mean_dice', performance, iter_num) - writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) - - if performance > best_performance: - best_performance = performance - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - logging.info( - 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}_labeled/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_entropy_minimization_3D.py b/code/train_entropy_minimization_3D.py deleted file mode 100644 index be5d316..0000000 --- a/code/train_entropy_minimization_3D.py +++ /dev/null @@ -1,246 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.brats2019 import (BraTS2019, CenterCrop, RandomCrop, - RandomRotFlip, ToTensor, - TwoStreamBatchSampler) -from networks.net_factory_3d import net_factory_3d -from utils import losses, metrics, ramps -from val_3D import test_all_case - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/BraTS2019', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='BraTs2019_Entropy_Minimization', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet_3D', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=4, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[96, 96, 96], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=2, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=25, - help='labeled data') - -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def update_ema_variables(model, ema_model, alpha, global_step): - # Use the true average until the exponential average is more correct - alpha = min(1 - 1 / (global_step + 1), alpha) - for ema_param, param in zip(ema_model.parameters(), model.parameters()): - ema_param.data.mul_(alpha).add_(1 - alpha, param.data) - - -def train(args, snapshot_path): - base_lr = args.base_lr - train_data_path = args.root_path - batch_size = args.batch_size - max_iterations = args.max_iterations - num_classes = 2 - - net = net_factory_3d(net_type=args.model, in_chns=1, class_num=num_classes) - model = net.cuda() - - db_train = BraTS2019(base_dir=train_data_path, - split='train', - num=None, - transform=transforms.Compose([ - RandomRotFlip(), - RandomCrop(args.patch_size), - ToTensor(), - ])) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - labeled_idxs = list(range(0, args.labeled_num)) - unlabeled_idxs = list(range(args.labeled_num, 250)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - model.train() - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - unlabeled_volume_batch = volume_batch[args.labeled_bs:] - - outputs = model(volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - - loss_ce = ce_loss(outputs[:args.labeled_bs], - label_batch[:args.labeled_bs]) - loss_dice = dice_loss( - outputs_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - supervised_loss = 0.5 * (loss_dice + loss_ce) - - consistency_weight = get_current_consistency_weight(iter_num//150) - consistency_loss = losses.entropy_loss(outputs_soft, C=2) - loss = supervised_loss + consistency_weight * consistency_loss - optimizer.zero_grad() - loss.backward() - optimizer.step() - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - - if iter_num % 20 == 0: - image = volume_batch[0, 0:1, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=True) - writer.add_image('train/Image', grid_image, iter_num) - - image = outputs_soft[0, 1:2, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Predicted_label', - grid_image, iter_num) - - image = label_batch[0, :, :, 20:61:10].unsqueeze( - 0).permute(3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Groundtruth_label', - grid_image, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - avg_metric = test_all_case( - model, args.root_path, test_list="val.txt", num_classes=2, patch_size=args.patch_size, - stride_xy=64, stride_z=64) - if avg_metric[:, 0].mean() > best_performance: - best_performance = avg_metric[:, 0].mean() - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - writer.add_scalar('info/val_dice_score', - avg_metric[0, 0], iter_num) - writer.add_scalar('info/val_hd95', - avg_metric[0, 1], iter_num) - logging.info( - 'iteration %d : dice_score : %f hd95 : %f' % (iter_num, avg_metric[0, 0].mean(), avg_metric[0, 1].mean())) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_fully_supervised_2D.py b/code/train_fully_supervised_2D.py deleted file mode 100644 index 588ddc4..0000000 --- a/code/train_fully_supervised_2D.py +++ /dev/null @@ -1,218 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.dataset import BaseDataSets, RandomGenerator -from networks.net_factory import net_factory -from utils import losses, metrics, ramps -from val_2D import test_single_volume, test_single_volume_ds - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/ACDC', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='ACDC/Fully_Supervised', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet', help='model_name') -parser.add_argument('--num_classes', type=int, default=4, - help='output channel of network') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=24, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[256, 256], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') -parser.add_argument('--labeled_num', type=int, default=50, - help='labeled data') -args = parser.parse_args() - - -def patients_to_slices(dataset, patiens_num): - ref_dict = None - if "ACDC" in dataset: - ref_dict = {"3": 68, "7": 136, - "14": 256, "21": 396, "28": 512, "35": 664, "140": 1312} - elif "Prostate": - ref_dict = {"2": 27, "4": 53, "8": 120, - "12": 179, "16": 256, "21": 312, "42": 623} - else: - print("Error") - return ref_dict[str(patiens_num)] - - -def train(args, snapshot_path): - base_lr = args.base_lr - num_classes = args.num_classes - batch_size = args.batch_size - max_iterations = args.max_iterations - - labeled_slice = patients_to_slices(args.root_path, args.labeled_num) - - model = net_factory(net_type=args.model, in_chns=1, class_num=num_classes) - db_train = BaseDataSets(base_dir=args.root_path, split="train", num=labeled_slice, transform=transforms.Compose([ - RandomGenerator(args.patch_size) - ])) - db_val = BaseDataSets(base_dir=args.root_path, split="val") - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - trainloader = DataLoader(db_train, batch_size=batch_size, shuffle=True, - num_workers=16, pin_memory=True, worker_init_fn=worker_init_fn) - valloader = DataLoader(db_val, batch_size=1, shuffle=False, - num_workers=1) - - model.train() - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - - outputs = model(volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - - loss_ce = ce_loss(outputs, label_batch[:].long()) - loss_dice = dice_loss(outputs_soft, label_batch.unsqueeze(1)) - loss = 0.5 * (loss_dice + loss_ce) - optimizer.zero_grad() - loss.backward() - optimizer.step() - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - - if iter_num % 20 == 0: - image = volume_batch[1, 0:1, :, :] - writer.add_image('train/Image', image, iter_num) - outputs = torch.argmax(torch.softmax( - outputs, dim=1), dim=1, keepdim=True) - writer.add_image('train/Prediction', - outputs[1, ...] * 50, iter_num) - labs = label_batch[1, ...].unsqueeze(0) * 50 - writer.add_image('train/GroundTruth', labs, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - metric_list = 0.0 - for i_batch, sampled_batch in enumerate(valloader): - metric_i = test_single_volume( - sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) - metric_list += np.array(metric_i) - metric_list = metric_list / len(db_val) - for class_i in range(num_classes-1): - writer.add_scalar('info/val_{}_dice'.format(class_i+1), - metric_list[class_i, 0], iter_num) - writer.add_scalar('info/val_{}_hd95'.format(class_i+1), - metric_list[class_i, 1], iter_num) - - performance = np.mean(metric_list, axis=0)[0] - - mean_hd95 = np.mean(metric_list, axis=0)[1] - writer.add_scalar('info/val_mean_dice', performance, iter_num) - writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) - - if performance > best_performance: - best_performance = performance - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - logging.info( - 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}_labeled/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_fully_supervised_3D.py b/code/train_fully_supervised_3D.py deleted file mode 100644 index 5f01b15..0000000 --- a/code/train_fully_supervised_3D.py +++ /dev/null @@ -1,203 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.brats2019 import (BraTS2019, CenterCrop, RandomCrop, - RandomRotFlip, ToTensor, - TwoStreamBatchSampler) -from networks.net_factory_3d import net_factory_3d -from utils import losses, metrics, ramps -from val_3D import test_all_case - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/BraTS2019', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='BraTs2019_Fully_Supervised', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet_3D', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=2, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[96, 96, 96], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') -parser.add_argument('--labeled_num', type=int, default=25, - help='labeled data') - -args = parser.parse_args() - - -def train(args, snapshot_path): - base_lr = args.base_lr - train_data_path = args.root_path - batch_size = args.batch_size - max_iterations = args.max_iterations - num_classes = 2 - model = net_factory_3d(net_type=args.model, in_chns=1, class_num=num_classes) - db_train = BraTS2019(base_dir=train_data_path, - split='train', - num=args.labeled_num, - transform=transforms.Compose([ - RandomRotFlip(), - RandomCrop(args.patch_size), - ToTensor(), - ])) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - trainloader = DataLoader(db_train, batch_size=batch_size, shuffle=True, - num_workers=16, pin_memory=True, worker_init_fn=worker_init_fn) - - model.train() - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(2) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - - outputs = model(volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - - loss_ce = ce_loss(outputs, label_batch) - loss_dice = dice_loss(outputs_soft, label_batch.unsqueeze(1)) - loss = 0.5 * (loss_dice + loss_ce) - optimizer.zero_grad() - loss.backward() - optimizer.step() - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - writer.add_scalar('loss/loss', loss, iter_num) - - if iter_num % 20 == 0: - image = volume_batch[0, 0:1, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=True) - writer.add_image('train/Image', grid_image, iter_num) - - image = outputs_soft[0, 1:2, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Predicted_label', - grid_image, iter_num) - - image = label_batch[0, :, :, 20:61:10].unsqueeze( - 0).permute(3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Groundtruth_label', - grid_image, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - avg_metric = test_all_case( - model, args.root_path, test_list="val.txt", num_classes=2, patch_size=args.patch_size, - stride_xy=64, stride_z=64) - if avg_metric[:, 0].mean() > best_performance: - best_performance = avg_metric[:, 0].mean() - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - writer.add_scalar('info/val_dice_score', - avg_metric[0, 0], iter_num) - writer.add_scalar('info/val_hd95', - avg_metric[0, 1], iter_num) - logging.info( - 'iteration %d : dice_score : %f hd95 : %f' % (iter_num, avg_metric[0, 0].mean(), avg_metric[0, 1].mean())) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}/{}".format(args.exp, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_interpolation_consistency_training_2D.py b/code/train_interpolation_consistency_training.py similarity index 60% rename from code/train_interpolation_consistency_training_2D.py rename to code/train_interpolation_consistency_training.py index bb1aed0..8f98d70 100644 --- a/code/train_interpolation_consistency_training_2D.py +++ b/code/train_interpolation_consistency_training.py @@ -1,301 +1,290 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.dataset import (BaseDataSets, RandomGenerator, - TwoStreamBatchSampler) -from networks.net_factory import net_factory -from utils import losses, metrics, ramps -from val_2D import test_single_volume - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/ACDC', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='ACDC/Interpolation_Consistency_Training', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=24, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[256, 256], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') -parser.add_argument('--num_classes', type=int, default=4, - help='output channel of network') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=12, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=300, - help='labeled data') -parser.add_argument('--ict_alpha', type=int, default=0.2, - help='ict_alpha') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() - - -def patients_to_slices(dataset, patiens_num): - ref_dict = None - if "ACDC" in dataset: - ref_dict = {"3": 68, "7": 136, - "14": 256, "21": 396, "28": 512, "35": 664, "140": 1312} - elif "Prostate": - ref_dict = {"2": 27, "4": 53, "8": 120, - "12": 179, "16": 256, "21": 312, "42": 623} - else: - print("Error") - return ref_dict[str(patiens_num)] - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def update_ema_variables(model, ema_model, alpha, global_step): - # Use the true average until the exponential average is more correct - alpha = min(1 - 1 / (global_step + 1), alpha) - for ema_param, param in zip(ema_model.parameters(), model.parameters()): - ema_param.data.mul_(alpha).add_(1 - alpha, param.data) - - -def train(args, snapshot_path): - base_lr = args.base_lr - num_classes = args.num_classes - batch_size = args.batch_size - max_iterations = args.max_iterations - - def create_model(ema=False): - # Network definition - model = net_factory(net_type=args.model, in_chns=1, - class_num=num_classes) - if ema: - for param in model.parameters(): - param.detach_() - return model - - model = create_model() - ema_model = create_model(ema=True) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - db_train = BaseDataSets(base_dir=args.root_path, split="train", num=None, transform=transforms.Compose([ - RandomGenerator(args.patch_size) - ])) - db_val = BaseDataSets(base_dir=args.root_path, split="val") - - total_slices = len(db_train) - labeled_slice = patients_to_slices(args.root_path, args.labeled_num) - print("Total silices is: {}, labeled slices is: {}".format( - total_slices, labeled_slice)) - labeled_idxs = list(range(0, labeled_slice)) - unlabeled_idxs = list(range(labeled_slice, total_slices)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - model.train() - - valloader = DataLoader(db_val, batch_size=1, shuffle=False, - num_workers=1) - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - unlabeled_volume_batch = volume_batch[args.labeled_bs:] - labeled_volume_batch = volume_batch[:args.labeled_bs] - - # ICT mix factors - ict_mix_factors = np.random.beta( - args.ict_alpha, args.ict_alpha, size=(args.labeled_bs//2, 1, 1, 1)) - ict_mix_factors = torch.tensor( - ict_mix_factors, dtype=torch.float).cuda() - unlabeled_volume_batch_0 = unlabeled_volume_batch[0:args.labeled_bs//2, ...] - unlabeled_volume_batch_1 = unlabeled_volume_batch[args.labeled_bs//2:, ...] - - # Mix images - batch_ux_mixed = unlabeled_volume_batch_0 * \ - (1.0 - ict_mix_factors) + \ - unlabeled_volume_batch_1 * ict_mix_factors - input_volume_batch = torch.cat( - [labeled_volume_batch, batch_ux_mixed], dim=0) - outputs = model(input_volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - with torch.no_grad(): - ema_output_ux0 = torch.softmax( - ema_model(unlabeled_volume_batch_0), dim=1) - ema_output_ux1 = torch.softmax( - ema_model(unlabeled_volume_batch_1), dim=1) - batch_pred_mixed = ema_output_ux0 * \ - (1.0 - ict_mix_factors) + ema_output_ux1 * ict_mix_factors - - loss_ce = ce_loss(outputs[:args.labeled_bs], - label_batch[:args.labeled_bs][:].long()) - loss_dice = dice_loss( - outputs_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - supervised_loss = 0.5 * (loss_dice + loss_ce) - consistency_weight = get_current_consistency_weight(iter_num//150) - consistency_loss = torch.mean( - (outputs_soft[args.labeled_bs:] - batch_pred_mixed) ** 2) - loss = supervised_loss + consistency_weight * consistency_loss - - optimizer.zero_grad() - loss.backward() - optimizer.step() - update_ema_variables(model, ema_model, args.ema_decay, iter_num) - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - - if iter_num % 20 == 0: - image = volume_batch[1, 0:1, :, :] - writer.add_image('train/Image', image, iter_num) - outputs = torch.argmax(torch.softmax( - outputs, dim=1), dim=1, keepdim=True) - writer.add_image('train/Prediction', - outputs[1, ...] * 50, iter_num) - image = batch_ux_mixed[1, 0:1, :, :] - writer.add_image('train/Mixed_Unlabeled', - image, iter_num) - labs = label_batch[1, ...].unsqueeze(0) * 50 - writer.add_image('train/GroundTruth', labs, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - metric_list = 0.0 - for i_batch, sampled_batch in enumerate(valloader): - metric_i = test_single_volume( - sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) - metric_list += np.array(metric_i) - metric_list = metric_list / len(db_val) - for class_i in range(num_classes-1): - writer.add_scalar('info/val_{}_dice'.format(class_i+1), - metric_list[class_i, 0], iter_num) - writer.add_scalar('info/val_{}_hd95'.format(class_i+1), - metric_list[class_i, 1], iter_num) - - performance = np.mean(metric_list, axis=0)[0] - - mean_hd95 = np.mean(metric_list, axis=0)[1] - writer.add_scalar('info/val_mean_dice', performance, iter_num) - writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) - - if performance > best_performance: - best_performance = performance - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - logging.info( - 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}_labeled/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) +import argparse +import logging +import os +import random +import shutil +import sys +import time +from itertools import cycle +import numpy as np +import torch +import torch.backends.cudnn as cudnn +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from tensorboardX import SummaryWriter +from torch.nn import BCEWithLogitsLoss +from torch.nn.modules.loss import CrossEntropyLoss +from torch.utils.data import DataLoader +from torchvision import transforms +from torchvision.utils import make_grid +from tqdm import tqdm + +from dataloaders.dataset import BaseDataSets, RandomGenerator +from networks.discriminator import FCDiscriminator +from networks.net_factory import net_factory +from utils import losses, metrics, ramps +from val_2D import test_single_volume + +parser = argparse.ArgumentParser() +parser.add_argument('--root_path', type=str, + default='../data/ACDC', help='Name of Experiment') +parser.add_argument('--exp', type=str, + default='ACDC/ICT', help='experiment_name') +parser.add_argument('--model', type=str, + default='unet', help='model_name') +parser.add_argument('--fold', type=int, + default=2, help='cross validation') +parser.add_argument('--max_iterations', type=int, + default=30000, help='maximum epoch number to train') +parser.add_argument('--batch_size', type=int, default=12, + help='batch_size per gpu') + +parser.add_argument('--deterministic', type=int, default=1, + help='whether use deterministic training') +parser.add_argument('--base_lr', type=float, default=0.01, + help='segmentation network learning rate') +parser.add_argument('--patch_size', type=list, default=[256, 256], + help='patch size of network input') +parser.add_argument('--seed', type=int, default=2022, help='random seed') +parser.add_argument('--num_classes', type=int, default=4, + help='output channel of network') + +# label and unlabel +parser.add_argument('--labeled_ratio', type=int, default=5, + help='1/labeled_ratio data is provided mask') +# costs +parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') +parser.add_argument('--consistency_type', type=str, + default="mse", help='consistency_type') +parser.add_argument('--consistency', type=float, + default=0.1, help='consistency') +parser.add_argument('--consistency_rampup', type=float, + default=200.0, help='consistency_rampup') +args = parser.parse_args() + + +def get_current_consistency_weight(epoch): + # Consistency ramp-up from https://arxiv.org/abs/1610.02242 + return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) + + +def update_ema_variables(model, ema_model, alpha, global_step): + # Use the true average until the exponential average is more correct + alpha = min(1 - 1 / (global_step + 1), alpha) + for ema_param, param in zip(ema_model.parameters(), model.parameters()): + ema_param.data.mul_(alpha).add_(1 - alpha, param.data) + + +def train(args, snapshot_path): + writer = SummaryWriter(snapshot_path + '/log') + base_lr = args.base_lr + num_classes = args.num_classes + max_iterations = args.max_iterations + + def worker_init_fn(worker_id): + random.seed(args.seed + worker_id) + + def create_model(ema=False): + # Network definition + model = net_factory(net_type=args.model, in_chns=1, + class_num=num_classes) + if ema: + for param in model.parameters(): + param.detach_() + return model + + model = create_model() + ema_model = create_model(ema=True) + + db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)])) + db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)])) + + trainloader_labeled = DataLoader( + db_train_labeled, batch_size=args.batch_size//2, shuffle=True) + trainloader_unlabeled = DataLoader( + db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) + logging.info("Labeled slices: {} ".format(len(db_train_labeled))) + logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) + + db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, + split="val", labeled_ratio=args.labeled_ratio) + valloader = DataLoader(db_val, batch_size=1, shuffle=False) + + model.train() + + optimizer = optim.SGD(model.parameters(), lr=base_lr, + momentum=0.9, weight_decay=0.0001) + + ce_loss = CrossEntropyLoss(ignore_index=4) + dice_loss = losses.DiceLoss(num_classes) + + logging.info("{} iterations per epoch".format(len(trainloader_labeled))) + + iter_num = 0 + max_epoch = max_iterations // len(trainloader_labeled) + 1 + best_performance = 0.0 + iterator = tqdm(range(max_epoch), ncols=70) + for epoch_num in iterator: + for i, data in enumerate(zip(cycle(trainloader_labeled), trainloader_unlabeled)): + sampled_batch_labeled, sampled_batch_unlabeled = data[0], data[1] + + volume_batch, label_batch = sampled_batch_labeled['image'], sampled_batch_labeled['label'] + volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() + unlabeled_volume_batch = sampled_batch_unlabeled['image'].cuda() + + outputs = model(volume_batch) + outputs_soft = torch.softmax(outputs, dim=1) + + supervised_loss = 0.5 * \ + (ce_loss(outputs, label_batch[:].long( + )) + dice_loss(outputs_soft, label_batch[:].unsqueeze(1))) + + if unlabeled_volume_batch.shape[0] != args.batch_size // 2: + loss = supervised_loss + consistency_weight = 0.0 + consistency_loss = 0.0 + else: + # ICT mix factors + ict_alpha = 0.2 + ict_mix_factors = np.random.beta( + ict_alpha, ict_alpha, size=(args.batch_size // 4, 1, 1, 1)) + ict_mix_factors = torch.tensor( + ict_mix_factors, dtype=torch.float).cuda() + unlabeled_volume_batch_0 = unlabeled_volume_batch[0:args.batch_size // 4, ...] + unlabeled_volume_batch_1 = unlabeled_volume_batch[args.batch_size // 4:, ...] + + # Mix images + batch_ux_mixed = unlabeled_volume_batch_0 * \ + (1.0 - ict_mix_factors) + \ + unlabeled_volume_batch_1 * ict_mix_factors + # input_volume_batch = torch.cat( + # [volume_batch, batch_ux_mixed], dim=0) + + outputs_unlabeled = model(batch_ux_mixed) + outputs_unlabeled_soft = torch.softmax(outputs_unlabeled, dim=1) + + with torch.no_grad(): + ema_output_ux0 = torch.softmax( + ema_model(unlabeled_volume_batch_0), dim=1) + ema_output_ux1 = torch.softmax( + ema_model(unlabeled_volume_batch_1), dim=1) + batch_pred_mixed = ema_output_ux0 * \ + (1.0 - ict_mix_factors) + ema_output_ux1 * ict_mix_factors + + consistency_weight = get_current_consistency_weight( + iter_num // 150) + consistency_loss = torch.mean( + (outputs_unlabeled_soft - batch_pred_mixed) ** 2) + loss = supervised_loss + consistency_weight * consistency_loss + + optimizer.zero_grad() + loss.backward() + optimizer.step() + update_ema_variables(model, ema_model, args.ema_decay, iter_num) + + lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 + for param_group in optimizer.param_groups: + param_group['lr'] = lr_ + + iter_num = iter_num + 1 + writer.add_scalar('info/lr', lr_, iter_num) + writer.add_scalar('info/total_loss', loss, iter_num) + writer.add_scalar('info/loss_ce', supervised_loss, iter_num) + writer.add_scalar('info/consistency_loss', + consistency_loss, iter_num) + writer.add_scalar('info/consistency_weight', + consistency_weight, iter_num) + + logging.info( + 'iteration %d : loss : %f, loss_ce: %f' % + (iter_num, loss.item(), supervised_loss.item())) + + if iter_num % 20 == 0: + image = volume_batch[0, 0:1, :, :] + writer.add_image('train/Image', image, iter_num) + outputs = torch.argmax(torch.softmax( + outputs, dim=1), dim=1, keepdim=True) + writer.add_image('train/Prediction', + outputs[0, ...] * 50, iter_num) + labs = label_batch[0, ...].unsqueeze(0) * 50 + writer.add_image('train/GroundTruth', labs, iter_num) + + if iter_num > 0 and iter_num % 200 == 0: + model.eval() + metric_list = 0.0 + for i_batch, sampled_batch in enumerate(valloader): + metric_i = test_single_volume( + sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) + metric_list += np.array(metric_i) + metric_list = metric_list / len(db_val) + for class_i in range(num_classes-1): + writer.add_scalar('info/val_{}_dice'.format(class_i+1), + metric_list[class_i, 0], iter_num) + writer.add_scalar('info/val_{}_hd95'.format(class_i+1), + metric_list[class_i, 1], iter_num) + + performance = np.mean(metric_list, axis=0)[0] + + mean_hd95 = np.mean(metric_list, axis=0)[1] + writer.add_scalar('info/val_mean_dice', performance, iter_num) + writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) + + if performance > best_performance: + best_performance = performance + save_mode_path = os.path.join(snapshot_path, + 'iter_{}_dice_{}.pth'.format( + iter_num, round(best_performance, 4))) + save_best = os.path.join(snapshot_path, + '{}_best_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_mode_path) + torch.save(model.state_dict(), save_best) + + logging.info( + 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) + model.train() + + if iter_num % 3000 == 0: + save_mode_path = os.path.join( + snapshot_path, 'iter_' + str(iter_num) + '.pth') + torch.save(model.state_dict(), save_mode_path) + logging.info("save model to {}".format(save_mode_path)) + + if iter_num >= max_iterations: + break + if iter_num >= max_iterations: + iterator.close() + break + writer.close() + return "Training Finished!" + + +if __name__ == "__main__": + if not args.deterministic: + cudnn.benchmark = True + cudnn.deterministic = False + else: + cudnn.benchmark = False + cudnn.deterministic = True + + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( + args.exp, args.labeled_ratio, args.fold) + if not os.path.exists(snapshot_path): + os.makedirs(snapshot_path) + if os.path.exists(snapshot_path + '/code'): + shutil.rmtree(snapshot_path + '/code') + shutil.copytree('.', snapshot_path + '/code', + shutil.ignore_patterns(['.git', '__pycache__'])) + + logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO, + format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') + logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) + logging.info(str(args)) + train(args, snapshot_path) diff --git a/code/train_interpolation_consistency_training_3D.py b/code/train_interpolation_consistency_training_3D.py deleted file mode 100644 index 689f76f..0000000 --- a/code/train_interpolation_consistency_training_3D.py +++ /dev/null @@ -1,284 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.brats2019 import (BraTS2019, CenterCrop, RandomCrop, - RandomRotFlip, ToTensor, - TwoStreamBatchSampler) -from networks.net_factory_3d import net_factory_3d -from utils import losses, metrics, ramps -from val_3D import test_all_case - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/BraTS2019', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='BraTS2019_Interpolation_Consistency_Training', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet_3D', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=4, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[96, 96, 96], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=2, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=14, - help='labeled data') -parser.add_argument('--total_labeled_num', type=int, default=140, - help='total labeled data') -parser.add_argument('--ict_alpha', type=int, default=0.2, - help='ict_alpha') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') - -args = parser.parse_args() - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def update_ema_variables(model, ema_model, alpha, global_step): - # Use the true average until the exponential average is more correct - alpha = min(1 - 1 / (global_step + 1), alpha) - for ema_param, param in zip(ema_model.parameters(), model.parameters()): - ema_param.data.mul_(alpha).add_(1 - alpha, param.data) - - -def train(args, snapshot_path): - base_lr = args.base_lr - train_data_path = args.root_path - batch_size = args.batch_size - max_iterations = args.max_iterations - num_classes = 2 - - def create_model(ema=False): - # Network definition - net = net_factory_3d(net_type=args.model, in_chns=1, class_num=num_classes) - model = net.cuda() - if ema: - for param in model.parameters(): - param.detach_() - return model - - model = create_model() - ema_model = create_model(ema=True) - - db_train = BraTS2019(base_dir=train_data_path, - split='train', - num=None, - transform=transforms.Compose([ - RandomRotFlip(), - RandomCrop(args.patch_size), - ToTensor(), - ])) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - labeled_idxs = list(range(0, args.labeled_num)) - unlabeled_idxs = list(range(args.labeled_num, args.total_labeled_num)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - model.train() - ema_model.train() - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(2) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - labeled_volume_batch = volume_batch[:args.labeled_bs] - unlabeled_volume_batch = volume_batch[args.labeled_bs:] - - # ICT mix factors - ict_mix_factors = np.random.beta( - args.ict_alpha, args.ict_alpha, size=(args.labeled_bs//2, 1, 1, 1, 1)) - ict_mix_factors = torch.tensor( - ict_mix_factors, dtype=torch.float).cuda() - unlabeled_volume_batch_0 = unlabeled_volume_batch[0:1, ...] - unlabeled_volume_batch_1 = unlabeled_volume_batch[1:2, ...] - - # Mix images - batch_ux_mixed = unlabeled_volume_batch_0 * \ - (1.0 - ict_mix_factors) + \ - unlabeled_volume_batch_1 * ict_mix_factors - input_volume_batch = torch.cat( - [labeled_volume_batch, batch_ux_mixed], dim=0) - outputs = model(input_volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - with torch.no_grad(): - ema_output_ux0 = torch.softmax( - ema_model(unlabeled_volume_batch_0), dim=1) - ema_output_ux1 = torch.softmax( - ema_model(unlabeled_volume_batch_1), dim=1) - batch_pred_mixed = ema_output_ux0 * \ - (1.0 - ict_mix_factors) + ema_output_ux1 * ict_mix_factors - - loss_ce = ce_loss(outputs[:args.labeled_bs], - label_batch[:args.labeled_bs][:]) - loss_dice = dice_loss( - outputs_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - supervised_loss = 0.5 * (loss_dice + loss_ce) - consistency_weight = get_current_consistency_weight(iter_num//150) - consistency_loss = torch.mean( - (outputs_soft[args.labeled_bs:] - batch_pred_mixed)**2) - loss = supervised_loss + consistency_weight * consistency_loss - optimizer.zero_grad() - loss.backward() - optimizer.step() - update_ema_variables(model, ema_model, args.ema_decay, iter_num) - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - writer.add_scalar('loss/loss', loss, iter_num) - - if iter_num % 20 == 0: - image = volume_batch[0, 0:1, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=True) - writer.add_image('train/Image', grid_image, iter_num) - - image = outputs_soft[0, 1:2, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Predicted_label', - grid_image, iter_num) - - image = label_batch[0, :, :, 20:61:10].unsqueeze( - 0).permute(3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Groundtruth_label', - grid_image, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - avg_metric = test_all_case( - model, args.root_path, test_list="val.txt", num_classes=2, patch_size=args.patch_size, - stride_xy=32, stride_z=32) - if avg_metric[:, 0].mean() > best_performance: - best_performance = avg_metric[:, 0].mean() - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - writer.add_scalar('info/val_dice_score', - avg_metric[0, 0], iter_num) - writer.add_scalar('info/val_hd95', - avg_metric[0, 1], iter_num) - logging.info( - 'iteration %d : dice_score : %f hd95 : %f' % (iter_num, avg_metric[0, 0].mean(), avg_metric[0, 1].mean())) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}_labeled/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_mean_teacher_2D.py b/code/train_mean_teacher.py similarity index 68% rename from code/train_mean_teacher_2D.py rename to code/train_mean_teacher.py index 591b5cc..eed285b 100644 --- a/code/train_mean_teacher_2D.py +++ b/code/train_mean_teacher.py @@ -1,283 +1,258 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.dataset import (BaseDataSets, RandomGenerator, - TwoStreamBatchSampler) -from networks.net_factory import net_factory -from utils import losses, metrics, ramps -from val_2D import test_single_volume - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/ACDC', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='ACDC/Mean_Teacher', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=24, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[256, 256], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') -parser.add_argument('--num_classes', type=int, default=4, - help='output channel of network') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=12, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=136, - help='labeled data') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() - - -def patients_to_slices(dataset, patiens_num): - ref_dict = None - if "ACDC" in dataset: - ref_dict = {"3": 68, "7": 136, - "14": 256, "21": 396, "28": 512, "35": 664, "140": 1312} - elif "Prostate": - ref_dict = {"2": 27, "4": 53, "8": 120, - "12": 179, "16": 256, "21": 312, "42": 623} - else: - print("Error") - return ref_dict[str(patiens_num)] - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def update_ema_variables(model, ema_model, alpha, global_step): - # Use the true average until the exponential average is more correct - alpha = min(1 - 1 / (global_step + 1), alpha) - for ema_param, param in zip(ema_model.parameters(), model.parameters()): - ema_param.data.mul_(alpha).add_(1 - alpha, param.data) - - -def train(args, snapshot_path): - base_lr = args.base_lr - num_classes = args.num_classes - batch_size = args.batch_size - max_iterations = args.max_iterations - - def create_model(ema=False): - # Network definition - model = net_factory(net_type=args.model, in_chns=1, - class_num=num_classes) - if ema: - for param in model.parameters(): - param.detach_() - return model - - model = create_model() - ema_model = create_model(ema=True) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - db_train = BaseDataSets(base_dir=args.root_path, split="train", num=None, transform=transforms.Compose([ - RandomGenerator(args.patch_size) - ])) - db_val = BaseDataSets(base_dir=args.root_path, split="val") - - total_slices = len(db_train) - labeled_slice = patients_to_slices(args.root_path, args.labeled_num) - print("Total silices is: {}, labeled slices is: {}".format( - total_slices, labeled_slice)) - labeled_idxs = list(range(0, labeled_slice)) - unlabeled_idxs = list(range(labeled_slice, total_slices)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - model.train() - - valloader = DataLoader(db_val, batch_size=1, shuffle=False, - num_workers=1) - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - unlabeled_volume_batch = volume_batch[args.labeled_bs:] - - noise = torch.clamp(torch.randn_like( - unlabeled_volume_batch) * 0.1, -0.2, 0.2) - ema_inputs = unlabeled_volume_batch + noise - - outputs = model(volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - with torch.no_grad(): - ema_output = ema_model(ema_inputs) - ema_output_soft = torch.softmax(ema_output, dim=1) - - loss_ce = ce_loss(outputs[:args.labeled_bs], - label_batch[:][:args.labeled_bs].long()) - loss_dice = dice_loss( - outputs_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - supervised_loss = 0.5 * (loss_dice + loss_ce) - consistency_weight = get_current_consistency_weight(iter_num//150) - if iter_num < 1000: - consistency_loss = 0.0 - else: - consistency_loss = torch.mean( - (outputs_soft[args.labeled_bs:]-ema_output_soft)**2) - loss = supervised_loss + consistency_weight * consistency_loss - optimizer.zero_grad() - loss.backward() - optimizer.step() - update_ema_variables(model, ema_model, args.ema_decay, iter_num) - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - - if iter_num % 20 == 0: - image = volume_batch[1, 0:1, :, :] - writer.add_image('train/Image', image, iter_num) - outputs = torch.argmax(torch.softmax( - outputs, dim=1), dim=1, keepdim=True) - writer.add_image('train/Prediction', - outputs[1, ...] * 50, iter_num) - labs = label_batch[1, ...].unsqueeze(0) * 50 - writer.add_image('train/GroundTruth', labs, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - metric_list = 0.0 - for i_batch, sampled_batch in enumerate(valloader): - metric_i = test_single_volume( - sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) - metric_list += np.array(metric_i) - metric_list = metric_list / len(db_val) - for class_i in range(num_classes-1): - writer.add_scalar('info/val_{}_dice'.format(class_i+1), - metric_list[class_i, 0], iter_num) - writer.add_scalar('info/val_{}_hd95'.format(class_i+1), - metric_list[class_i, 1], iter_num) - - performance = np.mean(metric_list, axis=0)[0] - - mean_hd95 = np.mean(metric_list, axis=0)[1] - writer.add_scalar('info/val_mean_dice', performance, iter_num) - writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) - - if performance > best_performance: - best_performance = performance - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - logging.info( - 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}_labeled/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) +import argparse +import logging +import os +import random +import shutil +import sys +import time +from itertools import cycle +import numpy as np +import torch +import torch.backends.cudnn as cudnn +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from tensorboardX import SummaryWriter +from torch.nn import BCEWithLogitsLoss +from torch.nn.modules.loss import CrossEntropyLoss +from torch.utils.data import DataLoader +from torchvision import transforms +from torchvision.utils import make_grid +from tqdm import tqdm + +from dataloaders.dataset import BaseDataSets, RandomGenerator +from networks.discriminator import FCDiscriminator +from networks.net_factory import net_factory +from utils import losses, metrics, ramps +from val_2D import test_single_volume + +parser = argparse.ArgumentParser() +parser.add_argument('--root_path', type=str, + default='../data/ProstateX', help='Name of Experiment') +parser.add_argument('--exp', type=str, + default='ProstateX/Mean_Teacher', help='experiment_name') +parser.add_argument('--model', type=str, + default='unet', help='model_name') +parser.add_argument('--fold', type=int, + default=3, help='cross validation') +parser.add_argument('--max_iterations', type=int, + default=30000, help='maximum epoch number to train') +parser.add_argument('--batch_size', type=int, default=16, + help='batch_size per gpu') + +parser.add_argument('--deterministic', type=int, default=1, + help='whether use deterministic training') +parser.add_argument('--base_lr', type=float, default=0.01, + help='segmentation network learning rate') +parser.add_argument('--patch_size', type=list, default=[256, 256], + help='patch size of network input') +parser.add_argument('--seed', type=int, default=2022, help='random seed') +parser.add_argument('--num_classes', type=int, default=3, + help='output channel of network') + +# label and unlabel +parser.add_argument('--labeled_ratio', type=int, default=8, + help='1/labeled_ratio data is provided mask') +# costs +parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') +parser.add_argument('--consistency_type', type=str, + default="mse", help='consistency_type') +parser.add_argument('--consistency', type=float, + default=0.1, help='consistency') +parser.add_argument('--consistency_rampup', type=float, + default=200.0, help='consistency_rampup') +args = parser.parse_args() + + +def get_current_consistency_weight(epoch): + # Consistency ramp-up from https://arxiv.org/abs/1610.02242 + return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) + +def update_ema_variables(model, ema_model, alpha, global_step): + # Use the true average until the exponential average is more correct + alpha = min(1 - 1 / (global_step + 1), alpha) + for ema_param, param in zip(ema_model.parameters(), model.parameters()): + ema_param.data.mul_(alpha).add_(1 - alpha, param.data) + + +def train(args, snapshot_path): + writer = SummaryWriter(snapshot_path + '/log') + base_lr = args.base_lr + num_classes = args.num_classes + max_iterations = args.max_iterations + + def worker_init_fn(worker_id): + random.seed(args.seed + worker_id) + + def create_model(ema=False): + # Network definition + model = net_factory(net_type=args.model, in_chns=1, + class_num=num_classes) + if ema: + for param in model.parameters(): + param.detach_() + return model + + model = create_model() + ema_model = create_model(ema=True) + + db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)])) + db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)])) + logging.info("Labeled slices: {} ".format(len(db_train_labeled))) + logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) + + trainloader_labeled = DataLoader(db_train_labeled, batch_size=args.batch_size//2, shuffle=True) + trainloader_unlabeled = DataLoader(db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) + + db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, split="val", labeled_ratio=args.labeled_ratio) + valloader = DataLoader(db_val, batch_size=1) + + model.train() + + optimizer = optim.SGD(model.parameters(), lr=base_lr, + momentum=0.9, weight_decay=0.0001) + + ce_loss = CrossEntropyLoss() + dice_loss = losses.DiceLoss(num_classes) + + logging.info("{} iterations per epoch".format(len(trainloader_labeled))) + + iter_num = 0 + max_epoch = max_iterations // len(trainloader_unlabeled) + 1 + best_performance = 0.0 + iterator = tqdm(range(max_epoch), ncols=70) + for epoch_num in iterator: + for i, (sampled_batch_labeled, sampled_batch_unlabeled) in enumerate(zip(cycle(trainloader_labeled), trainloader_unlabeled)): + volume_batch, label_batch = sampled_batch_labeled['image'], sampled_batch_labeled['label'] + volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() + unlabeled_volume_batch = sampled_batch_unlabeled['image'].cuda() + + noise = torch.clamp(torch.randn_like( + unlabeled_volume_batch) * 0.1, -0.2, 0.2) + ema_inputs = unlabeled_volume_batch + noise + + outputs = model(volume_batch) + outputs_soft = torch.softmax(outputs, dim=1) + + outputs_unlabeled = model(unlabeled_volume_batch) + outputs_unlabeled_soft = torch.softmax(outputs_unlabeled, dim=1) + + with torch.no_grad(): + ema_output = ema_model(ema_inputs) + ema_output_soft = torch.softmax(ema_output, dim=1) + + supervised_loss = 0.5*(ce_loss(outputs, label_batch[:].long()) + dice_loss(outputs_soft, label_batch[:].unsqueeze(1))) + consistency_weight = get_current_consistency_weight(iter_num // 150) + + consistency_loss = torch.mean((outputs_unlabeled_soft - ema_output_soft) ** 2) + loss = supervised_loss + consistency_weight * consistency_loss + optimizer.zero_grad() + loss.backward() + optimizer.step() + update_ema_variables(model, ema_model, args.ema_decay, iter_num) + + lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 + for param_group in optimizer.param_groups: + param_group['lr'] = lr_ + + iter_num = iter_num + 1 + writer.add_scalar('info/lr', lr_, iter_num) + writer.add_scalar('info/total_loss', loss, iter_num) + writer.add_scalar('info/loss_ce', supervised_loss, iter_num) + writer.add_scalar('info/consistency_loss', + consistency_loss, iter_num) + writer.add_scalar('info/consistency_weight', + consistency_weight, iter_num) + + logging.info( + 'iteration %d : loss : %f, loss_ce: %f' % + (iter_num, loss.item(), supervised_loss.item())) + + if iter_num % 20 == 0: + image = volume_batch[0, 0:1, :, :] + writer.add_image('train/Image', image, iter_num) + outputs = torch.argmax(torch.softmax( + outputs, dim=1), dim=1, keepdim=True) + writer.add_image('train/Prediction', + outputs[0, ...] * 50, iter_num) + labs = label_batch[0, ...].unsqueeze(0) * 50 + writer.add_image('train/GroundTruth', labs, iter_num) + + if iter_num > 0 and iter_num % 200 == 0: + model.eval() + metric_list = 0.0 + for i_batch, sampled_batch in enumerate(valloader): + metric_i = test_single_volume( + sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) + metric_list += np.array(metric_i) + metric_list = metric_list / len(db_val) + for class_i in range(num_classes-1): + writer.add_scalar('info/val_{}_dice'.format(class_i+1), + metric_list[class_i, 0], iter_num) + writer.add_scalar('info/val_{}_hd95'.format(class_i+1), + metric_list[class_i, 1], iter_num) + + performance = np.mean(metric_list, axis=0)[0] + + mean_hd95 = np.mean(metric_list, axis=0)[1] + writer.add_scalar('info/val_mean_dice', performance, iter_num) + writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) + + if performance > best_performance: + best_performance = performance + save_mode_path = os.path.join(snapshot_path, + 'iter_{}_dice_{}.pth'.format( + iter_num, round(best_performance, 4))) + save_best = os.path.join(snapshot_path, + '{}_best_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_mode_path) + torch.save(model.state_dict(), save_best) + + logging.info( + 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) + model.train() + + if iter_num % 3000 == 0: + save_mode_path = os.path.join( + snapshot_path, 'iter_' + str(iter_num) + '.pth') + torch.save(model.state_dict(), save_mode_path) + logging.info("save model to {}".format(save_mode_path)) + + if iter_num >= max_iterations: + break + if iter_num >= max_iterations: + iterator.close() + break + writer.close() + return "Training Finished!" + + +if __name__ == "__main__": + if not args.deterministic: + cudnn.benchmark = True + cudnn.deterministic = False + else: + cudnn.benchmark = False + cudnn.deterministic = True + + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( + args.exp, args.labeled_ratio, args.fold) + if not os.path.exists(snapshot_path): + os.makedirs(snapshot_path) + if os.path.exists(snapshot_path + '/code'): + shutil.rmtree(snapshot_path + '/code') + shutil.copytree('.', snapshot_path + '/code', + shutil.ignore_patterns(['.git', '__pycache__'])) + + logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO, + format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') + logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) + logging.info(str(args)) + train(args, snapshot_path) diff --git a/code/train_mean_teacher_3D.py b/code/train_mean_teacher_3D.py deleted file mode 100644 index e04fefc..0000000 --- a/code/train_mean_teacher_3D.py +++ /dev/null @@ -1,265 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.brats2019 import (BraTS2019, CenterCrop, RandomCrop, - RandomRotFlip, ToTensor, - TwoStreamBatchSampler) -from networks.net_factory_3d import net_factory_3d -from utils import losses, metrics, ramps -from val_3D import test_all_case - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/BraTS2019', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='BraTs2019_Mean_Teacher', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet_3D', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=4, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[96, 96, 96], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=2, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=25, - help='labeled data') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') - -args = parser.parse_args() - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def update_ema_variables(model, ema_model, alpha, global_step): - # Use the true average until the exponential average is more correct - alpha = min(1 - 1 / (global_step + 1), alpha) - for ema_param, param in zip(ema_model.parameters(), model.parameters()): - ema_param.data.mul_(alpha).add_(1 - alpha, param.data) - - -def train(args, snapshot_path): - base_lr = args.base_lr - train_data_path = args.root_path - batch_size = args.batch_size - max_iterations = args.max_iterations - num_classes = 2 - - def create_model(ema=False): - # Network definition - net = net_factory_3d(net_type=args.model, in_chns=1, class_num=num_classes) - model = net.cuda() - if ema: - for param in model.parameters(): - param.detach_() - return model - - model = create_model() - ema_model = create_model(ema=True) - - db_train = BraTS2019(base_dir=train_data_path, - split='train', - num=None, - transform=transforms.Compose([ - RandomRotFlip(), - RandomCrop(args.patch_size), - ToTensor(), - ])) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - labeled_idxs = list(range(0, args.labeled_num)) - unlabeled_idxs = list(range(args.labeled_num, 250)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - model.train() - ema_model.train() - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(2) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - unlabeled_volume_batch = volume_batch[args.labeled_bs:] - - noise = torch.clamp(torch.randn_like( - unlabeled_volume_batch) * 0.1, -0.2, 0.2) - ema_inputs = unlabeled_volume_batch + noise - - outputs = model(volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - with torch.no_grad(): - ema_output = ema_model(ema_inputs) - ema_output_soft = torch.softmax(ema_output, dim=1) - - loss_ce = ce_loss(outputs[:args.labeled_bs], - label_batch[:args.labeled_bs][:]) - loss_dice = dice_loss( - outputs_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - supervised_loss = 0.5 * (loss_dice + loss_ce) - consistency_weight = get_current_consistency_weight(iter_num//150) - consistency_loss = torch.mean( - (outputs_soft[args.labeled_bs:] - ema_output_soft)**2) - loss = supervised_loss + consistency_weight * consistency_loss - optimizer.zero_grad() - loss.backward() - optimizer.step() - update_ema_variables(model, ema_model, args.ema_decay, iter_num) - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - writer.add_scalar('loss/loss', loss, iter_num) - - if iter_num % 20 == 0: - image = volume_batch[0, 0:1, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=True) - writer.add_image('train/Image', grid_image, iter_num) - - image = outputs_soft[0, 1:2, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Predicted_label', - grid_image, iter_num) - - image = label_batch[0, :, :, 20:61:10].unsqueeze( - 0).permute(3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Groundtruth_label', - grid_image, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - avg_metric = test_all_case( - model, args.root_path, test_list="val.txt", num_classes=2, patch_size=args.patch_size, - stride_xy=64, stride_z=64) - if avg_metric[:, 0].mean() > best_performance: - best_performance = avg_metric[:, 0].mean() - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - writer.add_scalar('info/val_dice_score', - avg_metric[0, 0], iter_num) - writer.add_scalar('info/val_hd95', - avg_metric[0, 1], iter_num) - logging.info( - 'iteration %d : dice_score : %f hd95 : %f' % (iter_num, avg_metric[0, 0].mean(), avg_metric[0, 1].mean())) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_uncertainty_aware_mean_teacher_2D.py b/code/train_uncertainty_aware_mean_teacher.py similarity index 68% rename from code/train_uncertainty_aware_mean_teacher_2D.py rename to code/train_uncertainty_aware_mean_teacher.py index 2ec2e3b..caac5fb 100644 --- a/code/train_uncertainty_aware_mean_teacher_2D.py +++ b/code/train_uncertainty_aware_mean_teacher.py @@ -1,301 +1,280 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.dataset import (BaseDataSets, RandomGenerator, - TwoStreamBatchSampler) -from networks.net_factory import net_factory -from utils import losses, metrics, ramps -from val_2D import test_single_volume - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/ACDC', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='ACDC/Uncertainty_Aware_Mean_Teacher', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=24, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[256, 256], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') -parser.add_argument('--num_classes', type=int, default=4, - help='output channel of network') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=12, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=136, - help='labeled data') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() - - -def patients_to_slices(dataset, patiens_num): - ref_dict = None - if "ACDC" in dataset: - ref_dict = {"3": 68, "7": 136, - "14": 256, "21": 396, "28": 512, "35": 664, "140": 1312} - elif "Prostate": - ref_dict = {"2": 27, "4": 53, "8": 120, - "12": 179, "16": 256, "21": 312, "42": 623} - else: - print("Error") - return ref_dict[str(patiens_num)] - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def update_ema_variables(model, ema_model, alpha, global_step): - # Use the true average until the exponential average is more correct - alpha = min(1 - 1 / (global_step + 1), alpha) - for ema_param, param in zip(ema_model.parameters(), model.parameters()): - ema_param.data.mul_(alpha).add_(1 - alpha, param.data) - - -def train(args, snapshot_path): - base_lr = args.base_lr - num_classes = args.num_classes - batch_size = args.batch_size - max_iterations = args.max_iterations - - def create_model(ema=False): - # Network definition - model = net_factory(net_type=args.model, in_chns=1, - class_num=num_classes) - if ema: - for param in model.parameters(): - param.detach_() - return model - - model = create_model() - ema_model = create_model(ema=True) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - db_train = BaseDataSets(base_dir=args.root_path, split="train", num=None, transform=transforms.Compose([ - RandomGenerator(args.patch_size) - ])) - db_val = BaseDataSets(base_dir=args.root_path, split="val") - total_slices = len(db_train) - labeled_slice = patients_to_slices(args.root_path, args.labeled_num) - print("Total silices is: {}, labeled slices is: {}".format( - total_slices, labeled_slice)) - labeled_idxs = list(range(0, labeled_slice)) - unlabeled_idxs = list(range(labeled_slice, total_slices)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - model.train() - - valloader = DataLoader(db_val, batch_size=1, shuffle=False, - num_workers=1) - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - unlabeled_volume_batch = volume_batch[args.labeled_bs:] - - noise = torch.clamp(torch.randn_like( - unlabeled_volume_batch) * 0.1, -0.2, 0.2) - ema_inputs = unlabeled_volume_batch + noise - - outputs = model(volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - with torch.no_grad(): - ema_output = ema_model(ema_inputs) - T = 8 - _, _, w, h = unlabeled_volume_batch.shape - volume_batch_r = unlabeled_volume_batch.repeat(2, 1, 1, 1) - stride = volume_batch_r.shape[0] // 2 - preds = torch.zeros([stride * T, num_classes, w, h]).cuda() - for i in range(T//2): - ema_inputs = volume_batch_r + \ - torch.clamp(torch.randn_like( - volume_batch_r) * 0.1, -0.2, 0.2) - with torch.no_grad(): - preds[2 * stride * i:2 * stride * - (i + 1)] = ema_model(ema_inputs) - preds = F.softmax(preds, dim=1) - preds = preds.reshape(T, stride, num_classes, w, h) - preds = torch.mean(preds, dim=0) - uncertainty = -1.0 * \ - torch.sum(preds*torch.log(preds + 1e-6), dim=1, keepdim=True) - - loss_ce = ce_loss(outputs[:args.labeled_bs], - label_batch[:args.labeled_bs][:].long()) - loss_dice = dice_loss( - outputs_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - supervised_loss = 0.5 * (loss_dice + loss_ce) - consistency_weight = get_current_consistency_weight(iter_num//150) - consistency_dist = losses.softmax_mse_loss( - outputs[args.labeled_bs:], ema_output) # (batch, 2, 112,112,80) - threshold = (0.75+0.25*ramps.sigmoid_rampup(iter_num, - max_iterations))*np.log(2) - mask = (uncertainty < threshold).float() - consistency_loss = torch.sum( - mask*consistency_dist)/(2*torch.sum(mask)+1e-16) - - loss = supervised_loss + consistency_weight * consistency_loss - optimizer.zero_grad() - loss.backward() - optimizer.step() - update_ema_variables(model, ema_model, args.ema_decay, iter_num) - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - - if iter_num % 20 == 0: - image = volume_batch[1, 0:1, :, :] - writer.add_image('train/Image', image, iter_num) - outputs = torch.argmax(torch.softmax( - outputs, dim=1), dim=1, keepdim=True) - writer.add_image('train/Prediction', - outputs[1, ...] * 50, iter_num) - labs = label_batch[1, ...].unsqueeze(0) * 50 - writer.add_image('train/GroundTruth', labs, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - metric_list = 0.0 - for i_batch, sampled_batch in enumerate(valloader): - metric_i = test_single_volume( - sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) - metric_list += np.array(metric_i) - metric_list = metric_list / len(db_val) - for class_i in range(num_classes-1): - writer.add_scalar('info/val_{}_dice'.format(class_i+1), - metric_list[class_i, 0], iter_num) - writer.add_scalar('info/val_{}_hd95'.format(class_i+1), - metric_list[class_i, 1], iter_num) - - performance = np.mean(metric_list, axis=0)[0] - - mean_hd95 = np.mean(metric_list, axis=0)[1] - writer.add_scalar('info/val_mean_dice', performance, iter_num) - writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) - - if performance > best_performance: - best_performance = performance - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - logging.info( - 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}_labeled/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) +import argparse +import logging +import os +import random +import shutil +import sys +import time +from itertools import cycle +import numpy as np +import torch +import torch.backends.cudnn as cudnn +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from tensorboardX import SummaryWriter +from torch.nn import BCEWithLogitsLoss +from torch.nn.modules.loss import CrossEntropyLoss +from torch.utils.data import DataLoader +from torchvision import transforms +from torchvision.utils import make_grid +from tqdm import tqdm +from dataloaders.dataset import BaseDataSets, RandomGenerator +from networks.discriminator import FCDiscriminator +from networks.net_factory import net_factory +from utils import losses, metrics, ramps +from val_2D import test_single_volume + +parser = argparse.ArgumentParser() +parser.add_argument('--root_path', type=str, + default='../data/ACDC', help='Name of Experiment') +parser.add_argument('--exp', type=str, + default='ACDC/Uncertainty_Aware_Mean_Teacher', help='experiment_name') +parser.add_argument('--model', type=str, + default='unet', help='model_name') +parser.add_argument('--fold', type=int, + default=1, help='cross validation') +parser.add_argument('--max_iterations', type=int, + default=30000, help='maximum epoch number to train') +parser.add_argument('--batch_size', type=int, default=12, + help='batch_size per gpu') + +parser.add_argument('--deterministic', type=int, default=1, + help='whether use deterministic training') +parser.add_argument('--base_lr', type=float, default=0.01, + help='segmentation network learning rate') +parser.add_argument('--patch_size', type=list, default=[256, 256], + help='patch size of network input') +parser.add_argument('--seed', type=int, default=2022, help='random seed') +parser.add_argument('--num_classes', type=int, default=4, + help='output channel of network') + +# label and unlabel +parser.add_argument('--labeled_ratio', type=int, default=5, + help='1/labeled_ratio data is provided mask') +# costs +parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') +parser.add_argument('--consistency_type', type=str, + default="mse", help='consistency_type') +parser.add_argument('--consistency', type=float, + default=0.1, help='consistency') +parser.add_argument('--consistency_rampup', type=float, + default=200.0, help='consistency_rampup') +args = parser.parse_args() + + +def get_current_consistency_weight(epoch): + # Consistency ramp-up from https://arxiv.org/abs/1610.02242 + return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) + +def update_ema_variables(model, ema_model, alpha, global_step): + # Use the true average until the exponential average is more correct + alpha = min(1 - 1 / (global_step + 1), alpha) + for ema_param, param in zip(ema_model.parameters(), model.parameters()): + ema_param.data.mul_(alpha).add_(1 - alpha, param.data) + + +def train(args, snapshot_path): + writer = SummaryWriter(snapshot_path + '/log') + base_lr = args.base_lr + num_classes = args.num_classes + max_iterations = args.max_iterations + + def worker_init_fn(worker_id): + random.seed(args.seed + worker_id) + + def create_model(ema=False): + # Network definition + model = net_factory(net_type=args.model, in_chns=1, + class_num=num_classes) + if ema: + for param in model.parameters(): + param.detach_() + return model + + model = create_model() + ema_model = create_model(ema=True) + + db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)])) + db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)])) + + trainloader_labeled = DataLoader(db_train_labeled, batch_size=args.batch_size//2, shuffle=True) + trainloader_unlabeled = DataLoader(db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) + + logging.info("Labeled slices: {} ".format(len(db_train_labeled))) + logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) + + db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, split="val", labeled_ratio=args.labeled_ratio) + valloader = DataLoader(db_val, batch_size=1, shuffle=False) + model.train() + + optimizer = optim.SGD(model.parameters(), lr=base_lr, + momentum=0.9, weight_decay=0.0001) + + ce_loss = CrossEntropyLoss(ignore_index=4) + dice_loss = losses.DiceLoss(num_classes) + + iter_num = 0 + max_epoch = max_iterations // len(trainloader_labeled) + 1 + best_performance = 0.0 + iterator = tqdm(range(max_epoch), ncols=70) + for epoch_num in iterator: + for i, data in enumerate(zip(cycle(trainloader_labeled), trainloader_unlabeled)): + sampled_batch_labeled, sampled_batch_unlabeled = data[0], data[1] + + volume_batch, label_batch = sampled_batch_labeled['image'], sampled_batch_labeled['label'] + volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() + unlabeled_volume_batch = sampled_batch_unlabeled['image'].cuda() + + noise = torch.clamp(torch.randn_like( + unlabeled_volume_batch) * 0.1, -0.2, 0.2) + ema_inputs = unlabeled_volume_batch + noise + + outputs = model(volume_batch) + outputs_soft = torch.softmax(outputs, dim=1) + + outputs_unlabeled = model(unlabeled_volume_batch) + outputs_unlabeled_soft = torch.softmax(outputs_unlabeled, dim=1) + + with torch.no_grad(): + ema_output = ema_model(ema_inputs) + T = 8 + _, _, w, h = unlabeled_volume_batch.shape + volume_batch_r = unlabeled_volume_batch.repeat(2, 1, 1, 1) + stride = volume_batch_r.shape[0] // 2 + preds = torch.zeros([stride * T, num_classes, w, h]).cuda() + for i in range(T // 2): + ema_inputs = volume_batch_r + \ + torch.clamp(torch.randn_like( + volume_batch_r) * 0.1, -0.2, 0.2) + with torch.no_grad(): + preds[2 * stride * i:2 * stride * + (i + 1)] = ema_model(ema_inputs) + preds = F.softmax(preds, dim=1) + preds = preds.reshape(T, stride, num_classes, w, h) + preds = torch.mean(preds, dim=0) + uncertainty = -1.0 * \ + torch.sum(preds * torch.log(preds + 1e-6), dim=1, keepdim=True) + + loss_ce = ce_loss(outputs, label_batch[:].long()) + loss_dice = dice_loss(outputs_soft, label_batch.unsqueeze(1)) + supervised_loss = 0.5 * (loss_dice + loss_ce) + consistency_weight = get_current_consistency_weight(iter_num // 150) + consistency_dist = losses.softmax_mse_loss(outputs_unlabeled, ema_output) # (batch, 2, 112,112,80) + threshold = (0.75 + 0.25 * ramps.sigmoid_rampup(iter_num, + max_iterations)) * np.log(2) + mask = (uncertainty < threshold).float() + consistency_loss = torch.sum( + mask * consistency_dist) / (2 * torch.sum(mask) + 1e-16) + + loss = supervised_loss + consistency_weight * consistency_loss + optimizer.zero_grad() + loss.backward() + optimizer.step() + update_ema_variables(model, ema_model, args.ema_decay, iter_num) + + lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 + for param_group in optimizer.param_groups: + param_group['lr'] = lr_ + + iter_num = iter_num + 1 + writer.add_scalar('info/lr', lr_, iter_num) + writer.add_scalar('info/total_loss', loss, iter_num) + writer.add_scalar('info/loss_ce', supervised_loss, iter_num) + writer.add_scalar('info/consistency_loss', + consistency_loss, iter_num) + writer.add_scalar('info/consistency_weight', + consistency_weight, iter_num) + + logging.info( + 'iteration %d : loss : %f, loss_ce: %f' % + (iter_num, loss.item(), supervised_loss.item())) + + if iter_num % 20 == 0: + image = volume_batch[0, 0:1, :, :] + writer.add_image('train/Image', image, iter_num) + outputs = torch.argmax(torch.softmax( + outputs, dim=1), dim=1, keepdim=True) + writer.add_image('train/Prediction', + outputs[0, ...] * 50, iter_num) + labs = label_batch[0, ...].unsqueeze(0) * 50 + writer.add_image('train/GroundTruth', labs, iter_num) + + if iter_num > 0 and iter_num % 200 == 0: + model.eval() + metric_list = 0.0 + for i_batch, sampled_batch in enumerate(valloader): + metric_i = test_single_volume( + sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) + metric_list += np.array(metric_i) + metric_list = metric_list / len(db_val) + for class_i in range(num_classes-1): + writer.add_scalar('info/val_{}_dice'.format(class_i+1), + metric_list[class_i, 0], iter_num) + writer.add_scalar('info/val_{}_hd95'.format(class_i+1), + metric_list[class_i, 1], iter_num) + + performance = np.mean(metric_list, axis=0)[0] + + mean_hd95 = np.mean(metric_list, axis=0)[1] + writer.add_scalar('info/val_mean_dice', performance, iter_num) + writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) + + if performance > best_performance: + best_performance = performance + save_mode_path = os.path.join(snapshot_path, + 'iter_{}_dice_{}.pth'.format( + iter_num, round(best_performance, 4))) + save_best = os.path.join(snapshot_path, + '{}_best_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_mode_path) + torch.save(model.state_dict(), save_best) + + logging.info( + 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) + model.train() + + if iter_num % 3000 == 0: + save_mode_path = os.path.join( + snapshot_path, 'iter_' + str(iter_num) + '.pth') + torch.save(model.state_dict(), save_mode_path) + logging.info("save model to {}".format(save_mode_path)) + + if iter_num >= max_iterations: + break + if iter_num >= max_iterations: + iterator.close() + break + writer.close() + return "Training Finished!" + + +if __name__ == "__main__": + if not args.deterministic: + cudnn.benchmark = True + cudnn.deterministic = False + else: + cudnn.benchmark = False + cudnn.deterministic = True + + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( + args.exp, args.labeled_ratio, args.fold) + if not os.path.exists(snapshot_path): + os.makedirs(snapshot_path) + if os.path.exists(snapshot_path + '/new_code'): + shutil.rmtree(snapshot_path + '/new_code') + shutil.copytree('.', snapshot_path + '/new_code', + shutil.ignore_patterns(['.git', '__pycache__'])) + + logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO, + format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') + logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) + logging.info(str(args)) + train(args, snapshot_path) diff --git a/code/train_uncertainty_aware_mean_teacher_3D.py b/code/train_uncertainty_aware_mean_teacher_3D.py deleted file mode 100644 index ee5e2ce..0000000 --- a/code/train_uncertainty_aware_mean_teacher_3D.py +++ /dev/null @@ -1,288 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.brats2019 import (BraTS2019, CenterCrop, RandomCrop, - RandomRotFlip, ToTensor, - TwoStreamBatchSampler) -from networks.net_factory_3d import net_factory_3d -from utils import losses, metrics, ramps -from val_3D import test_all_case - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/BraTS2019', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='BraTs2019_Mean_Teacher', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet_3D', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=4, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[96, 96, 96], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=2, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=25, - help='labeled data') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') - -args = parser.parse_args() - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def update_ema_variables(model, ema_model, alpha, global_step): - # Use the true average until the exponential average is more correct - alpha = min(1 - 1 / (global_step + 1), alpha) - for ema_param, param in zip(ema_model.parameters(), model.parameters()): - ema_param.data.mul_(alpha).add_(1 - alpha, param.data) - - -def train(args, snapshot_path): - base_lr = args.base_lr - train_data_path = args.root_path - batch_size = args.batch_size - max_iterations = args.max_iterations - num_classes = 2 - - def create_model(ema=False): - # Network definition - net = net_factory_3d(net_type=args.model, - in_chns=1, class_num=num_classes) - model = net.cuda() - if ema: - for param in model.parameters(): - param.detach_() - return model - - model = create_model() - ema_model = create_model(ema=True) - - db_train = BraTS2019(base_dir=train_data_path, - split='train', - num=None, - transform=transforms.Compose([ - RandomRotFlip(), - RandomCrop(args.patch_size), - ToTensor(), - ])) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - labeled_idxs = list(range(0, args.labeled_num)) - unlabeled_idxs = list(range(args.labeled_num, 250)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - model.train() - ema_model.train() - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(2) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - unlabeled_volume_batch = volume_batch[args.labeled_bs:] - - noise = torch.clamp(torch.randn_like( - unlabeled_volume_batch) * 0.1, -0.2, 0.2) - ema_inputs = unlabeled_volume_batch + noise - - outputs = model(volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - with torch.no_grad(): - ema_output = ema_model(ema_inputs) - T = 8 - _, _, d, w, h = unlabeled_volume_batch.shape - volume_batch_r = unlabeled_volume_batch.repeat(2, 1, 1, 1, 1) - stride = volume_batch_r.shape[0] // 2 - preds = torch.zeros([stride * T, 2, d, w, h]).cuda() - for i in range(T//2): - ema_inputs = volume_batch_r + \ - torch.clamp(torch.randn_like( - volume_batch_r) * 0.1, -0.2, 0.2) - with torch.no_grad(): - preds[2 * stride * i:2 * stride * - (i + 1)] = ema_model(ema_inputs) - preds = torch.softmax(preds, dim=1) - preds = preds.reshape(T, stride, 2, d, w, h) - preds = torch.mean(preds, dim=0) - uncertainty = -1.0 * \ - torch.sum(preds*torch.log(preds + 1e-6), dim=1, keepdim=True) - - loss_ce = ce_loss(outputs[:args.labeled_bs], - label_batch[:args.labeled_bs]) - loss_dice = dice_loss( - outputs_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - supervised_loss = 0.5 * (loss_dice + loss_ce) - consistency_weight = get_current_consistency_weight(iter_num//150) - consistency_dist = losses.softmax_mse_loss( - outputs[args.labeled_bs:], ema_output) # (batch, 2, 112,112,80) - threshold = (0.75+0.25*ramps.sigmoid_rampup(iter_num, - max_iterations))*np.log(2) - mask = (uncertainty < threshold).float() - consistency_loss = torch.sum( - mask*consistency_dist)/(2*torch.sum(mask)+1e-16) - - loss = supervised_loss + consistency_weight * consistency_loss - optimizer.zero_grad() - loss.backward() - optimizer.step() - update_ema_variables(model, ema_model, args.ema_decay, iter_num) - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - writer.add_scalar('loss/loss', loss, iter_num) - - if iter_num % 20 == 0: - image = volume_batch[0, 0:1, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=True) - writer.add_image('train/Image', grid_image, iter_num) - - image = outputs_soft[0, 1:2, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Predicted_label', - grid_image, iter_num) - - image = label_batch[0, :, :, 20:61:10].unsqueeze( - 0).permute(3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Groundtruth_label', - grid_image, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - avg_metric = test_all_case( - model, args.root_path, test_list="val.txt", num_classes=2, patch_size=args.patch_size, - stride_xy=64, stride_z=64) - if avg_metric[:, 0].mean() > best_performance: - best_performance = avg_metric[:, 0].mean() - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - writer.add_scalar('info/val_dice_score', - avg_metric[0, 0], iter_num) - writer.add_scalar('info/val_hd95', - avg_metric[0, 1], iter_num) - logging.info( - 'iteration %d : dice_score : %f hd95 : %f' % (iter_num, avg_metric[0, 0].mean(), avg_metric[0, 1].mean())) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_uncertainty_rectified_pyramid_consistency_2D.py b/code/train_uncertainty_rectified_pyramid_consistency_2D.py deleted file mode 100644 index 660fdfd..0000000 --- a/code/train_uncertainty_rectified_pyramid_consistency_2D.py +++ /dev/null @@ -1,311 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.dataset import BaseDataSets, RandomGenerator, TwoStreamBatchSampler -from utils import losses, metrics, ramps -from val_2D import test_single_volume_ds -from networks.net_factory import net_factory - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/ACDC', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='ACDC/Uncertainty_Rectified_Pyramid_Consistency', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet_urpc', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=24, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[256, 256], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') -parser.add_argument('--num_classes', type=int, default=4, - help='output channel of network') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=12, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=7, - help='labeled data') -# costs -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=200.0, help='consistency_rampup') -args = parser.parse_args() - - -def patients_to_slices(dataset, patiens_num): - ref_dict = None - if "ACDC" in dataset: - ref_dict = {"3": 68, "7": 136, - "14": 256, "21": 396, "28": 512, "35": 664, "140": 1312} - elif "Prostate": - ref_dict = {"2": 27, "4": 53, "8": 120, - "12": 179, "16": 256, "21": 312, "42": 623} - else: - print("Error") - return ref_dict[str(patiens_num)] - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def train(args, snapshot_path): - base_lr = args.base_lr - num_classes = args.num_classes - batch_size = args.batch_size - max_iterations = args.max_iterations - - model = net_factory(net_type=args.model, in_chns=1, - class_num=num_classes) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - db_train = BaseDataSets(base_dir=args.root_path, split="train", num=None, transform=transforms.Compose([ - RandomGenerator(args.patch_size) - ])) - db_val = BaseDataSets(base_dir=args.root_path, split="val") - total_slices = len(db_train) - labeled_slice = patients_to_slices(args.root_path, args.labeled_num) - print("Total silices is: {}, labeled slices is: {}".format( - total_slices, labeled_slice)) - labeled_idxs = list(range(0, labeled_slice)) - unlabeled_idxs = list(range(labeled_slice, total_slices)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - model.train() - - valloader = DataLoader(db_val, batch_size=1, shuffle=False, - num_workers=1) - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - kl_distance = nn.KLDivLoss(reduction='none') - iterator = tqdm(range(max_epoch), ncols=70) - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - - outputs, outputs_aux1, outputs_aux2, outputs_aux3 = model( - volume_batch) - outputs_soft = torch.softmax(outputs, dim=1) - outputs_aux1_soft = torch.softmax(outputs_aux1, dim=1) - outputs_aux2_soft = torch.softmax(outputs_aux2, dim=1) - outputs_aux3_soft = torch.softmax(outputs_aux3, dim=1) - - loss_ce = ce_loss(outputs[:args.labeled_bs], - label_batch[:args.labeled_bs][:].long()) - loss_ce_aux1 = ce_loss(outputs_aux1[:args.labeled_bs], - label_batch[:args.labeled_bs][:].long()) - loss_ce_aux2 = ce_loss(outputs_aux2[:args.labeled_bs], - label_batch[:args.labeled_bs][:].long()) - loss_ce_aux3 = ce_loss(outputs_aux3[:args.labeled_bs], - label_batch[:args.labeled_bs][:].long()) - - loss_dice = dice_loss( - outputs_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - loss_dice_aux1 = dice_loss( - outputs_aux1_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - loss_dice_aux2 = dice_loss( - outputs_aux2_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - loss_dice_aux3 = dice_loss( - outputs_aux3_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - - supervised_loss = (loss_ce+loss_ce_aux1+loss_ce_aux2+loss_ce_aux3 + - loss_dice+loss_dice_aux1+loss_dice_aux2+loss_dice_aux3)/8 - - preds = (outputs_soft+outputs_aux1_soft + - outputs_aux2_soft+outputs_aux3_soft)/4 - - variance_main = torch.sum(kl_distance( - torch.log(outputs_soft[args.labeled_bs:]), preds[args.labeled_bs:]), dim=1, keepdim=True) - exp_variance_main = torch.exp(-variance_main) - - variance_aux1 = torch.sum(kl_distance( - torch.log(outputs_aux1_soft[args.labeled_bs:]), preds[args.labeled_bs:]), dim=1, keepdim=True) - exp_variance_aux1 = torch.exp(-variance_aux1) - - variance_aux2 = torch.sum(kl_distance( - torch.log(outputs_aux2_soft[args.labeled_bs:]), preds[args.labeled_bs:]), dim=1, keepdim=True) - exp_variance_aux2 = torch.exp(-variance_aux2) - - variance_aux3 = torch.sum(kl_distance( - torch.log(outputs_aux3_soft[args.labeled_bs:]), preds[args.labeled_bs:]), dim=1, keepdim=True) - exp_variance_aux3 = torch.exp(-variance_aux3) - - consistency_weight = get_current_consistency_weight(iter_num//150) - consistency_dist_main = ( - preds[args.labeled_bs:] - outputs_soft[args.labeled_bs:]) ** 2 - - consistency_loss_main = torch.mean( - consistency_dist_main * exp_variance_main) / (torch.mean(exp_variance_main) + 1e-8) + torch.mean(variance_main) - - consistency_dist_aux1 = ( - preds[args.labeled_bs:] - outputs_aux1_soft[args.labeled_bs:]) ** 2 - consistency_loss_aux1 = torch.mean( - consistency_dist_aux1 * exp_variance_aux1) / (torch.mean(exp_variance_aux1) + 1e-8) + torch.mean(variance_aux1) - - consistency_dist_aux2 = ( - preds[args.labeled_bs:] - outputs_aux2_soft[args.labeled_bs:]) ** 2 - consistency_loss_aux2 = torch.mean( - consistency_dist_aux2 * exp_variance_aux2) / (torch.mean(exp_variance_aux2) + 1e-8) + torch.mean(variance_aux2) - - consistency_dist_aux3 = ( - preds[args.labeled_bs:] - outputs_aux3_soft[args.labeled_bs:]) ** 2 - consistency_loss_aux3 = torch.mean( - consistency_dist_aux3 * exp_variance_aux3) / (torch.mean(exp_variance_aux3) + 1e-8) + torch.mean(variance_aux3) - - consistency_loss = (consistency_loss_main + consistency_loss_aux1 + - consistency_loss_aux2 + consistency_loss_aux3) / 4 - loss = supervised_loss + consistency_weight * consistency_loss - optimizer.zero_grad() - loss.backward() - optimizer.step() - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/loss_ce', loss_ce, iter_num) - writer.add_scalar('info/loss_dice', loss_dice, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - logging.info( - 'iteration %d : loss : %f, loss_ce: %f, loss_dice: %f' % - (iter_num, loss.item(), loss_ce.item(), loss_dice.item())) - - if iter_num % 20 == 0: - image = volume_batch[1, 0:1, :, :] - writer.add_image('train/Image', image, iter_num) - outputs = torch.argmax(torch.softmax( - outputs, dim=1), dim=1, keepdim=True) - writer.add_image('train/Prediction', - outputs[1, ...] * 50, iter_num) - labs = label_batch[1, ...].unsqueeze(0) * 50 - writer.add_image('train/GroundTruth', labs, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - metric_list = 0.0 - for i_batch, sampled_batch in enumerate(valloader): - metric_i = test_single_volume_ds( - sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) - metric_list += np.array(metric_i) - metric_list = metric_list / len(db_val) - for class_i in range(num_classes-1): - writer.add_scalar('info/val_{}_dice'.format(class_i+1), - metric_list[class_i, 0], iter_num) - writer.add_scalar('info/val_{}_hd95'.format(class_i+1), - metric_list[class_i, 1], iter_num) - - performance = np.mean(metric_list, axis=0)[0] - - mean_hd95 = np.mean(metric_list, axis=0)[1] - writer.add_scalar('info/val_mean_dice', performance, iter_num) - writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) - - if performance > best_performance: - best_performance = performance - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - - logging.info( - 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}_labeled/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/train_uncertainty_rectified_pyramid_consistency_3D.py b/code/train_uncertainty_rectified_pyramid_consistency_3D.py deleted file mode 100644 index 82c1751..0000000 --- a/code/train_uncertainty_rectified_pyramid_consistency_3D.py +++ /dev/null @@ -1,313 +0,0 @@ -import argparse -import logging -import os -import random -import shutil -import sys -import time - -import numpy as np -import torch -import torch.backends.cudnn as cudnn -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from tensorboardX import SummaryWriter -from torch.nn import BCEWithLogitsLoss -from torch.nn.modules.loss import CrossEntropyLoss -from torch.utils.data import DataLoader -from torchvision import transforms -from torchvision.utils import make_grid -from tqdm import tqdm - -from dataloaders import utils -from dataloaders.brats2019 import (BraTS2019, CenterCrop, RandomCrop, - RandomRotFlip, ToTensor, - TwoStreamBatchSampler) -from networks.unet_3D_dv_semi import unet_3D_dv_semi -from utils import losses, metrics, ramps -from val_urpc_util import test_all_case - -parser = argparse.ArgumentParser() -parser.add_argument('--root_path', type=str, - default='../data/GTV', help='Name of Experiment') -parser.add_argument('--exp', type=str, - default='GTV/uncertainty_rectified_pyramid_consistency', help='experiment_name') -parser.add_argument('--model', type=str, - default='unet_3D_dv_semi', help='model_name') -parser.add_argument('--max_iterations', type=int, - default=60000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=4, - help='batch_size per gpu') -parser.add_argument('--deterministic', type=int, default=1, - help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.1, - help='segmentation network learning rate') -parser.add_argument('--patch_size', type=list, default=[96, 96, 96], - help='patch size of network input') -parser.add_argument('--seed', type=int, default=1337, help='random seed') - -# label and unlabel -parser.add_argument('--labeled_bs', type=int, default=2, - help='labeled_batch_size per gpu') -parser.add_argument('--labeled_num', type=int, default=18, - help='labeled data') -parser.add_argument('--total_labeled_num', type=int, default=180, - help='total labeled data') -# costs -parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') -parser.add_argument('--consistency_type', type=str, - default="mse", help='consistency_type') -parser.add_argument('--consistency', type=float, - default=0.1, help='consistency') -parser.add_argument('--consistency_rampup', type=float, - default=400.0, help='consistency_rampup') -args = parser.parse_args() - - -def get_current_consistency_weight(epoch): - # Consistency ramp-up from https://arxiv.org/abs/1610.02242 - return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) - - -def update_ema_variables(model, ema_model, alpha, global_step): - # Use the true average until the exponential average is more correct - alpha = min(1 - 1 / (global_step + 1), alpha) - for ema_param, param in zip(ema_model.parameters(), model.parameters()): - ema_param.data.mul_(alpha).add_(1 - alpha, param.data) - - -def train(args, snapshot_path): - num_classes = 3 - base_lr = args.base_lr - train_data_path = args.root_path - batch_size = args.batch_size - max_iterations = args.max_iterations - - net = unet_3D_dv_semi(n_classes=num_classes, in_channels=1) - model = net.cuda() - - db_train = BraTS2019(base_dir=train_data_path, - split='train', - num=None, - transform=transforms.Compose([ - RandomRotFlip(), - RandomCrop(args.patch_size), - ToTensor(), - ])) - - def worker_init_fn(worker_id): - random.seed(args.seed + worker_id) - - labeled_idxs = list(range(0, args.labeled_num)) - unlabeled_idxs = list(range(args.labeled_num, args.total_labeled_num)) - batch_sampler = TwoStreamBatchSampler( - labeled_idxs, unlabeled_idxs, batch_size, batch_size-args.labeled_bs) - - trainloader = DataLoader(db_train, batch_sampler=batch_sampler, - num_workers=4, pin_memory=True, worker_init_fn=worker_init_fn) - - model.train() - - optimizer = optim.SGD(model.parameters(), lr=base_lr, - momentum=0.9, weight_decay=0.0001) - ce_loss = CrossEntropyLoss() - dice_loss = losses.DiceLoss(num_classes) - - writer = SummaryWriter(snapshot_path + '/log') - logging.info("{} iterations per epoch".format(len(trainloader))) - - iter_num = 0 - max_epoch = max_iterations // len(trainloader) + 1 - best_performance = 0.0 - iterator = tqdm(range(max_epoch), ncols=70) - kl_distance = nn.KLDivLoss(reduction='none') - for epoch_num in iterator: - for i_batch, sampled_batch in enumerate(trainloader): - - volume_batch, label_batch = sampled_batch['image'], sampled_batch['label'] - volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() - unlabeled_volume_batch = volume_batch[args.labeled_bs:] - - outputs_aux1, outputs_aux2, outputs_aux3, outputs_aux4, = model( - volume_batch) - outputs_aux1_soft = torch.softmax(outputs_aux1, dim=1) - outputs_aux2_soft = torch.softmax(outputs_aux2, dim=1) - outputs_aux3_soft = torch.softmax(outputs_aux3, dim=1) - outputs_aux4_soft = torch.softmax(outputs_aux4, dim=1) - - loss_ce_aux1 = ce_loss(outputs_aux1[:args.labeled_bs], - label_batch[:args.labeled_bs]) - loss_ce_aux2 = ce_loss(outputs_aux2[:args.labeled_bs], - label_batch[:args.labeled_bs]) - loss_ce_aux3 = ce_loss(outputs_aux3[:args.labeled_bs], - label_batch[:args.labeled_bs]) - loss_ce_aux4 = ce_loss(outputs_aux4[:args.labeled_bs], - label_batch[:args.labeled_bs]) - - loss_dice_aux1 = dice_loss( - outputs_aux1_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - loss_dice_aux2 = dice_loss( - outputs_aux2_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - loss_dice_aux3 = dice_loss( - outputs_aux3_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - loss_dice_aux4 = dice_loss( - outputs_aux4_soft[:args.labeled_bs], label_batch[:args.labeled_bs].unsqueeze(1)) - - supervised_loss = (loss_ce_aux1+loss_ce_aux2+loss_ce_aux3+loss_ce_aux4 + - loss_dice_aux1+loss_dice_aux2+loss_dice_aux3+loss_dice_aux4)/8 - - preds = (outputs_aux1_soft + - outputs_aux2_soft+outputs_aux3_soft+outputs_aux4_soft)/4 - - variance_aux1 = torch.sum(kl_distance( - torch.log(outputs_aux1_soft[args.labeled_bs:]), preds[args.labeled_bs:]), dim=1, keepdim=True) - exp_variance_aux1 = torch.exp(-variance_aux1) - - variance_aux2 = torch.sum(kl_distance( - torch.log(outputs_aux2_soft[args.labeled_bs:]), preds[args.labeled_bs:]), dim=1, keepdim=True) - exp_variance_aux2 = torch.exp(-variance_aux2) - - variance_aux3 = torch.sum(kl_distance( - torch.log(outputs_aux3_soft[args.labeled_bs:]), preds[args.labeled_bs:]), dim=1, keepdim=True) - exp_variance_aux3 = torch.exp(-variance_aux3) - - variance_aux4 = torch.sum(kl_distance( - torch.log(outputs_aux4_soft[args.labeled_bs:]), preds[args.labeled_bs:]), dim=1, keepdim=True) - exp_variance_aux4 = torch.exp(-variance_aux4) - - consistency_weight = get_current_consistency_weight(iter_num//150) - - consistency_dist_aux1 = ( - preds[args.labeled_bs:] - outputs_aux1_soft[args.labeled_bs:]) ** 2 - consistency_loss_aux1 = torch.mean( - consistency_dist_aux1 * exp_variance_aux1) / (torch.mean(exp_variance_aux1) + 1e-8) + torch.mean(variance_aux1) - - consistency_dist_aux2 = ( - preds[args.labeled_bs:] - outputs_aux2_soft[args.labeled_bs:]) ** 2 - consistency_loss_aux2 = torch.mean( - consistency_dist_aux2 * exp_variance_aux2) / (torch.mean(exp_variance_aux2) + 1e-8) + torch.mean(variance_aux2) - - consistency_dist_aux3 = ( - preds[args.labeled_bs:] - outputs_aux3_soft[args.labeled_bs:]) ** 2 - consistency_loss_aux3 = torch.mean( - consistency_dist_aux3 * exp_variance_aux3) / (torch.mean(exp_variance_aux3) + 1e-8) + torch.mean(variance_aux3) - - consistency_dist_aux4 = ( - preds[args.labeled_bs:] - outputs_aux4_soft[args.labeled_bs:]) ** 2 - consistency_loss_aux4 = torch.mean( - consistency_dist_aux4 * exp_variance_aux4) / (torch.mean(exp_variance_aux4) + 1e-8) + torch.mean(variance_aux4) - - consistency_loss = (consistency_loss_aux1 + - consistency_loss_aux2 + consistency_loss_aux3 + consistency_loss_aux4) / 4 - loss = supervised_loss + consistency_weight * consistency_loss - optimizer.zero_grad() - loss.backward() - optimizer.step() - - lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 - for param_group in optimizer.param_groups: - param_group['lr'] = lr_ - - iter_num = iter_num + 1 - writer.add_scalar('info/lr', lr_, iter_num) - writer.add_scalar('info/total_loss', loss, iter_num) - writer.add_scalar('info/supervised_loss', - supervised_loss, iter_num) - writer.add_scalar('info/consistency_loss', - consistency_loss, iter_num) - writer.add_scalar('info/consistency_weight', - consistency_weight, iter_num) - - logging.info( - 'iteration %d : loss : %f, supervised_loss: %f' % - (iter_num, loss.item(), supervised_loss.item())) - - if iter_num % 20 == 0: - image = volume_batch[0, 0:1, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) - grid_image = make_grid(image, 5, normalize=True) - writer.add_image('train/Image', grid_image, iter_num) - - image = torch.argmax(outputs_aux1_soft, dim=1, keepdim=True)[0, 0:1, :, :, 20:61:10].permute( - 3, 0, 1, 2).repeat(1, 3, 1, 1) * 100 - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Predicted_label', - grid_image, iter_num) - - image = label_batch[0, :, :, 20:61:10].unsqueeze( - 0).permute(3, 0, 1, 2).repeat(1, 3, 1, 1) * 100 - grid_image = make_grid(image, 5, normalize=False) - writer.add_image('train/Groundtruth_label', - grid_image, iter_num) - - if iter_num > 0 and iter_num % 200 == 0: - model.eval() - avg_metric = test_all_case( - model, args.root_path, test_list="val.txt", num_classes=num_classes, patch_size=args.patch_size, - stride_xy=64, stride_z=64) - if avg_metric[:, 0].mean() > best_performance: - best_performance = avg_metric[:, 0].mean() - save_mode_path = os.path.join(snapshot_path, - 'iter_{}_dice_{}.pth'.format( - iter_num, round(best_performance, 4))) - save_best = os.path.join(snapshot_path, - '{}_best_model.pth'.format(args.model)) - torch.save(model.state_dict(), save_mode_path) - torch.save(model.state_dict(), save_best) - for cls in range(1, num_classes): - writer.add_scalar('info/val_cls_{}_dice_score'.format(cls), - avg_metric[cls - 1, 0], iter_num) - writer.add_scalar('info/val_cls_{}_hd95'.format(cls), - avg_metric[cls - 1, 1], iter_num) - writer.add_scalar('info/val_mean_dice_score', - avg_metric[:, 0].mean(), iter_num) - writer.add_scalar('info/val_mean_hd95', - avg_metric[:, 1].mean(), iter_num) - logging.info( - 'iteration %d : dice_score : %f hd95 : %f' % ( - iter_num, avg_metric[:, 0].mean(), avg_metric[:, 1].mean())) - model.train() - - if iter_num % 3000 == 0: - save_mode_path = os.path.join( - snapshot_path, 'iter_' + str(iter_num) + '.pth') - torch.save(model.state_dict(), save_mode_path) - logging.info("save model to {}".format(save_mode_path)) - - if iter_num >= max_iterations: - break - if iter_num >= max_iterations: - iterator.close() - break - writer.close() - return "Training Finished!" - - -if __name__ == "__main__": - if not args.deterministic: - cudnn.benchmark = True - cudnn.deterministic = False - else: - cudnn.benchmark = False - cudnn.deterministic = True - - random.seed(args.seed) - np.random.seed(args.seed) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - snapshot_path = "../model/{}_{}/{}".format( - args.exp, args.labeled_num, args.model) - if not os.path.exists(snapshot_path): - os.makedirs(snapshot_path) - if os.path.exists(snapshot_path + '/code'): - shutil.rmtree(snapshot_path + '/code') - shutil.copytree('.', snapshot_path + '/code', - shutil.ignore_patterns(['.git', '__pycache__'])) - - logging.basicConfig(filename=snapshot_path+"/log.txt", level=logging.INFO, - format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') - logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) - logging.info(str(args)) - train(args, snapshot_path) diff --git a/code/utils/MIPloss.py b/code/utils/MIPloss.py new file mode 100755 index 0000000..de8f87c --- /dev/null +++ b/code/utils/MIPloss.py @@ -0,0 +1,125 @@ +import torch +from torch.nn import functional as F +import numpy as np +import torch.nn as nn +from torch.autograd import Variable +from random import uniform + + +class Rotated_MIP_Loss_Multiclass(nn.Module): + # 监督为softmax后的输出 + def __init__(self, num_rot=3, device='cuda'): + super(Rotated_MIP_Loss_Multiclass, self).__init__() + self.device = device + assert isinstance(num_rot, int) and num_rot > 0 + self.num_rot = num_rot + + def forward(self, batch_input, label_input): + rot_list = [uniform(0, 1.57) for _ in range(self.num_rot)] + + size = batch_input.shape + loss = self.max_project_loss(batch_input, label_input) + # print(loss) + for i in rot_list: + rot_mat = self.create_rot_matrix(size, i) + rot_grid = F.affine_grid(rot_mat, size).to(self.device) + inputs = F.grid_sample(batch_input, rot_grid) + labels = F.grid_sample(label_input, rot_grid) + loss_ = self.max_project_loss(inputs, labels) + loss += self.max_project_loss(inputs, labels) + # print(loss_) + return loss / (len(rot_list) + 1) + + def create_rot_matrix(self, size, angle): + b = size[0] + angle = torch.tensor(angle).to(self.device) + rotation_matrix = torch.zeros([b, 3, 3], dtype=torch.float) + rotation_matrix[:, 0, 0] = torch.cos(angle) + rotation_matrix[:, 0, 1] = -torch.sin(angle) + rotation_matrix[:, 1, 0] = torch.sin(angle) + rotation_matrix[:, 1, 1] = torch.cos(angle) + rotation_matrix[:, 2, 2] = 1.0 + return rotation_matrix[:, :2, :] + + def max_project_loss(self, score, target): + total_loss = 0.0 + for index, i in enumerate([-1, -2]): + new_target = torch.max(target, dim=i)[0].float() + new_score = torch.max(score, dim=i)[0].float() + smooth = 1e-5 + intersect = torch.sum(new_score * new_target, dim=-1) + y_sum = torch.sum(new_target, dim=-1) + z_sum = torch.sum(new_score, dim=-1) + loss = (2 * intersect + smooth) / (z_sum + y_sum + smooth) + loss = torch.mean(loss[:, 1:]) # 这里去掉了背景 + total_loss += (1.0 - loss) + return total_loss / 2 + + def mean_project_loss(self, score, target): + total_loss = 0.0 + for index, i in enumerate([-1, -2]): + new_target = torch.mean(target, dim=i, keepdim=True).float() + new_score = torch.mean(score, dim=i, keepdim=True).float() + new_target = new_target / new_target.max() + new_score = new_score / new_score.max() + loss = torch.nn.functional.mse_loss(new_score, new_target) + total_loss += loss + return total_loss / 2 + + +if __name__ == '__main__': + + import matplotlib.pyplot as plt + import matplotlib + + xs = np.asarray([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39]) /10. + ys = [0.035633184015750885, 0.04039973020553589, 0.04840249940752983, 0.049674104899168015, 0.05561297759413719, 0.06043963134288788, 0.06371262669563293, 0.06686598062515259, 0.07038842141628265, 0.07373663783073425, 0.07658594101667404, 0.07925296574831009, 0.08191808313131332, 0.08441270142793655, 0.08767103403806686, 0.09013877063989639, 0.09172293543815613, 0.09327001124620438, 0.09405311197042465, 0.0982089713215828, 0.0963851660490036, 0.10003411769866943, 0.10333031415939331, 0.10372356325387955, 0.10573377460241318, 0.1040252298116684, 0.10658720880746841, 0.10942418873310089, 0.10836256295442581, 0.1106153279542923, 0.11024092882871628, 0.11183327436447144, 0.11297039687633514, 0.11287706345319748, 0.11342803388834, 0.11346597969532013, 0.11305929720401764, 0.11496424674987793, 0.11747633665800095, 0.1148373931646347] + + zs = [0.0526] * 40 + ms = [0.0156] * 40 + + plt.rcParams["font.family"] = "Times New Roman" + plt.rcParams.update({'font.size': 16}) + # plt.rc('axes', titlesize=20) + plt.figure() + plt.plot(xs, ys, label='MPR loss') + plt.plot(xs, zs, label='Dice loss') + plt.plot(xs, ms, label='MSE loss') + # plt.xlabel("Position of disagreement region", fontsize=15) + # plt.ylabel("Loss value", fontsize=15) + plt.legend(fontsize=16) + plt.show() + # + # rmip = Rotated_MIP_Loss_Multiclass(num_rot=360, device='cpu') + # + # + # xs = [] + # mprs = [] + # for i in range(40): + # pred_a = torch.zeros((800, 800)) + # pred_b = torch.zeros((800, 800)) + # + # pred_a[200: 500, 0: 300] = 1 + # # pred_a[300: 400, 400: 500] = 1 + # + # pi = i * 10 + # pred_a[300: 400, 300+pi: 400+pi] = 1 + # + # pred_b[200: 500, 0: 300] = 1 + # + # pred_a = pred_a.unsqueeze(dim=0).unsqueeze(dim=0) + # pred_b = pred_b.unsqueeze(dim=0).unsqueeze(dim=0) + # + # + # print("dice: ", 1- 2 * (pred_a * pred_b).sum() / (pred_a.sum() + pred_b.sum())) + # + # print("mse: ", ((pred_a - pred_b) ** 2).mean()) + # + # mpr = rmip(pred_a, pred_b).item() + # print("mpr mean: ", mpr) + # + # xs.append(i) + # mprs.append(mpr) + # + # print(xs) + # print(mprs) diff --git a/code/utils/gate_crf_loss.py b/code/utils/gate_crf_loss.py new file mode 100755 index 0000000..a2a3e0a --- /dev/null +++ b/code/utils/gate_crf_loss.py @@ -0,0 +1,205 @@ +import torch +import torch.nn.functional as F + + +class ModelLossSemsegGatedCRF(torch.nn.Module): + """ + This module provides an implementation of the Gated CRF Loss for Weakly Supervised Semantic Image Segmentation. + This loss function promotes consistent label assignment guided by input features, such as RGBXY. + Please consider using the following bibtex for citation: + @article{obukhov2019gated, + author={Anton Obukhov and Stamatios Georgoulis and Dengxin Dai and Luc {Van Gool}}, + title={Gated {CRF} Loss for Weakly Supervised Semantic Image Segmentation}, + journal={CoRR}, + volume={abs/1906.04651}, + year={2019}, + url={http://arxiv.org/abs/1906.04651}, + } + """ + + def forward( + self, y_hat_softmax, kernels_desc, kernels_radius, sample, height_input, width_input, + mask_src=None, mask_dst=None, compatibility=None, custom_modality_downsamplers=None, out_kernels_vis=False + ): + """ + Performs the forward pass of the loss. + :param y_hat_softmax: A tensor of predicted per-pixel class probabilities of size NxCxHxW + :param kernels_desc: A list of dictionaries, each describing one Gaussian kernel composition from modalities. + The final kernel is a weighted sum of individual kernels. Following example is a composition of + RGBXY and XY kernels: + kernels_desc: [{ + 'weight': 0.9, # Weight of RGBXY kernel + 'xy': 6, # Sigma for XY + 'rgb': 0.1, # Sigma for RGB + },{ + 'weight': 0.1, # Weight of XY kernel + 'xy': 6, # Sigma for XY + }] + :param kernels_radius: Defines size of bounding box region around each pixel in which the kernel is constructed. + :param sample: A dictionary with modalities (except 'xy') used in kernels_desc parameter. Each of the provided + modalities is allowed to be larger than the shape of y_hat_softmax, in such case downsampling will be + invoked. Default downsampling method is area resize; this can be overriden by setting. + custom_modality_downsamplers parameter. + :param width_input, height_input: Dimensions of the full scale resolution of modalities + :param mask_src: (optional) Source mask. + :param mask_dst: (optional) Destination mask. + :param compatibility: (optional) Classes compatibility matrix, defaults to Potts model. + :param custom_modality_downsamplers: A dictionary of modality downsampling functions. + :param out_kernels_vis: Whether to return a tensor with kernels visualized with some step. + :return: Loss function value. + """ + assert y_hat_softmax.dim() == 4, 'Prediction must be a NCHW batch' + N, C, height_pred, width_pred = y_hat_softmax.shape + device = y_hat_softmax.device + + assert width_input % width_pred == 0 and height_input % height_pred == 0 and \ + width_input * height_pred == height_input * width_pred, \ + f'[{width_input}x{height_input}] !~= [{width_pred}x{height_pred}]' + + kernels = self._create_kernels( + kernels_desc, kernels_radius, sample, N, height_pred, width_pred, device, custom_modality_downsamplers + ) + + denom = N * height_pred * width_pred + + def resize_fix_mask(mask, name): + assert mask.dim() == 4 and mask.shape[:2] == (N, 1) and mask.dtype == torch.float32, \ + f'{name} mask must be a NCHW batch with C=1 and dtype float32' + if mask.shape[2:] != (height_pred, width_pred): + mask = ModelLossSemsegGatedCRF._downsample( + mask, 'mask', height_pred, width_pred, custom_modality_downsamplers + ) + mask[mask != mask] = 0.0 # handle NaN + # handle edges of mask after interpolation + mask[mask < 1.0] = 0.0 + return mask + + if mask_src is not None: + mask_src = resize_fix_mask(mask_src, 'Source') + denom = mask_src.sum().clamp(min=1) + mask_src = self._unfold(mask_src, kernels_radius) + kernels = kernels * mask_src + + if mask_dst is not None: + mask_dst = resize_fix_mask(mask_dst, 'Destination') + denom = mask_dst.sum().clamp(min=1) + mask_dst = mask_dst.view(N, 1, 1, 1, height_pred, width_pred) + kernels = kernels * mask_dst + + y_hat_unfolded = self._unfold(y_hat_softmax, kernels_radius) + + product_kernel_x_y_hat = (kernels * y_hat_unfolded) \ + .view(N, C, (kernels_radius * 2 + 1) ** 2, height_pred, width_pred) \ + .sum(dim=2, keepdim=False) + + if compatibility is None: + # Using shortcut for Pott's class compatibility model + loss = -(product_kernel_x_y_hat * y_hat_softmax).sum() + # comment out to save computation, total loss may go below 0 + loss = kernels.sum() + loss + else: + assert compatibility.shape == ( + C, C), f'Compatibility matrix expected shape [{C}x{C}]' + assert (compatibility < 0).int().sum( + ) == 0, 'Compatibility matrix must not have negative values' + assert compatibility.diag.sum() == 0, 'Compatibility matrix diagonal must be 0' + compat = (C-1) * \ + F.normalize(compatibility.float().to(device), p=1, dim=1) + y_hat_CxNHW = y_hat_softmax.permute( + 1, 0, 2, 3).contiguous().view(C, -1) + product_kernel_x_y_hat_NHWxC = product_kernel_x_y_hat.permute( + 0, 2, 3, 1).contiguous().view(-1, C) + product_CxC = torch.mm(y_hat_CxNHW, product_kernel_x_y_hat_NHWxC) + loss = (compat * product_CxC).sum() + + out = { + 'loss': loss / denom, + } + + if out_kernels_vis: + out['kernels_vis'] = self._visualize_kernels( + kernels, kernels_radius, height_input, width_input, height_pred, width_pred + ) + + return out + + @staticmethod + def _downsample(img, modality, height_dst, width_dst, custom_modality_downsamplers): + if custom_modality_downsamplers is not None and modality in custom_modality_downsamplers: + f_down = custom_modality_downsamplers[modality] + else: + f_down = F.adaptive_avg_pool2d + return f_down(img, (height_dst, width_dst)) + + @staticmethod + def _create_kernels( + kernels_desc, kernels_radius, sample, N, height_pred, width_pred, device, custom_modality_downsamplers + ): + kernels = None + for i, desc in enumerate(kernels_desc): + weight = desc['weight'] + features = [] + for modality, sigma in desc.items(): + if modality == 'weight': + continue + if modality == 'xy': + feature = ModelLossSemsegGatedCRF._get_mesh( + N, height_pred, width_pred, device) + else: + # assert modality in sample, 'Modality {} is listed in {}-th kernel descriptor, but not present in the sample'.format(modality, i) + feature = sample + feature = ModelLossSemsegGatedCRF._downsample( + feature, modality, height_pred, width_pred, custom_modality_downsamplers + ) + feature /= sigma + features.append(feature) + features = torch.cat(features, dim=1) + kernel = weight * \ + ModelLossSemsegGatedCRF._create_kernels_from_features( + features, kernels_radius) + kernels = kernel if kernels is None else kernel + kernels + return kernels + + @staticmethod + def _create_kernels_from_features(features, radius): + assert features.dim() == 4, 'Features must be a NCHW batch' + N, C, H, W = features.shape + kernels = ModelLossSemsegGatedCRF._unfold(features, radius) + kernels = kernels - kernels[:, :, radius, + radius, :, :].view(N, C, 1, 1, H, W) + kernels = (-0.5 * kernels ** 2).sum(dim=1, keepdim=True).exp() + kernels[:, :, radius, radius, :, :] = 0 + return kernels + + @staticmethod + def _get_mesh(N, H, W, device): + return torch.cat(( + torch.arange(0, W, 1, dtype=torch.float32, device=device).view( + 1, 1, 1, W).repeat(N, 1, H, 1), + torch.arange(0, H, 1, dtype=torch.float32, device=device).view( + 1, 1, H, 1).repeat(N, 1, 1, W) + ), 1) + + @staticmethod + def _unfold(img, radius): + assert img.dim() == 4, 'Unfolding requires NCHW batch' + N, C, H, W = img.shape + diameter = 2 * radius + 1 + return F.unfold(img, diameter, 1, radius).view(N, C, diameter, diameter, H, W) + + @staticmethod + def _visualize_kernels(kernels, radius, height_input, width_input, height_pred, width_pred): + diameter = 2 * radius + 1 + vis = kernels[:, :, :, :, radius::diameter, radius::diameter] + vis_nh, vis_nw = vis.shape[-2:] + vis = vis.permute(0, 1, 4, 2, 5, 3).contiguous().view( + kernels.shape[0], 1, diameter * vis_nh, diameter * vis_nw) + if vis.shape[2] > height_pred: + vis = vis[:, :, :height_pred, :] + if vis.shape[3] > width_pred: + vis = vis[:, :, :, :width_pred] + if vis.shape[2:] != (height_pred, width_pred): + vis = F.pad(vis, [0, width_pred-vis.shape[3], + 0, height_pred-vis.shape[2]]) + vis = F.interpolate(vis, (height_input, width_input), mode='nearest') + return vis diff --git a/code/utils/losses.py b/code/utils/losses.py old mode 100644 new mode 100755 index a44c535..e4d0103 --- a/code/utils/losses.py +++ b/code/utils/losses.py @@ -36,6 +36,35 @@ def entropy_loss(p, C=2): return ent +def nuclear_norm_maximum(pr): + ''' + + :param pr: outputs_soft [N * C * W * H * D] + :param C: + :return: + ''' + + N = pr.shape[0] + C = pr.shape[1] + + # L_FBNM = 0 + # for n in range(N): + # pr_batch = pr[n].view((C, -1)).t() + # list_svd, _ = torch.sort(torch.sqrt(torch.sum(torch.pow(pr_batch, 2), dim=0)), descending=True) + # nums = min(pr_batch.shape[0], pr_batch.shape[1]) + # L_FBNM += torch.sum(list_svd[:nums]) + # return L_FBNM / N + + L_BNM = 0 + for n in range(N): + pr_batch = pr[n].view((C, -1)).t() + L_BNM += -torch.norm(pr_batch, p='nuc') + + # S = torch.svdvals(A) + return L_BNM / N + + + def softmax_dice_loss(input_logits, target_logits): """Takes softmax on both sides and returns MSE loss diff --git a/code/utils/metrics.py b/code/utils/metrics.py old mode 100644 new mode 100755 diff --git a/code/utils/ramps.py b/code/utils/ramps.py old mode 100644 new mode 100755 diff --git a/code/utils/util.py b/code/utils/util.py old mode 100644 new mode 100755 diff --git a/code/val_2D.py b/code/val_2D.py old mode 100644 new mode 100755 index 1d37353..3c5b1c6 --- a/code/val_2D.py +++ b/code/val_2D.py @@ -7,58 +7,119 @@ def calculate_metric_percase(pred, gt): pred[pred > 0] = 1 gt[gt > 0] = 1 - if pred.sum() > 0: + if pred.sum() > 0 and gt.sum() > 0: dice = metric.binary.dc(pred, gt) - hd95 = metric.binary.hd95(pred, gt) - return dice, hd95 + hd95 = metric.binary.hd95(pred, gt, voxelspacing=[10, 1, 1]) + asd = metric.binary.asd(pred, gt, voxelspacing=[10, 1, 1]) + return dice, hd95, asd else: - return 0, 0 + return 0, 50, 10 def test_single_volume(image, label, net, classes, patch_size=[256, 256]): image, label = image.squeeze(0).cpu().detach( ).numpy(), label.squeeze(0).cpu().detach().numpy() - prediction = np.zeros_like(label) - for ind in range(image.shape[0]): - slice = image[ind, :, :] - x, y = slice.shape[0], slice.shape[1] - slice = zoom(slice, (patch_size[0] / x, patch_size[1] / y), order=0) - input = torch.from_numpy(slice).unsqueeze( + if len(image.shape) == 3: + prediction = np.zeros_like(label) + for ind in range(image.shape[0]): + slice = image[ind, :, :] + x, y = slice.shape[0], slice.shape[1] + slice = zoom( + slice, (patch_size[0] / x, patch_size[1] / y), order=0) + input = torch.from_numpy(slice).unsqueeze( + 0).unsqueeze(0).float().cuda() + net.eval() + with torch.no_grad(): + out = torch.argmax(torch.softmax( + net(input), dim=1), dim=1).squeeze(0) + out = out.cpu().detach().numpy() + pred = zoom( + out, (x / patch_size[0], y / patch_size[1]), order=0) + prediction[ind] = pred + else: + input = torch.from_numpy(image).unsqueeze( 0).unsqueeze(0).float().cuda() net.eval() with torch.no_grad(): out = torch.argmax(torch.softmax( net(input), dim=1), dim=1).squeeze(0) - out = out.cpu().detach().numpy() - pred = zoom(out, (x / patch_size[0], y / patch_size[1]), order=0) - prediction[ind] = pred + prediction = out.cpu().detach().numpy() metric_list = [] for i in range(1, classes): metric_list.append(calculate_metric_percase( prediction == i, label == i)) - return metric_list + # return metric_list, image, prediction, label + return metric_list -def test_single_volume_ds(image, label, net, classes, patch_size=[256, 256]): +# +# def test_single_volume_ds(image, label, net, classes, patch_size=[256, 256]): +# image, label = image.squeeze(0).cpu().detach( +# ).numpy(), label.squeeze(0).cpu().detach().numpy() +# if len(image.shape) == 3: +# prediction = np.zeros_like(label) +# for ind in range(image.shape[0]): +# slice = image[ind, :, :] +# x, y = slice.shape[0], slice.shape[1] +# slice = zoom( +# slice, (patch_size[0] / x, patch_size[1] / y), order=0) +# input = torch.from_numpy(slice).unsqueeze( +# 0).unsqueeze(0).float().cuda() +# net.eval() +# with torch.no_grad(): +# output_main, _, _, _ = net(input) +# out = torch.argmax(torch.softmax( +# output_main, dim=1), dim=1).squeeze(0) +# out = out.cpu().detach().numpy() +# pred = zoom( +# out, (x / patch_size[0], y / patch_size[1]), order=0) +# prediction[ind] = pred +# else: +# input = torch.from_numpy(image).unsqueeze( +# 0).unsqueeze(0).float().cuda() +# net.eval() +# with torch.no_grad(): +# output_main, _, _, _ = net(input) +# out = torch.argmax(torch.softmax( +# output_main, dim=1), dim=1).squeeze(0) +# prediction = out.cpu().detach().numpy() +# metric_list = [] +# for i in range(1, classes): +# metric_list.append(calculate_metric_percase( +# prediction == i, label == i)) +# return metric_list +# +# +def test_single_volume_multitask(image, label, net, classes, patch_size=[256, 256]): image, label = image.squeeze(0).cpu().detach( ).numpy(), label.squeeze(0).cpu().detach().numpy() - prediction = np.zeros_like(label) - for ind in range(image.shape[0]): - slice = image[ind, :, :] - x, y = slice.shape[0], slice.shape[1] - slice = zoom(slice, (patch_size[0] / x, patch_size[1] / y), order=0) - input = torch.from_numpy(slice).unsqueeze( + if len(image.shape) == 3: + prediction = np.zeros_like(label) + for ind in range(image.shape[0]): + slice = image[ind, :, :] + x, y = slice.shape[0], slice.shape[1] + slice = zoom( + slice, (patch_size[0] / x, patch_size[1] / y), order=0) + input = torch.from_numpy(slice).unsqueeze( + 0).unsqueeze(0).float().cuda() + net.eval() + with torch.no_grad(): + out = torch.argmax(torch.softmax( + net(input)[0], dim=1), dim=1).squeeze(0) + out = out.cpu().detach().numpy() + pred = zoom( + out, (x / patch_size[0], y / patch_size[1]), order=0) + prediction[ind] = pred + else: + input = torch.from_numpy(image).unsqueeze( 0).unsqueeze(0).float().cuda() net.eval() with torch.no_grad(): - output_main, _, _, _ = net(input) out = torch.argmax(torch.softmax( - output_main, dim=1), dim=1).squeeze(0) - out = out.cpu().detach().numpy() - pred = zoom(out, (x / patch_size[0], y / patch_size[1]), order=0) - prediction[ind] = pred + net(input)[0], dim=1), dim=1).squeeze(0) + prediction = out.cpu().detach().numpy() metric_list = [] for i in range(1, classes): - metric_list.append(calculate_metric_percase( - prediction == i, label == i)) + metric_list.append(calculate_metric_percase(prediction == i, label == i)) + # return metric_list, image, prediction, label return metric_list diff --git a/code/val_3D.py b/code/val_3D.py deleted file mode 100644 index 4befb87..0000000 --- a/code/val_3D.py +++ /dev/null @@ -1,107 +0,0 @@ -import math -from glob import glob - -import h5py -import nibabel as nib -import numpy as np -import SimpleITK as sitk -import torch -import torch.nn.functional as F -from medpy import metric -from tqdm import tqdm - - -def test_single_case(net, image, stride_xy, stride_z, patch_size, num_classes=1): - w, h, d = image.shape - - # if the size of image is less than patch_size, then padding it - add_pad = False - if w < patch_size[0]: - w_pad = patch_size[0]-w - add_pad = True - else: - w_pad = 0 - if h < patch_size[1]: - h_pad = patch_size[1]-h - add_pad = True - else: - h_pad = 0 - if d < patch_size[2]: - d_pad = patch_size[2]-d - add_pad = True - else: - d_pad = 0 - wl_pad, wr_pad = w_pad//2, w_pad-w_pad//2 - hl_pad, hr_pad = h_pad//2, h_pad-h_pad//2 - dl_pad, dr_pad = d_pad//2, d_pad-d_pad//2 - if add_pad: - image = np.pad(image, [(wl_pad, wr_pad), (hl_pad, hr_pad), - (dl_pad, dr_pad)], mode='constant', constant_values=0) - ww, hh, dd = image.shape - - sx = math.ceil((ww - patch_size[0]) / stride_xy) + 1 - sy = math.ceil((hh - patch_size[1]) / stride_xy) + 1 - sz = math.ceil((dd - patch_size[2]) / stride_z) + 1 - # print("{}, {}, {}".format(sx, sy, sz)) - score_map = np.zeros((num_classes, ) + image.shape).astype(np.float32) - cnt = np.zeros(image.shape).astype(np.float32) - - for x in range(0, sx): - xs = min(stride_xy*x, ww-patch_size[0]) - for y in range(0, sy): - ys = min(stride_xy * y, hh-patch_size[1]) - for z in range(0, sz): - zs = min(stride_z * z, dd-patch_size[2]) - test_patch = image[xs:xs+patch_size[0], - ys:ys+patch_size[1], zs:zs+patch_size[2]] - test_patch = np.expand_dims(np.expand_dims( - test_patch, axis=0), axis=0).astype(np.float32) - test_patch = torch.from_numpy(test_patch).cuda() - - with torch.no_grad(): - y1 = net(test_patch) - # ensemble - y = torch.softmax(y1, dim=1) - y = y.cpu().data.numpy() - y = y[0, :, :, :, :] - score_map[:, xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] \ - = score_map[:, xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] + y - cnt[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] \ - = cnt[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] + 1 - score_map = score_map/np.expand_dims(cnt, axis=0) - label_map = np.argmax(score_map, axis=0) - - if add_pad: - label_map = label_map[wl_pad:wl_pad+w, - hl_pad:hl_pad+h, dl_pad:dl_pad+d] - score_map = score_map[:, wl_pad:wl_pad + - w, hl_pad:hl_pad+h, dl_pad:dl_pad+d] - return label_map - - -def cal_metric(gt, pred): - if pred.sum() > 0 and gt.sum() > 0: - dice = metric.binary.dc(pred, gt) - hd95 = metric.binary.hd95(pred, gt) - return np.array([dice, hd95]) - else: - return np.zeros(2) - - -def test_all_case(net, base_dir, test_list="full_test.list", num_classes=4, patch_size=(48, 160, 160), stride_xy=32, stride_z=24): - with open(base_dir + '/{}'.format(test_list), 'r') as f: - image_list = f.readlines() - image_list = [base_dir + "/data/{}.h5".format( - item.replace('\n', '').split(",")[0]) for item in image_list] - total_metric = np.zeros((num_classes-1, 2)) - print("Validation begin") - for image_path in tqdm(image_list): - h5f = h5py.File(image_path, 'r') - image = h5f['image'][:] - label = h5f['label'][:] - prediction = test_single_case( - net, image, stride_xy, stride_z, patch_size, num_classes=num_classes) - for i in range(1, num_classes): - total_metric[i-1, :] += cal_metric(label == i, prediction == i) - print("Validation end") - return total_metric / len(image_list) diff --git a/code/val_urpc_util.py b/code/val_urpc_util.py deleted file mode 100644 index a4257d1..0000000 --- a/code/val_urpc_util.py +++ /dev/null @@ -1,107 +0,0 @@ -import math -from glob import glob - -import h5py -import nibabel as nib -import numpy as np -import SimpleITK as sitk -import torch -import torch.nn.functional as F -from medpy import metric -from tqdm import tqdm - - -def test_single_case(net, image, stride_xy, stride_z, patch_size, num_classes=1): - w, h, d = image.shape - - # if the size of image is less than patch_size, then padding it - add_pad = False - if w < patch_size[0]: - w_pad = patch_size[0]-w - add_pad = True - else: - w_pad = 0 - if h < patch_size[1]: - h_pad = patch_size[1]-h - add_pad = True - else: - h_pad = 0 - if d < patch_size[2]: - d_pad = patch_size[2]-d - add_pad = True - else: - d_pad = 0 - wl_pad, wr_pad = w_pad//2, w_pad-w_pad//2 - hl_pad, hr_pad = h_pad//2, h_pad-h_pad//2 - dl_pad, dr_pad = d_pad//2, d_pad-d_pad//2 - if add_pad: - image = np.pad(image, [(wl_pad, wr_pad), (hl_pad, hr_pad), - (dl_pad, dr_pad)], mode='constant', constant_values=0) - ww, hh, dd = image.shape - - sx = math.ceil((ww - patch_size[0]) / stride_xy) + 1 - sy = math.ceil((hh - patch_size[1]) / stride_xy) + 1 - sz = math.ceil((dd - patch_size[2]) / stride_z) + 1 - # print("{}, {}, {}".format(sx, sy, sz)) - score_map = np.zeros((num_classes, ) + image.shape).astype(np.float32) - cnt = np.zeros(image.shape).astype(np.float32) - - for x in range(0, sx): - xs = min(stride_xy*x, ww-patch_size[0]) - for y in range(0, sy): - ys = min(stride_xy * y, hh-patch_size[1]) - for z in range(0, sz): - zs = min(stride_z * z, dd-patch_size[2]) - test_patch = image[xs:xs+patch_size[0], - ys:ys+patch_size[1], zs:zs+patch_size[2]] - test_patch = np.expand_dims(np.expand_dims( - test_patch, axis=0), axis=0).astype(np.float32) - test_patch = torch.from_numpy(test_patch).cuda() - - with torch.no_grad(): - y1, _, _, _ = net(test_patch) - # ensemble - y = torch.softmax(y1, dim=1) - y = y.cpu().data.numpy() - y = y[0, :, :, :, :] - score_map[:, xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] \ - = score_map[:, xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] + y - cnt[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] \ - = cnt[xs:xs+patch_size[0], ys:ys+patch_size[1], zs:zs+patch_size[2]] + 1 - score_map = score_map/np.expand_dims(cnt, axis=0) - label_map = np.argmax(score_map, axis=0) - - if add_pad: - label_map = label_map[wl_pad:wl_pad+w, - hl_pad:hl_pad+h, dl_pad:dl_pad+d] - score_map = score_map[:, wl_pad:wl_pad + - w, hl_pad:hl_pad+h, dl_pad:dl_pad+d] - return label_map - - -def cal_metric(gt, pred): - if pred.sum() > 0 and gt.sum() > 0: - dice = metric.binary.dc(pred, gt) - hd95 = metric.binary.hd95(pred, gt) - return np.array([dice, hd95]) - else: - return np.zeros(2) - - -def test_all_case(net, base_dir, test_list="val.list", num_classes=4, patch_size=(48, 160, 160), stride_xy=32, stride_z=24): - with open(base_dir + '/{}'.format(test_list), 'r') as f: - image_list = f.readlines() - image_list = [base_dir + "/data/{}.h5".format( - item.replace('\n', '').split(",")[0]) for item in image_list] - total_metric = np.zeros((num_classes-1, 2)) - print("Validation begin") - for image_path in tqdm(image_list): - h5f = h5py.File(image_path, 'r') - image = h5f['image'][:] - label = h5f['label'][:] - prediction = test_single_case( - net, image, stride_xy, stride_z, patch_size, num_classes=num_classes) - for i in range(1, num_classes): - total_metric[i-1, :] += cal_metric(label == i, prediction == i) - print("Validation end") - return total_metric / len(image_list) From f4f8879c5b9e91db67aede8bdfbe9ca5fff1c63b Mon Sep 17 00:00:00 2001 From: Xiangde Luo <19658103+Luoxd1996@users.noreply.github.com> Date: Sun, 6 Mar 2022 15:30:21 +0800 Subject: [PATCH 2/7] Create README.md --- data/README.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 data/README.md diff --git a/data/README.md b/data/README.md new file mode 100644 index 0000000..caa4c88 --- /dev/null +++ b/data/README.md @@ -0,0 +1,5 @@ +There are three datasets: +(1) [ACDC](https://acdc.creatis.insa-lyon.fr) with 200 3D MRI volumes; +(2) [ProstateX](https://github.com/ykl-ucla/prostate_zonal_seg) with 201 3D MRI volumes; +(3) [CHAOS](https://chaos.grand-challenge.org) with 20 MRI-T2 volumes. +The preprocessed datasets are provided in [Here](https://drive.google.com/file/d/1BtT4mEtRPBJb2F6buGc9HqBZ60eUKa33/view?usp=sharing). From 8f267b0cd72841c1d5d3b486b9c713ae49d0450c Mon Sep 17 00:00:00 2001 From: Xiangde Luo <19658103+Luoxd1996@users.noreply.github.com> Date: Sun, 6 Mar 2022 17:26:43 +0800 Subject: [PATCH 3/7] Update README.md --- data/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/data/README.md b/data/README.md index caa4c88..3cf4884 100644 --- a/data/README.md +++ b/data/README.md @@ -2,4 +2,4 @@ There are three datasets: (1) [ACDC](https://acdc.creatis.insa-lyon.fr) with 200 3D MRI volumes; (2) [ProstateX](https://github.com/ykl-ucla/prostate_zonal_seg) with 201 3D MRI volumes; (3) [CHAOS](https://chaos.grand-challenge.org) with 20 MRI-T2 volumes. -The preprocessed datasets are provided in [Here](https://drive.google.com/file/d/1BtT4mEtRPBJb2F6buGc9HqBZ60eUKa33/view?usp=sharing). +The preprocessed datasets are provided in [Google Drive](https://drive.google.com/file/d/1BtT4mEtRPBJb2F6buGc9HqBZ60eUKa33/view?usp=sharing) or [Pan.Baidu]https://pan.baidu.com/s/15J5G8hw8ATy5VTnLOxZPLA), pwd: ```data```. From e97c8ba8c8efb559956e2ad9bec65e146a3fe2cc Mon Sep 17 00:00:00 2001 From: Xiangde Luo <19658103+Luoxd1996@users.noreply.github.com> Date: Sun, 6 Mar 2022 17:26:59 +0800 Subject: [PATCH 4/7] Update README.md --- data/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/data/README.md b/data/README.md index 3cf4884..33717d0 100644 --- a/data/README.md +++ b/data/README.md @@ -2,4 +2,4 @@ There are three datasets: (1) [ACDC](https://acdc.creatis.insa-lyon.fr) with 200 3D MRI volumes; (2) [ProstateX](https://github.com/ykl-ucla/prostate_zonal_seg) with 201 3D MRI volumes; (3) [CHAOS](https://chaos.grand-challenge.org) with 20 MRI-T2 volumes. -The preprocessed datasets are provided in [Google Drive](https://drive.google.com/file/d/1BtT4mEtRPBJb2F6buGc9HqBZ60eUKa33/view?usp=sharing) or [Pan.Baidu]https://pan.baidu.com/s/15J5G8hw8ATy5VTnLOxZPLA), pwd: ```data```. +The preprocessed datasets are provided in [Google Drive](https://drive.google.com/file/d/1BtT4mEtRPBJb2F6buGc9HqBZ60eUKa33/view?usp=sharing) or [Pan.Baidu](https://pan.baidu.com/s/15J5G8hw8ATy5VTnLOxZPLA), pwd: ```data```. From 126084dd856eae0969333b8abe9ab540279c3bc7 Mon Sep 17 00:00:00 2001 From: luoxd Date: Mon, 7 Mar 2022 08:58:37 +0800 Subject: [PATCH 5/7] update 5fold cross validation code --- code/dataloaders/dataset.py | 66 +- .../prostate_dataset_preprocessing.py | 100 -- code/networks/unet.py | 12 +- code/test.py | 90 +- code/train_cross_pseudo_supervision.py | 314 ++++ code/train_deep_adversarial_network.py | 271 ++++ code/train_entropy_minimization.py | 238 +++ code/train_fullsup_exp_fold0.sh | 75 + code/train_fully_supervised.py | 219 +++ ...rain_interpolation_consistency_training.py | 12 +- code/train_mean_teacher.py | 24 +- code/train_uncertainty_aware_mean_teacher.py | 26 +- code/val_2D.py | 93 +- data/ACDC/README.md | 2 - data/ACDC/test.list | 40 - data/ACDC/train.list | 140 -- data/ACDC/train_slices.list | 1312 ----------------- data/ACDC/val.list | 20 - data/BraTS2019/README.md | 1 - data/BraTS2019/test.txt | 60 - data/BraTS2019/train.txt | 250 ---- data/BraTS2019/val.txt | 25 - 22 files changed, 1374 insertions(+), 2016 deletions(-) delete mode 100644 code/dataloaders/prostate_dataset_preprocessing.py create mode 100644 code/train_cross_pseudo_supervision.py create mode 100644 code/train_deep_adversarial_network.py create mode 100644 code/train_entropy_minimization.py create mode 100644 code/train_fullsup_exp_fold0.sh create mode 100644 code/train_fully_supervised.py delete mode 100644 data/ACDC/README.md delete mode 100755 data/ACDC/test.list delete mode 100755 data/ACDC/train.list delete mode 100644 data/ACDC/train_slices.list delete mode 100755 data/ACDC/val.list delete mode 100644 data/BraTS2019/README.md delete mode 100755 data/BraTS2019/test.txt delete mode 100755 data/BraTS2019/train.txt delete mode 100755 data/BraTS2019/val.txt diff --git a/code/dataloaders/dataset.py b/code/dataloaders/dataset.py index 1724e00..29365f8 100644 --- a/code/dataloaders/dataset.py +++ b/code/dataloaders/dataset.py @@ -10,8 +10,13 @@ import torch from scipy import ndimage from scipy.ndimage.interpolation import zoom -from torch.utils.data import Dataset from sklearn.model_selection import KFold +from torch.utils.data import Dataset + +try: # SciPy >= 0.19 + from scipy.special import comb +except ImportError: + from scipy.misc import comb class BaseDataSets(Dataset): @@ -86,7 +91,7 @@ def __getitem__(self, idx): return sample -def random_rot_flip(image, label): +def random_flip(image, label): k = np.random.randint(0, 4) image = np.rot90(image, k) label = np.rot90(label, k) @@ -112,6 +117,57 @@ def random_noise(image, label, mu=0, sigma=0.1): return image, label +def bernstein_poly(i, n, t): + """ + The Bernstein polynomial of n, i as a function of t + """ + return comb(n, i) * (t**(n-i)) * (1 - t)**i + + +def bezier_curve(points, nTimes=1000): + """ + Given a set of control points, return the + bezier curve defined by the control points. + Control points should be a list of lists, or list of tuples + such as [ [1,1], + [2,3], + [4,5], ..[Xn, Yn] ] + nTimes is the number of time steps, defaults to 1000 + See http://processingjs.nihongoresources.com/bezierinfo/ + """ + + nPoints = len(points) + xPoints = np.array([p[0] for p in points]) + yPoints = np.array([p[1] for p in points]) + + t = np.linspace(0.0, 1.0, nTimes) + + polynomial_array = np.array( + [bernstein_poly(i, nPoints-1, t) for i in range(0, nPoints)]) + + xvals = np.dot(xPoints, polynomial_array) + yvals = np.dot(yPoints, polynomial_array) + + return xvals, yvals + + +def nonlinear_transformation(x, label, prob=0.5): + if random.random() >= prob: + return x, label + points = [[0, 0], [random.random(), random.random()], [ + random.random(), random.random()], [1, 1]] + xpoints = [p[0] for p in points] + ypoints = [p[1] for p in points] + xvals, yvals = bezier_curve(points, nTimes=100000) + if random.random() < 0.5: + # Half change to get flip + xvals = np.sort(xvals) + else: + xvals, yvals = np.sort(xvals), np.sort(yvals) + nonlinear_x = np.interp(x, xvals, yvals) + return nonlinear_x, label + + class RandomGenerator(object): def __init__(self, output_size): self.output_size = output_size @@ -119,11 +175,13 @@ def __init__(self, output_size): def __call__(self, sample): image, label = sample['image'], sample['label'] if random.random() > 0.5: - image, label = random_rot_flip(image, label) + image, label = random_flip(image, label) if random.random() > 0.5: image, label = random_rotate(image, label, cval=0) if random.random() > 0.5: image, label = random_noise(image, label) + if random.random() > 0.5: + image, label = nonlinear_transformation(image, label) x, y = image.shape image = zoom( image, (self.output_size[0] / x, self.output_size[1] / y), order=0) @@ -131,6 +189,6 @@ def __call__(self, sample): label, (self.output_size[0] / x, self.output_size[1] / y), order=0) image = torch.from_numpy( image.astype(np.float32)).unsqueeze(0) - label = torch.from_numpy(label.astype(np.uint8)) + label = torch.from_numpy(label.astype(np.int16)) sample = {'image': image, 'label': label} return sample diff --git a/code/dataloaders/prostate_dataset_preprocessing.py b/code/dataloaders/prostate_dataset_preprocessing.py deleted file mode 100644 index a1af761..0000000 --- a/code/dataloaders/prostate_dataset_preprocessing.py +++ /dev/null @@ -1,100 +0,0 @@ -# save images in slice level -import glob -import os - -import h5py -import numpy as np -import SimpleITK as sitk - - -class MedicalImageDeal(object): - def __init__(self, img, percent=1): - self.img = img - self.percent = percent - - @property - def valid_img(self): - from skimage import exposure - cdf = exposure.cumulative_distribution(self.img) - watershed = cdf[1][cdf[0] >= self.percent][0] - return np.clip(self.img, self.img.min(), watershed) - - @property - def norm_img(self): - return (self.img - self.img.min()) / (self.img.max() - self.img.min()) - -# slice_num = 0 -# mask_path = sorted( -# glob.glob("/home/SENSETIME/luoxiangde.vendor/Desktop/SSL4MIS_5Fold/data/prostate_zonal_nii/*_lab.nii.gz")) -# for image_path in mask_path: -# image_itk = sitk.ReadImage(image_path.replace("_lab", "")) -# image = sitk.GetArrayFromImage(image_itk) - -# image = MedicalImageDeal(image, percent=0.99).valid_img -# image = (image - image.min()) / (image.max() - image.min()) -# norm_img_itk = sitk.GetImageFromArray(image) -# norm_img_itk.CopyInformation(image_itk) -# sitk.WriteImage(norm_img_itk, image_path.replace("_lab", "")) - - -# saving images in slice level - -slice_num = 0 -mask_path = sorted( - glob.glob("/home/SENSETIME/luoxiangde.vendor/Desktop/SSL4MIS_5Fold/data/CHAOS_NII/label/*.nii.gz")) -for case in mask_path: - label_itk = sitk.ReadImage(case) - label = sitk.GetArrayFromImage(label_itk) - - image_path = case.replace("/label/", "/image/") - image_itk = sitk.ReadImage(image_path) - image = sitk.GetArrayFromImage(image_itk) - spacing = image_itk.GetSpacing() - - image = MedicalImageDeal(image, percent=0.99).valid_img - image = (image - image.min()) / (image.max() - image.min()) - print(image.shape) - image = image.astype(np.float32) - item = case.split("/")[-1].split(".")[0].replace("_gt", "") - if image.shape != label.shape: - print("Error") - print(item) - - f = h5py.File( - '/home/SENSETIME/luoxiangde.vendor/Desktop/SSL4MIS_5Fold/data/CHAOS/all_volumes/{}.h5'.format(item), 'w') - f.create_dataset( - 'image', data=image, compression="gzip") - f.create_dataset('label', data=label, compression="gzip") - f.create_dataset('spacing', data=np.array(spacing), compression="gzip") - f.close() -print("Converted all ACDC volumes to 2D slices") -print("Total {} slices".format(slice_num)) -# # saving images in volume level - -# slice_num = 0 -# mask_path = sorted( -# glob.glob("/home/SENSETIME/luoxiangde.vendor/Desktop/SSL4MIS_5Fold/data/prostate_zonal_nii/*_lab.nii.gz")) -# for case in mask_path: -# label_itk = sitk.ReadImage(case) -# label = sitk.GetArrayFromImage(label_itk) - -# image_path = case.replace("_lab", "") -# image_itk = sitk.ReadImage(image_path) -# image = sitk.GetArrayFromImage(image_itk) -# spacing = image_itk.GetSpacing() - -# image = image.astype(np.float32) -# item = case.split("/")[-1].split(".")[0].replace("_lab", "") -# if image.shape != label.shape: -# print("Error") -# print(item) - -# f = h5py.File( -# '/home/SENSETIME/luoxiangde.vendor/Desktop/SSL4MIS_5Fold/data/ProstateX/all_volumes/{}.h5'.format(item), 'w') -# f.create_dataset( -# 'image', data=image, compression="gzip") -# f.create_dataset('label', data=label, compression="gzip") -# f.create_dataset('spacing', data=np.array(spacing), compression="gzip") -# f.close() -# print("Converted all Prostate volumes to 2D slices") -# print("Total {} slices".format(slice_num)) \ No newline at end of file diff --git a/code/networks/unet.py b/code/networks/unet.py index 4ab4b9d..102d821 100755 --- a/code/networks/unet.py +++ b/code/networks/unet.py @@ -48,7 +48,7 @@ class UpBlock(nn.Module): """Upssampling followed by ConvBlock""" def __init__(self, in_channels1, in_channels2, out_channels, dropout_p, - bilinear=True): + bilinear=False): super(UpBlock, self).__init__() self.bilinear = bilinear if bilinear: @@ -109,13 +109,13 @@ def __init__(self, params): assert (len(self.ft_chns) == 5) self.up1 = UpBlock( - self.ft_chns[4], self.ft_chns[3], self.ft_chns[3], dropout_p=0.0) + self.ft_chns[4], self.ft_chns[3], self.ft_chns[3], dropout_p=0.0, bilinear=self.bilinear) self.up2 = UpBlock( - self.ft_chns[3], self.ft_chns[2], self.ft_chns[2], dropout_p=0.0) + self.ft_chns[3], self.ft_chns[2], self.ft_chns[2], dropout_p=0.0, bilinear=self.bilinear) self.up3 = UpBlock( - self.ft_chns[2], self.ft_chns[1], self.ft_chns[1], dropout_p=0.0) + self.ft_chns[2], self.ft_chns[1], self.ft_chns[1], dropout_p=0.0, bilinear=self.bilinear) self.up4 = UpBlock( - self.ft_chns[1], self.ft_chns[0], self.ft_chns[0], dropout_p=0.0) + self.ft_chns[1], self.ft_chns[0], self.ft_chns[0], dropout_p=0.0, bilinear=self.bilinear) self.out_conv = nn.Conv2d(self.ft_chns[0], self.n_class, kernel_size=3, padding=1) @@ -291,7 +291,7 @@ def __init__(self, in_chns, class_num): 'feature_chns': [16, 32, 64, 128, 256], 'dropout': [0.05, 0.1, 0.2, 0.3, 0.5], 'class_num': class_num, - 'bilinear': False, + 'bilinear': True, 'acti_func': 'relu'} self.encoder = Encoder(params) diff --git a/code/test.py b/code/test.py index e3566a4..71170ea 100644 --- a/code/test.py +++ b/code/test.py @@ -4,16 +4,17 @@ import shutil import h5py -from matplotlib.pyplot import axis import nibabel as nib import numpy as np import SimpleITK as sitk import torch +from matplotlib.pyplot import axis from medpy import metric from scipy.ndimage import zoom from scipy.ndimage.interpolation import zoom from sklearn.model_selection import KFold from tqdm import tqdm + from networks.net_factory import net_factory parser = argparse.ArgumentParser() @@ -26,7 +27,7 @@ parser.add_argument('--labeled_ratio', type=int, default=8, help='1/labeled_ratio data is provided mask') parser.add_argument('--fold', type=int, - default=1, help='fold') + default=3, help='fold') parser.add_argument('--patch_size', type=list, default=[256, 256], help='patch size of network input') parser.add_argument('--num_classes', type=int, default=3, @@ -59,27 +60,88 @@ def calculate_metric_percase(pred, gt, spacing): return dice, hd95, asd -def test_single_volume(case, net, test_save_path, FLAGS): +# def test_single_volume(case, net, test_save_path, FLAGS): +# h5f = h5py.File(FLAGS.root_path + +# "/all_volumes/{}".format(case), 'r') +# image = h5f['image'][:] +# label = h5f['label'][:] +# spacing = h5f['spacing'][:] +# prediction = np.zeros_like(label) +# for ind in range(image.shape[0]): +# slice = image[ind, :, :] +# x, y = slice.shape[0], slice.shape[1] +# slice = zoom(slice, (FLAGS.patch_size / x, FLAGS.patch_size / y), order=0) +# input = torch.from_numpy(slice).unsqueeze( +# 0).unsqueeze(0).float().cuda() +# net.eval() +# with torch.no_grad(): +# out_main = net(input) +# out = torch.argmax(torch.softmax( +# out_main, dim=1), dim=1).squeeze(0) +# out = out.cpu().detach().numpy() +# pred = zoom(out, (x / FLAGS.patch_size, y / FLAGS.patch_size), order=0) +# prediction[ind] = pred +# case = case.replace(".h5", "") + +# metric_list = [] +# for i in range(1, FLAGS.num_classes): +# metric_list.append(calculate_metric_percase( +# prediction == i, label == i, spacing=(spacing[2], spacing[0], spacing[1]))) +# img_itk = sitk.GetImageFromArray(image.astype(np.float32)) +# img_itk.SetSpacing(spacing) +# prd_itk = sitk.GetImageFromArray(prediction.astype(np.float32)) +# prd_itk.SetSpacing(spacing) +# lab_itk = sitk.GetImageFromArray(label.astype(np.float32)) +# lab_itk.SetSpacing(spacing) +# sitk.WriteImage(prd_itk, test_save_path + case + "_pred.nii.gz") +# sitk.WriteImage(img_itk, test_save_path + case + "_img.nii.gz") +# sitk.WriteImage(lab_itk, test_save_path + case + "_gt.nii.gz") +# return np.array(metric_list) + + +def test_single_volume(case, net, test_save_path, FLAGS, batch_size=12): h5f = h5py.File(FLAGS.root_path + "/all_volumes/{}".format(case), 'r') image = h5f['image'][:] label = h5f['label'][:] spacing = h5f['spacing'][:] - prediction = np.zeros_like(label) - for ind in range(image.shape[0]): - slice = image[ind, :, :] - x, y = slice.shape[0], slice.shape[1] - slice = zoom(slice, (FLAGS.patch_size / x, FLAGS.patch_size / y), order=0) - input = torch.from_numpy(slice).unsqueeze( + if len(image.shape) == 3: + prediction = np.zeros_like(label) + ind_x = np.array([i for i in range(image.shape[0])]) + for ind in ind_x[::batch_size]: + if ind + batch_size < image.shape[0]: + slice = image[ind:ind + batch_size, ...] + thickness, x, y = slice.shape[0], slice.shape[1], slice.shape[2] + slice = zoom(slice, (1, FLAGS.patch_size[0] / x, FLAGS.patch_size[1] / y), order=0) + input = torch.from_numpy(slice).unsqueeze(1).float().cuda() + net.eval() + with torch.no_grad(): + out = torch.argmax(torch.softmax( + net(input), dim=1), dim=1) + out = out.cpu().detach().numpy() + pred = zoom(out, (1, x / FLAGS.patch_size[0], y / FLAGS.patch_size[1]), order=0) + prediction[ind:ind + batch_size, ...] = pred + else: + slice = image[ind:, ...] + thickness, x, y = slice.shape[0], slice.shape[1], slice.shape[2] + slice = zoom(slice, (1, FLAGS.patch_size[0] / x, FLAGS.patch_size[1] / y), order=0) + input = torch.from_numpy(slice).unsqueeze(1).float().cuda() + net.eval() + with torch.no_grad(): + out = torch.argmax(torch.softmax( + net(input), dim=1), dim=1) + out = out.cpu().detach().numpy() + pred = zoom(out, (1, x / FLAGS.patch_size[0], y / FLAGS.patch_size[1]), order=0) + prediction[ind:, ...] = pred + else: + input = torch.from_numpy(image).unsqueeze( 0).unsqueeze(0).float().cuda() net.eval() with torch.no_grad(): - out_main = net(input) out = torch.argmax(torch.softmax( - out_main, dim=1), dim=1).squeeze(0) - out = out.cpu().detach().numpy() - pred = zoom(out, (x / FLAGS.patch_size, y / FLAGS.patch_size), order=0) - prediction[ind] = pred + net(input), dim=1), dim=1).squeeze(0) + prediction = out.cpu().detach().numpy() + case = case.replace(".h5", "") metric_list = [] diff --git a/code/train_cross_pseudo_supervision.py b/code/train_cross_pseudo_supervision.py new file mode 100644 index 0000000..ca1d085 --- /dev/null +++ b/code/train_cross_pseudo_supervision.py @@ -0,0 +1,314 @@ +import argparse +import logging +import os +import random +import shutil +import sys +import time +from itertools import cycle + +import numpy as np +import torch +import torch.backends.cudnn as cudnn +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from tensorboardX import SummaryWriter +from torch.nn import BCEWithLogitsLoss +from torch.nn.modules.loss import CrossEntropyLoss +from torch.utils.data import DataLoader +from torchvision import transforms +from torchvision.utils import make_grid +from tqdm import tqdm + +from dataloaders.dataset import BaseDataSets, RandomGenerator +from networks.discriminator import FCDiscriminator +from networks.net_factory import net_factory +from utils import losses, metrics, ramps +from val_2D import test_single_volume + +parser = argparse.ArgumentParser() +parser.add_argument('--root_path', type=str, + default='../data/ProstateX', help='Name of Experiment') +parser.add_argument('--exp', type=str, + default='ProstateX/CPS', help='experiment_name') +parser.add_argument('--model', type=str, + default='unet', help='model_name') +parser.add_argument('--fold', type=int, + default=1, help='cross validation') +parser.add_argument('--max_iterations', type=int, + default=30000, help='maximum epoch number to train') +parser.add_argument('--batch_size', type=int, default=16, + help='batch_size per gpu') + +parser.add_argument('--deterministic', type=int, default=1, + help='whether use deterministic training') +parser.add_argument('--base_lr', type=float, default=0.03, + help='segmentation network learning rate') +parser.add_argument('--patch_size', type=list, default=[256, 256], + help='patch size of network input') +parser.add_argument('--seed', type=int, default=2022, help='random seed') +parser.add_argument('--num_classes', type=int, default=3, + help='output channel of network') + +# label and unlabel +parser.add_argument('--labeled_ratio', type=int, default=8, + help='1/labeled_ratio data is provided mask') +# costs +parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') +parser.add_argument('--consistency_type', type=str, + default="mse", help='consistency_type') +parser.add_argument('--consistency', type=float, + default=0.1, help='consistency') +parser.add_argument('--consistency_rampup', type=float, + default=200.0, help='consistency_rampup') +args = parser.parse_args() + + +def get_current_consistency_weight(epoch): + # Consistency ramp-up from https://arxiv.org/abs/1610.02242 + return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) + + +def train(args, snapshot_path): + writer = SummaryWriter(snapshot_path + '/log') + base_lr = args.base_lr + num_classes = args.num_classes + max_iterations = args.max_iterations + + model1 = net_factory(net_type=args.model, in_chns=1, + class_num=num_classes) + model2 = net_factory(net_type=args.model, in_chns=1, + class_num=num_classes) + + db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)])) + db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)])) + logging.info("Labeled slices: {} ".format(len(db_train_labeled))) + logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) + + trainloader_labeled = DataLoader( + db_train_labeled, batch_size=args.batch_size//2, shuffle=True) + trainloader_unlabeled = DataLoader( + db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) + + db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, + split="val", labeled_ratio=args.labeled_ratio) + valloader = DataLoader(db_val, batch_size=1) + + model1.train() + model2.train() + + optimizer1 = optim.SGD(model1.parameters(), lr=base_lr, + momentum=0.9, weight_decay=0.0001) + optimizer2 = optim.SGD(model2.parameters(), lr=base_lr, + momentum=0.9, weight_decay=0.0001) + ce_loss = CrossEntropyLoss() + dice_loss = losses.DiceLoss(num_classes) + + logging.info("{} iterations per epoch".format(len(trainloader_unlabeled))) + + iter_num = 0 + max_epoch = max_iterations // len(trainloader_unlabeled) + 1 + best_performance1 = 0.0 + best_performance2 = 0.0 + iterator = tqdm(range(max_epoch), ncols=70) + for epoch_num in iterator: + for i, (sampled_batch_labeled, sampled_batch_unlabeled) in enumerate(zip(cycle(trainloader_labeled), trainloader_unlabeled)): + volume_batch, label_batch = sampled_batch_labeled['image'], sampled_batch_labeled['label'] + volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() + unlabeled_volume_batch = sampled_batch_unlabeled['image'].cuda() + + outputs1 = model1(volume_batch) + outputs1_soft = torch.softmax(outputs1, dim=1) + + outputs1_unlabeled = model1(unlabeled_volume_batch) + outputs1_unlabeled_soft = torch.softmax(outputs1_unlabeled, dim=1) + + outputs2 = model2(volume_batch) + outputs2_soft = torch.softmax(outputs2, dim=1) + + outputs2_unlabeled = model2(unlabeled_volume_batch) + outputs2_unlabeled_soft = torch.softmax(outputs2_unlabeled, dim=1) + + supervised_loss1 = 0.5 * \ + (ce_loss(outputs1, label_batch[:].long( + )) + dice_loss(outputs1_soft, label_batch[:].unsqueeze(1))) + supervised_loss2 = 0.5 * \ + (ce_loss(outputs2, label_batch[:].long( + )) + dice_loss(outputs2_soft, label_batch[:].unsqueeze(1))) + + pseudo_outputs1 = torch.argmax( + outputs1_unlabeled_soft.detach(), dim=1, keepdim=False) + pseudo_outputs2 = torch.argmax( + outputs2_unlabeled_soft.detach(), dim=1, keepdim=False) + + pseudo_supervision1 = ce_loss(outputs1_unlabeled, pseudo_outputs2) + pseudo_supervision2 = ce_loss(outputs2_unlabeled, pseudo_outputs1) + + consistency_weight = get_current_consistency_weight( + iter_num // (args.max_iterations/args.consistency_rampup)) + + model1_loss = supervised_loss1 + consistency_weight * pseudo_supervision1 + model2_loss = supervised_loss2 + consistency_weight * pseudo_supervision2 + + loss = model1_loss + model2_loss + + optimizer1.zero_grad() + optimizer2.zero_grad() + + loss.backward() + + optimizer1.step() + optimizer2.step() + + iter_num = iter_num + 1 + + lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 + for param_group in optimizer1.param_groups: + param_group['lr'] = lr_ + for param_group in optimizer2.param_groups: + param_group['lr'] = lr_ + + writer.add_scalar('lr', lr_, iter_num) + writer.add_scalar( + 'consistency_weight/consistency_weight', consistency_weight, iter_num) + writer.add_scalar('loss/model1_loss', + model1_loss, iter_num) + writer.add_scalar('loss/model2_loss', + model2_loss, iter_num) + logging.info('iteration %d : model1 loss : %f model2 loss : %f' % ( + iter_num, model1_loss.item(), model2_loss.item())) + if iter_num % 20 == 0: + image = volume_batch[0, 0:1, :, :] + writer.add_image('train/Image', image, iter_num) + outputs = torch.argmax(torch.softmax( + outputs1, dim=1), dim=1, keepdim=True) + writer.add_image('train/model1_Prediction', + outputs[0, ...] * 50, iter_num) + outputs = torch.argmax(torch.softmax( + outputs2, dim=1), dim=1, keepdim=True) + writer.add_image('train/model2_Prediction', + outputs[0, ...] * 50, iter_num) + labs = label_batch[0, ...].unsqueeze(0) * 50 + writer.add_image('train/GroundTruth', labs, iter_num) + + if iter_num > 0 and iter_num % 200 == 0: + model1.eval() + metric_list = 0.0 + for i_batch, sampled_batch in enumerate(valloader): + metric_i = test_single_volume( + sampled_batch["image"], sampled_batch["label"], model1, classes=num_classes) + metric_list += np.array(metric_i) + metric_list = metric_list / len(db_val) + for class_i in range(num_classes-1): + writer.add_scalar('info/model1_val_{}_dice'.format(class_i+1), + metric_list[class_i, 0], iter_num) + writer.add_scalar('info/model1_val_{}_hd95'.format(class_i+1), + metric_list[class_i, 1], iter_num) + + performance1 = np.mean(metric_list, axis=0)[0] + + mean_hd951 = np.mean(metric_list, axis=0)[1] + writer.add_scalar('info/model1_val_mean_dice', + performance1, iter_num) + writer.add_scalar('info/model1_val_mean_hd95', + mean_hd951, iter_num) + + if performance1 > best_performance1: + best_performance1 = performance1 + save_mode_path = os.path.join(snapshot_path, + 'model1_iter_{}_dice_{}.pth'.format( + iter_num, round(best_performance1, 4))) + save_best = os.path.join(snapshot_path, + '{}_best_model1.pth'.format(args.model)) + torch.save(model1.state_dict(), save_mode_path) + torch.save(model1.state_dict(), save_best) + + logging.info( + 'iteration %d : model1_mean_dice : %f model1_mean_hd95 : %f' % (iter_num, performance1, mean_hd951)) + model1.train() + + model2.eval() + metric_list = 0.0 + for i_batch, sampled_batch in enumerate(valloader): + metric_i = test_single_volume( + sampled_batch["image"], sampled_batch["label"], model2, classes=num_classes) + metric_list += np.array(metric_i) + metric_list = metric_list / len(db_val) + for class_i in range(num_classes-1): + writer.add_scalar('info/model2_val_{}_dice'.format(class_i+1), + metric_list[class_i, 0], iter_num) + writer.add_scalar('info/model2_val_{}_hd95'.format(class_i+1), + metric_list[class_i, 1], iter_num) + + performance2 = np.mean(metric_list, axis=0)[0] + + mean_hd952 = np.mean(metric_list, axis=0)[1] + writer.add_scalar('info/model2_val_mean_dice', + performance2, iter_num) + writer.add_scalar('info/model2_val_mean_hd95', + mean_hd952, iter_num) + + if performance2 > best_performance2: + best_performance2 = performance2 + save_mode_path = os.path.join(snapshot_path, + 'model2_iter_{}_dice_{}.pth'.format( + iter_num, round(best_performance2, 4))) + save_best = os.path.join(snapshot_path, + '{}_best_model2.pth'.format(args.model)) + torch.save(model2.state_dict(), save_mode_path) + torch.save(model2.state_dict(), save_best) + + logging.info( + 'iteration %d : model2_mean_dice : %f model2_mean_hd95 : %f' % (iter_num, performance2, mean_hd952)) + model2.train() + + if iter_num % 3000 == 0: + save_mode_path = os.path.join( + snapshot_path, 'model1_iter_' + str(iter_num) + '.pth') + torch.save(model1.state_dict(), save_mode_path) + logging.info("save model1 to {}".format(save_mode_path)) + + save_mode_path = os.path.join( + snapshot_path, 'model2_iter_' + str(iter_num) + '.pth') + torch.save(model2.state_dict(), save_mode_path) + logging.info("save model2 to {}".format(save_mode_path)) + + if iter_num >= max_iterations: + break + if iter_num >= max_iterations: + iterator.close() + break + writer.close() + return "Training Finished!" + + +if __name__ == "__main__": + if not args.deterministic: + cudnn.benchmark = True + cudnn.deterministic = False + else: + cudnn.benchmark = False + cudnn.deterministic = True + + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( + args.exp, args.labeled_ratio, args.fold) + if not os.path.exists(snapshot_path): + os.makedirs(snapshot_path) + if os.path.exists(snapshot_path + '/code'): + shutil.rmtree(snapshot_path + '/code') + shutil.copytree('.', snapshot_path + '/code', + shutil.ignore_patterns(['.git', '__pycache__'])) + + logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO, + format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') + logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) + logging.info(str(args)) + train(args, snapshot_path) diff --git a/code/train_deep_adversarial_network.py b/code/train_deep_adversarial_network.py new file mode 100644 index 0000000..c99fd54 --- /dev/null +++ b/code/train_deep_adversarial_network.py @@ -0,0 +1,271 @@ +import argparse +import logging +import os +import random +import shutil +import sys +import time +from itertools import cycle +import numpy as np +import torch +import torch.backends.cudnn as cudnn +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from tensorboardX import SummaryWriter +from torch.nn import BCEWithLogitsLoss +from torch.nn.modules.loss import CrossEntropyLoss +from torch.utils.data import DataLoader +from torchvision import transforms +from torchvision.utils import make_grid +from tqdm import tqdm + +from dataloaders.dataset import BaseDataSets, RandomGenerator +from networks.discriminator import FCDiscriminator +from networks.net_factory import net_factory +from utils import losses, metrics, ramps +from val_2D import test_single_volume + +parser = argparse.ArgumentParser() +parser.add_argument('--root_path', type=str, + default='../data/ProstateX', help='Name of Experiment') +parser.add_argument('--exp', type=str, + default='ProstateX/DAN', help='experiment_name') +parser.add_argument('--model', type=str, + default='unet', help='model_name') +parser.add_argument('--fold', type=int, + default=3, help='cross validation') +parser.add_argument('--max_iterations', type=int, + default=30000, help='maximum epoch number to train') +parser.add_argument('--batch_size', type=int, default=16, + help='batch_size per gpu') + +parser.add_argument('--deterministic', type=int, default=1, + help='whether use deterministic training') +parser.add_argument('--base_lr', type=float, default=0.03, + help='segmentation network learning rate') +parser.add_argument('--patch_size', type=list, default=[256, 256], + help='patch size of network input') +parser.add_argument('--seed', type=int, default=2022, help='random seed') +parser.add_argument('--num_classes', type=int, default=3, + help='output channel of network') + +# label and unlabel +parser.add_argument('--labeled_ratio', type=int, default=8, + help='1/labeled_ratio data is provided mask') +# costs +parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') +parser.add_argument('--consistency_type', type=str, + default="mse", help='consistency_type') +parser.add_argument('--consistency', type=float, + default=0.1, help='consistency') +parser.add_argument('--consistency_rampup', type=float, + default=200.0, help='consistency_rampup') +args = parser.parse_args() + + +def get_current_consistency_weight(epoch): + # Consistency ramp-up from https://arxiv.org/abs/1610.02242 + return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) + + +def train(args, snapshot_path): + writer = SummaryWriter(snapshot_path + '/log') + base_lr = args.base_lr + num_classes = args.num_classes + max_iterations = args.max_iterations + + model = net_factory(net_type=args.model, in_chns=1, + class_num=num_classes) + DAN = FCDiscriminator(num_classes=num_classes) + DAN = DAN.cuda() + + db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)])) + db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)])) + logging.info("Labeled slices: {} ".format(len(db_train_labeled))) + logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) + + trainloader_labeled = DataLoader( + db_train_labeled, batch_size=args.batch_size//2, shuffle=True) + trainloader_unlabeled = DataLoader( + db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) + + db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, + split="val", labeled_ratio=args.labeled_ratio) + valloader = DataLoader(db_val, batch_size=1) + + model.train() + + optimizer = optim.SGD(model.parameters(), lr=base_lr, + momentum=0.9, weight_decay=0.0001) + + DAN_optimizer = optim.Adam( + DAN.parameters(), lr=1e-4, betas=(0.9, 0.99)) + + ce_loss = CrossEntropyLoss() + dice_loss = losses.DiceLoss(num_classes) + + logging.info("{} iterations per epoch".format(len(trainloader_labeled))) + + iter_num = 0 + max_epoch = max_iterations // len(trainloader_unlabeled) + 1 + best_performance = 0.0 + iterator = tqdm(range(max_epoch), ncols=70) + for epoch_num in iterator: + for i, (sampled_batch_labeled, sampled_batch_unlabeled) in enumerate(zip(cycle(trainloader_labeled), trainloader_unlabeled)): + volume_batch, label_batch = sampled_batch_labeled['image'], sampled_batch_labeled['label'] + volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() + unlabeled_volume_batch = sampled_batch_unlabeled['image'].cuda() + + DAN_target = torch.tensor( + [0] * (unlabeled_volume_batch.shape[0] + volume_batch.shape[0])).cuda() + DAN_target[:volume_batch.shape[0]] = 1 + model.train() + DAN.eval() + + outputs = model(volume_batch) + outputs_soft = torch.softmax(outputs, dim=1) + + outputs_unlabeled = model(unlabeled_volume_batch) + outputs_unlabeled_soft = torch.softmax(outputs_unlabeled, dim=1) + + supervised_loss = 0.5 * \ + (ce_loss(outputs, label_batch[:].long( + )) + dice_loss(outputs_soft, label_batch[:].unsqueeze(1))) + consistency_weight = get_current_consistency_weight( + iter_num // (args.max_iterations/args.consistency_rampup)) + + DAN_outputs = DAN(outputs_unlabeled_soft, unlabeled_volume_batch) + + DAN_target_unlabeled = torch.tensor( + [1] * unlabeled_volume_batch.shape[0]).cuda() + + consistency_loss = F.cross_entropy( + DAN_outputs, DAN_target_unlabeled.long()) + loss = supervised_loss + consistency_weight * consistency_loss + optimizer.zero_grad() + loss.backward() + optimizer.step() + + model.eval() + DAN.train() + with torch.no_grad(): + outputs = model(volume_batch) + outputs_soft = torch.softmax(outputs, dim=1) + + outputs_unlabeled = model(unlabeled_volume_batch) + outputs_unlabeled_soft = torch.softmax( + outputs_unlabeled, dim=1) + + DAN_outputs = DAN(torch.cat([outputs_soft, outputs_unlabeled_soft], dim=0), torch.cat( + [volume_batch, unlabeled_volume_batch], dim=0)) + DAN_loss = F.cross_entropy(DAN_outputs, DAN_target.long()) + DAN_optimizer.zero_grad() + DAN_loss.backward() + DAN_optimizer.step() + + lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 + for param_group in optimizer.param_groups: + param_group['lr'] = lr_ + + iter_num = iter_num + 1 + writer.add_scalar('info/lr', lr_, iter_num) + writer.add_scalar('info/total_loss', loss, iter_num) + writer.add_scalar('info/loss_ce', supervised_loss, iter_num) + writer.add_scalar('info/consistency_loss', + consistency_loss, iter_num) + writer.add_scalar('info/consistency_weight', + consistency_weight, iter_num) + + logging.info( + 'iteration %d : loss : %f, loss_ce: %f' % + (iter_num, loss.item(), supervised_loss.item())) + + if iter_num % 20 == 0: + image = volume_batch[0, 0:1, :, :] + writer.add_image('train/Image', image, iter_num) + outputs = torch.argmax(torch.softmax( + outputs, dim=1), dim=1, keepdim=True) + writer.add_image('train/Prediction', + outputs[0, ...] * 50, iter_num) + labs = label_batch[0, ...].unsqueeze(0) * 50 + writer.add_image('train/GroundTruth', labs, iter_num) + + if iter_num > 0 and iter_num % 200 == 0: + model.eval() + metric_list = 0.0 + for i_batch, sampled_batch in enumerate(valloader): + metric_i = test_single_volume( + sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) + metric_list += np.array(metric_i) + metric_list = metric_list / len(db_val) + for class_i in range(num_classes-1): + writer.add_scalar('info/val_{}_dice'.format(class_i+1), + metric_list[class_i, 0], iter_num) + writer.add_scalar('info/val_{}_hd95'.format(class_i+1), + metric_list[class_i, 1], iter_num) + + performance = np.mean(metric_list, axis=0)[0] + + mean_hd95 = np.mean(metric_list, axis=0)[1] + writer.add_scalar('info/val_mean_dice', performance, iter_num) + writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) + + if performance > best_performance: + best_performance = performance + save_mode_path = os.path.join(snapshot_path, + 'iter_{}_dice_{}.pth'.format( + iter_num, round(best_performance, 4))) + save_best = os.path.join(snapshot_path, + '{}_best_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_mode_path) + torch.save(model.state_dict(), save_best) + + logging.info( + 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) + model.train() + + if iter_num % 3000 == 0: + save_mode_path = os.path.join( + snapshot_path, 'iter_' + str(iter_num) + '.pth') + torch.save(model.state_dict(), save_mode_path) + logging.info("save model to {}".format(save_mode_path)) + + if iter_num >= max_iterations: + break + if iter_num >= max_iterations: + iterator.close() + break + writer.close() + return "Training Finished!" + + +if __name__ == "__main__": + if not args.deterministic: + cudnn.benchmark = True + cudnn.deterministic = False + else: + cudnn.benchmark = False + cudnn.deterministic = True + + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( + args.exp, args.labeled_ratio, args.fold) + if not os.path.exists(snapshot_path): + os.makedirs(snapshot_path) + if os.path.exists(snapshot_path + '/code'): + shutil.rmtree(snapshot_path + '/code') + shutil.copytree('.', snapshot_path + '/code', + shutil.ignore_patterns(['.git', '__pycache__'])) + + logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO, + format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') + logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) + logging.info(str(args)) + train(args, snapshot_path) diff --git a/code/train_entropy_minimization.py b/code/train_entropy_minimization.py new file mode 100644 index 0000000..a1eb291 --- /dev/null +++ b/code/train_entropy_minimization.py @@ -0,0 +1,238 @@ +import argparse +import logging +import os +import random +import shutil +import sys +import time +from itertools import cycle +import numpy as np +import torch +import torch.backends.cudnn as cudnn +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from tensorboardX import SummaryWriter +from torch.nn import BCEWithLogitsLoss +from torch.nn.modules.loss import CrossEntropyLoss +from torch.utils.data import DataLoader +from torchvision import transforms +from torchvision.utils import make_grid +from tqdm import tqdm + +from dataloaders.dataset import BaseDataSets, RandomGenerator +from networks.discriminator import FCDiscriminator +from networks.net_factory import net_factory +from utils import losses, metrics, ramps +from val_2D import test_single_volume + +parser = argparse.ArgumentParser() +parser.add_argument('--root_path', type=str, + default='../data/ProstateX', help='Name of Experiment') +parser.add_argument('--exp', type=str, + default='ProstateX/Mean_Teacher', help='experiment_name') +parser.add_argument('--model', type=str, + default='unet', help='model_name') +parser.add_argument('--fold', type=int, + default=3, help='cross validation') +parser.add_argument('--max_iterations', type=int, + default=30000, help='maximum epoch number to train') +parser.add_argument('--batch_size', type=int, default=16, + help='batch_size per gpu') + +parser.add_argument('--deterministic', type=int, default=1, + help='whether use deterministic training') +parser.add_argument('--base_lr', type=float, default=0.03, + help='segmentation network learning rate') +parser.add_argument('--patch_size', type=list, default=[256, 256], + help='patch size of network input') +parser.add_argument('--seed', type=int, default=2022, help='random seed') +parser.add_argument('--num_classes', type=int, default=3, + help='output channel of network') + +# label and unlabel +parser.add_argument('--labeled_ratio', type=int, default=8, + help='1/labeled_ratio data is provided mask') +# costs +parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') +parser.add_argument('--consistency_type', type=str, + default="mse", help='consistency_type') +parser.add_argument('--consistency', type=float, + default=0.1, help='consistency') +parser.add_argument('--consistency_rampup', type=float, + default=200.0, help='consistency_rampup') +args = parser.parse_args() + + +def get_current_consistency_weight(epoch): + # Consistency ramp-up from https://arxiv.org/abs/1610.02242 + return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) + + +def train(args, snapshot_path): + writer = SummaryWriter(snapshot_path + '/log') + base_lr = args.base_lr + num_classes = args.num_classes + max_iterations = args.max_iterations + + model = net_factory(net_type=args.model, in_chns=1, + class_num=num_classes) + + db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)])) + db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)])) + logging.info("Labeled slices: {} ".format(len(db_train_labeled))) + logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) + + trainloader_labeled = DataLoader( + db_train_labeled, batch_size=args.batch_size//2, shuffle=True) + trainloader_unlabeled = DataLoader( + db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) + + db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, + split="val", labeled_ratio=args.labeled_ratio) + valloader = DataLoader(db_val, batch_size=1) + + model.train() + + optimizer = optim.SGD(model.parameters(), lr=base_lr, + momentum=0.9, weight_decay=0.0001) + + ce_loss = CrossEntropyLoss() + dice_loss = losses.DiceLoss(num_classes) + + logging.info("{} iterations per epoch".format(len(trainloader_labeled))) + + iter_num = 0 + max_epoch = max_iterations // len(trainloader_unlabeled) + 1 + best_performance = 0.0 + iterator = tqdm(range(max_epoch), ncols=70) + for epoch_num in iterator: + for i, (sampled_batch_labeled, sampled_batch_unlabeled) in enumerate(zip(cycle(trainloader_labeled), trainloader_unlabeled)): + volume_batch, label_batch = sampled_batch_labeled['image'], sampled_batch_labeled['label'] + volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() + unlabeled_volume_batch = sampled_batch_unlabeled['image'].cuda() + + outputs = model(volume_batch) + outputs_soft = torch.softmax(outputs, dim=1) + + outputs_unlabeled = model(unlabeled_volume_batch) + outputs_unlabeled_soft = torch.softmax(outputs_unlabeled, dim=1) + + supervised_loss = 0.5 * \ + (ce_loss(outputs, label_batch[:].long( + )) + dice_loss(outputs_soft, label_batch[:].unsqueeze(1))) + consistency_weight = get_current_consistency_weight( + iter_num // (args.max_iterations/args.consistency_rampup)) + + ent_loss = losses.entropy_loss( + outputs_unlabeled_soft, C=args.num_classes) + loss = supervised_loss + consistency_weight * ent_loss + optimizer.zero_grad() + loss.backward() + optimizer.step() + + lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 + for param_group in optimizer.param_groups: + param_group['lr'] = lr_ + + iter_num = iter_num + 1 + writer.add_scalar('info/lr', lr_, iter_num) + writer.add_scalar('info/total_loss', loss, iter_num) + writer.add_scalar('info/loss_ce', supervised_loss, iter_num) + writer.add_scalar('info/consistency_loss', + ent_loss, iter_num) + writer.add_scalar('info/consistency_weight', + consistency_weight, iter_num) + + logging.info( + 'iteration %d : loss : %f, loss_ce: %f' % + (iter_num, loss.item(), supervised_loss.item())) + + if iter_num % 20 == 0: + image = volume_batch[0, 0:1, :, :] + writer.add_image('train/Image', image, iter_num) + outputs = torch.argmax(torch.softmax( + outputs, dim=1), dim=1, keepdim=True) + writer.add_image('train/Prediction', + outputs[0, ...] * 50, iter_num) + labs = label_batch[0, ...].unsqueeze(0) * 50 + writer.add_image('train/GroundTruth', labs, iter_num) + + if iter_num > 0 and iter_num % 200 == 0: + model.eval() + metric_list = 0.0 + for i_batch, sampled_batch in enumerate(valloader): + metric_i = test_single_volume( + sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) + metric_list += np.array(metric_i) + metric_list = metric_list / len(db_val) + for class_i in range(num_classes-1): + writer.add_scalar('info/val_{}_dice'.format(class_i+1), + metric_list[class_i, 0], iter_num) + writer.add_scalar('info/val_{}_hd95'.format(class_i+1), + metric_list[class_i, 1], iter_num) + + performance = np.mean(metric_list, axis=0)[0] + + mean_hd95 = np.mean(metric_list, axis=0)[1] + writer.add_scalar('info/val_mean_dice', performance, iter_num) + writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) + + if performance > best_performance: + best_performance = performance + save_mode_path = os.path.join(snapshot_path, + 'iter_{}_dice_{}.pth'.format( + iter_num, round(best_performance, 4))) + save_best = os.path.join(snapshot_path, + '{}_best_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_mode_path) + torch.save(model.state_dict(), save_best) + + logging.info( + 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) + model.train() + + if iter_num % 3000 == 0: + save_mode_path = os.path.join( + snapshot_path, 'iter_' + str(iter_num) + '.pth') + torch.save(model.state_dict(), save_mode_path) + logging.info("save model to {}".format(save_mode_path)) + + if iter_num >= max_iterations: + break + if iter_num >= max_iterations: + iterator.close() + break + writer.close() + return "Training Finished!" + + +if __name__ == "__main__": + if not args.deterministic: + cudnn.benchmark = True + cudnn.deterministic = False + else: + cudnn.benchmark = False + cudnn.deterministic = True + + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( + args.exp, args.labeled_ratio, args.fold) + if not os.path.exists(snapshot_path): + os.makedirs(snapshot_path) + if os.path.exists(snapshot_path + '/code'): + shutil.rmtree(snapshot_path + '/code') + shutil.copytree('.', snapshot_path + '/code', + shutil.ignore_patterns(['.git', '__pycache__'])) + + logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO, + format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') + logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) + logging.info(str(args)) + train(args, snapshot_path) diff --git a/code/train_fullsup_exp_fold0.sh b/code/train_fullsup_exp_fold0.sh new file mode 100644 index 0000000..a40fff5 --- /dev/null +++ b/code/train_fullsup_exp_fold0.sh @@ -0,0 +1,75 @@ +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ProstateX --exp ProstateX/FullSup --batch_size 16 --num_classes 3 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 1 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 2 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 4 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 1 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 2 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 3 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/CHAOS --exp CHAOS/FullSup --batch_size 16 --num_classes 5 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 0 \ No newline at end of file diff --git a/code/train_fully_supervised.py b/code/train_fully_supervised.py new file mode 100644 index 0000000..802e7c5 --- /dev/null +++ b/code/train_fully_supervised.py @@ -0,0 +1,219 @@ +import argparse +import logging +import os +import random +import shutil +import sys +import time +from itertools import cycle + +import numpy as np +import torch +import torch.backends.cudnn as cudnn +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from tensorboardX import SummaryWriter +from torch.nn import BCEWithLogitsLoss +from torch.nn.modules.loss import CrossEntropyLoss +from torch.utils.data import DataLoader +from torchvision import transforms +from torchvision.utils import make_grid +from tqdm import tqdm + +from dataloaders.dataset import BaseDataSets, RandomGenerator +from networks.discriminator import FCDiscriminator +from networks.net_factory import net_factory +from utils import losses, metrics, ramps +from val_2D import test_single_volume + +parser = argparse.ArgumentParser() +parser.add_argument('--root_path', type=str, + default='../data/ACDC', help='Name of Experiment') +parser.add_argument('--exp', type=str, + default='ACDC/FullSup', help='experiment_name') +parser.add_argument('--model', type=str, + default='unet', help='model_name') +parser.add_argument('--fold', type=int, + default=5, help='cross validation') +parser.add_argument('--max_iterations', type=int, + default=30000, help='maximum epoch number to train') +parser.add_argument('--batch_size', type=int, default=12, + help='batch_size per gpu') + +parser.add_argument('--deterministic', type=int, default=1, + help='whether use deterministic training') +parser.add_argument('--base_lr', type=float, default=0.03, + help='segmentation network learning rate') +parser.add_argument('--patch_size', type=list, default=[256, 256], + help='patch size of network input') +parser.add_argument('--seed', type=int, default=2022, help='random seed') +parser.add_argument('--num_classes', type=int, default=4, + help='output channel of network') + +# label and unlabel +parser.add_argument('--labeled_ratio', type=int, default=8, + help='1/labeled_ratio data is provided mask') +# costs +parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') +parser.add_argument('--consistency_type', type=str, + default="mse", help='consistency_type') +parser.add_argument('--consistency', type=float, + default=0.1, help='consistency') +parser.add_argument('--consistency_rampup', type=float, + default=200.0, help='consistency_rampup') +args = parser.parse_args() + + +def get_current_consistency_weight(epoch): + # Consistency ramp-up from https://arxiv.org/abs/1610.02242 + return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) + + +def train(args, snapshot_path): + writer = SummaryWriter(snapshot_path + '/log') + base_lr = args.base_lr + num_classes = args.num_classes + max_iterations = args.max_iterations + + model = net_factory(net_type=args.model, in_chns=1, + class_num=num_classes) + db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)])) + + logging.info("Labeled slices: {} ".format(len(db_train_labeled))) + + trainloader_labeled = DataLoader( + db_train_labeled, batch_size=args.batch_size, shuffle=True) + + db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, + split="val", labeled_ratio=args.labeled_ratio) + valloader = DataLoader(db_val, batch_size=1) + + model.train() + + optimizer = optim.SGD(model.parameters(), lr=base_lr, + momentum=0.9, weight_decay=0.0001) + + ce_loss = CrossEntropyLoss() + dice_loss = losses.DiceLoss(num_classes) + + logging.info("{} iterations per epoch".format(len(trainloader_labeled))) + + iter_num = 0 + max_epoch = max_iterations // len(trainloader_labeled) + 1 + best_performance = 0.0 + iterator = tqdm(range(max_epoch), ncols=70) + for epoch_num in iterator: + for i, sampled_batch_labeled in enumerate(trainloader_labeled): + volume_batch, label_batch = sampled_batch_labeled['image'], sampled_batch_labeled['label'] + volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() + + outputs = model(volume_batch) + outputs_soft = torch.softmax(outputs, dim=1) + + loss = 0.5 * \ + (ce_loss(outputs, label_batch[:].long( + )) + dice_loss(outputs_soft, label_batch[:].unsqueeze(1))) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 + for param_group in optimizer.param_groups: + param_group['lr'] = lr_ + + iter_num = iter_num + 1 + writer.add_scalar('info/lr', lr_, iter_num) + writer.add_scalar('info/total_loss', loss, iter_num) + logging.info( + 'iteration %d : loss : %f, loss_ce: %f' % + (iter_num, loss.item(), loss.item())) + + if iter_num % 20 == 0: + image = volume_batch[0, 0:1, :, :] + writer.add_image('train/Image', image, iter_num) + outputs = torch.argmax(torch.softmax( + outputs, dim=1), dim=1, keepdim=True) + writer.add_image('train/Prediction', + outputs[0, ...] * 50, iter_num) + labs = label_batch[0, ...].unsqueeze(0) * 50 + writer.add_image('train/GroundTruth', labs, iter_num) + + if iter_num > 0 and iter_num % 200 == 0: + model.eval() + metric_list = 0.0 + for i_batch, sampled_batch in enumerate(valloader): + metric_i = test_single_volume( + sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) + metric_list += np.array(metric_i) + metric_list = metric_list / len(db_val) + for class_i in range(num_classes-1): + writer.add_scalar('info/val_{}_dice'.format(class_i+1), + metric_list[class_i, 0], iter_num) + writer.add_scalar('info/val_{}_hd95'.format(class_i+1), + metric_list[class_i, 1], iter_num) + + performance = np.mean(metric_list, axis=0)[0] + + mean_hd95 = np.mean(metric_list, axis=0)[1] + writer.add_scalar('info/val_mean_dice', performance, iter_num) + writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) + + if performance > best_performance: + best_performance = performance + save_mode_path = os.path.join(snapshot_path, + 'iter_{}_dice_{}.pth'.format( + iter_num, round(best_performance, 4))) + save_best = os.path.join(snapshot_path, + '{}_best_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_mode_path) + torch.save(model.state_dict(), save_best) + + logging.info( + 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) + model.train() + + if iter_num % 3000 == 0: + save_mode_path = os.path.join( + snapshot_path, 'iter_' + str(iter_num) + '.pth') + torch.save(model.state_dict(), save_mode_path) + logging.info("save model to {}".format(save_mode_path)) + + if iter_num >= max_iterations: + break + if iter_num >= max_iterations: + iterator.close() + break + writer.close() + return "Training Finished!" + + +if __name__ == "__main__": + if not args.deterministic: + cudnn.benchmark = True + cudnn.deterministic = False + else: + cudnn.benchmark = False + cudnn.deterministic = True + + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( + args.exp, args.labeled_ratio, args.fold) + if not os.path.exists(snapshot_path): + os.makedirs(snapshot_path) + if os.path.exists(snapshot_path + '/code'): + shutil.rmtree(snapshot_path + '/code') + shutil.copytree('.', snapshot_path + '/code', + shutil.ignore_patterns(['.git', '__pycache__'])) + + logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO, + format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') + logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) + logging.info(str(args)) + train(args, snapshot_path) diff --git a/code/train_interpolation_consistency_training.py b/code/train_interpolation_consistency_training.py index 8f98d70..114eb77 100644 --- a/code/train_interpolation_consistency_training.py +++ b/code/train_interpolation_consistency_training.py @@ -42,7 +42,7 @@ parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, +parser.add_argument('--base_lr', type=float, default=0.03, help='segmentation network learning rate') parser.add_argument('--patch_size', type=list, default=[256, 256], help='patch size of network input') @@ -141,7 +141,7 @@ def create_model(ema=False): supervised_loss = 0.5 * \ (ce_loss(outputs, label_batch[:].long( )) + dice_loss(outputs_soft, label_batch[:].unsqueeze(1))) - + if unlabeled_volume_batch.shape[0] != args.batch_size // 2: loss = supervised_loss consistency_weight = 0.0 @@ -164,7 +164,8 @@ def create_model(ema=False): # [volume_batch, batch_ux_mixed], dim=0) outputs_unlabeled = model(batch_ux_mixed) - outputs_unlabeled_soft = torch.softmax(outputs_unlabeled, dim=1) + outputs_unlabeled_soft = torch.softmax( + outputs_unlabeled, dim=1) with torch.no_grad(): ema_output_ux0 = torch.softmax( @@ -172,10 +173,11 @@ def create_model(ema=False): ema_output_ux1 = torch.softmax( ema_model(unlabeled_volume_batch_1), dim=1) batch_pred_mixed = ema_output_ux0 * \ - (1.0 - ict_mix_factors) + ema_output_ux1 * ict_mix_factors + (1.0 - ict_mix_factors) + \ + ema_output_ux1 * ict_mix_factors consistency_weight = get_current_consistency_weight( - iter_num // 150) + iter_num // (args.max_iterations/args.consistency_rampup)) consistency_loss = torch.mean( (outputs_unlabeled_soft - batch_pred_mixed) ** 2) loss = supervised_loss + consistency_weight * consistency_loss diff --git a/code/train_mean_teacher.py b/code/train_mean_teacher.py index eed285b..d968714 100644 --- a/code/train_mean_teacher.py +++ b/code/train_mean_teacher.py @@ -42,7 +42,7 @@ parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, +parser.add_argument('--base_lr', type=float, default=0.03, help='segmentation network learning rate') parser.add_argument('--patch_size', type=list, default=[256, 256], help='patch size of network input') @@ -68,6 +68,7 @@ def get_current_consistency_weight(epoch): # Consistency ramp-up from https://arxiv.org/abs/1610.02242 return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) + def update_ema_variables(model, ema_model, alpha, global_step): # Use the true average until the exponential average is more correct alpha = min(1 - 1 / (global_step + 1), alpha) @@ -103,10 +104,13 @@ def create_model(ema=False): logging.info("Labeled slices: {} ".format(len(db_train_labeled))) logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) - trainloader_labeled = DataLoader(db_train_labeled, batch_size=args.batch_size//2, shuffle=True) - trainloader_unlabeled = DataLoader(db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) + trainloader_labeled = DataLoader( + db_train_labeled, batch_size=args.batch_size//2, shuffle=True) + trainloader_unlabeled = DataLoader( + db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) - db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, split="val", labeled_ratio=args.labeled_ratio) + db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, + split="val", labeled_ratio=args.labeled_ratio) valloader = DataLoader(db_val, batch_size=1) model.train() @@ -128,7 +132,7 @@ def create_model(ema=False): volume_batch, label_batch = sampled_batch_labeled['image'], sampled_batch_labeled['label'] volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() unlabeled_volume_batch = sampled_batch_unlabeled['image'].cuda() - + noise = torch.clamp(torch.randn_like( unlabeled_volume_batch) * 0.1, -0.2, 0.2) ema_inputs = unlabeled_volume_batch + noise @@ -143,10 +147,14 @@ def create_model(ema=False): ema_output = ema_model(ema_inputs) ema_output_soft = torch.softmax(ema_output, dim=1) - supervised_loss = 0.5*(ce_loss(outputs, label_batch[:].long()) + dice_loss(outputs_soft, label_batch[:].unsqueeze(1))) - consistency_weight = get_current_consistency_weight(iter_num // 150) + supervised_loss = 0.5 * \ + (ce_loss(outputs, label_batch[:].long( + )) + dice_loss(outputs_soft, label_batch[:].unsqueeze(1))) + consistency_weight = get_current_consistency_weight( + iter_num // (args.max_iterations/args.consistency_rampup)) - consistency_loss = torch.mean((outputs_unlabeled_soft - ema_output_soft) ** 2) + consistency_loss = torch.mean( + (outputs_unlabeled_soft - ema_output_soft) ** 2) loss = supervised_loss + consistency_weight * consistency_loss optimizer.zero_grad() loss.backward() diff --git a/code/train_uncertainty_aware_mean_teacher.py b/code/train_uncertainty_aware_mean_teacher.py index caac5fb..febf888 100644 --- a/code/train_uncertainty_aware_mean_teacher.py +++ b/code/train_uncertainty_aware_mean_teacher.py @@ -41,7 +41,7 @@ parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training') -parser.add_argument('--base_lr', type=float, default=0.01, +parser.add_argument('--base_lr', type=float, default=0.03, help='segmentation network learning rate') parser.add_argument('--patch_size', type=list, default=[256, 256], help='patch size of network input') @@ -67,6 +67,7 @@ def get_current_consistency_weight(epoch): # Consistency ramp-up from https://arxiv.org/abs/1610.02242 return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) + def update_ema_variables(model, ema_model, alpha, global_step): # Use the true average until the exponential average is more correct alpha = min(1 - 1 / (global_step + 1), alpha) @@ -100,13 +101,16 @@ def create_model(ema=False): db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ RandomGenerator(args.patch_size)])) - trainloader_labeled = DataLoader(db_train_labeled, batch_size=args.batch_size//2, shuffle=True) - trainloader_unlabeled = DataLoader(db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) + trainloader_labeled = DataLoader( + db_train_labeled, batch_size=args.batch_size//2, shuffle=True) + trainloader_unlabeled = DataLoader( + db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) logging.info("Labeled slices: {} ".format(len(db_train_labeled))) logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) - db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, split="val", labeled_ratio=args.labeled_ratio) + db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, + split="val", labeled_ratio=args.labeled_ratio) valloader = DataLoader(db_val, batch_size=1, shuffle=False) model.train() @@ -147,22 +151,24 @@ def create_model(ema=False): preds = torch.zeros([stride * T, num_classes, w, h]).cuda() for i in range(T // 2): ema_inputs = volume_batch_r + \ - torch.clamp(torch.randn_like( - volume_batch_r) * 0.1, -0.2, 0.2) + torch.clamp(torch.randn_like( + volume_batch_r) * 0.1, -0.2, 0.2) with torch.no_grad(): preds[2 * stride * i:2 * stride * - (i + 1)] = ema_model(ema_inputs) + (i + 1)] = ema_model(ema_inputs) preds = F.softmax(preds, dim=1) preds = preds.reshape(T, stride, num_classes, w, h) preds = torch.mean(preds, dim=0) uncertainty = -1.0 * \ - torch.sum(preds * torch.log(preds + 1e-6), dim=1, keepdim=True) + torch.sum(preds * torch.log(preds + 1e-6), dim=1, keepdim=True) loss_ce = ce_loss(outputs, label_batch[:].long()) loss_dice = dice_loss(outputs_soft, label_batch.unsqueeze(1)) supervised_loss = 0.5 * (loss_dice + loss_ce) - consistency_weight = get_current_consistency_weight(iter_num // 150) - consistency_dist = losses.softmax_mse_loss(outputs_unlabeled, ema_output) # (batch, 2, 112,112,80) + consistency_weight = get_current_consistency_weight( + iter_num // (args.max_iterations/args.consistency_rampup)) + consistency_dist = losses.softmax_mse_loss( + outputs_unlabeled, ema_output) # (batch, 2, 112,112,80) threshold = (0.75 + 0.25 * ramps.sigmoid_rampup(iter_num, max_iterations)) * np.log(2) mask = (uncertainty < threshold).float() diff --git a/code/val_2D.py b/code/val_2D.py index 3c5b1c6..89a9a30 100755 --- a/code/val_2D.py +++ b/code/val_2D.py @@ -16,26 +16,41 @@ def calculate_metric_percase(pred, gt): return 0, 50, 10 -def test_single_volume(image, label, net, classes, patch_size=[256, 256]): +def test_single_volume(image, label, net, classes, patch_size=[256, 256], batch_size=8): image, label = image.squeeze(0).cpu().detach( ).numpy(), label.squeeze(0).cpu().detach().numpy() if len(image.shape) == 3: prediction = np.zeros_like(label) - for ind in range(image.shape[0]): - slice = image[ind, :, :] - x, y = slice.shape[0], slice.shape[1] - slice = zoom( - slice, (patch_size[0] / x, patch_size[1] / y), order=0) - input = torch.from_numpy(slice).unsqueeze( - 0).unsqueeze(0).float().cuda() - net.eval() - with torch.no_grad(): - out = torch.argmax(torch.softmax( - net(input), dim=1), dim=1).squeeze(0) - out = out.cpu().detach().numpy() - pred = zoom( - out, (x / patch_size[0], y / patch_size[1]), order=0) - prediction[ind] = pred + ind_x = np.array([i for i in range(image.shape[0])]) + for ind in ind_x[::batch_size]: + if ind + batch_size < image.shape[0]: + slice = image[ind:ind + batch_size, ...] + thickness, x, y = slice.shape[0], slice.shape[1], slice.shape[2] + slice = zoom( + slice, (1, patch_size[0] / x, patch_size[1] / y), order=0) + input = torch.from_numpy(slice).unsqueeze(1).float().cuda() + net.eval() + with torch.no_grad(): + out = torch.argmax(torch.softmax( + net(input), dim=1), dim=1) + out = out.cpu().detach().numpy() + pred = zoom( + out, (1, x / patch_size[0], y / patch_size[1]), order=0) + prediction[ind:ind + batch_size, ...] = pred + else: + slice = image[ind:, ...] + thickness, x, y = slice.shape[0], slice.shape[1], slice.shape[2] + slice = zoom( + slice, (1, patch_size[0] / x, patch_size[1] / y), order=0) + input = torch.from_numpy(slice).unsqueeze(1).float().cuda() + net.eval() + with torch.no_grad(): + out = torch.argmax(torch.softmax( + net(input), dim=1), dim=1) + out = out.cpu().detach().numpy() + pred = zoom( + out, (1, x / patch_size[0], y / patch_size[1]), order=0) + prediction[ind:, ...] = pred else: input = torch.from_numpy(image).unsqueeze( 0).unsqueeze(0).float().cuda() @@ -48,10 +63,47 @@ def test_single_volume(image, label, net, classes, patch_size=[256, 256]): for i in range(1, classes): metric_list.append(calculate_metric_percase( prediction == i, label == i)) - - # return metric_list, image, prediction, label return metric_list + +# def test_single_volume(image, label, net, classes, patch_size=[256, 256]): +# image, label = image.squeeze(0).cpu().detach( +# ).numpy(), label.squeeze(0).cpu().detach().numpy() +# if len(image.shape) == 3: +# prediction = np.zeros_like(label) +# ind_x = np.array([i for i in range(image.shape[0])]) + +# for ind in range(image.shape[0]): +# slice = image[ind, :, :] +# x, y = slice.shape[0], slice.shape[1] +# slice = zoom( +# slice, (patch_size[0] / x, patch_size[1] / y), order=0) +# input = torch.from_numpy(slice).unsqueeze( +# 0).unsqueeze(0).float().cuda() +# net.eval() +# with torch.no_grad(): +# out = torch.argmax(torch.softmax( +# net(input), dim=1), dim=1).squeeze(0) +# out = out.cpu().detach().numpy() +# pred = zoom( +# out, (x / patch_size[0], y / patch_size[1]), order=0) +# prediction[ind] = pred +# else: +# input = torch.from_numpy(image).unsqueeze( +# 0).unsqueeze(0).float().cuda() +# net.eval() +# with torch.no_grad(): +# out = torch.argmax(torch.softmax( +# net(input), dim=1), dim=1).squeeze(0) +# prediction = out.cpu().detach().numpy() +# metric_list = [] +# for i in range(1, classes): +# metric_list.append(calculate_metric_percase( +# prediction == i, label == i)) + +# # return metric_list, image, prediction, label +# return metric_list + # # def test_single_volume_ds(image, label, net, classes, patch_size=[256, 256]): # image, label = image.squeeze(0).cpu().detach( @@ -90,6 +142,8 @@ def test_single_volume(image, label, net, classes, patch_size=[256, 256]): # return metric_list # # + + def test_single_volume_multitask(image, label, net, classes, patch_size=[256, 256]): image, label = image.squeeze(0).cpu().detach( ).numpy(), label.squeeze(0).cpu().detach().numpy() @@ -120,6 +174,7 @@ def test_single_volume_multitask(image, label, net, classes, patch_size=[256, 25 prediction = out.cpu().detach().numpy() metric_list = [] for i in range(1, classes): - metric_list.append(calculate_metric_percase(prediction == i, label == i)) + metric_list.append(calculate_metric_percase( + prediction == i, label == i)) # return metric_list, image, prediction, label return metric_list diff --git a/data/ACDC/README.md b/data/ACDC/README.md deleted file mode 100644 index 4ad05b5..0000000 --- a/data/ACDC/README.md +++ /dev/null @@ -1,2 +0,0 @@ -- Download the processed ACDC data from [BaiduDisk](https://pan.baidu.com/s/1d0cFhj3LU029oHajNni8KQ), the password is *code*, and decompress the zip file to [data/ACDC](https://github.com/Luoxd1996/SSL4MIS/edit/master/data/ACDC). More details of this dataset can be found at: https://www.creatis.insa-lyon.fr/Challenge/acdc/databases.html. -- If you want to use the [ACDC dataset](https://www.creatis.insa-lyon.fr/Challenge/acdc/databases.html) in your paper, please cite the original paper [TMI2018](https://ieeexplore.ieee.org/document/8360453). diff --git a/data/ACDC/test.list b/data/ACDC/test.list deleted file mode 100755 index 021039c..0000000 --- a/data/ACDC/test.list +++ /dev/null @@ -1,40 +0,0 @@ -patient011_frame01 -patient011_frame02 -patient013_frame01 -patient013_frame02 -patient084_frame01 -patient084_frame02 -patient033_frame01 -patient033_frame02 -patient093_frame01 -patient093_frame02 -patient022_frame01 -patient022_frame02 -patient068_frame01 -patient068_frame02 -patient024_frame01 -patient024_frame02 -patient083_frame01 -patient083_frame02 -patient081_frame01 -patient081_frame02 -patient080_frame01 -patient080_frame02 -patient001_frame01 -patient001_frame02 -patient007_frame01 -patient007_frame02 -patient066_frame01 -patient066_frame02 -patient008_frame01 -patient008_frame02 -patient065_frame01 -patient065_frame02 -patient075_frame01 -patient075_frame02 -patient064_frame01 -patient064_frame02 -patient059_frame01 -patient059_frame02 -patient052_frame01 -patient052_frame02 diff --git a/data/ACDC/train.list b/data/ACDC/train.list deleted file mode 100755 index 8dfe643..0000000 --- a/data/ACDC/train.list +++ /dev/null @@ -1,140 +0,0 @@ -patient099_frame01 -patient099_frame02 -patient038_frame01 -patient038_frame02 -patient050_frame01 -patient050_frame02 -patient100_frame01 -patient100_frame02 -patient058_frame01 -patient058_frame02 -patient021_frame01 -patient021_frame02 -patient049_frame01 -patient049_frame02 -patient020_frame01 -patient020_frame02 -patient072_frame01 -patient072_frame02 -patient040_frame01 -patient040_frame02 -patient060_frame01 -patient060_frame02 -patient089_frame01 -patient089_frame02 -patient004_frame01 -patient004_frame02 -patient056_frame01 -patient056_frame02 -patient098_frame01 -patient098_frame02 -patient096_frame01 -patient096_frame02 -patient031_frame01 -patient031_frame02 -patient018_frame01 -patient018_frame02 -patient094_frame01 -patient094_frame02 -patient047_frame01 -patient047_frame02 -patient048_frame01 -patient048_frame02 -patient055_frame01 -patient055_frame02 -patient097_frame01 -patient097_frame02 -patient074_frame01 -patient074_frame02 -patient043_frame01 -patient043_frame02 -patient041_frame01 -patient041_frame02 -patient063_frame01 -patient063_frame02 -patient037_frame01 -patient037_frame02 -patient095_frame01 -patient095_frame02 -patient054_frame01 -patient054_frame02 -patient026_frame01 -patient026_frame02 -patient088_frame01 -patient088_frame02 -patient032_frame01 -patient032_frame02 -patient069_frame01 -patient069_frame02 -patient006_frame01 -patient006_frame02 -patient071_frame01 -patient071_frame02 -patient012_frame01 -patient012_frame02 -patient073_frame01 -patient073_frame02 -patient061_frame01 -patient061_frame02 -patient017_frame01 -patient017_frame02 -patient025_frame01 -patient025_frame02 -patient010_frame01 -patient010_frame02 -patient057_frame01 -patient057_frame02 -patient029_frame01 -patient029_frame02 -patient051_frame01 -patient051_frame02 -patient005_frame01 -patient005_frame02 -patient036_frame01 -patient036_frame02 -patient046_frame01 -patient046_frame02 -patient062_frame01 -patient062_frame02 -patient034_frame01 -patient034_frame02 -patient076_frame01 -patient076_frame02 -patient092_frame01 -patient092_frame02 -patient070_frame01 -patient070_frame02 -patient077_frame01 -patient077_frame02 -patient067_frame01 -patient067_frame02 -patient003_frame01 -patient003_frame02 -patient091_frame01 -patient091_frame02 -patient016_frame01 -patient016_frame02 -patient014_frame01 -patient014_frame02 -patient044_frame01 -patient044_frame02 -patient042_frame01 -patient042_frame02 -patient090_frame01 -patient090_frame02 -patient053_frame01 -patient053_frame02 -patient027_frame01 -patient027_frame02 -patient035_frame01 -patient035_frame02 -patient086_frame01 -patient086_frame02 -patient023_frame01 -patient023_frame02 -patient009_frame01 -patient009_frame02 -patient079_frame01 -patient079_frame02 -patient015_frame01 -patient015_frame02 diff --git a/data/ACDC/train_slices.list b/data/ACDC/train_slices.list deleted file mode 100644 index 944bdcb..0000000 --- a/data/ACDC/train_slices.list +++ /dev/null @@ -1,1312 +0,0 @@ -patient099_frame01_slice_15 -patient099_frame01_slice_11 -patient099_frame01_slice_1 -patient099_frame01_slice_16 -patient099_frame01_slice_9 -patient099_frame01_slice_7 -patient099_frame01_slice_4 -patient099_frame01_slice_13 -patient099_frame01_slice_8 -patient099_frame01_slice_5 -patient099_frame01_slice_14 -patient099_frame01_slice_6 -patient099_frame01_slice_12 -patient099_frame01_slice_3 -patient099_frame01_slice_10 -patient099_frame01_slice_2 -patient099_frame02_slice_5 -patient099_frame02_slice_9 -patient099_frame02_slice_11 -patient099_frame02_slice_2 -patient099_frame02_slice_3 -patient099_frame02_slice_10 -patient099_frame02_slice_1 -patient099_frame02_slice_15 -patient099_frame02_slice_14 -patient099_frame02_slice_4 -patient099_frame02_slice_6 -patient099_frame02_slice_16 -patient099_frame02_slice_13 -patient099_frame02_slice_8 -patient099_frame02_slice_12 -patient099_frame02_slice_7 -patient038_frame01_slice_5 -patient038_frame01_slice_3 -patient038_frame01_slice_4 -patient038_frame01_slice_7 -patient038_frame01_slice_6 -patient038_frame01_slice_2 -patient038_frame01_slice_1 -patient038_frame01_slice_8 -patient038_frame02_slice_8 -patient038_frame02_slice_5 -patient038_frame02_slice_1 -patient038_frame02_slice_7 -patient038_frame02_slice_4 -patient038_frame02_slice_3 -patient038_frame02_slice_6 -patient038_frame02_slice_2 -patient050_frame01_slice_8 -patient050_frame01_slice_3 -patient050_frame01_slice_5 -patient050_frame01_slice_10 -patient050_frame01_slice_6 -patient050_frame01_slice_4 -patient050_frame01_slice_7 -patient050_frame01_slice_2 -patient050_frame01_slice_1 -patient050_frame01_slice_9 -patient050_frame02_slice_10 -patient050_frame02_slice_8 -patient050_frame02_slice_5 -patient050_frame02_slice_3 -patient050_frame02_slice_7 -patient050_frame02_slice_2 -patient050_frame02_slice_1 -patient050_frame02_slice_4 -patient050_frame02_slice_9 -patient050_frame02_slice_6 -patient100_frame01_slice_3 -patient100_frame01_slice_1 -patient100_frame01_slice_5 -patient100_frame01_slice_2 -patient100_frame01_slice_8 -patient100_frame01_slice_7 -patient100_frame01_slice_4 -patient100_frame01_slice_6 -patient100_frame02_slice_1 -patient100_frame02_slice_5 -patient100_frame02_slice_6 -patient100_frame02_slice_8 -patient100_frame02_slice_4 -patient100_frame02_slice_2 -patient100_frame02_slice_7 -patient100_frame02_slice_3 -patient058_frame01_slice_2 -patient058_frame01_slice_3 -patient058_frame01_slice_9 -patient058_frame01_slice_6 -patient058_frame01_slice_8 -patient058_frame01_slice_1 -patient058_frame01_slice_7 -patient058_frame01_slice_5 -patient058_frame01_slice_4 -patient058_frame02_slice_2 -patient058_frame02_slice_9 -patient058_frame02_slice_3 -patient058_frame02_slice_7 -patient058_frame02_slice_1 -patient058_frame02_slice_4 -patient058_frame02_slice_5 -patient058_frame02_slice_6 -patient058_frame02_slice_8 -patient021_frame01_slice_2 -patient021_frame01_slice_4 -patient021_frame01_slice_10 -patient021_frame01_slice_3 -patient021_frame01_slice_1 -patient021_frame01_slice_7 -patient021_frame01_slice_6 -patient021_frame01_slice_9 -patient021_frame01_slice_8 -patient021_frame01_slice_5 -patient021_frame02_slice_9 -patient021_frame02_slice_6 -patient021_frame02_slice_8 -patient021_frame02_slice_7 -patient021_frame02_slice_10 -patient021_frame02_slice_5 -patient021_frame02_slice_2 -patient021_frame02_slice_1 -patient021_frame02_slice_3 -patient021_frame02_slice_4 -patient049_frame01_slice_1 -patient049_frame01_slice_3 -patient049_frame01_slice_5 -patient049_frame01_slice_2 -patient049_frame01_slice_7 -patient049_frame01_slice_4 -patient049_frame01_slice_6 -patient049_frame02_slice_4 -patient049_frame02_slice_6 -patient049_frame02_slice_1 -patient049_frame02_slice_7 -patient049_frame02_slice_3 -patient049_frame02_slice_2 -patient049_frame02_slice_5 -patient020_frame01_slice_8 -patient020_frame01_slice_6 -patient020_frame01_slice_2 -patient020_frame01_slice_1 -patient020_frame01_slice_7 -patient020_frame01_slice_3 -patient020_frame01_slice_4 -patient020_frame01_slice_5 -patient020_frame02_slice_7 -patient020_frame02_slice_5 -patient020_frame02_slice_2 -patient020_frame02_slice_6 -patient020_frame02_slice_8 -patient020_frame02_slice_3 -patient020_frame02_slice_4 -patient020_frame02_slice_1 -patient072_frame01_slice_7 -patient072_frame01_slice_4 -patient072_frame01_slice_2 -patient072_frame01_slice_1 -patient072_frame01_slice_3 -patient072_frame01_slice_8 -patient072_frame01_slice_6 -patient072_frame01_slice_5 -patient072_frame02_slice_7 -patient072_frame02_slice_3 -patient072_frame02_slice_5 -patient072_frame02_slice_1 -patient072_frame02_slice_4 -patient072_frame02_slice_8 -patient072_frame02_slice_6 -patient072_frame02_slice_2 -patient040_frame01_slice_2 -patient040_frame01_slice_5 -patient040_frame01_slice_9 -patient040_frame01_slice_1 -patient040_frame01_slice_3 -patient040_frame01_slice_6 -patient040_frame01_slice_4 -patient040_frame01_slice_7 -patient040_frame01_slice_10 -patient040_frame01_slice_8 -patient040_frame02_slice_8 -patient040_frame02_slice_3 -patient040_frame02_slice_2 -patient040_frame02_slice_9 -patient040_frame02_slice_5 -patient040_frame02_slice_4 -patient040_frame02_slice_6 -patient040_frame02_slice_10 -patient040_frame02_slice_7 -patient040_frame02_slice_1 -patient060_frame01_slice_1 -patient060_frame01_slice_8 -patient060_frame01_slice_7 -patient060_frame01_slice_4 -patient060_frame01_slice_5 -patient060_frame01_slice_6 -patient060_frame01_slice_3 -patient060_frame01_slice_2 -patient060_frame01_slice_9 -patient060_frame02_slice_3 -patient060_frame02_slice_9 -patient060_frame02_slice_8 -patient060_frame02_slice_6 -patient060_frame02_slice_5 -patient060_frame02_slice_7 -patient060_frame02_slice_2 -patient060_frame02_slice_1 -patient060_frame02_slice_4 -patient089_frame01_slice_6 -patient089_frame01_slice_3 -patient089_frame01_slice_2 -patient089_frame01_slice_5 -patient089_frame01_slice_4 -patient089_frame01_slice_1 -patient089_frame02_slice_2 -patient089_frame02_slice_6 -patient089_frame02_slice_3 -patient089_frame02_slice_5 -patient089_frame02_slice_4 -patient089_frame02_slice_1 -patient004_frame01_slice_5 -patient004_frame01_slice_8 -patient004_frame01_slice_2 -patient004_frame01_slice_7 -patient004_frame01_slice_4 -patient004_frame01_slice_1 -patient004_frame01_slice_9 -patient004_frame01_slice_10 -patient004_frame01_slice_6 -patient004_frame01_slice_3 -patient004_frame02_slice_9 -patient004_frame02_slice_4 -patient004_frame02_slice_3 -patient004_frame02_slice_6 -patient004_frame02_slice_7 -patient004_frame02_slice_8 -patient004_frame02_slice_2 -patient004_frame02_slice_1 -patient004_frame02_slice_5 -patient004_frame02_slice_10 -patient056_frame01_slice_2 -patient056_frame01_slice_8 -patient056_frame01_slice_4 -patient056_frame01_slice_3 -patient056_frame01_slice_9 -patient056_frame01_slice_6 -patient056_frame01_slice_1 -patient056_frame01_slice_5 -patient056_frame01_slice_7 -patient056_frame02_slice_4 -patient056_frame02_slice_8 -patient056_frame02_slice_3 -patient056_frame02_slice_6 -patient056_frame02_slice_7 -patient056_frame02_slice_5 -patient056_frame02_slice_1 -patient056_frame02_slice_9 -patient056_frame02_slice_2 -patient098_frame01_slice_2 -patient098_frame01_slice_6 -patient098_frame01_slice_4 -patient098_frame01_slice_7 -patient098_frame01_slice_5 -patient098_frame01_slice_1 -patient098_frame01_slice_3 -patient098_frame02_slice_4 -patient098_frame02_slice_3 -patient098_frame02_slice_5 -patient098_frame02_slice_1 -patient098_frame02_slice_2 -patient098_frame02_slice_6 -patient098_frame02_slice_7 -patient096_frame01_slice_15 -patient096_frame01_slice_12 -patient096_frame01_slice_4 -patient096_frame01_slice_3 -patient096_frame01_slice_9 -patient096_frame01_slice_2 -patient096_frame01_slice_16 -patient096_frame01_slice_17 -patient096_frame01_slice_8 -patient096_frame01_slice_18 -patient096_frame01_slice_11 -patient096_frame01_slice_7 -patient096_frame01_slice_1 -patient096_frame01_slice_10 -patient096_frame01_slice_6 -patient096_frame01_slice_14 -patient096_frame01_slice_5 -patient096_frame01_slice_13 -patient096_frame02_slice_17 -patient096_frame02_slice_8 -patient096_frame02_slice_4 -patient096_frame02_slice_15 -patient096_frame02_slice_14 -patient096_frame02_slice_12 -patient096_frame02_slice_1 -patient096_frame02_slice_16 -patient096_frame02_slice_3 -patient096_frame02_slice_6 -patient096_frame02_slice_2 -patient096_frame02_slice_11 -patient096_frame02_slice_13 -patient096_frame02_slice_10 -patient096_frame02_slice_7 -patient096_frame02_slice_9 -patient096_frame02_slice_18 -patient096_frame02_slice_5 -patient031_frame01_slice_3 -patient031_frame01_slice_2 -patient031_frame01_slice_6 -patient031_frame01_slice_9 -patient031_frame01_slice_1 -patient031_frame01_slice_5 -patient031_frame01_slice_8 -patient031_frame01_slice_10 -patient031_frame01_slice_4 -patient031_frame01_slice_7 -patient031_frame02_slice_10 -patient031_frame02_slice_5 -patient031_frame02_slice_9 -patient031_frame02_slice_4 -patient031_frame02_slice_7 -patient031_frame02_slice_2 -patient031_frame02_slice_1 -patient031_frame02_slice_6 -patient031_frame02_slice_8 -patient031_frame02_slice_3 -patient018_frame01_slice_5 -patient018_frame01_slice_1 -patient018_frame01_slice_6 -patient018_frame01_slice_4 -patient018_frame01_slice_2 -patient018_frame01_slice_3 -patient018_frame01_slice_7 -patient018_frame01_slice_8 -patient018_frame02_slice_5 -patient018_frame02_slice_3 -patient018_frame02_slice_1 -patient018_frame02_slice_6 -patient018_frame02_slice_7 -patient018_frame02_slice_2 -patient018_frame02_slice_8 -patient018_frame02_slice_4 -patient094_frame01_slice_10 -patient094_frame01_slice_9 -patient094_frame01_slice_3 -patient094_frame01_slice_5 -patient094_frame01_slice_1 -patient094_frame01_slice_7 -patient094_frame01_slice_4 -patient094_frame01_slice_8 -patient094_frame01_slice_2 -patient094_frame01_slice_6 -patient094_frame02_slice_8 -patient094_frame02_slice_1 -patient094_frame02_slice_3 -patient094_frame02_slice_10 -patient094_frame02_slice_2 -patient094_frame02_slice_6 -patient094_frame02_slice_5 -patient094_frame02_slice_9 -patient094_frame02_slice_7 -patient094_frame02_slice_4 -patient047_frame01_slice_3 -patient047_frame01_slice_4 -patient047_frame01_slice_9 -patient047_frame01_slice_8 -patient047_frame01_slice_2 -patient047_frame01_slice_7 -patient047_frame01_slice_5 -patient047_frame01_slice_1 -patient047_frame01_slice_6 -patient047_frame02_slice_8 -patient047_frame02_slice_3 -patient047_frame02_slice_7 -patient047_frame02_slice_2 -patient047_frame02_slice_6 -patient047_frame02_slice_9 -patient047_frame02_slice_5 -patient047_frame02_slice_1 -patient047_frame02_slice_4 -patient048_frame01_slice_2 -patient048_frame01_slice_8 -patient048_frame01_slice_3 -patient048_frame01_slice_5 -patient048_frame01_slice_6 -patient048_frame01_slice_4 -patient048_frame01_slice_7 -patient048_frame01_slice_1 -patient048_frame02_slice_5 -patient048_frame02_slice_3 -patient048_frame02_slice_2 -patient048_frame02_slice_1 -patient048_frame02_slice_4 -patient048_frame02_slice_7 -patient048_frame02_slice_8 -patient048_frame02_slice_6 -patient055_frame01_slice_5 -patient055_frame01_slice_8 -patient055_frame01_slice_1 -patient055_frame01_slice_2 -patient055_frame01_slice_6 -patient055_frame01_slice_4 -patient055_frame01_slice_9 -patient055_frame01_slice_3 -patient055_frame01_slice_7 -patient055_frame02_slice_6 -patient055_frame02_slice_9 -patient055_frame02_slice_1 -patient055_frame02_slice_2 -patient055_frame02_slice_7 -patient055_frame02_slice_5 -patient055_frame02_slice_4 -patient055_frame02_slice_8 -patient055_frame02_slice_3 -patient097_frame01_slice_8 -patient097_frame01_slice_4 -patient097_frame01_slice_2 -patient097_frame01_slice_7 -patient097_frame01_slice_6 -patient097_frame01_slice_1 -patient097_frame01_slice_5 -patient097_frame01_slice_3 -patient097_frame02_slice_8 -patient097_frame02_slice_2 -patient097_frame02_slice_6 -patient097_frame02_slice_3 -patient097_frame02_slice_7 -patient097_frame02_slice_4 -patient097_frame02_slice_5 -patient097_frame02_slice_1 -patient074_frame01_slice_5 -patient074_frame01_slice_3 -patient074_frame01_slice_2 -patient074_frame01_slice_6 -patient074_frame01_slice_1 -patient074_frame01_slice_8 -patient074_frame01_slice_4 -patient074_frame01_slice_7 -patient074_frame02_slice_7 -patient074_frame02_slice_6 -patient074_frame02_slice_1 -patient074_frame02_slice_8 -patient074_frame02_slice_5 -patient074_frame02_slice_3 -patient074_frame02_slice_2 -patient074_frame02_slice_4 -patient043_frame01_slice_9 -patient043_frame01_slice_3 -patient043_frame01_slice_10 -patient043_frame01_slice_11 -patient043_frame01_slice_6 -patient043_frame01_slice_12 -patient043_frame01_slice_5 -patient043_frame01_slice_1 -patient043_frame01_slice_4 -patient043_frame01_slice_2 -patient043_frame01_slice_7 -patient043_frame01_slice_8 -patient043_frame02_slice_6 -patient043_frame02_slice_1 -patient043_frame02_slice_5 -patient043_frame02_slice_7 -patient043_frame02_slice_8 -patient043_frame02_slice_11 -patient043_frame02_slice_2 -patient043_frame02_slice_9 -patient043_frame02_slice_4 -patient043_frame02_slice_10 -patient043_frame02_slice_12 -patient043_frame02_slice_3 -patient041_frame01_slice_1 -patient041_frame01_slice_5 -patient041_frame01_slice_2 -patient041_frame01_slice_3 -patient041_frame01_slice_4 -patient041_frame01_slice_6 -patient041_frame02_slice_1 -patient041_frame02_slice_3 -patient041_frame02_slice_2 -patient041_frame02_slice_4 -patient041_frame02_slice_6 -patient041_frame02_slice_5 -patient063_frame01_slice_3 -patient063_frame01_slice_7 -patient063_frame01_slice_1 -patient063_frame01_slice_8 -patient063_frame01_slice_4 -patient063_frame01_slice_2 -patient063_frame01_slice_5 -patient063_frame01_slice_6 -patient063_frame02_slice_8 -patient063_frame02_slice_6 -patient063_frame02_slice_3 -patient063_frame02_slice_7 -patient063_frame02_slice_5 -patient063_frame02_slice_4 -patient063_frame02_slice_1 -patient063_frame02_slice_2 -patient037_frame01_slice_2 -patient037_frame01_slice_6 -patient037_frame01_slice_7 -patient037_frame01_slice_1 -patient037_frame01_slice_4 -patient037_frame01_slice_5 -patient037_frame01_slice_3 -patient037_frame02_slice_6 -patient037_frame02_slice_2 -patient037_frame02_slice_1 -patient037_frame02_slice_7 -patient037_frame02_slice_4 -patient037_frame02_slice_5 -patient037_frame02_slice_3 -patient095_frame01_slice_5 -patient095_frame01_slice_10 -patient095_frame01_slice_12 -patient095_frame01_slice_2 -patient095_frame01_slice_13 -patient095_frame01_slice_8 -patient095_frame01_slice_4 -patient095_frame01_slice_7 -patient095_frame01_slice_1 -patient095_frame01_slice_6 -patient095_frame01_slice_11 -patient095_frame01_slice_3 -patient095_frame01_slice_14 -patient095_frame01_slice_9 -patient095_frame02_slice_12 -patient095_frame02_slice_2 -patient095_frame02_slice_10 -patient095_frame02_slice_6 -patient095_frame02_slice_14 -patient095_frame02_slice_7 -patient095_frame02_slice_3 -patient095_frame02_slice_13 -patient095_frame02_slice_5 -patient095_frame02_slice_9 -patient095_frame02_slice_4 -patient095_frame02_slice_8 -patient095_frame02_slice_11 -patient095_frame02_slice_1 -patient054_frame01_slice_7 -patient054_frame01_slice_1 -patient054_frame01_slice_6 -patient054_frame01_slice_8 -patient054_frame01_slice_2 -patient054_frame01_slice_3 -patient054_frame01_slice_4 -patient054_frame01_slice_5 -patient054_frame02_slice_8 -patient054_frame02_slice_7 -patient054_frame02_slice_4 -patient054_frame02_slice_5 -patient054_frame02_slice_2 -patient054_frame02_slice_6 -patient054_frame02_slice_3 -patient054_frame02_slice_1 -patient026_frame01_slice_6 -patient026_frame01_slice_8 -patient026_frame01_slice_3 -patient026_frame01_slice_7 -patient026_frame01_slice_10 -patient026_frame01_slice_2 -patient026_frame01_slice_9 -patient026_frame01_slice_5 -patient026_frame01_slice_4 -patient026_frame01_slice_1 -patient026_frame02_slice_6 -patient026_frame02_slice_1 -patient026_frame02_slice_3 -patient026_frame02_slice_7 -patient026_frame02_slice_10 -patient026_frame02_slice_9 -patient026_frame02_slice_2 -patient026_frame02_slice_4 -patient026_frame02_slice_5 -patient026_frame02_slice_8 -patient088_frame01_slice_6 -patient088_frame01_slice_15 -patient088_frame01_slice_1 -patient088_frame01_slice_9 -patient088_frame01_slice_8 -patient088_frame01_slice_12 -patient088_frame01_slice_11 -patient088_frame01_slice_13 -patient088_frame01_slice_3 -patient088_frame01_slice_10 -patient088_frame01_slice_2 -patient088_frame01_slice_7 -patient088_frame01_slice_4 -patient088_frame01_slice_16 -patient088_frame01_slice_5 -patient088_frame01_slice_14 -patient088_frame02_slice_4 -patient088_frame02_slice_12 -patient088_frame02_slice_3 -patient088_frame02_slice_5 -patient088_frame02_slice_2 -patient088_frame02_slice_13 -patient088_frame02_slice_6 -patient088_frame02_slice_15 -patient088_frame02_slice_10 -patient088_frame02_slice_11 -patient088_frame02_slice_1 -patient088_frame02_slice_16 -patient088_frame02_slice_7 -patient088_frame02_slice_14 -patient088_frame02_slice_9 -patient088_frame02_slice_8 -patient032_frame01_slice_5 -patient032_frame01_slice_3 -patient032_frame01_slice_4 -patient032_frame01_slice_7 -patient032_frame01_slice_10 -patient032_frame01_slice_8 -patient032_frame01_slice_6 -patient032_frame01_slice_9 -patient032_frame01_slice_1 -patient032_frame01_slice_2 -patient032_frame02_slice_8 -patient032_frame02_slice_2 -patient032_frame02_slice_5 -patient032_frame02_slice_3 -patient032_frame02_slice_10 -patient032_frame02_slice_4 -patient032_frame02_slice_7 -patient032_frame02_slice_6 -patient032_frame02_slice_1 -patient032_frame02_slice_9 -patient069_frame01_slice_1 -patient069_frame01_slice_3 -patient069_frame01_slice_6 -patient069_frame01_slice_4 -patient069_frame01_slice_7 -patient069_frame01_slice_2 -patient069_frame01_slice_5 -patient069_frame02_slice_5 -patient069_frame02_slice_3 -patient069_frame02_slice_4 -patient069_frame02_slice_2 -patient069_frame02_slice_7 -patient069_frame02_slice_1 -patient069_frame02_slice_6 -patient006_frame01_slice_6 -patient006_frame01_slice_10 -patient006_frame01_slice_11 -patient006_frame01_slice_3 -patient006_frame01_slice_8 -patient006_frame01_slice_9 -patient006_frame01_slice_2 -patient006_frame01_slice_5 -patient006_frame01_slice_1 -patient006_frame01_slice_4 -patient006_frame01_slice_7 -patient006_frame02_slice_9 -patient006_frame02_slice_10 -patient006_frame02_slice_4 -patient006_frame02_slice_5 -patient006_frame02_slice_2 -patient006_frame02_slice_3 -patient006_frame02_slice_7 -patient006_frame02_slice_1 -patient006_frame02_slice_6 -patient006_frame02_slice_11 -patient006_frame02_slice_8 -patient071_frame01_slice_8 -patient071_frame01_slice_9 -patient071_frame01_slice_7 -patient071_frame01_slice_5 -patient071_frame01_slice_6 -patient071_frame01_slice_1 -patient071_frame01_slice_10 -patient071_frame01_slice_3 -patient071_frame01_slice_2 -patient071_frame01_slice_4 -patient071_frame02_slice_8 -patient071_frame02_slice_2 -patient071_frame02_slice_4 -patient071_frame02_slice_7 -patient071_frame02_slice_1 -patient071_frame02_slice_3 -patient071_frame02_slice_9 -patient071_frame02_slice_6 -patient071_frame02_slice_5 -patient071_frame02_slice_10 -patient012_frame01_slice_5 -patient012_frame01_slice_1 -patient012_frame01_slice_8 -patient012_frame01_slice_7 -patient012_frame01_slice_3 -patient012_frame01_slice_10 -patient012_frame01_slice_2 -patient012_frame01_slice_6 -patient012_frame01_slice_4 -patient012_frame01_slice_9 -patient012_frame02_slice_8 -patient012_frame02_slice_4 -patient012_frame02_slice_9 -patient012_frame02_slice_7 -patient012_frame02_slice_6 -patient012_frame02_slice_2 -patient012_frame02_slice_3 -patient012_frame02_slice_5 -patient012_frame02_slice_10 -patient012_frame02_slice_1 -patient073_frame01_slice_3 -patient073_frame01_slice_7 -patient073_frame01_slice_5 -patient073_frame01_slice_1 -patient073_frame01_slice_4 -patient073_frame01_slice_2 -patient073_frame01_slice_6 -patient073_frame02_slice_5 -patient073_frame02_slice_7 -patient073_frame02_slice_4 -patient073_frame02_slice_1 -patient073_frame02_slice_3 -patient073_frame02_slice_2 -patient073_frame02_slice_6 -patient061_frame01_slice_4 -patient061_frame01_slice_7 -patient061_frame01_slice_2 -patient061_frame01_slice_5 -patient061_frame01_slice_8 -patient061_frame01_slice_3 -patient061_frame01_slice_1 -patient061_frame01_slice_9 -patient061_frame01_slice_6 -patient061_frame02_slice_8 -patient061_frame02_slice_4 -patient061_frame02_slice_2 -patient061_frame02_slice_3 -patient061_frame02_slice_9 -patient061_frame02_slice_7 -patient061_frame02_slice_5 -patient061_frame02_slice_6 -patient061_frame02_slice_1 -patient017_frame01_slice_4 -patient017_frame01_slice_8 -patient017_frame01_slice_6 -patient017_frame01_slice_2 -patient017_frame01_slice_3 -patient017_frame01_slice_5 -patient017_frame01_slice_9 -patient017_frame01_slice_7 -patient017_frame01_slice_1 -patient017_frame02_slice_3 -patient017_frame02_slice_5 -patient017_frame02_slice_4 -patient017_frame02_slice_8 -patient017_frame02_slice_6 -patient017_frame02_slice_7 -patient017_frame02_slice_1 -patient017_frame02_slice_2 -patient017_frame02_slice_9 -patient025_frame01_slice_7 -patient025_frame01_slice_8 -patient025_frame01_slice_6 -patient025_frame01_slice_1 -patient025_frame01_slice_5 -patient025_frame01_slice_4 -patient025_frame01_slice_3 -patient025_frame01_slice_9 -patient025_frame01_slice_2 -patient025_frame02_slice_3 -patient025_frame02_slice_6 -patient025_frame02_slice_2 -patient025_frame02_slice_1 -patient025_frame02_slice_5 -patient025_frame02_slice_8 -patient025_frame02_slice_7 -patient025_frame02_slice_9 -patient025_frame02_slice_4 -patient010_frame01_slice_9 -patient010_frame01_slice_10 -patient010_frame01_slice_1 -patient010_frame01_slice_2 -patient010_frame01_slice_7 -patient010_frame01_slice_3 -patient010_frame01_slice_8 -patient010_frame01_slice_6 -patient010_frame01_slice_4 -patient010_frame01_slice_5 -patient010_frame02_slice_8 -patient010_frame02_slice_5 -patient010_frame02_slice_9 -patient010_frame02_slice_4 -patient010_frame02_slice_10 -patient010_frame02_slice_2 -patient010_frame02_slice_7 -patient010_frame02_slice_1 -patient010_frame02_slice_3 -patient010_frame02_slice_6 -patient057_frame01_slice_1 -patient057_frame01_slice_7 -patient057_frame01_slice_5 -patient057_frame01_slice_2 -patient057_frame01_slice_3 -patient057_frame01_slice_8 -patient057_frame01_slice_4 -patient057_frame01_slice_6 -patient057_frame02_slice_7 -patient057_frame02_slice_4 -patient057_frame02_slice_1 -patient057_frame02_slice_5 -patient057_frame02_slice_6 -patient057_frame02_slice_3 -patient057_frame02_slice_8 -patient057_frame02_slice_2 -patient029_frame01_slice_9 -patient029_frame01_slice_1 -patient029_frame01_slice_5 -patient029_frame01_slice_7 -patient029_frame01_slice_11 -patient029_frame01_slice_3 -patient029_frame01_slice_4 -patient029_frame01_slice_6 -patient029_frame01_slice_10 -patient029_frame01_slice_8 -patient029_frame01_slice_2 -patient029_frame02_slice_6 -patient029_frame02_slice_8 -patient029_frame02_slice_2 -patient029_frame02_slice_3 -patient029_frame02_slice_10 -patient029_frame02_slice_1 -patient029_frame02_slice_11 -patient029_frame02_slice_4 -patient029_frame02_slice_9 -patient029_frame02_slice_7 -patient029_frame02_slice_5 -patient051_frame01_slice_7 -patient051_frame01_slice_9 -patient051_frame01_slice_3 -patient051_frame01_slice_4 -patient051_frame01_slice_10 -patient051_frame01_slice_8 -patient051_frame01_slice_2 -patient051_frame01_slice_1 -patient051_frame01_slice_6 -patient051_frame01_slice_5 -patient051_frame02_slice_5 -patient051_frame02_slice_10 -patient051_frame02_slice_1 -patient051_frame02_slice_4 -patient051_frame02_slice_3 -patient051_frame02_slice_8 -patient051_frame02_slice_9 -patient051_frame02_slice_7 -patient051_frame02_slice_6 -patient051_frame02_slice_2 -patient005_frame01_slice_3 -patient005_frame01_slice_8 -patient005_frame01_slice_9 -patient005_frame01_slice_1 -patient005_frame01_slice_5 -patient005_frame01_slice_2 -patient005_frame01_slice_4 -patient005_frame01_slice_6 -patient005_frame01_slice_10 -patient005_frame01_slice_7 -patient005_frame02_slice_4 -patient005_frame02_slice_8 -patient005_frame02_slice_3 -patient005_frame02_slice_2 -patient005_frame02_slice_7 -patient005_frame02_slice_10 -patient005_frame02_slice_9 -patient005_frame02_slice_6 -patient005_frame02_slice_5 -patient005_frame02_slice_1 -patient036_frame01_slice_5 -patient036_frame01_slice_2 -patient036_frame01_slice_8 -patient036_frame01_slice_7 -patient036_frame01_slice_1 -patient036_frame01_slice_4 -patient036_frame01_slice_6 -patient036_frame01_slice_3 -patient036_frame02_slice_6 -patient036_frame02_slice_1 -patient036_frame02_slice_8 -patient036_frame02_slice_5 -patient036_frame02_slice_7 -patient036_frame02_slice_3 -patient036_frame02_slice_4 -patient036_frame02_slice_2 -patient046_frame01_slice_6 -patient046_frame01_slice_1 -patient046_frame01_slice_3 -patient046_frame01_slice_7 -patient046_frame01_slice_8 -patient046_frame01_slice_5 -patient046_frame01_slice_4 -patient046_frame01_slice_2 -patient046_frame01_slice_9 -patient046_frame02_slice_1 -patient046_frame02_slice_7 -patient046_frame02_slice_9 -patient046_frame02_slice_6 -patient046_frame02_slice_3 -patient046_frame02_slice_8 -patient046_frame02_slice_5 -patient046_frame02_slice_2 -patient046_frame02_slice_4 -patient062_frame01_slice_1 -patient062_frame01_slice_10 -patient062_frame01_slice_4 -patient062_frame01_slice_9 -patient062_frame01_slice_8 -patient062_frame01_slice_7 -patient062_frame01_slice_2 -patient062_frame01_slice_3 -patient062_frame01_slice_5 -patient062_frame01_slice_6 -patient062_frame02_slice_9 -patient062_frame02_slice_2 -patient062_frame02_slice_7 -patient062_frame02_slice_6 -patient062_frame02_slice_5 -patient062_frame02_slice_10 -patient062_frame02_slice_3 -patient062_frame02_slice_4 -patient062_frame02_slice_1 -patient062_frame02_slice_8 -patient034_frame01_slice_6 -patient034_frame01_slice_2 -patient034_frame01_slice_7 -patient034_frame01_slice_8 -patient034_frame01_slice_9 -patient034_frame01_slice_10 -patient034_frame01_slice_1 -patient034_frame01_slice_4 -patient034_frame01_slice_5 -patient034_frame01_slice_3 -patient034_frame02_slice_1 -patient034_frame02_slice_3 -patient034_frame02_slice_5 -patient034_frame02_slice_10 -patient034_frame02_slice_9 -patient034_frame02_slice_4 -patient034_frame02_slice_8 -patient034_frame02_slice_6 -patient034_frame02_slice_7 -patient034_frame02_slice_2 -patient076_frame01_slice_1 -patient076_frame01_slice_5 -patient076_frame01_slice_6 -patient076_frame01_slice_3 -patient076_frame01_slice_7 -patient076_frame01_slice_2 -patient076_frame01_slice_4 -patient076_frame01_slice_8 -patient076_frame02_slice_6 -patient076_frame02_slice_5 -patient076_frame02_slice_4 -patient076_frame02_slice_3 -patient076_frame02_slice_2 -patient076_frame02_slice_1 -patient076_frame02_slice_8 -patient076_frame02_slice_7 -patient092_frame01_slice_3 -patient092_frame01_slice_9 -patient092_frame01_slice_12 -patient092_frame01_slice_4 -patient092_frame01_slice_6 -patient092_frame01_slice_7 -patient092_frame01_slice_2 -patient092_frame01_slice_11 -patient092_frame01_slice_5 -patient092_frame01_slice_10 -patient092_frame01_slice_13 -patient092_frame01_slice_14 -patient092_frame01_slice_1 -patient092_frame01_slice_8 -patient092_frame01_slice_15 -patient092_frame02_slice_8 -patient092_frame02_slice_1 -patient092_frame02_slice_5 -patient092_frame02_slice_12 -patient092_frame02_slice_14 -patient092_frame02_slice_9 -patient092_frame02_slice_3 -patient092_frame02_slice_6 -patient092_frame02_slice_13 -patient092_frame02_slice_15 -patient092_frame02_slice_11 -patient092_frame02_slice_7 -patient092_frame02_slice_10 -patient092_frame02_slice_4 -patient092_frame02_slice_2 -patient070_frame01_slice_3 -patient070_frame01_slice_4 -patient070_frame01_slice_1 -patient070_frame01_slice_6 -patient070_frame01_slice_5 -patient070_frame01_slice_2 -patient070_frame02_slice_4 -patient070_frame02_slice_2 -patient070_frame02_slice_1 -patient070_frame02_slice_6 -patient070_frame02_slice_3 -patient070_frame02_slice_5 -patient077_frame01_slice_6 -patient077_frame01_slice_4 -patient077_frame01_slice_5 -patient077_frame01_slice_7 -patient077_frame01_slice_8 -patient077_frame01_slice_1 -patient077_frame01_slice_2 -patient077_frame01_slice_3 -patient077_frame02_slice_3 -patient077_frame02_slice_2 -patient077_frame02_slice_7 -patient077_frame02_slice_4 -patient077_frame02_slice_1 -patient077_frame02_slice_8 -patient077_frame02_slice_5 -patient077_frame02_slice_6 -patient067_frame01_slice_8 -patient067_frame01_slice_2 -patient067_frame01_slice_1 -patient067_frame01_slice_3 -patient067_frame01_slice_6 -patient067_frame01_slice_7 -patient067_frame01_slice_4 -patient067_frame01_slice_10 -patient067_frame01_slice_9 -patient067_frame01_slice_5 -patient067_frame02_slice_10 -patient067_frame02_slice_7 -patient067_frame02_slice_6 -patient067_frame02_slice_8 -patient067_frame02_slice_9 -patient067_frame02_slice_4 -patient067_frame02_slice_3 -patient067_frame02_slice_1 -patient067_frame02_slice_5 -patient067_frame02_slice_2 -patient003_frame01_slice_5 -patient003_frame01_slice_3 -patient003_frame01_slice_8 -patient003_frame01_slice_4 -patient003_frame01_slice_2 -patient003_frame01_slice_6 -patient003_frame01_slice_9 -patient003_frame01_slice_10 -patient003_frame01_slice_7 -patient003_frame01_slice_1 -patient003_frame02_slice_2 -patient003_frame02_slice_5 -patient003_frame02_slice_9 -patient003_frame02_slice_6 -patient003_frame02_slice_8 -patient003_frame02_slice_3 -patient003_frame02_slice_7 -patient003_frame02_slice_1 -patient003_frame02_slice_4 -patient003_frame02_slice_10 -patient091_frame01_slice_2 -patient091_frame01_slice_8 -patient091_frame01_slice_6 -patient091_frame01_slice_5 -patient091_frame01_slice_7 -patient091_frame01_slice_4 -patient091_frame01_slice_1 -patient091_frame01_slice_3 -patient091_frame02_slice_6 -patient091_frame02_slice_7 -patient091_frame02_slice_4 -patient091_frame02_slice_1 -patient091_frame02_slice_2 -patient091_frame02_slice_5 -patient091_frame02_slice_3 -patient091_frame02_slice_8 -patient016_frame01_slice_9 -patient016_frame01_slice_5 -patient016_frame01_slice_1 -patient016_frame01_slice_2 -patient016_frame01_slice_8 -patient016_frame01_slice_10 -patient016_frame01_slice_6 -patient016_frame01_slice_3 -patient016_frame01_slice_7 -patient016_frame01_slice_4 -patient016_frame02_slice_2 -patient016_frame02_slice_3 -patient016_frame02_slice_8 -patient016_frame02_slice_5 -patient016_frame02_slice_1 -patient016_frame02_slice_6 -patient016_frame02_slice_4 -patient016_frame02_slice_9 -patient016_frame02_slice_7 -patient016_frame02_slice_10 -patient014_frame01_slice_5 -patient014_frame01_slice_9 -patient014_frame01_slice_1 -patient014_frame01_slice_8 -patient014_frame01_slice_6 -patient014_frame01_slice_2 -patient014_frame01_slice_3 -patient014_frame01_slice_4 -patient014_frame01_slice_7 -patient014_frame01_slice_10 -patient014_frame02_slice_6 -patient014_frame02_slice_3 -patient014_frame02_slice_5 -patient014_frame02_slice_10 -patient014_frame02_slice_7 -patient014_frame02_slice_9 -patient014_frame02_slice_8 -patient014_frame02_slice_2 -patient014_frame02_slice_1 -patient014_frame02_slice_4 -patient044_frame01_slice_4 -patient044_frame01_slice_6 -patient044_frame01_slice_9 -patient044_frame01_slice_5 -patient044_frame01_slice_7 -patient044_frame01_slice_1 -patient044_frame01_slice_8 -patient044_frame01_slice_2 -patient044_frame01_slice_3 -patient044_frame02_slice_3 -patient044_frame02_slice_5 -patient044_frame02_slice_1 -patient044_frame02_slice_7 -patient044_frame02_slice_2 -patient044_frame02_slice_9 -patient044_frame02_slice_6 -patient044_frame02_slice_4 -patient044_frame02_slice_8 -patient042_frame01_slice_2 -patient042_frame01_slice_6 -patient042_frame01_slice_3 -patient042_frame01_slice_8 -patient042_frame01_slice_7 -patient042_frame01_slice_5 -patient042_frame01_slice_4 -patient042_frame01_slice_9 -patient042_frame01_slice_1 -patient042_frame02_slice_4 -patient042_frame02_slice_2 -patient042_frame02_slice_5 -patient042_frame02_slice_9 -patient042_frame02_slice_1 -patient042_frame02_slice_6 -patient042_frame02_slice_8 -patient042_frame02_slice_3 -patient042_frame02_slice_7 -patient090_frame01_slice_4 -patient090_frame01_slice_3 -patient090_frame01_slice_7 -patient090_frame01_slice_6 -patient090_frame01_slice_2 -patient090_frame01_slice_5 -patient090_frame01_slice_1 -patient090_frame02_slice_1 -patient090_frame02_slice_4 -patient090_frame02_slice_2 -patient090_frame02_slice_7 -patient090_frame02_slice_5 -patient090_frame02_slice_3 -patient090_frame02_slice_6 -patient053_frame01_slice_1 -patient053_frame01_slice_7 -patient053_frame01_slice_5 -patient053_frame01_slice_6 -patient053_frame01_slice_3 -patient053_frame01_slice_4 -patient053_frame01_slice_2 -patient053_frame02_slice_2 -patient053_frame02_slice_6 -patient053_frame02_slice_4 -patient053_frame02_slice_5 -patient053_frame02_slice_7 -patient053_frame02_slice_1 -patient053_frame02_slice_3 -patient027_frame01_slice_2 -patient027_frame01_slice_5 -patient027_frame01_slice_6 -patient027_frame01_slice_10 -patient027_frame01_slice_1 -patient027_frame01_slice_4 -patient027_frame01_slice_7 -patient027_frame01_slice_8 -patient027_frame01_slice_9 -patient027_frame01_slice_3 -patient027_frame02_slice_5 -patient027_frame02_slice_9 -patient027_frame02_slice_7 -patient027_frame02_slice_6 -patient027_frame02_slice_1 -patient027_frame02_slice_8 -patient027_frame02_slice_4 -patient027_frame02_slice_3 -patient027_frame02_slice_2 -patient027_frame02_slice_10 -patient035_frame01_slice_8 -patient035_frame01_slice_6 -patient035_frame01_slice_11 -patient035_frame01_slice_5 -patient035_frame01_slice_2 -patient035_frame01_slice_3 -patient035_frame01_slice_9 -patient035_frame01_slice_12 -patient035_frame01_slice_10 -patient035_frame01_slice_7 -patient035_frame01_slice_13 -patient035_frame01_slice_1 -patient035_frame01_slice_4 -patient035_frame02_slice_9 -patient035_frame02_slice_4 -patient035_frame02_slice_3 -patient035_frame02_slice_13 -patient035_frame02_slice_12 -patient035_frame02_slice_10 -patient035_frame02_slice_8 -patient035_frame02_slice_6 -patient035_frame02_slice_1 -patient035_frame02_slice_2 -patient035_frame02_slice_5 -patient035_frame02_slice_11 -patient035_frame02_slice_7 -patient086_frame01_slice_6 -patient086_frame01_slice_5 -patient086_frame01_slice_3 -patient086_frame01_slice_4 -patient086_frame01_slice_7 -patient086_frame01_slice_2 -patient086_frame01_slice_1 -patient086_frame02_slice_1 -patient086_frame02_slice_5 -patient086_frame02_slice_3 -patient086_frame02_slice_2 -patient086_frame02_slice_7 -patient086_frame02_slice_6 -patient086_frame02_slice_4 -patient023_frame01_slice_7 -patient023_frame01_slice_1 -patient023_frame01_slice_5 -patient023_frame01_slice_2 -patient023_frame01_slice_4 -patient023_frame01_slice_8 -patient023_frame01_slice_9 -patient023_frame01_slice_3 -patient023_frame01_slice_6 -patient023_frame02_slice_7 -patient023_frame02_slice_6 -patient023_frame02_slice_9 -patient023_frame02_slice_2 -patient023_frame02_slice_5 -patient023_frame02_slice_1 -patient023_frame02_slice_3 -patient023_frame02_slice_8 -patient023_frame02_slice_4 -patient009_frame01_slice_4 -patient009_frame01_slice_10 -patient009_frame01_slice_7 -patient009_frame01_slice_5 -patient009_frame01_slice_2 -patient009_frame01_slice_6 -patient009_frame01_slice_8 -patient009_frame01_slice_3 -patient009_frame01_slice_1 -patient009_frame01_slice_9 -patient009_frame02_slice_2 -patient009_frame02_slice_4 -patient009_frame02_slice_7 -patient009_frame02_slice_10 -patient009_frame02_slice_9 -patient009_frame02_slice_1 -patient009_frame02_slice_5 -patient009_frame02_slice_8 -patient009_frame02_slice_6 -patient009_frame02_slice_3 -patient079_frame01_slice_5 -patient079_frame01_slice_4 -patient079_frame01_slice_1 -patient079_frame01_slice_2 -patient079_frame01_slice_8 -patient079_frame01_slice_9 -patient079_frame01_slice_6 -patient079_frame01_slice_3 -patient079_frame01_slice_7 -patient079_frame02_slice_6 -patient079_frame02_slice_5 -patient079_frame02_slice_1 -patient079_frame02_slice_4 -patient079_frame02_slice_8 -patient079_frame02_slice_9 -patient079_frame02_slice_7 -patient079_frame02_slice_2 -patient079_frame02_slice_3 -patient015_frame01_slice_6 -patient015_frame01_slice_7 -patient015_frame01_slice_9 -patient015_frame01_slice_1 -patient015_frame01_slice_5 -patient015_frame01_slice_8 -patient015_frame01_slice_4 -patient015_frame01_slice_3 -patient015_frame01_slice_2 -patient015_frame02_slice_5 -patient015_frame02_slice_1 -patient015_frame02_slice_6 -patient015_frame02_slice_3 -patient015_frame02_slice_2 -patient015_frame02_slice_7 -patient015_frame02_slice_4 -patient015_frame02_slice_9 -patient015_frame02_slice_8 diff --git a/data/ACDC/val.list b/data/ACDC/val.list deleted file mode 100755 index bda8dad..0000000 --- a/data/ACDC/val.list +++ /dev/null @@ -1,20 +0,0 @@ -patient028_frame01 -patient028_frame02 -patient085_frame01 -patient085_frame02 -patient082_frame01 -patient082_frame02 -patient087_frame01 -patient087_frame02 -patient019_frame01 -patient019_frame02 -patient030_frame01 -patient030_frame02 -patient078_frame01 -patient078_frame02 -patient045_frame01 -patient045_frame02 -patient002_frame01 -patient002_frame02 -patient039_frame01 -patient039_frame02 diff --git a/data/BraTS2019/README.md b/data/BraTS2019/README.md deleted file mode 100644 index 07fabc4..0000000 --- a/data/BraTS2019/README.md +++ /dev/null @@ -1 +0,0 @@ -- Download the processed BraTS2019 data (we just used the Flair images for whole tumor segmentation) from [BaiduDisk](https://pan.baidu.com/s/1CrMNP8hUExGuQNrHPuGb7w), the password is *code*, and decompress the zip file to [data/BraTS2019](https://github.com/Luoxd1996/SSL4MIS/edit/master/data/BraTS2019). diff --git a/data/BraTS2019/test.txt b/data/BraTS2019/test.txt deleted file mode 100755 index 5be2024..0000000 --- a/data/BraTS2019/test.txt +++ /dev/null @@ -1,60 +0,0 @@ -BraTS19_TCIA02_309_1 -BraTS19_CBICA_AYA_1 -BraTS19_CBICA_AYG_1 -BraTS19_CBICA_ANV_1 -BraTS19_CBICA_BAN_1 -BraTS19_TCIA10_408_1 -BraTS19_TCIA01_448_1 -BraTS19_TCIA10_261_1 -BraTS19_CBICA_ALU_1 -BraTS19_CBICA_AWH_1 -BraTS19_TCIA10_420_1 -BraTS19_2013_11_1 -BraTS19_TCIA01_460_1 -BraTS19_TCIA10_639_1 -BraTS19_TCIA10_130_1 -BraTS19_TCIA09_254_1 -BraTS19_CBICA_BLJ_1 -BraTS19_TCIA01_390_1 -BraTS19_2013_19_1 -BraTS19_CBICA_BDK_1 -BraTS19_TCIA13_630_1 -BraTS19_TCIA02_430_1 -BraTS19_TCIA06_165_1 -BraTS19_CBICA_AOS_1 -BraTS19_CBICA_AZH_1 -BraTS19_CBICA_BGX_1 -BraTS19_TCIA02_455_1 -BraTS19_TCIA10_330_1 -BraTS19_TMC_12866_1 -BraTS19_TCIA02_321_1 -BraTS19_TCIA03_257_1 -BraTS19_CBICA_AWI_1 -BraTS19_CBICA_AQZ_1 -BraTS19_TCIA10_152_1 -BraTS19_TMC_27374_1 -BraTS19_TCIA01_378_1 -BraTS19_TCIA08_218_1 -BraTS19_CBICA_ASY_1 -BraTS19_TCIA02_168_1 -BraTS19_CBICA_BJY_1 -BraTS19_2013_18_1 -BraTS19_TCIA03_121_1 -BraTS19_CBICA_BIC_1 -BraTS19_TCIA03_133_1 -BraTS19_TCIA02_171_1 -BraTS19_TMC_06290_1 -BraTS19_TCIA10_175_1 -BraTS19_CBICA_BGR_1 -BraTS19_TCIA10_276_1 -BraTS19_TCIA09_402_1 -BraTS19_TCIA10_442_1 -BraTS19_CBICA_AUN_1 -BraTS19_2013_20_1 -BraTS19_TCIA10_490_1 -BraTS19_TCIA06_409_1 -BraTS19_TMC_21360_1 -BraTS19_CBICA_AQQ_1 -BraTS19_TCIA10_307_1 -BraTS19_TCIA04_437_1 -BraTS19_CBICA_AQR_1 diff --git a/data/BraTS2019/train.txt b/data/BraTS2019/train.txt deleted file mode 100755 index a24855c..0000000 --- a/data/BraTS2019/train.txt +++ /dev/null @@ -1,250 +0,0 @@ -BraTS19_TCIA02_370_1 -BraTS19_CBICA_ASA_1 -BraTS19_TCIA12_470_1 -BraTS19_2013_8_1 -BraTS19_TCIA01_429_1 -BraTS19_TCIA08_234_1 -BraTS19_TCIA10_266_1 -BraTS19_TCIA13_633_1 -BraTS19_TCIA03_199_1 -BraTS19_TCIA10_629_1 -BraTS19_CBICA_ATB_1 -BraTS19_CBICA_BCL_1 -BraTS19_TCIA08_469_1 -BraTS19_TCIA04_343_1 -BraTS19_CBICA_AOO_1 -BraTS19_CBICA_ASF_1 -BraTS19_TCIA13_618_1 -BraTS19_CBICA_AVG_1 -BraTS19_TCIA09_255_1 -BraTS19_TCIA08_280_1 -BraTS19_2013_9_1 -BraTS19_TCIA10_202_1 -BraTS19_TCIA04_149_1 -BraTS19_CBICA_AZD_1 -BraTS19_TCIA01_131_1 -BraTS19_CBICA_BHZ_1 -BraTS19_TCIA02_179_1 -BraTS19_CBICA_BEM_1 -BraTS19_TCIA02_300_1 -BraTS19_CBICA_ARF_1 -BraTS19_CBICA_ABY_1 -BraTS19_TCIA02_608_1 -BraTS19_2013_17_1 -BraTS19_CBICA_BHV_1 -BraTS19_TCIA02_117_1 -BraTS19_TCIA12_249_1 -BraTS19_TCIA08_162_1 -BraTS19_TCIA03_498_1 -BraTS19_TCIA01_235_1 -BraTS19_2013_15_1 -BraTS19_CBICA_AOP_1 -BraTS19_CBICA_AUA_1 -BraTS19_CBICA_AAB_1 -BraTS19_CBICA_BFB_1 -BraTS19_TCIA09_451_1 -BraTS19_TCIA02_322_1 -BraTS19_CBICA_ATV_1 -BraTS19_CBICA_BCF_1 -BraTS19_CBICA_AQJ_1 -BraTS19_CBICA_AVV_1 -BraTS19_CBICA_ASU_1 -BraTS19_CBICA_AYW_1 -BraTS19_CBICA_AUR_1 -BraTS19_CBICA_AYC_1 -BraTS19_TCIA02_607_1 -BraTS19_TCIA08_436_1 -BraTS19_TCIA02_471_1 -BraTS19_CBICA_AQG_1 -BraTS19_TCIA10_387_1 -BraTS19_TCIA02_606_1 -BraTS19_CBICA_ASR_1 -BraTS19_CBICA_ASO_1 -BraTS19_TCIA08_167_1 -BraTS19_CBICA_AXN_1 -BraTS19_CBICA_ABB_1 -BraTS19_CBICA_AWG_1 -BraTS19_TCIA13_653_1 -BraTS19_2013_23_1 -BraTS19_2013_14_1 -BraTS19_CBICA_APY_1 -BraTS19_CBICA_ATX_1 -BraTS19_CBICA_ATF_1 -BraTS19_CBICA_AQV_1 -BraTS19_CBICA_ASW_1 -BraTS19_TCIA10_449_1 -BraTS19_CBICA_AQO_1 -BraTS19_TCIA06_247_1 -BraTS19_CBICA_AXQ_1 -BraTS19_TCIA02_290_1 -BraTS19_CBICA_APZ_1 -BraTS19_TCIA10_640_1 -BraTS19_TCIA02_118_1 -BraTS19_TCIA02_151_1 -BraTS19_CBICA_BGE_1 -BraTS19_CBICA_AOD_1 -BraTS19_TCIA01_401_1 -BraTS19_2013_27_1 -BraTS19_TCIA01_180_1 -BraTS19_TCIA01_231_1 -BraTS19_TCIA02_491_1 -BraTS19_TCIA05_444_1 -BraTS19_CBICA_AQT_1 -BraTS19_CBICA_AQU_1 -BraTS19_CBICA_AME_1 -BraTS19_TCIA12_298_1 -BraTS19_CBICA_AMH_1 -BraTS19_CBICA_ANI_1 -BraTS19_2013_28_1 -BraTS19_CBICA_BGN_1 -BraTS19_TCIA13_650_1 -BraTS19_TCIA13_634_1 -BraTS19_CBICA_APK_1 -BraTS19_TCIA13_624_1 -BraTS19_TCIA10_628_1 -BraTS19_TCIA01_221_1 -BraTS19_TCIA06_332_1 -BraTS19_CBICA_BHK_1 -BraTS19_TCIA09_428_1 -BraTS19_2013_3_1 -BraTS19_CBICA_BAX_1 -BraTS19_TCIA10_299_1 -BraTS19_TCIA10_310_1 -BraTS19_TCIA02_331_1 -BraTS19_TCIA03_419_1 -BraTS19_CBICA_AAG_1 -BraTS19_TCIA12_480_1 -BraTS19_CBICA_BGO_1 -BraTS19_TCIA01_425_1 -BraTS19_TCIA10_637_1 -BraTS19_2013_2_1 -BraTS19_CBICA_AWX_1 -BraTS19_TCIA13_642_1 -BraTS19_TCIA04_328_1 -BraTS19_CBICA_AWV_1 -BraTS19_CBICA_AAL_1 -BraTS19_CBICA_AVF_1 -BraTS19_TCIA08_406_1 -BraTS19_TCIA03_296_1 -BraTS19_CBICA_ASK_1 -BraTS19_TCIA05_396_1 -BraTS19_TCIA02_368_1 -BraTS19_CBICA_AUQ_1 -BraTS19_2013_21_1 -BraTS19_CBICA_BGG_1 -BraTS19_TCIA13_654_1 -BraTS19_CBICA_AUW_1 -BraTS19_CBICA_BHM_1 -BraTS19_2013_29_1 -BraTS19_CBICA_AOZ_1 -BraTS19_2013_0_1 -BraTS19_CBICA_ASH_1 -BraTS19_CBICA_ANZ_1 -BraTS19_2013_26_1 -BraTS19_TCIA02_283_1 -BraTS19_TCIA02_473_1 -BraTS19_TCIA09_141_1 -BraTS19_TCIA08_278_1 -BraTS19_TCIA03_375_1 -BraTS19_2013_10_1 -BraTS19_2013_7_1 -BraTS19_TCIA12_101_1 -BraTS19_TCIA04_192_1 -BraTS19_CBICA_ASN_1 -BraTS19_TCIA08_242_1 -BraTS19_TCIA02_394_1 -BraTS19_CBICA_AXW_1 -BraTS19_TCIA04_361_1 -BraTS19_CBICA_AQY_1 -BraTS19_TCIA01_412_1 -BraTS19_CBICA_AVT_1 -BraTS19_TCIA01_186_1 -BraTS19_TCIA02_377_1 -BraTS19_TCIA01_411_1 -BraTS19_TCIA02_198_1 -BraTS19_CBICA_ASE_1 -BraTS19_TCIA10_413_1 -BraTS19_CBICA_APR_1 -BraTS19_CBICA_ALN_1 -BraTS19_TCIA10_393_1 -BraTS19_TMC_09043_1 -BraTS19_2013_5_1 -BraTS19_TCIA01_201_1 -BraTS19_CBICA_ABM_1 -BraTS19_2013_4_1 -BraTS19_TCIA13_621_1 -BraTS19_TCIA09_620_1 -BraTS19_TCIA10_325_1 -BraTS19_CBICA_ANP_1 -BraTS19_TCIA01_335_1 -BraTS19_CBICA_AXM_1 -BraTS19_TCIA08_105_1 -BraTS19_TCIA02_226_1 -BraTS19_CBICA_ANG_1 -BraTS19_TCIA03_338_1 -BraTS19_TCIA08_113_1 -BraTS19_TCIA02_222_1 -BraTS19_TCIA10_625_1 -BraTS19_TCIA09_493_1 -BraTS19_CBICA_AOH_1 -BraTS19_CBICA_BAP_1 -BraTS19_TCIA02_605_1 -BraTS19_TCIA01_190_1 -BraTS19_TCIA10_109_1 -BraTS19_CBICA_AUX_1 -BraTS19_2013_13_1 -BraTS19_TMC_06643_1 -BraTS19_TCIA10_241_1 -BraTS19_CBICA_AXL_1 -BraTS19_CBICA_AXO_1 -BraTS19_TCIA10_632_1 -BraTS19_TCIA13_645_1 -BraTS19_TCIA02_314_1 -BraTS19_TCIA05_277_1 -BraTS19_TCIA03_474_1 -BraTS19_TCIA02_374_1 -BraTS19_CBICA_AXJ_1 -BraTS19_TCIA10_351_1 -BraTS19_CBICA_AQD_1 -BraTS19_TCIA09_462_1 -BraTS19_CBICA_AYI_1 -BraTS19_2013_25_1 -BraTS19_CBICA_AVJ_1 -BraTS19_2013_12_1 -BraTS19_TCIA06_372_1 -BraTS19_2013_24_1 -BraTS19_TCIA04_479_1 -BraTS19_CBICA_ASV_1 -BraTS19_CBICA_ABN_1 -BraTS19_CBICA_ATN_1 -BraTS19_CBICA_ABE_1 -BraTS19_TCIA08_319_1 -BraTS19_CBICA_AQP_1 -BraTS19_TCIA03_138_1 -BraTS19_TCIA10_410_1 -BraTS19_CBICA_BGW_1 -BraTS19_CBICA_BNR_1 -BraTS19_CBICA_BHB_1 -BraTS19_CBICA_BFP_1 -BraTS19_CBICA_BGT_1 -BraTS19_CBICA_AYU_1 -BraTS19_CBICA_ATD_1 -BraTS19_CBICA_ATP_1 -BraTS19_CBICA_BBG_1 -BraTS19_CBICA_ARZ_1 -BraTS19_TCIA06_211_1 -BraTS19_2013_6_1 -BraTS19_CBICA_ALX_1 -BraTS19_TCIA05_478_1 -BraTS19_TCIA06_603_1 -BraTS19_TCIA03_265_1 -BraTS19_TCIA13_623_1 -BraTS19_TCIA04_111_1 -BraTS19_CBICA_ASG_1 -BraTS19_TCIA01_499_1 -BraTS19_TCIA12_466_1 -BraTS19_TCIA01_203_1 -BraTS19_2013_22_1 -BraTS19_TCIA02_274_1 -BraTS19_TCIA01_150_1 -BraTS19_TCIA10_346_1 diff --git a/data/BraTS2019/val.txt b/data/BraTS2019/val.txt deleted file mode 100755 index 7684793..0000000 --- a/data/BraTS2019/val.txt +++ /dev/null @@ -1,25 +0,0 @@ -BraTS19_CBICA_AAP_1 -BraTS19_TCIA10_282_1 -BraTS19_TCIA01_147_1 -BraTS19_TCIA13_615_1 -BraTS19_CBICA_AQA_1 -BraTS19_TCIA08_205_1 -BraTS19_TCIA09_177_1 -BraTS19_TCIA02_208_1 -BraTS19_CBICA_AQN_1 -BraTS19_2013_16_1 -BraTS19_CBICA_BHQ_1 -BraTS19_TCIA02_135_1 -BraTS19_CBICA_AVB_1 -BraTS19_TCIA10_644_1 -BraTS19_CBICA_BKV_1 -BraTS19_TMC_15477_1 -BraTS19_2013_1_1 -BraTS19_CBICA_ARW_1 -BraTS19_TMC_11964_1 -BraTS19_TCIA09_312_1 -BraTS19_TCIA06_184_1 -BraTS19_TMC_30014_1 -BraTS19_TCIA10_103_1 -BraTS19_CBICA_AOC_1 -BraTS19_CBICA_ABO_1 From a47bbe497bd79d6fb8b224d257d157de21e91ebf Mon Sep 17 00:00:00 2001 From: luoxd Date: Thu, 10 Mar 2022 13:07:16 +0800 Subject: [PATCH 6/7] upload fixmatch code --- code/dataloaders/dataset.py | 85 ++++- code/networks/unet.py | 4 +- code/test.py | 76 ++-- code/train_acdc_ent_mini.sh | 10 + code/train_acdc_ict_mini.sh | 10 + code/train_cross_pseudo_supervision.py | 35 +- code/train_deep_adversarial_network.py | 23 +- code/train_deep_co_training.py | 324 ++++++++++++++++++ code/train_entropy_minimization.py | 24 +- code/train_exp1.sh | 11 + code/train_fixmatch.py | 281 +++++++++++++++ code/train_fully_supervised.py | 23 +- ...rain_interpolation_consistency_training.py | 24 +- code/train_mean_teacher.py | 24 +- code/train_uncertainty_aware_mean_teacher.py | 24 +- code/utils/losses.py | 32 +- 16 files changed, 923 insertions(+), 87 deletions(-) create mode 100644 code/train_acdc_ent_mini.sh create mode 100644 code/train_acdc_ict_mini.sh create mode 100644 code/train_deep_co_training.py create mode 100644 code/train_exp1.sh create mode 100644 code/train_fixmatch.py diff --git a/code/dataloaders/dataset.py b/code/dataloaders/dataset.py index 29365f8..36002d3 100644 --- a/code/dataloaders/dataset.py +++ b/code/dataloaders/dataset.py @@ -10,8 +10,10 @@ import torch from scipy import ndimage from scipy.ndimage.interpolation import zoom -from sklearn.model_selection import KFold +from skimage import exposure +from sklearn.model_selection import KFold, train_test_split from torch.utils.data import Dataset +import copy try: # SciPy >= 0.19 from scipy.special import comb @@ -20,14 +22,18 @@ class BaseDataSets(Dataset): - def __init__(self, base_dir=None, labeled_type="labeled", labeled_ratio=10, split='train', transform=None, fold=1): + def __init__(self, base_dir=None, labeled_type="labeled", labeled_ratio=10, split='train', transform=None, fold=1, cross_val=True): self._base_dir = base_dir self.sample_list = [] self.split = split self.transform = transform self.labeled_type = labeled_type self.all_volumes = sorted(os.listdir(self._base_dir + "/all_volumes")) - train_ids, test_ids = self._get_fold_ids(fold) + if cross_val: # 5-fold cross validation + train_ids, val_ids = self._get_fold_ids(fold) + else: + train_ids, val_ids, _ = self._get_split_ids() + train_ids = sorted(train_ids) all_labeled_ids = train_ids[::labeled_ratio] if self.split == 'train': self.all_slices = os.listdir(self._base_dir + "/all_slices") @@ -50,11 +56,9 @@ def __init__(self, base_dir=None, labeled_type="labeled", labeled_ratio=10, spli print("total unlabeled {} samples".format(len(self.sample_list))) elif self.split == 'val': - print("test_ids", test_ids) - self.all_volumes = os.listdir( - self._base_dir + "/all_volumes") + print("val_ids", val_ids) self.sample_list = [] - for ids in test_ids: + for ids in val_ids: new_data_list = list(filter(lambda x: re.match( '{}.*'.format(ids.replace(".h5", "")), x) != None, self.all_volumes)) self.sample_list.extend(new_data_list) @@ -65,7 +69,18 @@ def _get_fold_ids(self, fold): k_fold_data = [] for trn_idx, val_idx in folds.split(all_cases): k_fold_data.append([all_cases[trn_idx], all_cases[val_idx]]) - return k_fold_data[fold][0], k_fold_data[fold][1] + train_set = k_fold_data[fold][0] + test_set = k_fold_data[fold][1] + return train_set, test_set + + def _get_split_ids(self): + all_cases = np.array(self.all_volumes) + rest_set, test_set = train_test_split(all_cases, test_size=int( + len(self.all_volumes)*0.2), shuffle=True, random_state=1234) + train_set, val_set = train_test_split(rest_set, test_size=int( + len(self.all_volumes)*0.1), shuffle=True, random_state=1234) + print("test_set", sorted(test_set)) + return train_set, val_set, test_set def __len__(self): return len(self.sample_list) @@ -168,7 +183,17 @@ def nonlinear_transformation(x, label, prob=0.5): return nonlinear_x, label -class RandomGenerator(object): +def random_rescale_intensity(image, label): + image = exposure.rescale_intensity(image) + return image, label + + +def random_equalize_hist(image, label): + image = exposure.equalize_hist(image) + return image, label + + +class RandomGenerator_Strong_Weak(object): def __init__(self, output_size): self.output_size = output_size @@ -180,8 +205,50 @@ def __call__(self, sample): image, label = random_rotate(image, label, cval=0) if random.random() > 0.5: image, label = random_noise(image, label) + + x, y = image.shape + image_w = copy.deepcopy(image) + image_w = zoom( + image_w, (self.output_size[0] / x, self.output_size[1] / y), order=0) + + if random.random() > 0.33: + image, label = nonlinear_transformation(image, label) + elif random.random() < 0.66 and random.random() > 0.33: + image, label = random_rescale_intensity(image, label) + else: + image, label = random_equalize_hist(image, label) + image_s = image + image_s = zoom( + image_s, (self.output_size[0] / x, self.output_size[1] / y), order=0) + label = zoom( + label, (self.output_size[0] / x, self.output_size[1] / y), order=0) + image_w = torch.from_numpy( + image_w.astype(np.float32)).unsqueeze(0) + image_s = torch.from_numpy( + image_s.astype(np.float32)).unsqueeze(0) + label = torch.from_numpy(label.astype(np.int16)) + sample = {'image_w': image_w, 'image_s': image_s, 'label': label} + return sample + + +class RandomGenerator(object): + def __init__(self, output_size): + self.output_size = output_size + + def __call__(self, sample): + image, label = sample['image'], sample['label'] + if random.random() > 0.5: + image, label = random_flip(image, label) + if random.random() > 0.5: + image, label = random_rotate(image, label, cval=0) if random.random() > 0.5: + image, label = random_noise(image, label) + if random.random() > 0.33: image, label = nonlinear_transformation(image, label) + elif random.random() < 0.66 and random.random() > 0.33: + image, label = random_rescale_intensity(image, label) + elif random.random() > 0.66: + image, label = random_equalize_hist(image, label) x, y = image.shape image = zoom( image, (self.output_size[0] / x, self.output_size[1] / y), order=0) diff --git a/code/networks/unet.py b/code/networks/unet.py index 102d821..f691230 100755 --- a/code/networks/unet.py +++ b/code/networks/unet.py @@ -291,7 +291,7 @@ def __init__(self, in_chns, class_num): 'feature_chns': [16, 32, 64, 128, 256], 'dropout': [0.05, 0.1, 0.2, 0.3, 0.5], 'class_num': class_num, - 'bilinear': True, + 'bilinear': False, 'acti_func': 'relu'} self.encoder = Encoder(params) @@ -311,7 +311,7 @@ def __init__(self, in_chns, class_num): 'feature_chns': [16, 32, 64, 128, 256], 'dropout': [0.05, 0.1, 0.2, 0.3, 0.5], 'class_num': class_num, - 'bilinear': False, + 'bilinear': True, 'acti_func': 'relu'} self.encoder = Encoder(params) self.decoder = Decoder_DS(params) diff --git a/code/test.py b/code/test.py index 71170ea..ac49f1f 100644 --- a/code/test.py +++ b/code/test.py @@ -12,39 +12,51 @@ from medpy import metric from scipy.ndimage import zoom from scipy.ndimage.interpolation import zoom -from sklearn.model_selection import KFold +from sklearn.model_selection import KFold, train_test_split from tqdm import tqdm from networks.net_factory import net_factory parser = argparse.ArgumentParser() parser.add_argument('--root_path', type=str, - default='../data/ProstateX', help='Name of Experiment') + default='../data/ACDC', help='Name of Experiment') parser.add_argument('--exp', type=str, - default='ProstateX/Mean_Teacher', help='experiment_name') + default='ACDC/FullSup', help='experiment_name') parser.add_argument('--model', type=str, default='unet', help='model_name') -parser.add_argument('--labeled_ratio', type=int, default=8, +parser.add_argument('--best_model', type=int, + default=1, help='the best or latest checkpoints') +parser.add_argument('--labeled_ratio', type=int, default=10, help='1/labeled_ratio data is provided mask') parser.add_argument('--fold', type=int, default=3, help='fold') +parser.add_argument('--cross_val', type=int, + default=0, help='5-fold cross validation or random split 7/1/2 for training/validation/testing') parser.add_argument('--patch_size', type=list, default=[256, 256], help='patch size of network input') -parser.add_argument('--num_classes', type=int, default=3, +parser.add_argument('--num_classes', type=int, default=4, help='output channel of network') parser.add_argument('--sup_type', type=str, default="label", help='label') def get_fold_ids(FLAGS): - all_volumes = sorted(os.listdir(FLAGS.root_path + "/all_volumes")) + all_volumes = os.listdir(FLAGS.root_path + "/all_volumes") folds = KFold(n_splits=5, shuffle=False) all_cases = np.array(all_volumes) k_fold_data = [] for trn_idx, val_idx in folds.split(all_cases): k_fold_data.append([all_cases[trn_idx], all_cases[val_idx]]) - return k_fold_data[FLAGS.fold][0], k_fold_data[FLAGS.fold][1] + return sorted(k_fold_data[FLAGS.fold][0]), sorted(k_fold_data[FLAGS.fold][1]) + +def get_split_ids(FLAGS): + all_cases = sorted(np.array(os.listdir(FLAGS.root_path + "/all_volumes"))) + rest_set, test_set = train_test_split(all_cases, test_size=int( + len(all_cases)*0.2), shuffle=True, random_state=1234) + train_set, val_set = train_test_split(rest_set, test_size=int( + len(all_cases)*0.1), shuffle=True, random_state=1234) + return sorted(train_set), sorted(val_set), sorted(test_set) def calculate_metric_percase(pred, gt, spacing): if pred.sum() > 0 and gt.sum() > 0: @@ -161,27 +173,39 @@ def test_single_volume(case, net, test_save_path, FLAGS, batch_size=12): def Inference(FLAGS): - train_ids, test_ids = get_fold_ids(FLAGS) - all_volumes = os.listdir( - FLAGS.root_path + "/all_volumes") + if FLAGS.cross_val: + _, test_ids = get_fold_ids(FLAGS) + else: + _, _, test_ids = get_split_ids(FLAGS) + print(test_ids) + all_volumes = sorted(os.listdir(FLAGS.root_path + "/all_volumes")) image_list = [] for ids in test_ids: new_data_list = list(filter(lambda x: re.match( '{}.*'.format(ids), x) != None, all_volumes)) image_list.extend(new_data_list) - snapshot_path = snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( - FLAGS.exp, FLAGS.labeled_ratio, FLAGS.fold) - test_save_path = "../model/{}/1_of_{}_labeled/fold{}/prediction/".format( - FLAGS.exp, FLAGS.labeled_ratio, FLAGS.fold) + + if FLAGS.cross_val: + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}/{}".format( + FLAGS.exp, FLAGS.labeled_ratio, FLAGS.fold, FLAGS.model) + test_save_path = "../model/{}/1_of_{}_labeled/fold{}/{}/prediction/".format( + FLAGS.exp, FLAGS.labeled_ratio, FLAGS.fold, FLAGS.model) + + else: + snapshot_path = "../model/{}/1_of_{}_labeled/{}".format( + FLAGS.exp, FLAGS.labeled_ratio, FLAGS.model) + test_save_path = "../model/{}/1_of_{}_labeled/{}/prediction/".format( + FLAGS.exp, FLAGS.labeled_ratio, FLAGS.model) + if os.path.exists(test_save_path): shutil.rmtree(test_save_path) os.makedirs(test_save_path) net = net_factory(net_type=FLAGS.model, in_chns=1, class_num=FLAGS.num_classes) - save_mode_path = os.path.join( - snapshot_path, '{}_best_model.pth'.format(FLAGS.model)) - # save_mode_path = os.path.join( - # snapshot_path, 'iter_60000.pth') + if FLAGS.best_model: + save_mode_path = os.path.join(snapshot_path, '{}_best_model.pth'.format(FLAGS.model)) + else: + save_mode_path = os.path.join(snapshot_path, '{}_latest_model.pth'.format(FLAGS.model)) net.load_state_dict(torch.load(save_mode_path)) print("init weight from {}".format(save_mode_path)) net.eval() @@ -193,19 +217,15 @@ def Inference(FLAGS): case, net, test_save_path, FLAGS) print(cases_metric) metric_array[ind, ...] = cases_metric - np.save("../model/{}/1_of_{}_labeled/fold{}/prediction/Results.npy".format( - FLAGS.exp, FLAGS.labeled_ratio, FLAGS.fold), metric_array) + np.save(test_save_path+"/Results.npy", metric_array) return metric_array if __name__ == '__main__': FLAGS = parser.parse_args() - total = 0.0 - for i in [3]: - FLAGS.fold = i - print("Inference fold{}".format(i)) - metric_array = Inference(FLAGS) - print("mean class results:", np.mean(metric_array, axis=0)) - print("mean case results:", np.mean(metric_array, axis=0).mean(axis=0)) - print(total/1) + print("Inference fold{}".format(FLAGS.fold)) + metric_array = Inference(FLAGS) + print("mean class results:", np.mean(metric_array, axis=0)) + print("mean case results:", np.mean(metric_array, axis=0).mean(axis=0)) + \ No newline at end of file diff --git a/code/train_acdc_ent_mini.sh b/code/train_acdc_ent_mini.sh new file mode 100644 index 0000000..d09df77 --- /dev/null +++ b/code/train_acdc_ent_mini.sh @@ -0,0 +1,10 @@ +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_entropy_minimization.py --root_path ../data/ACDC --exp ACDC/EntMini --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 1 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_entropy_minimization.py --root_path ../data/ACDC --exp ACDC/EntMini --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 2 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_entropy_minimization.py --root_path ../data/ACDC --exp ACDC/EntMini --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 3 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_entropy_minimization.py --root_path ../data/ACDC --exp ACDC/EntMini --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_entropy_minimization.py --root_path ../data/ACDC --exp ACDC/EntMini --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 0 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_entropy_minimization.py --root_path ../data/ACDC --exp ACDC/EntMini --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 1 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_entropy_minimization.py --root_path ../data/ACDC --exp ACDC/EntMini --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 2 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_entropy_minimization.py --root_path ../data/ACDC --exp ACDC/EntMini --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 3 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_entropy_minimization.py --root_path ../data/ACDC --exp ACDC/EntMini --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_entropy_minimization.py --root_path ../data/ACDC --exp ACDC/EntMini --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 0 \ No newline at end of file diff --git a/code/train_acdc_ict_mini.sh b/code/train_acdc_ict_mini.sh new file mode 100644 index 0000000..fbb99bc --- /dev/null +++ b/code/train_acdc_ict_mini.sh @@ -0,0 +1,10 @@ +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_interpolation_consistency_training.py --root_path ../data/ACDC --exp ACDC/ICT --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 1 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_interpolation_consistency_training.py --root_path ../data/ACDC --exp ACDC/ICT --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 2 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_interpolation_consistency_training.py --root_path ../data/ACDC --exp ACDC/ICT --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 3 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_interpolation_consistency_training.py --root_path ../data/ACDC --exp ACDC/ICT --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_interpolation_consistency_training.py --root_path ../data/ACDC --exp ACDC/ICT --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 8 --fold 0 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_interpolation_consistency_training.py --root_path ../data/ACDC --exp ACDC/ICT --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 1 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_interpolation_consistency_training.py --root_path ../data/ACDC --exp ACDC/ICT --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 2 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_interpolation_consistency_training.py --root_path ../data/ACDC --exp ACDC/ICT --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 3 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_interpolation_consistency_training.py --root_path ../data/ACDC --exp ACDC/ICT --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 4 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_interpolation_consistency_training.py --root_path ../data/ACDC --exp ACDC/ICT --batch_size 16 --num_classes 4 --max_iterations 60000 --base_lr 0.03 --labeled_ratio 16 --fold 0 \ No newline at end of file diff --git a/code/train_cross_pseudo_supervision.py b/code/train_cross_pseudo_supervision.py index ca1d085..0aeca14 100644 --- a/code/train_cross_pseudo_supervision.py +++ b/code/train_cross_pseudo_supervision.py @@ -29,18 +29,19 @@ parser = argparse.ArgumentParser() parser.add_argument('--root_path', type=str, - default='../data/ProstateX', help='Name of Experiment') + default='../data/ACDC', help='Name of Experiment') parser.add_argument('--exp', type=str, - default='ProstateX/CPS', help='experiment_name') + default='ACDC/CPS', help='experiment_name') parser.add_argument('--model', type=str, default='unet', help='model_name') parser.add_argument('--fold', type=int, default=1, help='cross validation') parser.add_argument('--max_iterations', type=int, default=30000, help='maximum epoch number to train') -parser.add_argument('--batch_size', type=int, default=16, +parser.add_argument('--batch_size', type=int, default=12, help='batch_size per gpu') - +parser.add_argument('--cross_val', type=int, + default=0, help='5-fold cross validation or random split 7/1/2 for training/validation/testing') parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training') parser.add_argument('--base_lr', type=float, default=0.03, @@ -48,11 +49,11 @@ parser.add_argument('--patch_size', type=list, default=[256, 256], help='patch size of network input') parser.add_argument('--seed', type=int, default=2022, help='random seed') -parser.add_argument('--num_classes', type=int, default=3, +parser.add_argument('--num_classes', type=int, default=4, help='output channel of network') # label and unlabel -parser.add_argument('--labeled_ratio', type=int, default=8, +parser.add_argument('--labeled_ratio', type=int, default=10, help='1/labeled_ratio data is provided mask') # costs parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') @@ -82,9 +83,9 @@ def train(args, snapshot_path): class_num=num_classes) db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ - RandomGenerator(args.patch_size)])) + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ - RandomGenerator(args.patch_size)])) + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) logging.info("Labeled slices: {} ".format(len(db_train_labeled))) logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) @@ -94,7 +95,7 @@ def train(args, snapshot_path): db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, - split="val", labeled_ratio=args.labeled_ratio) + split="val", labeled_ratio=args.labeled_ratio, cross_val=args.cross_val) valloader = DataLoader(db_val, batch_size=1) model1.train() @@ -278,6 +279,14 @@ def train(args, snapshot_path): if iter_num >= max_iterations: break + + save_latest = os.path.join( + snapshot_path, '{}_latest_model1.pth'.format(args.model)) + torch.save(model1.state_dict(), save_latest) + save_latest = os.path.join( + snapshot_path, '{}_latest_model2.pth'.format(args.model)) + torch.save(model2.state_dict(), save_latest) + if iter_num >= max_iterations: iterator.close() break @@ -298,8 +307,12 @@ def train(args, snapshot_path): torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) - snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( - args.exp, args.labeled_ratio, args.fold) + if args.cross_val: + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}/{}".format( + args.exp, args.labeled_ratio, args.fold, args.model) + else: + snapshot_path = "../model/{}/1_of_{}_labeled/{}".format( + args.exp, args.labeled_ratio, args.model) if not os.path.exists(snapshot_path): os.makedirs(snapshot_path) if os.path.exists(snapshot_path + '/code'): diff --git a/code/train_deep_adversarial_network.py b/code/train_deep_adversarial_network.py index c99fd54..4d5bab4 100644 --- a/code/train_deep_adversarial_network.py +++ b/code/train_deep_adversarial_network.py @@ -6,6 +6,7 @@ import sys import time from itertools import cycle + import numpy as np import torch import torch.backends.cudnn as cudnn @@ -39,7 +40,8 @@ default=30000, help='maximum epoch number to train') parser.add_argument('--batch_size', type=int, default=16, help='batch_size per gpu') - +parser.add_argument('--cross_val', type=bool, + default=True, help='5-fold cross validation or random split 7/1/2 for training/validation/testing') parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training') parser.add_argument('--base_lr', type=float, default=0.03, @@ -81,9 +83,9 @@ def train(args, snapshot_path): DAN = DAN.cuda() db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ - RandomGenerator(args.patch_size)])) + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ - RandomGenerator(args.patch_size)])) + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) logging.info("Labeled slices: {} ".format(len(db_train_labeled))) logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) @@ -93,7 +95,7 @@ def train(args, snapshot_path): db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, - split="val", labeled_ratio=args.labeled_ratio) + split="val", labeled_ratio=args.labeled_ratio, cross_val=args.cross_val) valloader = DataLoader(db_val, batch_size=1) model.train() @@ -235,6 +237,11 @@ def train(args, snapshot_path): if iter_num >= max_iterations: break + + save_latest = os.path.join( + snapshot_path, '{}_latest_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_latest) + if iter_num >= max_iterations: iterator.close() break @@ -255,8 +262,12 @@ def train(args, snapshot_path): torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) - snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( - args.exp, args.labeled_ratio, args.fold) + if args.cross_val: + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}/{}".format( + args.exp, args.labeled_ratio, args.fold, args.model) + else: + snapshot_path = "../model/{}/1_of_{}_labeled/{}".format( + args.exp, args.labeled_ratio, args.model) if not os.path.exists(snapshot_path): os.makedirs(snapshot_path) if os.path.exists(snapshot_path + '/code'): diff --git a/code/train_deep_co_training.py b/code/train_deep_co_training.py new file mode 100644 index 0000000..6926a6f --- /dev/null +++ b/code/train_deep_co_training.py @@ -0,0 +1,324 @@ +import argparse +import logging +import os +import random +import shutil +import sys +import time +from itertools import cycle + +import numpy as np +import torch +import torch.backends.cudnn as cudnn +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from tensorboardX import SummaryWriter +from torch.nn import BCEWithLogitsLoss +from torch.nn.modules.loss import CrossEntropyLoss +from torch.utils.data import DataLoader +from torchvision import transforms +from torchvision.utils import make_grid +from tqdm import tqdm + +from dataloaders.dataset import BaseDataSets, RandomGenerator +from networks.discriminator import FCDiscriminator +from networks.net_factory import net_factory +from utils import losses, metrics, ramps +from val_2D import test_single_volume + +parser = argparse.ArgumentParser() +parser.add_argument('--root_path', type=str, + default='../data/ProstateX', help='Name of Experiment') +parser.add_argument('--exp', type=str, + default='ProstateX/DCT', help='experiment_name') +parser.add_argument('--model', type=str, + default='unet', help='model_name') +parser.add_argument('--fold', type=int, + default=1, help='cross validation') +parser.add_argument('--max_iterations', type=int, + default=30000, help='maximum epoch number to train') +parser.add_argument('--batch_size', type=int, default=16, + help='batch_size per gpu') +parser.add_argument('--cross_val', type=bool, + default=True, help='5-fold cross validation or random split 7/1/2 for training/validation/testing') +parser.add_argument('--deterministic', type=int, default=1, + help='whether use deterministic training') +parser.add_argument('--base_lr', type=float, default=0.03, + help='segmentation network learning rate') +parser.add_argument('--patch_size', type=list, default=[256, 256], + help='patch size of network input') +parser.add_argument('--seed', type=int, default=2022, help='random seed') +parser.add_argument('--num_classes', type=int, default=3, + help='output channel of network') + +# label and unlabel +parser.add_argument('--labeled_ratio', type=int, default=8, + help='1/labeled_ratio data is provided mask') +# costs +parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') +parser.add_argument('--consistency_type', type=str, + default="mse", help='consistency_type') +parser.add_argument('--consistency', type=float, + default=0.1, help='consistency') +parser.add_argument('--consistency_rampup', type=float, + default=200.0, help='consistency_rampup') +args = parser.parse_args() + + +def get_current_consistency_weight(epoch): + # Consistency ramp-up from https://arxiv.org/abs/1610.02242 + return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) + + +def train(args, snapshot_path): + writer = SummaryWriter(snapshot_path + '/log') + base_lr = args.base_lr + num_classes = args.num_classes + max_iterations = args.max_iterations + + model1 = net_factory(net_type=args.model, in_chns=1, + class_num=num_classes) + model2 = net_factory(net_type=args.model, in_chns=1, + class_num=num_classes) + + db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) + db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) + logging.info("Labeled slices: {} ".format(len(db_train_labeled))) + logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) + + trainloader_labeled = DataLoader( + db_train_labeled, batch_size=args.batch_size//2, shuffle=True) + trainloader_unlabeled = DataLoader( + db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) + + db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, + split="val", labeled_ratio=args.labeled_ratio, cross_val=args.cross_val) + valloader = DataLoader(db_val, batch_size=1) + + model1.train() + model2.train() + + optimizer1 = optim.SGD(model1.parameters(), lr=base_lr, + momentum=0.9, weight_decay=0.0001) + optimizer2 = optim.SGD(model2.parameters(), lr=base_lr, + momentum=0.9, weight_decay=0.0001) + ce_loss = CrossEntropyLoss() + dice_loss = losses.DiceLoss(num_classes) + + logging.info("{} iterations per epoch".format(len(trainloader_unlabeled))) + + iter_num = 0 + max_epoch = max_iterations // len(trainloader_unlabeled) + 1 + best_performance1 = 0.0 + best_performance2 = 0.0 + iterator = tqdm(range(max_epoch), ncols=70) + for epoch_num in iterator: + for i, (sampled_batch_labeled, sampled_batch_unlabeled) in enumerate(zip(cycle(trainloader_labeled), trainloader_unlabeled)): + volume_batch, label_batch = sampled_batch_labeled['image'], sampled_batch_labeled['label'] + volume_batch, label_batch = volume_batch.cuda(), label_batch.cuda() + unlabeled_volume_batch = sampled_batch_unlabeled['image'].cuda() + + outputs1 = model1(volume_batch) + outputs1_soft = torch.softmax(outputs1, dim=1) + + outputs1_unlabeled = model1(unlabeled_volume_batch) + outputs1_unlabeled_soft = torch.softmax(outputs1_unlabeled, dim=1) + + outputs2 = model2(volume_batch) + outputs2_soft = torch.softmax(outputs2, dim=1) + + outputs2_unlabeled = model2(unlabeled_volume_batch) + outputs2_unlabeled_soft = torch.softmax(outputs2_unlabeled, dim=1) + + supervised_loss1 = 0.5 * \ + (ce_loss(outputs1, label_batch[:].long( + )) + dice_loss(outputs1_soft, label_batch[:].unsqueeze(1))) + supervised_loss2 = 0.5 * \ + (ce_loss(outputs2, label_batch[:].long( + )) + dice_loss(outputs2_soft, label_batch[:].unsqueeze(1))) + + pseudo_outputs1 = torch.argmax( + outputs1_unlabeled_soft.detach(), dim=1, keepdim=False) + pseudo_outputs2 = torch.argmax( + outputs2_unlabeled_soft.detach(), dim=1, keepdim=False) + + pseudo_supervision1 = ce_loss(outputs1_unlabeled, pseudo_outputs2) + pseudo_supervision2 = ce_loss(outputs2_unlabeled, pseudo_outputs1) + + consistency_weight = get_current_consistency_weight( + iter_num // (args.max_iterations/args.consistency_rampup)) + + model1_loss = supervised_loss1 + consistency_weight * pseudo_supervision1 + model2_loss = supervised_loss2 + consistency_weight * pseudo_supervision2 + + loss = model1_loss + model2_loss + + optimizer1.zero_grad() + optimizer2.zero_grad() + + loss.backward() + + optimizer1.step() + optimizer2.step() + + iter_num = iter_num + 1 + + lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 + for param_group in optimizer1.param_groups: + param_group['lr'] = lr_ + for param_group in optimizer2.param_groups: + param_group['lr'] = lr_ + + writer.add_scalar('lr', lr_, iter_num) + writer.add_scalar( + 'consistency_weight/consistency_weight', consistency_weight, iter_num) + writer.add_scalar('loss/model1_loss', + model1_loss, iter_num) + writer.add_scalar('loss/model2_loss', + model2_loss, iter_num) + logging.info('iteration %d : model1 loss : %f model2 loss : %f' % ( + iter_num, model1_loss.item(), model2_loss.item())) + if iter_num % 20 == 0: + image = volume_batch[0, 0:1, :, :] + writer.add_image('train/Image', image, iter_num) + outputs = torch.argmax(torch.softmax( + outputs1, dim=1), dim=1, keepdim=True) + writer.add_image('train/model1_Prediction', + outputs[0, ...] * 50, iter_num) + outputs = torch.argmax(torch.softmax( + outputs2, dim=1), dim=1, keepdim=True) + writer.add_image('train/model2_Prediction', + outputs[0, ...] * 50, iter_num) + labs = label_batch[0, ...].unsqueeze(0) * 50 + writer.add_image('train/GroundTruth', labs, iter_num) + + if iter_num > 0 and iter_num % 200 == 0: + model1.eval() + metric_list = 0.0 + for i_batch, sampled_batch in enumerate(valloader): + metric_i = test_single_volume( + sampled_batch["image"], sampled_batch["label"], model1, classes=num_classes) + metric_list += np.array(metric_i) + metric_list = metric_list / len(db_val) + for class_i in range(num_classes-1): + writer.add_scalar('info/model1_val_{}_dice'.format(class_i+1), + metric_list[class_i, 0], iter_num) + writer.add_scalar('info/model1_val_{}_hd95'.format(class_i+1), + metric_list[class_i, 1], iter_num) + + performance1 = np.mean(metric_list, axis=0)[0] + + mean_hd951 = np.mean(metric_list, axis=0)[1] + writer.add_scalar('info/model1_val_mean_dice', + performance1, iter_num) + writer.add_scalar('info/model1_val_mean_hd95', + mean_hd951, iter_num) + + if performance1 > best_performance1: + best_performance1 = performance1 + save_mode_path = os.path.join(snapshot_path, + 'model1_iter_{}_dice_{}.pth'.format( + iter_num, round(best_performance1, 4))) + save_best = os.path.join(snapshot_path, + '{}_best_model1.pth'.format(args.model)) + torch.save(model1.state_dict(), save_mode_path) + torch.save(model1.state_dict(), save_best) + + logging.info( + 'iteration %d : model1_mean_dice : %f model1_mean_hd95 : %f' % (iter_num, performance1, mean_hd951)) + model1.train() + + model2.eval() + metric_list = 0.0 + for i_batch, sampled_batch in enumerate(valloader): + metric_i = test_single_volume( + sampled_batch["image"], sampled_batch["label"], model2, classes=num_classes) + metric_list += np.array(metric_i) + metric_list = metric_list / len(db_val) + for class_i in range(num_classes-1): + writer.add_scalar('info/model2_val_{}_dice'.format(class_i+1), + metric_list[class_i, 0], iter_num) + writer.add_scalar('info/model2_val_{}_hd95'.format(class_i+1), + metric_list[class_i, 1], iter_num) + + performance2 = np.mean(metric_list, axis=0)[0] + + mean_hd952 = np.mean(metric_list, axis=0)[1] + writer.add_scalar('info/model2_val_mean_dice', + performance2, iter_num) + writer.add_scalar('info/model2_val_mean_hd95', + mean_hd952, iter_num) + + if performance2 > best_performance2: + best_performance2 = performance2 + save_mode_path = os.path.join(snapshot_path, + 'model2_iter_{}_dice_{}.pth'.format( + iter_num, round(best_performance2, 4))) + save_best = os.path.join(snapshot_path, + '{}_best_model2.pth'.format(args.model)) + torch.save(model2.state_dict(), save_mode_path) + torch.save(model2.state_dict(), save_best) + + logging.info( + 'iteration %d : model2_mean_dice : %f model2_mean_hd95 : %f' % (iter_num, performance2, mean_hd952)) + model2.train() + + if iter_num % 3000 == 0: + save_mode_path = os.path.join( + snapshot_path, 'model1_iter_' + str(iter_num) + '.pth') + torch.save(model1.state_dict(), save_mode_path) + logging.info("save model1 to {}".format(save_mode_path)) + + save_mode_path = os.path.join( + snapshot_path, 'model2_iter_' + str(iter_num) + '.pth') + torch.save(model2.state_dict(), save_mode_path) + logging.info("save model2 to {}".format(save_mode_path)) + + if iter_num >= max_iterations: + break + save_latest = os.path.join( + snapshot_path, '{}_latest_model.pth'.format(args.model)) + torch.save(model1.state_dict(), save_latest) + + if iter_num >= max_iterations: + iterator.close() + break + writer.close() + return "Training Finished!" + + +if __name__ == "__main__": + if not args.deterministic: + cudnn.benchmark = True + cudnn.deterministic = False + else: + cudnn.benchmark = False + cudnn.deterministic = True + + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + if args.cross_val: + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}/{}".format( + args.exp, args.labeled_ratio, args.fold, args.model) + else: + snapshot_path = "../model/{}/1_of_{}_labeled/{}".format( + args.exp, args.labeled_ratio, args.model) + + if not os.path.exists(snapshot_path): + os.makedirs(snapshot_path) + if os.path.exists(snapshot_path + '/code'): + shutil.rmtree(snapshot_path + '/code') + shutil.copytree('.', snapshot_path + '/code', + shutil.ignore_patterns(['.git', '__pycache__'])) + + logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO, + format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') + logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) + logging.info(str(args)) + train(args, snapshot_path) diff --git a/code/train_entropy_minimization.py b/code/train_entropy_minimization.py index a1eb291..6a1e40d 100644 --- a/code/train_entropy_minimization.py +++ b/code/train_entropy_minimization.py @@ -6,6 +6,7 @@ import sys import time from itertools import cycle + import numpy as np import torch import torch.backends.cudnn as cudnn @@ -39,7 +40,8 @@ default=30000, help='maximum epoch number to train') parser.add_argument('--batch_size', type=int, default=16, help='batch_size per gpu') - +parser.add_argument('--cross_val', type=bool, + default=True, help='5-fold cross validation or random split 7/1/2 for training/validation/testing') parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training') parser.add_argument('--base_lr', type=float, default=0.03, @@ -79,9 +81,9 @@ def train(args, snapshot_path): class_num=num_classes) db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ - RandomGenerator(args.patch_size)])) + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ - RandomGenerator(args.patch_size)])) + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) logging.info("Labeled slices: {} ".format(len(db_train_labeled))) logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) @@ -91,7 +93,7 @@ def train(args, snapshot_path): db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, - split="val", labeled_ratio=args.labeled_ratio) + split="val", labeled_ratio=args.labeled_ratio, cross_val=args.cross_val) valloader = DataLoader(db_val, batch_size=1) model.train() @@ -202,6 +204,11 @@ def train(args, snapshot_path): if iter_num >= max_iterations: break + + save_latest = os.path.join( + snapshot_path, '{}_latest_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_latest) + if iter_num >= max_iterations: iterator.close() break @@ -222,8 +229,13 @@ def train(args, snapshot_path): torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) - snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( - args.exp, args.labeled_ratio, args.fold) + if args.cross_val: + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}/{}".format( + args.exp, args.labeled_ratio, args.fold, args.model) + else: + snapshot_path = "../model/{}/1_of_{}_labeled/{}".format( + args.exp, args.labeled_ratio, args.model) + if not os.path.exists(snapshot_path): os.makedirs(snapshot_path) if os.path.exists(snapshot_path + '/code'): diff --git a/code/train_exp1.sh b/code/train_exp1.sh new file mode 100644 index 0000000..b815077 --- /dev/null +++ b/code/train_exp1.sh @@ -0,0 +1,11 @@ +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_cross_pseudo_supervision.py --root_path ../data/ACDC --exp ACDC/CPS --batch_size 16 --num_classes 4 --max_iterations 30000 --base_lr 0.03 --labeled_ratio 10 --cross_val 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_cross_pseudo_supervision.py --root_path ../data/ACDC --exp ACDC/CPS --batch_size 16 --num_classes 4 --max_iterations 30000 --base_lr 0.03 --labeled_ratio 20 --cross_val 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_cross_pseudo_supervision.py --root_path ../data/ACDC --exp ACDC/CPS --batch_size 16 --num_classes 4 --max_iterations 30000 --base_lr 0.03 --labeled_ratio 40 --cross_val 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 30000 --base_lr 0.03 --labeled_ratio 10 --cross_val 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 30000 --base_lr 0.03 --labeled_ratio 20 --cross_val 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 30000 --base_lr 0.03 --labeled_ratio 40 --cross_val 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fixmatch.py --root_path ../data/ACDC --exp ACDC/FixMatch --batch_size 16 --num_classes 4 --max_iterations 30000 --base_lr 0.03 --labeled_ratio 10 --cross_val 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fixmatch.py --root_path ../data/ACDC --exp ACDC/FixMatch --batch_size 16 --num_classes 4 --max_iterations 30000 --base_lr 0.03 --labeled_ratio 20 --cross_val 0 & +# srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fixmatch.py --root_path ../data/ACDC --exp ACDC/FixMatch --batch_size 16 --num_classes 4 --max_iterations 30000 --base_lr 0.03 --labeled_ratio 40 --cross_val 0 +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fully_supervised.py --root_path ../data/ACDC --exp ACDC/FullSup --batch_size 16 --num_classes 4 --max_iterations 30000 --base_lr 0.03 --labeled_ratio 1 --cross_val 0 & +srun -p MIA -n 1 -c 4 --mpi=pmi2 --gres=gpu:1 python -u train_fixmatch.py --root_path ../data/ACDC --exp ACDC/FixMatch --batch_size 16 --num_classes 4 --max_iterations 30000 --base_lr 0.03 --labeled_ratio 20 --cross_val 0 \ No newline at end of file diff --git a/code/train_fixmatch.py b/code/train_fixmatch.py new file mode 100644 index 0000000..cf1ce8d --- /dev/null +++ b/code/train_fixmatch.py @@ -0,0 +1,281 @@ +import argparse +import logging +import os +import random +import shutil +import sys +import time +from itertools import cycle + +import numpy as np +import torch +import torch.backends.cudnn as cudnn +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from tensorboardX import SummaryWriter +from torch.nn import BCEWithLogitsLoss +from torch.nn.modules.loss import CrossEntropyLoss +from torch.utils.data import DataLoader +from torchvision import transforms +from torchvision.utils import make_grid +from tqdm import tqdm + +from dataloaders.dataset import BaseDataSets, RandomGenerator_Strong_Weak +from networks.discriminator import FCDiscriminator +from networks.net_factory import net_factory +from utils import losses, metrics, ramps +from val_2D import test_single_volume + +parser = argparse.ArgumentParser() +parser.add_argument('--root_path', type=str, + default='../data/ACDC', help='Name of Experiment') +parser.add_argument('--exp', type=str, + default='ACDC/FixMatch', help='experiment_name') +parser.add_argument('--model', type=str, + default='unet', help='model_name') +parser.add_argument('--fold', type=int, + default=3, help='cross validation') +parser.add_argument('--max_iterations', type=int, + default=30000, help='maximum epoch number to train') +parser.add_argument('--batch_size', type=int, default=12, + help='batch_size per gpu') +parser.add_argument('--cross_val', type=int, + default=0, help='5-fold cross validation or random split 7/1/2 for training/validation/testing') +parser.add_argument('--deterministic', type=int, default=1, + help='whether use deterministic training') +parser.add_argument('--base_lr', type=float, default=0.03, + help='segmentation network learning rate') +parser.add_argument('--patch_size', type=list, default=[256, 256], + help='patch size of network input') +parser.add_argument('--seed', type=int, default=2022, help='random seed') +parser.add_argument('--num_classes', type=int, default=4, + help='output channel of network') + +# label and unlabel +parser.add_argument('--labeled_ratio', type=int, default=10, + help='1/labeled_ratio data is provided mask') +# costs +parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') +parser.add_argument('--consistency_type', type=str, + default="mse", help='consistency_type') +parser.add_argument('--consistency', type=float, + default=0.1, help='consistency') +parser.add_argument('--consistency_rampup', type=float, + default=200.0, help='consistency_rampup') +args = parser.parse_args() + + +def get_current_consistency_weight(epoch): + # Consistency ramp-up from https://arxiv.org/abs/1610.02242 + return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup) + + +def update_ema_variables(model, ema_model, alpha, global_step): + # Use the true average until the exponential average is more correct + alpha = min(1 - 1 / (global_step + 1), alpha) + for ema_param, param in zip(ema_model.parameters(), model.parameters()): + ema_param.data.mul_(alpha).add_(1 - alpha, param.data) + + +def train(args, snapshot_path): + writer = SummaryWriter(snapshot_path + '/log') + base_lr = args.base_lr + num_classes = args.num_classes + max_iterations = args.max_iterations + + def worker_init_fn(worker_id): + random.seed(args.seed + worker_id) + + def create_model(ema=False): + # Network definition + model = net_factory(net_type=args.model, in_chns=1, + class_num=num_classes) + if ema: + for param in model.parameters(): + param.detach_() + return model + + model = create_model() + ema_model = create_model(ema=True) + + db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator_Strong_Weak(args.patch_size)]), cross_val=args.cross_val) + db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ + RandomGenerator_Strong_Weak(args.patch_size)]), cross_val=args.cross_val) + logging.info("Labeled slices: {} ".format(len(db_train_labeled))) + logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) + + trainloader_labeled = DataLoader( + db_train_labeled, batch_size=args.batch_size//2, shuffle=True) + trainloader_unlabeled = DataLoader( + db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) + + db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, + split="val", labeled_ratio=args.labeled_ratio, cross_val=args.cross_val) + valloader = DataLoader(db_val, batch_size=1) + + model.train() + + optimizer = optim.SGD(model.parameters(), lr=base_lr, + momentum=0.9, weight_decay=0.0001) + + ce_loss = CrossEntropyLoss() + dice_loss = losses.DiceLoss(num_classes) + + logging.info("{} iterations per epoch".format(len(trainloader_labeled))) + + iter_num = 0 + max_epoch = max_iterations // len(trainloader_unlabeled) + 1 + best_performance = 0.0 + iterator = tqdm(range(max_epoch), ncols=70) + for epoch_num in iterator: + for i, (sampled_batch_labeled, sampled_batch_unlabeled) in enumerate(zip(cycle(trainloader_labeled), trainloader_unlabeled)): + volume_batch_sa, volume_batch_wa, label_batch = sampled_batch_labeled['image_s'], sampled_batch_labeled['image_w'], sampled_batch_labeled['label'] + volume_batch_sa, volume_batch_wa, label_batch = volume_batch_sa.cuda(), volume_batch_wa.cuda(), label_batch.cuda() + unlabeled_volume_batch_sa, unlabeled_volume_batch_wa = sampled_batch_unlabeled['image_s'].cuda(), sampled_batch_unlabeled['image_w'].cuda() + + outputs = model(volume_batch_sa) + outputs_soft = torch.softmax(outputs, dim=1) + + outputs_unlabeled = model(unlabeled_volume_batch_sa) + outputs_unlabeled_soft = torch.softmax(outputs_unlabeled, dim=1) + + T = 1 + threshold = 0.95 + + with torch.no_grad(): + ema_output = ema_model(unlabeled_volume_batch_wa) + pseudo_label = torch.softmax(ema_output.detach()/T, dim=1) + max_probs, targets_u = torch.max(pseudo_label, dim=1) + mask = max_probs.ge(threshold).float() + + supervised_loss = 0.5 * \ + (ce_loss(outputs, label_batch[:].long( + )) + dice_loss(outputs_soft, label_batch[:].unsqueeze(1))) + + consistency_weight = get_current_consistency_weight( + iter_num // (args.max_iterations/args.consistency_rampup)) + + unsupervised_loss = (F.cross_entropy(outputs_unlabeled, targets_u, + reduction='none') * mask).mean() + + loss = supervised_loss + consistency_weight * unsupervised_loss + optimizer.zero_grad() + loss.backward() + optimizer.step() + update_ema_variables(model, ema_model, args.ema_decay, iter_num) + + lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 + for param_group in optimizer.param_groups: + param_group['lr'] = lr_ + + iter_num = iter_num + 1 + writer.add_scalar('info/lr', lr_, iter_num) + writer.add_scalar('info/total_loss', loss, iter_num) + writer.add_scalar('info/loss_ce', supervised_loss, iter_num) + writer.add_scalar('info/unsupervised_loss', + unsupervised_loss, iter_num) + writer.add_scalar('info/consistency_weight', + consistency_weight, iter_num) + + logging.info( + 'iteration %d : loss : %f, loss_ce: %f' % + (iter_num, loss.item(), supervised_loss.item())) + + if iter_num % 20 == 0: + image = volume_batch_sa[0, 0:1, :, :] + writer.add_image('train/Image', image, iter_num) + outputs = torch.argmax(torch.softmax( + outputs, dim=1), dim=1, keepdim=True) + writer.add_image('train/Prediction', + outputs[0, ...] * 50, iter_num) + labs = label_batch[0, ...].unsqueeze(0) * 50 + writer.add_image('train/GroundTruth', labs, iter_num) + + if iter_num > 0 and iter_num % 200 == 0: + model.eval() + metric_list = 0.0 + for i_batch, sampled_batch in enumerate(valloader): + metric_i = test_single_volume( + sampled_batch["image"], sampled_batch["label"], model, classes=num_classes) + metric_list += np.array(metric_i) + metric_list = metric_list / len(db_val) + for class_i in range(num_classes-1): + writer.add_scalar('info/val_{}_dice'.format(class_i+1), + metric_list[class_i, 0], iter_num) + writer.add_scalar('info/val_{}_hd95'.format(class_i+1), + metric_list[class_i, 1], iter_num) + + performance = np.mean(metric_list, axis=0)[0] + + mean_hd95 = np.mean(metric_list, axis=0)[1] + writer.add_scalar('info/val_mean_dice', performance, iter_num) + writer.add_scalar('info/val_mean_hd95', mean_hd95, iter_num) + + if performance > best_performance: + best_performance = performance + save_mode_path = os.path.join(snapshot_path, + 'iter_{}_dice_{}.pth'.format( + iter_num, round(best_performance, 4))) + save_best = os.path.join(snapshot_path, + '{}_best_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_mode_path) + torch.save(model.state_dict(), save_best) + + logging.info( + 'iteration %d : mean_dice : %f mean_hd95 : %f' % (iter_num, performance, mean_hd95)) + model.train() + + if iter_num % 3000 == 0: + save_mode_path = os.path.join( + snapshot_path, 'iter_' + str(iter_num) + '.pth') + torch.save(model.state_dict(), save_mode_path) + logging.info("save model to {}".format(save_mode_path)) + + if iter_num >= max_iterations: + break + + save_latest = os.path.join( + snapshot_path, '{}_latest_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_latest) + + if iter_num >= max_iterations: + iterator.close() + break + writer.close() + return "Training Finished!" + + +if __name__ == "__main__": + if not args.deterministic: + cudnn.benchmark = True + cudnn.deterministic = False + else: + cudnn.benchmark = False + cudnn.deterministic = True + + random.seed(args.seed) + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + if args.cross_val: + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}/{}".format( + args.exp, args.labeled_ratio, args.fold, args.model) + else: + snapshot_path = "../model/{}/1_of_{}_labeled/{}".format( + args.exp, args.labeled_ratio, args.model) + + if not os.path.exists(snapshot_path): + os.makedirs(snapshot_path) + if os.path.exists(snapshot_path + '/code'): + shutil.rmtree(snapshot_path + '/code') + shutil.copytree('.', snapshot_path + '/code', + shutil.ignore_patterns(['.git', '__pycache__'])) + + logging.basicConfig(filename=snapshot_path + "/log.txt", level=logging.INFO, + format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S') + logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) + logging.info(str(args)) + train(args, snapshot_path) diff --git a/code/train_fully_supervised.py b/code/train_fully_supervised.py index 802e7c5..c57564d 100644 --- a/code/train_fully_supervised.py +++ b/code/train_fully_supervised.py @@ -40,7 +40,8 @@ default=30000, help='maximum epoch number to train') parser.add_argument('--batch_size', type=int, default=12, help='batch_size per gpu') - +parser.add_argument('--cross_val', type=int, + default=0, help='5-fold cross validation or random split 7/1/2 for training/validation/testing') parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training') parser.add_argument('--base_lr', type=float, default=0.03, @@ -52,7 +53,7 @@ help='output channel of network') # label and unlabel -parser.add_argument('--labeled_ratio', type=int, default=8, +parser.add_argument('--labeled_ratio', type=int, default=10, help='1/labeled_ratio data is provided mask') # costs parser.add_argument('--ema_decay', type=float, default=0.99, help='ema_decay') @@ -79,7 +80,7 @@ def train(args, snapshot_path): model = net_factory(net_type=args.model, in_chns=1, class_num=num_classes) db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ - RandomGenerator(args.patch_size)])) + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) logging.info("Labeled slices: {} ".format(len(db_train_labeled))) @@ -87,7 +88,7 @@ def train(args, snapshot_path): db_train_labeled, batch_size=args.batch_size, shuffle=True) db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, - split="val", labeled_ratio=args.labeled_ratio) + split="val", labeled_ratio=args.labeled_ratio, cross_val=args.cross_val) valloader = DataLoader(db_val, batch_size=1) model.train() @@ -183,6 +184,11 @@ def train(args, snapshot_path): if iter_num >= max_iterations: break + + save_latest = os.path.join( + snapshot_path, '{}_latest_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_latest) + if iter_num >= max_iterations: iterator.close() break @@ -203,8 +209,13 @@ def train(args, snapshot_path): torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) - snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( - args.exp, args.labeled_ratio, args.fold) + if args.cross_val: + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}/{}".format( + args.exp, args.labeled_ratio, args.fold, args.model) + else: + snapshot_path = "../model/{}/1_of_{}_labeled/{}".format( + args.exp, args.labeled_ratio, args.model) + if not os.path.exists(snapshot_path): os.makedirs(snapshot_path) if os.path.exists(snapshot_path + '/code'): diff --git a/code/train_interpolation_consistency_training.py b/code/train_interpolation_consistency_training.py index 114eb77..672cbed 100644 --- a/code/train_interpolation_consistency_training.py +++ b/code/train_interpolation_consistency_training.py @@ -6,6 +6,7 @@ import sys import time from itertools import cycle + import numpy as np import torch import torch.backends.cudnn as cudnn @@ -39,7 +40,8 @@ default=30000, help='maximum epoch number to train') parser.add_argument('--batch_size', type=int, default=12, help='batch_size per gpu') - +parser.add_argument('--cross_val', type=bool, + default=True, help='5-fold cross validation or random split 7/1/2 for training/validation/testing') parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training') parser.add_argument('--base_lr', type=float, default=0.03, @@ -98,9 +100,9 @@ def create_model(ema=False): ema_model = create_model(ema=True) db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ - RandomGenerator(args.patch_size)])) + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ - RandomGenerator(args.patch_size)])) + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) trainloader_labeled = DataLoader( db_train_labeled, batch_size=args.batch_size//2, shuffle=True) @@ -110,7 +112,7 @@ def create_model(ema=False): logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, - split="val", labeled_ratio=args.labeled_ratio) + split="val", labeled_ratio=args.labeled_ratio, cross_val=args.cross_val) valloader = DataLoader(db_val, batch_size=1, shuffle=False) model.train() @@ -256,6 +258,11 @@ def create_model(ema=False): if iter_num >= max_iterations: break + + save_latest = os.path.join( + snapshot_path, '{}_latest_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_latest) + if iter_num >= max_iterations: iterator.close() break @@ -276,8 +283,13 @@ def create_model(ema=False): torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) - snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( - args.exp, args.labeled_ratio, args.fold) + if args.cross_val: + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}/{}".format( + args.exp, args.labeled_ratio, args.fold, args.model) + else: + snapshot_path = "../model/{}/1_of_{}_labeled/{}".format( + args.exp, args.labeled_ratio, args.model) + if not os.path.exists(snapshot_path): os.makedirs(snapshot_path) if os.path.exists(snapshot_path + '/code'): diff --git a/code/train_mean_teacher.py b/code/train_mean_teacher.py index d968714..a0e3584 100644 --- a/code/train_mean_teacher.py +++ b/code/train_mean_teacher.py @@ -6,6 +6,7 @@ import sys import time from itertools import cycle + import numpy as np import torch import torch.backends.cudnn as cudnn @@ -39,7 +40,8 @@ default=30000, help='maximum epoch number to train') parser.add_argument('--batch_size', type=int, default=16, help='batch_size per gpu') - +parser.add_argument('--cross_val', type=bool, + default=True, help='5-fold cross validation or random split 7/1/2 for training/validation/testing') parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training') parser.add_argument('--base_lr', type=float, default=0.03, @@ -98,9 +100,9 @@ def create_model(ema=False): ema_model = create_model(ema=True) db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ - RandomGenerator(args.patch_size)])) + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ - RandomGenerator(args.patch_size)])) + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) logging.info("Labeled slices: {} ".format(len(db_train_labeled))) logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) @@ -110,7 +112,7 @@ def create_model(ema=False): db_train_unlabeled, batch_size=args.batch_size//2, shuffle=True) db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, - split="val", labeled_ratio=args.labeled_ratio) + split="val", labeled_ratio=args.labeled_ratio, cross_val=args.cross_val) valloader = DataLoader(db_val, batch_size=1) model.train() @@ -230,6 +232,11 @@ def create_model(ema=False): if iter_num >= max_iterations: break + + save_latest = os.path.join( + snapshot_path, '{}_latest_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_latest) + if iter_num >= max_iterations: iterator.close() break @@ -250,8 +257,13 @@ def create_model(ema=False): torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) - snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( - args.exp, args.labeled_ratio, args.fold) + if args.cross_val: + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}/{}".format( + args.exp, args.labeled_ratio, args.fold, args.model) + else: + snapshot_path = "../model/{}/1_of_{}_labeled/{}".format( + args.exp, args.labeled_ratio, args.model) + if not os.path.exists(snapshot_path): os.makedirs(snapshot_path) if os.path.exists(snapshot_path + '/code'): diff --git a/code/train_uncertainty_aware_mean_teacher.py b/code/train_uncertainty_aware_mean_teacher.py index febf888..2a46446 100644 --- a/code/train_uncertainty_aware_mean_teacher.py +++ b/code/train_uncertainty_aware_mean_teacher.py @@ -6,6 +6,7 @@ import sys import time from itertools import cycle + import numpy as np import torch import torch.backends.cudnn as cudnn @@ -19,6 +20,7 @@ from torchvision import transforms from torchvision.utils import make_grid from tqdm import tqdm + from dataloaders.dataset import BaseDataSets, RandomGenerator from networks.discriminator import FCDiscriminator from networks.net_factory import net_factory @@ -38,7 +40,8 @@ default=30000, help='maximum epoch number to train') parser.add_argument('--batch_size', type=int, default=12, help='batch_size per gpu') - +parser.add_argument('--cross_val', type=bool, + default=True, help='5-fold cross validation or random split 7/1/2 for training/validation/testing') parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training') parser.add_argument('--base_lr', type=float, default=0.03, @@ -97,9 +100,9 @@ def create_model(ema=False): ema_model = create_model(ema=True) db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ - RandomGenerator(args.patch_size)])) + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) db_train_unlabeled = BaseDataSets(base_dir=args.root_path, labeled_type="unlabeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ - RandomGenerator(args.patch_size)])) + RandomGenerator(args.patch_size)]), cross_val=args.cross_val) trainloader_labeled = DataLoader( db_train_labeled, batch_size=args.batch_size//2, shuffle=True) @@ -110,7 +113,7 @@ def create_model(ema=False): logging.info("Unlabeled slices: {} ".format(len(db_train_unlabeled))) db_val = BaseDataSets(base_dir=args.root_path, fold=args.fold, - split="val", labeled_ratio=args.labeled_ratio) + split="val", labeled_ratio=args.labeled_ratio, cross_val=args.cross_val) valloader = DataLoader(db_val, batch_size=1, shuffle=False) model.train() @@ -250,6 +253,10 @@ def create_model(ema=False): if iter_num >= max_iterations: break + save_latest = os.path.join( + snapshot_path, '{}_latest_model.pth'.format(args.model)) + torch.save(model.state_dict(), save_latest) + if iter_num >= max_iterations: iterator.close() break @@ -270,8 +277,13 @@ def create_model(ema=False): torch.manual_seed(args.seed) torch.cuda.manual_seed(args.seed) - snapshot_path = "../model/{}/1_of_{}_labeled/fold{}".format( - args.exp, args.labeled_ratio, args.fold) + if args.cross_val: + snapshot_path = "../model/{}/1_of_{}_labeled/fold{}/{}".format( + args.exp, args.labeled_ratio, args.fold, args.model) + else: + snapshot_path = "../model/{}/1_of_{}_labeled/{}".format( + args.exp, args.labeled_ratio, args.model) + if not os.path.exists(snapshot_path): os.makedirs(snapshot_path) if os.path.exists(snapshot_path + '/new_code'): diff --git a/code/utils/losses.py b/code/utils/losses.py index e4d0103..5466aef 100755 --- a/code/utils/losses.py +++ b/code/utils/losses.py @@ -64,7 +64,6 @@ def nuclear_norm_maximum(pr): return L_BNM / N - def softmax_dice_loss(input_logits, target_logits): """Takes softmax on both sides and returns MSE loss @@ -232,3 +231,34 @@ def entropy_map(p): ent_map = -1*torch.sum(p * torch.log(p + 1e-6), dim=1, keepdim=True) return ent_map + + +def js_loss(p1, p2): + # the Jensen-Shannon divergence between p1(x) and p2(x) + a1 = 0.5 * (p1 + p2) + loss1 = a1 * torch.log(a1) + loss1 = -torch.sum(loss1) + loss2 = p1 * torch.log(p1) + loss2 = -torch.sum(loss2) + loss3 = p2 * torch.log(p2) + loss3 = -torch.sum(loss3) + return (loss1 - 0.5 * (loss2 + loss3)) / p1.shape[0] + + +# def loss_diff(logit_S1, logit_S2, perturbed_logit_S1, perturbed_logit_S2, logit_U1, logit_U2, perturbed_logit_U1, perturbed_logit_U2): +# S = nn.Softmax(dim=1) +# LS = nn.LogSoftmax(dim=1) + +# a = S(logit_S2) * LS(perturbed_logit_S1) +# a = torch.sum(a) + +# b = S(logit_S1) * LS(perturbed_logit_S2) +# b = torch.sum(b) + +# c = S(logit_U2) * LS(perturbed_logit_U1) +# c = torch.sum(c) + +# d = S(logit_U1) * LS(perturbed_logit_U2) +# d = torch.sum(d) + +# return -(a+b+c+d)/batch_size From a18871784dad1821f7c594725c824f21c8511f49 Mon Sep 17 00:00:00 2001 From: luoxd Date: Thu, 10 Mar 2022 14:29:29 +0800 Subject: [PATCH 7/7] upload fixmatch code --- code/train_fixmatch.py | 45 ++++++++++++++++++++++++++++-------------- 1 file changed, 30 insertions(+), 15 deletions(-) diff --git a/code/train_fixmatch.py b/code/train_fixmatch.py index cf1ce8d..3e8cce9 100644 --- a/code/train_fixmatch.py +++ b/code/train_fixmatch.py @@ -51,6 +51,8 @@ parser.add_argument('--seed', type=int, default=2022, help='random seed') parser.add_argument('--num_classes', type=int, default=4, help='output channel of network') +parser.add_argument('--ema', type=int, default=0, + help='ema') # label and unlabel parser.add_argument('--labeled_ratio', type=int, default=10, @@ -97,7 +99,8 @@ def create_model(ema=False): return model model = create_model() - ema_model = create_model(ema=True) + if args.ema: + ema_model = create_model(ema=args.ema) db_train_labeled = BaseDataSets(base_dir=args.root_path, labeled_type="labeled", labeled_ratio=args.labeled_ratio, fold=args.fold, split="train", transform=transforms.Compose([ RandomGenerator_Strong_Weak(args.patch_size)]), cross_val=args.cross_val) @@ -131,9 +134,12 @@ def create_model(ema=False): iterator = tqdm(range(max_epoch), ncols=70) for epoch_num in iterator: for i, (sampled_batch_labeled, sampled_batch_unlabeled) in enumerate(zip(cycle(trainloader_labeled), trainloader_unlabeled)): - volume_batch_sa, volume_batch_wa, label_batch = sampled_batch_labeled['image_s'], sampled_batch_labeled['image_w'], sampled_batch_labeled['label'] - volume_batch_sa, volume_batch_wa, label_batch = volume_batch_sa.cuda(), volume_batch_wa.cuda(), label_batch.cuda() - unlabeled_volume_batch_sa, unlabeled_volume_batch_wa = sampled_batch_unlabeled['image_s'].cuda(), sampled_batch_unlabeled['image_w'].cuda() + volume_batch_sa, volume_batch_wa, label_batch = sampled_batch_labeled[ + 'image_s'], sampled_batch_labeled['image_w'], sampled_batch_labeled['label'] + volume_batch_sa, volume_batch_wa, label_batch = volume_batch_sa.cuda( + ), volume_batch_wa.cuda(), label_batch.cuda() + unlabeled_volume_batch_sa, unlabeled_volume_batch_wa = sampled_batch_unlabeled['image_s'].cuda( + ), sampled_batch_unlabeled['image_w'].cuda() outputs = model(volume_batch_sa) outputs_soft = torch.softmax(outputs, dim=1) @@ -143,10 +149,15 @@ def create_model(ema=False): T = 1 threshold = 0.95 - - with torch.no_grad(): - ema_output = ema_model(unlabeled_volume_batch_wa) - pseudo_label = torch.softmax(ema_output.detach()/T, dim=1) + if args.ema: + with torch.no_grad(): + ema_output = ema_model(unlabeled_volume_batch_wa) + pseudo_label = torch.softmax(ema_output.detach()/T, dim=1) + max_probs, targets_u = torch.max(pseudo_label, dim=1) + mask = max_probs.ge(threshold).float() + else: + output_wa = model(unlabeled_volume_batch_wa) + pseudo_label = torch.softmax(output_wa.detach()/T, dim=1) max_probs, targets_u = torch.max(pseudo_label, dim=1) mask = max_probs.ge(threshold).float() @@ -158,13 +169,16 @@ def create_model(ema=False): iter_num // (args.max_iterations/args.consistency_rampup)) unsupervised_loss = (F.cross_entropy(outputs_unlabeled, targets_u, - reduction='none') * mask).mean() + reduction='none') * mask).mean() loss = supervised_loss + consistency_weight * unsupervised_loss optimizer.zero_grad() loss.backward() optimizer.step() - update_ema_variables(model, ema_model, args.ema_decay, iter_num) + if args.ema: + update_ema_variables(model, ema_model, args.ema_decay, iter_num) + else: + pass lr_ = base_lr * (1.0 - iter_num / max_iterations) ** 0.9 for param_group in optimizer.param_groups: @@ -235,7 +249,7 @@ def create_model(ema=False): if iter_num >= max_iterations: break - + save_latest = os.path.join( snapshot_path, '{}_latest_model.pth'.format(args.model)) torch.save(model.state_dict(), save_latest) @@ -261,11 +275,11 @@ def create_model(ema=False): torch.cuda.manual_seed(args.seed) if args.cross_val: - snapshot_path = "../model/{}/1_of_{}_labeled/fold{}/{}".format( - args.exp, args.labeled_ratio, args.fold, args.model) + snapshot_path = "../model/{}_ema_{}/1_of_{}_labeled/fold{}/{}".format( + args.exp, args.ema, args.labeled_ratio, args.fold, args.model) else: - snapshot_path = "../model/{}/1_of_{}_labeled/{}".format( - args.exp, args.labeled_ratio, args.model) + snapshot_path = "../model/{}_ema_{}/1_of_{}_labeled/{}".format( + args.exp, args.ema, args.labeled_ratio, args.model) if not os.path.exists(snapshot_path): os.makedirs(snapshot_path) @@ -279,3 +293,4 @@ def create_model(ema=False): logging.getLogger().addHandler(logging.StreamHandler(sys.stdout)) logging.info(str(args)) train(args, snapshot_path) +