325 lines
No EOL
13 KiB
Python
325 lines
No EOL
13 KiB
Python
import os
|
|
import time
|
|
import numpy as np
|
|
import SimpleITK as sitk
|
|
from tqdm import tqdm
|
|
import json
|
|
|
|
def load_nifti(file_path):
|
|
"""Load a NIfTI file"""
|
|
return sitk.GetArrayFromImage(sitk.ReadImage(file_path))
|
|
|
|
def compute_dice(pred, gt):
|
|
"""Compute Dice coefficient"""
|
|
intersection = np.sum(pred & gt)
|
|
sum_pred = np.sum(pred)
|
|
sum_gt = np.sum(gt)
|
|
|
|
if sum_pred == 0 and sum_gt == 0:
|
|
return 1.0
|
|
elif sum_pred == 0 or sum_gt == 0:
|
|
return 0.0
|
|
|
|
return 2.0 * intersection / (sum_pred + sum_gt)
|
|
|
|
def preprocess_binary_mask(mask, apply_closing=True):
|
|
"""Preprocess binary mask with morphological operations"""
|
|
mask_sitk = sitk.GetImageFromArray(mask)
|
|
|
|
if apply_closing:
|
|
closing_filter = sitk.BinaryMorphologicalClosingImageFilter()
|
|
closing_filter.SetKernelRadius(2)
|
|
mask = sitk.GetArrayFromImage(closing_filter.Execute(mask_sitk))
|
|
|
|
return mask
|
|
|
|
def separate_structures(combined_mask, structure_name=""):
|
|
"""Separate left and right structures"""
|
|
sitk_image = sitk.GetImageFromArray(combined_mask)
|
|
labeled_components = sitk.ConnectedComponent(sitk_image)
|
|
labeled_array = sitk.GetArrayFromImage(labeled_components)
|
|
|
|
unique, counts = np.unique(labeled_array, return_counts=True)
|
|
if len(unique) < 3:
|
|
print(f"Warning: Not enough connected regions found for {structure_name}")
|
|
return None, None
|
|
|
|
region_sizes = [(i, count) for i, count in zip(unique[1:], counts[1:])]
|
|
sorted_regions = sorted(region_sizes, key=lambda x: x[1], reverse=True)
|
|
|
|
if len(sorted_regions) < 2:
|
|
print(f"Warning: Only one major region found for {structure_name}")
|
|
return None, None
|
|
|
|
region1_mask = (labeled_array == sorted_regions[0][0])
|
|
region2_mask = (labeled_array == sorted_regions[1][0])
|
|
|
|
def get_centroid(mask):
|
|
indices = np.where(mask)
|
|
return np.mean(indices[2])
|
|
|
|
centroid1 = get_centroid(region1_mask)
|
|
centroid2 = get_centroid(region2_mask)
|
|
|
|
if centroid1 > centroid2:
|
|
right = region1_mask
|
|
left = region2_mask
|
|
else:
|
|
right = region2_mask
|
|
left = region1_mask
|
|
|
|
return right.astype(np.uint8), left.astype(np.uint8)
|
|
|
|
def evaluate_case(gt_path, pred_path, labels):
|
|
"""Evaluate a single case"""
|
|
print(f"\nTest file: {os.path.basename(pred_path)}")
|
|
|
|
try:
|
|
start_time = time.time()
|
|
gt_img = load_nifti(gt_path)
|
|
pred_img = load_nifti(pred_path)
|
|
load_time = time.time() - start_time
|
|
|
|
print(f"Image size: {gt_img.shape}")
|
|
print(f"Loading time: {load_time:.2f}s")
|
|
|
|
results = {}
|
|
total_time = 0
|
|
|
|
# Process eyes
|
|
start_time = time.time()
|
|
combined_gt_eye = ((gt_img == 2) | (gt_img == 3)).astype(np.uint8)
|
|
combined_pred_eye = ((pred_img == 2) | (pred_img == 3)).astype(np.uint8)
|
|
|
|
processed_pred_eye = preprocess_binary_mask(combined_pred_eye)
|
|
processed_gt_eye = preprocess_binary_mask(combined_gt_eye)
|
|
|
|
eye_dice = compute_dice(processed_pred_eye, processed_gt_eye)
|
|
eye_time = time.time() - start_time
|
|
total_time += eye_time
|
|
results['combined_eye'] = eye_dice
|
|
print(f"Eyes combined computation time: {eye_time:.2f}s, Dice: {eye_dice:.4f}")
|
|
|
|
right_pred_eye, left_pred_eye = separate_structures(processed_pred_eye, "Eyes")
|
|
right_gt_eye, left_gt_eye = separate_structures(processed_gt_eye, "Eyes")
|
|
|
|
# Process optic nerves
|
|
start_time = time.time()
|
|
combined_gt_nerve = ((gt_img == 5) | (gt_img == 6)).astype(np.uint8)
|
|
combined_pred_nerve = ((pred_img == 5) | (pred_img == 6)).astype(np.uint8)
|
|
|
|
processed_pred_nerve = preprocess_binary_mask(combined_pred_nerve)
|
|
processed_gt_nerve = preprocess_binary_mask(combined_gt_nerve)
|
|
|
|
nerve_dice = compute_dice(processed_pred_nerve, processed_gt_nerve)
|
|
nerve_time = time.time() - start_time
|
|
total_time += nerve_time
|
|
results['combined_nerve'] = nerve_dice
|
|
print(f"Optic nerves combined computation time: {nerve_time:.2f}s, Dice: {nerve_dice:.4f}")
|
|
|
|
right_pred_nerve, left_pred_nerve = separate_structures(processed_pred_nerve, "Optic Nerves")
|
|
right_gt_nerve, left_gt_nerve = separate_structures(processed_gt_nerve, "Optic Nerves")
|
|
|
|
for label in labels:
|
|
start_time = time.time()
|
|
|
|
if label == 1: # Brainstem
|
|
gt_binary = (gt_img == label).astype(np.uint8)
|
|
pred_binary = (pred_img == label).astype(np.uint8)
|
|
apply_closing = True
|
|
elif label == 2 and right_pred_eye is not None: # Right Eye
|
|
gt_binary = right_gt_eye
|
|
pred_binary = right_pred_eye
|
|
apply_closing = False # Already processed
|
|
elif label == 3 and left_pred_eye is not None: # Left Eye
|
|
gt_binary = left_gt_eye
|
|
pred_binary = left_pred_eye
|
|
apply_closing = False # Already processed
|
|
elif label == 4: # Optic Chiasm
|
|
gt_binary = (gt_img == label).astype(np.uint8)
|
|
pred_binary = (pred_img == label).astype(np.uint8)
|
|
apply_closing = True
|
|
elif label == 5 and right_pred_nerve is not None: # Right Optic Nerve
|
|
gt_binary = right_gt_nerve
|
|
pred_binary = right_pred_nerve
|
|
apply_closing = False # Already processed
|
|
elif label == 6 and left_pred_nerve is not None: # Left Optic Nerve
|
|
gt_binary = left_gt_nerve
|
|
pred_binary = left_pred_nerve
|
|
apply_closing = False # Already processed
|
|
elif label == 7: # TV
|
|
gt_binary = (gt_img == label).astype(np.uint8)
|
|
pred_binary = (pred_img == label).astype(np.uint8)
|
|
apply_closing = True
|
|
else:
|
|
continue
|
|
|
|
try:
|
|
if apply_closing:
|
|
pred_binary = preprocess_binary_mask(pred_binary)
|
|
gt_binary = preprocess_binary_mask(gt_binary)
|
|
|
|
dice = compute_dice(pred_binary, gt_binary)
|
|
compute_time = time.time() - start_time
|
|
total_time += compute_time
|
|
results[label] = dice
|
|
print(f"Label {label} computation time: {compute_time:.2f}s, Dice: {dice:.4f}")
|
|
except Exception as e:
|
|
print(f"Error processing label {label}: {str(e)}")
|
|
results[label] = np.nan
|
|
continue
|
|
|
|
print(f"Total computation time for this case: {total_time:.2f}s")
|
|
return results
|
|
except Exception as e:
|
|
print(f"Error evaluating case: {str(e)}")
|
|
return None
|
|
|
|
def evaluate_validation_fold(gt_base_folder, validation_folder, labels):
|
|
"""Evaluate the validation set for a single fold"""
|
|
fold_results = {label: [] for label in labels}
|
|
fold_results['combined_eye'] = []
|
|
fold_results['combined_nerve'] = []
|
|
|
|
pred_files = sorted([f for f in os.listdir(validation_folder) if f.endswith('.nii.gz')])
|
|
|
|
for pred_file in tqdm(pred_files, desc=f"Processing {os.path.basename(validation_folder)}"):
|
|
patient_id = pred_file.split('.nii.gz')[0]
|
|
gt_path = os.path.join(gt_base_folder, f"{patient_id}.nii.gz")
|
|
pred_path = os.path.join(validation_folder, pred_file)
|
|
|
|
if not os.path.exists(gt_path):
|
|
print(f"Could not find corresponding ground truth file: {gt_path}")
|
|
continue
|
|
|
|
case_results = evaluate_case(gt_path, pred_path, labels)
|
|
if case_results is not None:
|
|
for key in case_results:
|
|
if not np.isnan(case_results[key]):
|
|
if key in fold_results:
|
|
fold_results[key].append(case_results[key])
|
|
|
|
return fold_results
|
|
|
|
def evaluate_model(gt_folder, model_path, model_name, labels, organ_names, output_dir):
|
|
"""Evaluate the validation set for all folds of a model"""
|
|
print(f"\nStarting evaluation of model: {model_name}")
|
|
all_results = {label: [] for label in labels}
|
|
all_results['combined_eye'] = []
|
|
all_results['combined_nerve'] = []
|
|
fold_results = {}
|
|
|
|
for fold in range(5):
|
|
validation_folder = os.path.join(model_path, f"fold_{fold}", "validation")
|
|
|
|
if not os.path.exists(validation_folder):
|
|
print(f"Could not find fold_{fold} validation folder: {validation_folder}")
|
|
continue
|
|
|
|
print(f"\nProcessing {model_name} - Fold {fold} validation set")
|
|
results = evaluate_validation_fold(gt_folder, validation_folder, labels)
|
|
fold_results[f"fold_{fold}"] = results
|
|
|
|
for key in results:
|
|
all_results[key].extend(results[key])
|
|
|
|
summary = {}
|
|
for key in all_results:
|
|
if all_results[key]:
|
|
mean_dice = np.mean(all_results[key])
|
|
std_dice = np.std(all_results[key])
|
|
summary[key] = {
|
|
'mean_dice': float(mean_dice),
|
|
'std_dice': float(std_dice),
|
|
'n_samples': len(all_results[key])
|
|
}
|
|
else:
|
|
summary[key] = {
|
|
'mean_dice': np.nan,
|
|
'std_dice': np.nan,
|
|
'n_samples': 0
|
|
}
|
|
|
|
model_results = {
|
|
'summary': summary,
|
|
'fold_results': fold_results
|
|
}
|
|
|
|
output_path = os.path.join(output_dir, 'T1C_OAR_TV', 'Dice', 'validation')
|
|
os.makedirs(output_path, exist_ok=True)
|
|
|
|
filename = os.path.join(output_path, f'dice_results_{model_name}.json')
|
|
with open(filename, 'w') as f:
|
|
json.dump(model_results, f, indent=4)
|
|
print(f"\nValidation set results for {model_name} saved to {filename}")
|
|
|
|
print(f"\nValidation set Dice results for {model_name} (average across all folds):")
|
|
|
|
# Print individual organ results
|
|
print("\nOrgan-specific results:")
|
|
for label in labels:
|
|
if label in summary:
|
|
mean = summary[label]['mean_dice']
|
|
std = summary[label]['std_dice']
|
|
n = summary[label]['n_samples']
|
|
organ_name = organ_names.get(label, f"Label_{label}")
|
|
if not np.isnan(mean):
|
|
print(f"{organ_name}: {mean:.4f} ± {std:.4f} (n={n})")
|
|
else:
|
|
print(f"{organ_name}: Unable to compute (n={n})")
|
|
|
|
# Print combined structure results
|
|
print("\nCombined structure results:")
|
|
for key in ['combined_eye', 'combined_nerve']:
|
|
if key in summary:
|
|
mean = summary[key]['mean_dice']
|
|
std = summary[key]['std_dice']
|
|
n = summary[key]['n_samples']
|
|
name = "Combined Eyes" if key == 'combined_eye' else "Combined Optic Nerves"
|
|
if not np.isnan(mean):
|
|
print(f"{name}: {mean:.4f} ± {std:.4f} (n={n})")
|
|
else:
|
|
print(f"{name}: Unable to compute (n={n})")
|
|
|
|
return summary
|
|
|
|
if __name__ == "__main__":
|
|
base_dir = "/mnt/1248/onlylian"
|
|
gt_folder = os.path.join(base_dir, "nnUNet/nnUNet_raw/Dataset012_OAR_TV/labelsTr")
|
|
output_dir = base_dir
|
|
|
|
models = {
|
|
'2D': os.path.join(base_dir, "nnUNet/nnUNet_results/Dataset012_OAR_TV/nnUNetTrainer__nnUNetPlans__2d"),
|
|
'3D_fullres': os.path.join(base_dir, "nnUNet/nnUNet_results/Dataset012_OAR_TV/nnUNetTrainer__nnUNetPlans__3d_fullres"),
|
|
'3D_lowres': os.path.join(base_dir, "nnUNet/nnUNet_results/Dataset012_OAR_TV/nnUNetTrainer__nnUNetPlans__3d_lowres")
|
|
}
|
|
|
|
labels = [1, 2, 3, 4, 5, 6, 7]
|
|
organ_names = {
|
|
1: "Brainstem",
|
|
2: "Right_Eye",
|
|
3: "Left_Eye",
|
|
4: "Optic_Chiasm",
|
|
5: "Right_Optic_Nerve",
|
|
6: "Left_Optic_Nerve",
|
|
7: "TV"
|
|
}
|
|
|
|
print("Available models:")
|
|
for i, model_name in enumerate(models.keys(), 1):
|
|
print(f"{i}. {model_name}")
|
|
|
|
try:
|
|
choice = input("\nPlease select a model to run (enter 1-3): ")
|
|
model_index = int(choice) - 1
|
|
model_names = list(models.keys())
|
|
selected_model = model_names[model_index]
|
|
evaluate_model(gt_folder, models[selected_model], selected_model, labels, organ_names, output_dir)
|
|
except (ValueError, IndexError):
|
|
print("Invalid choice! Please enter a number between 1 and 3.")
|
|
except KeyboardInterrupt:
|
|
print("\nProgram interrupted by user")
|
|
print("You can run the program again to continue processing")
|
|
except Exception as e:
|
|
print(f"\nAn unexpected error occurred: {str(e)}")
|
|
print("Please check the error message and try again") |