|
- # Implementation of this model is borrowed and modified
- # (from torch to paddle) from here:
- # https://github.com/MIC-DKFZ/nnUNet
-
- # Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
- #
- # 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.
-
- import argparse
- import os
- import shutil
- import paddle
- import numpy as np
- import pickle
- import sys
-
- parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 2)))
- sys.path.insert(0, parent_path)
-
- from medicalseg.cvlibs import Config
- from medicalseg.utils import get_sys_env, logger, config_check, utils
-
- from nnunet.utils import DynamicPredictor, save_segmentation_nifti_from_softmax, aggregate_scores, determine_postprocessing, resample_and_save, predict_next_stage
-
-
- def to_one_hot(seg, all_seg_labels=None):
- if all_seg_labels is None:
- all_seg_labels = np.unique(seg)
- result = np.zeros((len(all_seg_labels), *seg.shape), dtype=seg.dtype)
- for i, l in enumerate(all_seg_labels):
- result[i][seg == l] = 1
- return result
-
-
- def evaluate(model, eval_dataset, args):
- """
- Launch evalution.
- Args:
- model(nn.Layer): A sementic segmentation model.
- eval_dataset (paddle.io.Dataset): Used to read and process validation datasets.
- args(str, optional): configura args.
- """
- model.eval()
- predictor = DynamicPredictor(model)
- plans = model.plans
- if 'segmentation_export_params' in plans.keys():
- force_separate_z = plans['segmentation_export_params'][
- 'force_separate_z']
- interpolation_order = plans['segmentation_export_params'][
- 'interpolation_order']
- interpolation_order_z = plans['segmentation_export_params'][
- 'interpolation_order_z']
- else:
- force_separate_z = None
- interpolation_order = 1
- interpolation_order_z = 0
- output_base = args.val_save_folder
- output_folder = os.path.join(output_base,
- 'fold_{}'.format(eval_dataset.fold))
- validation_raw_folder = os.path.join(output_folder, 'validation_raw')
-
- os.makedirs(validation_raw_folder, exist_ok=True)
- if not eval_dataset.data_aug_params['do_mirror']:
- raise RuntimeError(
- "We did not train with mirroring so you cannot do inference with mirroring enabled"
- )
- mirror_axes = eval_dataset.data_aug_params['mirror_axes']
-
- if args.predict_next_stage:
- next_stage_output_folder = eval_dataset.folder_with_segs_from_prev_stage
- os.makedirs(next_stage_output_folder, exist_ok=True)
-
- print('start evaluating...')
- pred_gt_tuples = []
- gt_niftis_folder = os.path.join(eval_dataset.preprocessed_dir,
- "gt_segmentations") # gt dir
- for k in eval_dataset.dataset_val.keys():
- with open(eval_dataset.dataset[k]['properties_file'], 'rb') as f:
- properties = pickle.load(f)
- fname = properties['list_of_data_files'][0].split("/")[-1][:-12]
- pred_gt_tuples.append([
- os.path.join(validation_raw_folder, fname + '.nii.gz'),
- os.path.join(gt_niftis_folder, fname + '.nii.gz'),
- ])
- if os.path.exists(
- os.path.join(validation_raw_folder, fname + '.nii.gz')):
- print('{} already exists, skip.'.format(
- os.path.join(validation_raw_folder, fname + '.nii.gz')))
- continue
- data = np.load(eval_dataset.dataset[k]['data_file'])['data']
- print(k, data.shape)
- data[-1][data[-1] == -1] = 0
- data = data[:-1]
-
- if eval_dataset.stage == 1 and eval_dataset.cascade:
- seg_pre_path = os.path.join(
- eval_dataset.folder_with_segs_from_prev_stage,
- k + "_segFromPrevStage.npz")
- if not os.path.exists(seg_pre_path):
- raise UserWarning(
- 'cannot find stage 1 segmentation result for {}.'.format(
- seg_pre_path))
- seg_from_prev_stage = np.load(seg_pre_path)['data'][None]
- data = np.concatenate(
- (data, to_one_hot(seg_from_prev_stage[0],
- range(1, eval_dataset.num_classes))))
-
- argmax_pred, softmax_pred = predictor.predict_3D(
- data,
- do_mirroring=args.do_mirroring,
- mirror_axes=mirror_axes,
- use_sliding_window=args.use_sliding_window,
- step_size=args.step_size,
- patch_size=eval_dataset.patch_size,
- regions_class_order=None,
- use_gaussian=args.use_gaussian,
- pad_border_mode='constant',
- pad_kwargs=None,
- verbose=args.verbose,
- mixed_precision=args.precision == 'fp16')
-
- softmax_pred = softmax_pred.transpose(
- [0] + [i + 1 for i in eval_dataset.transpose_backward])
- softmax_fname = os.path.join(validation_raw_folder, fname + '.npz')
-
- if np.prod(softmax_pred.shape) > (2e9 / 4 * 0.85):
- np.save(
- os.path.join(validation_raw_folder, fname + ".npy"),
- softmax_pred)
- softmax_pred = os.path.join(validation_raw_folder, fname + ".npy")
-
- save_segmentation_nifti_from_softmax(
- softmax_pred,
- os.path.join(validation_raw_folder, fname + '.nii.gz'), properties,
- interpolation_order, None, None, None, softmax_fname, None,
- force_separate_z, interpolation_order_z)
-
- _ = aggregate_scores(
- pred_gt_tuples,
- labels=list(range(eval_dataset.num_classes)),
- json_output_file=os.path.join(validation_raw_folder, "summary.json"),
- json_name=" val tiled %s" % (str(args.use_sliding_window)),
- json_author="medicalseg",
- json_task='task',
- num_threads=4)
-
- determine_postprocessing(
- output_folder,
- gt_niftis_folder,
- 'validation_raw',
- final_subf_name='validation_raw' + "_postprocessed",
- debug=False)
-
- gt_nifti_folder = os.path.join(output_base, "gt_niftis")
-
- if not os.path.exists(gt_nifti_folder):
- os.makedirs(gt_nifti_folder, exist_ok=True)
- print('copy gt from {} to {}.'.format(gt_niftis_folder, gt_nifti_folder))
-
- gt_files = []
- for f_name in os.listdir(gt_niftis_folder):
- if os.path.isfile(os.path.join(gt_niftis_folder,
- f_name)) and f_name.endswith('.nii.gz'):
- gt_files.append(os.path.join(gt_niftis_folder, f_name))
-
- for f in gt_files:
- success = False
- attempts = 0
- e = None
- while not success and attempts < 10:
- try:
- shutil.copy(f, gt_nifti_folder)
- success = True
- except OSError as e:
- attempts += 1
- if not success:
- print("Could not copy gt nifti file %s into folder %s" %
- (f, gt_nifti_folder))
- if e is not None:
- raise e
-
- if args.predict_next_stage:
- predict_next_stage(
- model,
- plans,
- eval_dataset,
- eval_dataset.preprocessed_dir,
- os.path.join(eval_dataset.preprocessed_dir,
- plans['data_identifier'] + "_stage%d" % 1),
- args.precision == 'fp16',
- num_threads=4)
-
-
- def parse_args():
- parser = argparse.ArgumentParser(description='Model evaluation')
-
- # params of evaluate
- parser.add_argument(
- "--config", dest="cfg", help="The config file.", default=None, type=str)
-
- parser.add_argument(
- '--model_path',
- dest='model_path',
- help='The path of model for evaluation',
- type=str,
- default="saved_model/vnet_lung_coronavirus_128_128_128_15k/best_model/model.pdparams"
- )
-
- parser.add_argument(
- '--predict_next_stage',
- action='store_true',
- default=False,
- help='whether predict stage 2 training data.')
-
- parser.add_argument(
- '--val_save_folder',
- dest='val_save_folder',
- help='The path to val data predicted result',
- type=str,
- default="val_save_folder")
-
- parser.add_argument(
- "--precision",
- default="fp32",
- type=str,
- choices=["fp32", "fp16"],
- help="Use AMP (Auto mixed precision) if precision='fp16'. If precision='fp32', the training is normal."
- )
-
- parser.add_argument(
- "--do_mirroring",
- dest='do_mirroring',
- type=bool,
- default=True,
- help="Whether use mirroring when inference.")
-
- parser.add_argument(
- "--use_sliding_window",
- dest='use_sliding_window',
- type=bool,
- default=True,
- help="Whether use sliding window when inference.")
-
- parser.add_argument(
- "--step_size",
- dest='step_size',
- type=float,
- default=0.5,
- help="step size when predict.")
-
- parser.add_argument(
- "--use_gaussian",
- dest='use_gaussian',
- type=bool,
- default=True,
- help="Whether use gaussian.")
-
- parser.add_argument(
- "--verbose",
- dest='verbose',
- type=bool,
- default=True,
- help="Whether print log.")
-
- return parser.parse_args()
-
-
- def main(args):
- env_info = get_sys_env()
- place = 'gpu' if env_info['Paddle compiled with cuda'] and env_info[
- 'GPUs used'] else 'cpu'
-
- paddle.set_device(place)
- if not args.cfg:
- raise RuntimeError('No configuration file specified.')
-
- cfg = Config(args.cfg)
-
- val_dataset = cfg.val_dataset
- if val_dataset is None:
- raise RuntimeError(
- 'The verification dataset is not specified in the configuration file.'
- )
- elif len(val_dataset) == 0:
- raise ValueError(
- 'The length of val_dataset is 0. Please check if your dataset is valid'
- )
-
- msg = '\n---------------Config Information---------------\n'
- msg += str(cfg)
- msg += '------------------------------------------------'
- logger.info(msg)
-
- model = cfg.model
- if args.model_path:
- utils.load_entire_model(model, args.model_path)
- logger.info('Loaded trained params of model successfully')
-
- config_check(cfg, val_dataset=val_dataset)
-
- evaluate(model, val_dataset, args)
-
-
- if __name__ == '__main__':
- args = parse_args()
- main(args)
|