""" author: Clément Zotti (clement.zotti@usherbrooke.ca) date: April 2017 DESCRIPTION : The script provide helpers functions to handle nifti image format: - load_nii() - save_nii() to generate metrics for two images: - metrics() And it is callable from the command line (see below). Each function provided in this script has comments to understand how they works. HOW-TO: This script was tested for python 3.4. First, you need to install the required packages with pip install -r requirements.txt After the installation, you have two ways of running this script: 1) python metrics.py ground_truth/patient001_ED.nii.gz prediction/patient001_ED.nii.gz 2) python metrics.py ground_truth/ prediction/ The first option will print in the console the dice and volume of each class for the given image. The second option wiil ouput a csv file where each images will have the dice and volume of each class. Link: http://acdc.creatis.insa-lyon.fr """ import os from glob import glob import time import re import argparse import nibabel as nib import pandas as pd from medpy.metric.binary import hd, dc import numpy as np HEADER = ["Name", "Dice LV", "Volume LV", "Err LV(ml)", "Dice RV", "Volume RV", "Err RV(ml)", "Dice MYO", "Volume MYO", "Err MYO(ml)"] # # Utils functions used to sort strings into a natural order # def conv_int(i): return int(i) if i.isdigit() else i def natural_order(sord): """ Sort a (list,tuple) of strings into natural order. Ex: ['1','10','2'] -> ['1','2','10'] ['abc1def','ab10d','b2c','ab1d'] -> ['ab1d','ab10d', 'abc1def', 'b2c'] """ if isinstance(sord, tuple): sord = sord[0] return [conv_int(c) for c in re.split(r'(\d+)', sord)] # # Utils function to load and save nifti files with the nibabel package # def load_nii(img_path): """ Function to load a 'nii' or 'nii.gz' file, The function returns everyting needed to save another 'nii' or 'nii.gz' in the same dimensional space, i.e. the affine matrix and the header Parameters ---------- img_path: string String with the path of the 'nii' or 'nii.gz' image file name. Returns ------- Three element, the first is a numpy array of the image values, the second is the affine transformation of the image, and the last one is the header of the image. """ nimg = nib.load(img_path) return nimg.get_data(), nimg.affine, nimg.header def save_nii(img_path, data, affine, header): """ Function to save a 'nii' or 'nii.gz' file. Parameters ---------- img_path: string Path to save the image should be ending with '.nii' or '.nii.gz'. data: np.array Numpy array of the image data. affine: list of list or np.array The affine transformation to save with the image. header: nib.Nifti1Header The header that define everything about the data (pleasecheck nibabel documentation). """ nimg = nib.Nifti1Image(data, affine=affine, header=header) nimg.to_filename(img_path) # # Functions to process files, directories and metrics # def metrics(img_gt, img_pred, voxel_size): """ Function to compute the metrics between two segmentation maps given as input. Parameters ---------- img_gt: np.array Array of the ground truth segmentation map. img_pred: np.array Array of the predicted segmentation map. voxel_size: list, tuple or np.array The size of a voxel of the images used to compute the volumes. Return ------ A list of metrics in this order, [Dice LV, Volume LV, Err LV(ml), Dice RV, Volume RV, Err RV(ml), Dice MYO, Volume MYO, Err MYO(ml)] """ if img_gt.ndim != img_pred.ndim: raise ValueError("The arrays 'img_gt' and 'img_pred' should have the " "same dimension, {} against {}".format(img_gt.ndim, img_pred.ndim)) res = [] # Loop on each classes of the input images for c in [3, 1, 2]: # Copy the gt image to not alterate the input gt_c_i = np.copy(img_gt) gt_c_i[gt_c_i != c] = 0 # Copy the pred image to not alterate the input pred_c_i = np.copy(img_pred) pred_c_i[pred_c_i != c] = 0 # Clip the value to compute the volumes gt_c_i = np.clip(gt_c_i, 0, 1) pred_c_i = np.clip(pred_c_i, 0, 1) # Compute the Dice dice = dc(gt_c_i, pred_c_i) # Compute volume volpred = pred_c_i.sum() * np.prod(voxel_size) / 1000. volgt = gt_c_i.sum() * np.prod(voxel_size) / 1000. res += [dice, volpred, volpred-volgt] return res def compute_metrics_on_files(path_gt, path_pred): """ Function to give the metrics for two files Parameters ---------- path_gt: string Path of the ground truth image. path_pred: string Path of the predicted image. """ gt, _, header = load_nii(path_gt) pred, _, _ = load_nii(path_pred) zooms = header.get_zooms() name = os.path.basename(path_gt) name = name.split('.')[0] res = metrics(gt, pred, zooms) res = ["{:.3f}".format(r) for r in res] formatting = "{:>14}, {:>7}, {:>9}, {:>10}, {:>7}, {:>9}, {:>10}, {:>8}, {:>10}, {:>11}" print(formatting.format(*HEADER)) print(formatting.format(name, *res)) def compute_metrics_on_directories(dir_gt, dir_pred): """ Function to generate a csv file for each images of two directories. Parameters ---------- path_gt: string Directory of the ground truth segmentation maps. path_pred: string Directory of the predicted segmentation maps. """ lst_gt = sorted(glob(os.path.join(dir_gt, '*')), key=natural_order) lst_pred = sorted(glob(os.path.join(dir_pred, '*')), key=natural_order) res = [] for p_gt, p_pred in zip(lst_gt, lst_pred): if os.path.basename(p_gt) != os.path.basename(p_pred): raise ValueError("The two files don't have the same name" " {}, {}.".format(os.path.basename(p_gt), os.path.basename(p_pred))) gt, _, header = load_nii(p_gt) pred, _, _ = load_nii(p_pred) zooms = header.get_zooms() res.append(metrics(gt, pred, zooms)) lst_name_gt = [os.path.basename(gt).split(".")[0] for gt in lst_gt] res = [[n,] + r for r, n in zip(res, lst_name_gt)] df = pd.DataFrame(res, columns=HEADER) df.to_csv("results_{}.csv".format(time.strftime("%Y%m%d_%H%M%S")), index=False) def main(path_gt, path_pred): """ Main function to select which method to apply on the input parameters. """ if os.path.isfile(path_gt) and os.path.isfile(path_pred): compute_metrics_on_files(path_gt, path_pred) elif os.path.isdir(path_gt) and os.path.isdir(path_pred): compute_metrics_on_directories(path_gt, path_pred) else: raise ValueError( "The paths given needs to be two directories or two files.") if __name__ == "__main__": parser = argparse.ArgumentParser( description="Script to compute ACDC challenge metrics.") parser.add_argument("GT_IMG", type=str, help="Ground Truth image") parser.add_argument("PRED_IMG", type=str, help="Predicted image") args = parser.parse_args() main(args.GT_IMG, args.PRED_IMG)