| | from copy import deepcopy |
| | from multiprocessing.pool import Pool |
| |
|
| | from batchgenerators.utilities.file_and_folder_operations import * |
| | from medpy import metric |
| | import SimpleITK as sitk |
| | import numpy as np |
| | from nnunet.configuration import default_num_threads |
| | from nnunet.postprocessing.consolidate_postprocessing import collect_cv_niftis |
| |
|
| |
|
| | def get_brats_regions(): |
| | """ |
| | this is only valid for the brats data in here where the labels are 1, 2, and 3. The original brats data have a |
| | different labeling convention! |
| | :return: |
| | """ |
| | regions = { |
| | "whole tumor": (1, 2, 3), |
| | "tumor core": (2, 3), |
| | "enhancing tumor": (3,) |
| | } |
| | return regions |
| |
|
| |
|
| | def get_KiTS_regions(): |
| | regions = { |
| | "kidney incl tumor": (1, 2), |
| | "tumor": (2,) |
| | } |
| | return regions |
| |
|
| |
|
| | def create_region_from_mask(mask, join_labels: tuple): |
| | mask_new = np.zeros_like(mask, dtype=np.uint8) |
| | for l in join_labels: |
| | mask_new[mask == l] = 1 |
| | return mask_new |
| |
|
| |
|
| | def evaluate_case(file_pred: str, file_gt: str, regions): |
| | image_gt = sitk.GetArrayFromImage(sitk.ReadImage(file_gt)) |
| | image_pred = sitk.GetArrayFromImage(sitk.ReadImage(file_pred)) |
| | results = [] |
| | for r in regions: |
| | mask_pred = create_region_from_mask(image_pred, r) |
| | mask_gt = create_region_from_mask(image_gt, r) |
| | dc = np.nan if np.sum(mask_gt) == 0 and np.sum(mask_pred) == 0 else metric.dc(mask_pred, mask_gt) |
| | results.append(dc) |
| | return results |
| |
|
| |
|
| | def evaluate_regions(folder_predicted: str, folder_gt: str, regions: dict, processes=default_num_threads): |
| | region_names = list(regions.keys()) |
| | files_in_pred = subfiles(folder_predicted, suffix='.nii.gz', join=False) |
| | files_in_gt = subfiles(folder_gt, suffix='.nii.gz', join=False) |
| | have_no_gt = [i for i in files_in_pred if i not in files_in_gt] |
| | assert len(have_no_gt) == 0, "Some files in folder_predicted have not ground truth in folder_gt" |
| | have_no_pred = [i for i in files_in_gt if i not in files_in_pred] |
| | if len(have_no_pred) > 0: |
| | print("WARNING! Some files in folder_gt were not predicted (not present in folder_predicted)!") |
| |
|
| | files_in_gt.sort() |
| | files_in_pred.sort() |
| |
|
| | |
| | full_filenames_gt = [join(folder_gt, i) for i in files_in_pred] |
| | full_filenames_pred = [join(folder_predicted, i) for i in files_in_pred] |
| |
|
| | p = Pool(processes) |
| | res = p.starmap(evaluate_case, zip(full_filenames_pred, full_filenames_gt, [list(regions.values())] * len(files_in_gt))) |
| | p.close() |
| | p.join() |
| |
|
| | all_results = {r: [] for r in region_names} |
| | with open(join(folder_predicted, 'summary.csv'), 'w') as f: |
| | f.write("casename") |
| | for r in region_names: |
| | f.write(",%s" % r) |
| | f.write("\n") |
| | for i in range(len(files_in_pred)): |
| | f.write(files_in_pred[i][:-7]) |
| | result_here = res[i] |
| | for k, r in enumerate(region_names): |
| | dc = result_here[k] |
| | f.write(",%02.4f" % dc) |
| | all_results[r].append(dc) |
| | f.write("\n") |
| |
|
| | f.write('mean') |
| | for r in region_names: |
| | f.write(",%02.4f" % np.nanmean(all_results[r])) |
| | f.write("\n") |
| | f.write('median') |
| | for r in region_names: |
| | f.write(",%02.4f" % np.nanmedian(all_results[r])) |
| | f.write("\n") |
| |
|
| | f.write('mean (nan is 1)') |
| | for r in region_names: |
| | tmp = np.array(all_results[r]) |
| | tmp[np.isnan(tmp)] = 1 |
| | f.write(",%02.4f" % np.mean(tmp)) |
| | f.write("\n") |
| | f.write('median (nan is 1)') |
| | for r in region_names: |
| | tmp = np.array(all_results[r]) |
| | tmp[np.isnan(tmp)] = 1 |
| | f.write(",%02.4f" % np.median(tmp)) |
| | f.write("\n") |
| |
|
| |
|
| | if __name__ == '__main__': |
| | collect_cv_niftis('./', './cv_niftis') |
| | evaluate_regions('./cv_niftis/', './gt_niftis/', get_brats_regions()) |
| |
|
