forked from MIC-DKFZ/nnUNet
-
Notifications
You must be signed in to change notification settings - Fork 0
/
COVID-19-Seg-Evaluation.py
83 lines (65 loc) · 3.5 KB
/
COVID-19-Seg-Evaluation.py
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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
import numpy as np
import nibabel as nb
import os
from collections import OrderedDict
import pandas as pd
from SurfaceDice import compute_surface_distances, compute_surface_dice_at_tolerance, compute_dice_coefficient, compute_iou_coefficient
join = os.path.join
seg_path = 'output'
gt_path = 'save/nnUNet_raw/nnUNet_raw_data/Task101_Infection/labelsTr'
save_path = 'output'
save_name = 'Task_num_SegMetric.csv'
filenames = os.listdir(seg_path)
filenames.sort()
num_labels = np.max(nb.load(join(gt_path, filenames[0])).get_fdata())
seg_metrics = OrderedDict()
seg_metrics['Name'] = list()
if num_labels == 1:
seg_metrics['LesionDSC'] = list()
seg_metrics['LesionIOU'] = list()
seg_metrics['LesionNSD-3mm'] = list()
elif num_labels == 2:
seg_metrics['L-lungDSC'] = list()
seg_metrics['L-lung-1mm'] = list()
seg_metrics['R-lungDSC'] = list()
seg_metrics['R-lung-1mm'] = list()
else:
seg_metrics['L-lungDSC'] = list()
seg_metrics['L-lung-1mm'] = list()
seg_metrics['R-lungDSC'] = list()
seg_metrics['R-lung-1mm'] = list()
seg_metrics['LesionDSC'] = list()
seg_metrics['LesionNSD-3mm'] = list()
for name in filenames:
seg_metrics['Name'].append(name)
# load grond truth and segmentation
gt_nii = nb.load(join(gt_path, name))
case_spacing = gt_nii.header.get_zooms()
gt_data = np.uint8(gt_nii.get_fdata())
seg_data = nb.load(join(seg_path, name)).get_fdata()
labels = np.unique(gt_data)[1:]
labels = labels.tolist()
if num_labels == 1: # Lesion
surface_distances = compute_surface_distances(gt_data == 1, seg_data == 1, case_spacing)
seg_metrics['LesionDSC'].append(compute_dice_coefficient(gt_data == 1, seg_data == 1))
seg_metrics['LesionIOU'].append(compute_iou_coefficient(gt_data == 1, seg_data == 1))
seg_metrics['LesionNSD-3mm'].append(compute_surface_dice_at_tolerance(surface_distances, 3))
elif num_labels == 2: # left and right lung
surface_distances = compute_surface_distances(gt_data == 1, seg_data == 1, case_spacing)
seg_metrics['L-lungDSC'].append(compute_dice_coefficient(gt_data == 1, seg_data == 1))
seg_metrics['L-lung-1mm'].append(compute_surface_dice_at_tolerance(surface_distances, 1))
surface_distances = compute_surface_distances(gt_data == 2, seg_data == 2, case_spacing)
seg_metrics['R-lungDSC'].append(compute_dice_coefficient(gt_data == 2, seg_data == 2))
seg_metrics['R-lung-1mm'].append(compute_surface_dice_at_tolerance(surface_distances, 1))
else: # left lung, right lung and infections
surface_distances = compute_surface_distances(gt_data == 1, seg_data == 1, case_spacing)
seg_metrics['L-lungDSC'].append(compute_dice_coefficient(gt_data == 1, seg_data == 1))
seg_metrics['L-lung-1mm'].append(compute_surface_dice_at_tolerance(surface_distances, 1))
surface_distances = compute_surface_distances(gt_data == 2, seg_data == 2, case_spacing)
seg_metrics['R-lungDSC'].append(compute_dice_coefficient(gt_data == 2, seg_data == 2))
seg_metrics['R-lung-1mm'].append(compute_surface_dice_at_tolerance(surface_distances, 1))
surface_distances = compute_surface_distances(gt_data == 3, seg_data == 3, case_spacing)
seg_metrics['LesionDSC'].append(compute_dice_coefficient(gt_data == 3, seg_data == 3))
seg_metrics['LesionNSD-3mm'].append(compute_surface_dice_at_tolerance(surface_distances, 3))
dataframe = pd.DataFrame(seg_metrics)
dataframe.to_csv(join(save_path, save_name), index=False)