306 lines
No EOL
11 KiB
Python
306 lines
No EOL
11 KiB
Python
import json
|
||
import os
|
||
import numpy as np
|
||
import SimpleITK as sitk
|
||
from tqdm import tqdm
|
||
import time
|
||
|
||
def load_nifti(file_path):
|
||
"""加載 NIfTI 文件"""
|
||
return sitk.GetArrayFromImage(sitk.ReadImage(file_path))
|
||
|
||
def compute_dice(pred, gt):
|
||
"""計算Dice係數"""
|
||
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 separate_structures(combined_mask, structure_name=""):
|
||
"""分離左右結構"""
|
||
# 使用SimpleITK的ConnectedComponent
|
||
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: # 背景算一個，所以至少需要3個
|
||
print(f"警告：{structure_name} 未找到足夠的連通區域")
|
||
return None, None
|
||
|
||
# 排除背景（標籤0）並按大小排序
|
||
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"警告：{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]) # 使用x軸（第三維度）的平均值
|
||
|
||
centroid1 = get_centroid(region1_mask)
|
||
centroid2 = get_centroid(region2_mask)
|
||
|
||
# 根據x軸位置確定左右
|
||
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):
|
||
"""評估單個案例"""
|
||
print(f"\n測試文件: {os.path.basename(pred_path)}")
|
||
|
||
# 加載圖像
|
||
start_time = time.time()
|
||
gt_img = load_nifti(gt_path)
|
||
pred_img = load_nifti(pred_path)
|
||
load_time = time.time() - start_time
|
||
|
||
print(f"圖像大小: {gt_img.shape}")
|
||
print(f"加載時間: {load_time:.2f}秒")
|
||
|
||
results = {}
|
||
total_time = 0
|
||
|
||
# 處理眼睛
|
||
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)
|
||
|
||
# 形態學處理眼睛
|
||
combined_pred_eye_sitk = sitk.GetImageFromArray(combined_pred_eye)
|
||
closing_filter = sitk.BinaryMorphologicalClosingImageFilter()
|
||
closing_filter.SetKernelRadius(2)
|
||
processed_pred_eye = sitk.GetArrayFromImage(closing_filter.Execute(combined_pred_eye_sitk))
|
||
|
||
# 計算眼睛的Dice
|
||
eye_dice = compute_dice(processed_pred_eye, combined_gt_eye)
|
||
eye_time = time.time() - start_time
|
||
total_time += eye_time
|
||
results['combined_eye'] = eye_dice
|
||
print(f"眼睛合併計算時間: {eye_time:.2f}秒, Dice: {eye_dice:.4f}")
|
||
|
||
# 分離眼睛
|
||
right_pred_eye, left_pred_eye = separate_structures(processed_pred_eye, "眼睛")
|
||
right_gt_eye, left_gt_eye = separate_structures(combined_gt_eye, "眼睛")
|
||
|
||
# 處理視神經
|
||
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)
|
||
|
||
# 形態學處理視神經
|
||
combined_pred_nerve_sitk = sitk.GetImageFromArray(combined_pred_nerve)
|
||
processed_pred_nerve = sitk.GetArrayFromImage(closing_filter.Execute(combined_pred_nerve_sitk))
|
||
|
||
# 計算視神經的Dice
|
||
nerve_dice = compute_dice(processed_pred_nerve, combined_gt_nerve)
|
||
nerve_time = time.time() - start_time
|
||
total_time += nerve_time
|
||
results['combined_nerve'] = nerve_dice
|
||
print(f"視神經合併計算時間: {nerve_time:.2f}秒, Dice: {nerve_dice:.4f}")
|
||
|
||
# 分離視神經
|
||
right_pred_nerve, left_pred_nerve = separate_structures(processed_pred_nerve, "視神經")
|
||
right_gt_nerve, left_gt_nerve = separate_structures(combined_gt_nerve, "視神經")
|
||
|
||
# 計算各個標籤的Dice
|
||
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)
|
||
elif label == 2 and right_pred_eye is not None: # Right Eye
|
||
gt_binary = right_gt_eye
|
||
pred_binary = right_pred_eye
|
||
elif label == 3 and left_pred_eye is not None: # Left Eye
|
||
gt_binary = left_gt_eye
|
||
pred_binary = left_pred_eye
|
||
elif label == 4: # Optic Chiasm
|
||
gt_binary = (gt_img == label).astype(np.uint8)
|
||
pred_binary = (pred_img == label).astype(np.uint8)
|
||
elif label == 5 and right_pred_nerve is not None: # Right Optic Nerve
|
||
gt_binary = right_gt_nerve
|
||
pred_binary = right_pred_nerve
|
||
elif label == 6 and left_pred_nerve is not None: # Left Optic Nerve
|
||
gt_binary = left_gt_nerve
|
||
pred_binary = left_pred_nerve
|
||
else:
|
||
continue
|
||
|
||
try:
|
||
dice = compute_dice(pred_binary, gt_binary)
|
||
compute_time = time.time() - start_time
|
||
total_time += compute_time
|
||
results[label] = dice
|
||
print(f"標籤 {label} 計算時間: {compute_time:.2f}秒, Dice: {dice:.4f}")
|
||
except Exception as e:
|
||
print(f"處理標籤 {label} 時出錯：{str(e)}")
|
||
results[label] = np.nan
|
||
continue
|
||
|
||
print(f"單個案例總計算時間: {total_time:.2f}秒")
|
||
return results
|
||
|
||
def evaluate_fold(gt_folder, pred_folder, labels):
|
||
"""評估單個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(pred_folder) if f.endswith('.nii.gz')])
|
||
|
||
for pred_file in tqdm(pred_files, desc=f"Processing {os.path.basename(pred_folder)}"):
|
||
gt_path = os.path.join(gt_folder, pred_file)
|
||
pred_path = os.path.join(pred_folder, pred_file)
|
||
|
||
if not os.path.exists(gt_path):
|
||
print(f"找不到對應的真實標籤文件：{gt_path}")
|
||
continue
|
||
|
||
case_results = evaluate_case(gt_path, pred_path, labels)
|
||
|
||
# 收集結果
|
||
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, base_pred_folder, model_name, labels, organ_names):
|
||
"""評估單個模型的所有fold"""
|
||
print(f"\n開始評估模型：{model_name}")
|
||
all_results = {label: [] for label in labels}
|
||
all_results['combined_eye'] = []
|
||
all_results['combined_nerve'] = []
|
||
fold_results = {}
|
||
|
||
# 處理 fold_0 到 fold_4
|
||
for fold in range(5):
|
||
print(f"\n處理 {model_name} - Fold {fold}")
|
||
pred_folder = os.path.join(base_pred_folder, f"fold_{fold}")
|
||
if not os.path.exists(pred_folder):
|
||
print(f"找不到預測文件夾：{pred_folder}")
|
||
continue
|
||
|
||
results = evaluate_fold(gt_folder, pred_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
|
||
}
|
||
|
||
# 保存結果到指定路径
|
||
save_dir = "/mnt/1248/onlylian/CTC/Dice/test"
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
filename = os.path.join(save_dir, f'CTC_test_dice_results_{model_name}.json')
|
||
|
||
model_results = {
|
||
'summary': summary,
|
||
'fold_results': fold_results
|
||
}
|
||
|
||
with open(filename, 'w') as f:
|
||
json.dump(model_results, f, indent=4)
|
||
print(f"\n{model_name}結果已保存到 {filename}")
|
||
|
||
# 打印結果
|
||
print(f"\n{model_name} 模型的 Dice 結果（所有fold的平均值）:")
|
||
|
||
# 打印器官結果
|
||
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}: 無法計算 (n={n})")
|
||
|
||
# 打印合併結果
|
||
print("\n合併結構的結果:")
|
||
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 = "合併眼睛" if key == 'combined_eye' else "合併視神經"
|
||
if not np.isnan(mean):
|
||
print(f"{name}: {mean:.4f} ± {std:.4f} (n={n})")
|
||
else:
|
||
print(f"{name}: 無法計算 (n={n})")
|
||
|
||
return summary
|
||
|
||
if __name__ == "__main__":
|
||
# 設置路徑和標籤
|
||
gt_folder = "/mnt/1248/onlylian/nnUNet/nnUNet_raw/Dataset666/labelsTs"
|
||
models = {
|
||
'2D': '/mnt/1248/onlylian/CTC/output_predictions/2d',
|
||
'3D_fullres': '/mnt/1248/onlylian/CTC/output_predictions/3d_fullres',
|
||
'3D_lowres': '/mnt/1248/onlylian/CTC/output_predictions/3d_lowres'
|
||
}
|
||
labels = [1, 2, 3, 4, 5, 6]
|
||
organ_names = {
|
||
1: "Brainstem",
|
||
2: "Right_Eye",
|
||
3: "Left_Eye",
|
||
4: "Optic_Chiasm",
|
||
5: "Right_Optic_Nerve",
|
||
6: "Left_Optic_Nerve"
|
||
}
|
||
|
||
# 讓用戶選擇要運行哪個模型
|
||
print("可用的模型：")
|
||
for i, model_name in enumerate(models.keys(), 1):
|
||
print(f"{i}. {model_name}")
|
||
|
||
choice = input("\n請選擇要運行的模型 (輸入數字1-3): ")
|
||
try:
|
||
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)
|
||
except (ValueError, IndexError):
|
||
print("無效的選擇！請輸入1-3之間的數字。") |