A Python scipt for calculating the dice scores

dice_calculator.py

#!/usr/bin/env python3
import numpy
import nibabel
import argparse


def calculate_dice_coefficients(file1, file2):
    input1_nii = nibabel.load(file1)
    input2_nii = nibabel.load(file2)

    voxel_data_1 = input1_nii.get_fdata()
    voxel_data_2 = input2_nii.get_fdata()

    voxel_data_1[numpy.isnan(voxel_data_1)] = 0
    voxel_data_2[numpy.isnan(voxel_data_2)] = 0

    mask_data_1 = numpy.sign(voxel_data_1)
    mask_data_2 = numpy.sign(voxel_data_2)

    n_1 = numpy.sum(mask_data_1)
    n_2 = numpy.sum(mask_data_2)

    w_1 = numpy.sum(voxel_data_1)
    w_2 = numpy.sum(voxel_data_2)

    mask_intersection = numpy.logical_and(mask_data_1, mask_data_2)

    n_ = numpy.sum(mask_intersection)

    w_1_ = numpy.sum(voxel_data_1[mask_intersection])
    w_2_ = numpy.sum(voxel_data_2[mask_intersection])

    dice_standard = 2 * n_ / (n_1 + n_2)
    dice_weighted = (w_1_ + w_2_) / (w_1 + w_2)

    return dice_weighted, dice_standard


def main():
    parser = argparse.ArgumentParser(
        description="Calculate Dice Coefficients.",
        epilog="Original logic created by Fan Zhang and converted by Ryan Zurrin"
        )
    parser.add_argument("file1", help="Path to the first .nii or .nii.gz file")
    parser.add_argument(
        "file2", help="Path to the second .nii or .nii.gz file")
    args = parser.parse_args()

    dice_weighted, dice_standard = calculate_dice_coefficients(
        args.file1, args.file2)
    print(f" ** wDice = {dice_weighted:.2f}, sDice = {dice_standard:.2f}")


if __name__ == "__main__":
    main()