318 lines
13 KiB
Python
318 lines
13 KiB
Python
|
import json
|
||
|
|
import os
|
||
|
|
import numpy as np
|
||
|
|
import SimpleITK as sitk
|
||
|
|
import pandas as pd
|
||
|
|
import matplotlib.pyplot as plt
|
||
|
|
import seaborn as sns
|
||
|
|
from tqdm import tqdm
|
||
|
|
import time
|
||
|
|
from datetime import datetime
|
||
|
|
from scipy.ndimage import distance_transform_edt
|
||
|
|
import multiprocessing
|
||
|
|
from joblib import Parallel, delayed
|
||
|
|
|
||
|
|
class IntegratedAnalyzer:
|
||
|
|
def __init__(self, base_path):
|
||
|
|
"""初始化分析器"""
|
||
|
|
self.base_path = base_path
|
||
|
|
self.organ_names = {
|
||
|
|
1: "Brainstem",
|
||
|
|
2: "Right_Eye",
|
||
|
|
3: "Left_Eye",
|
||
|
|
4: "Optic_Chiasm",
|
||
|
|
5: "Right_Optic_Nerve",
|
||
|
|
6: "Left_Optic_Nerve",
|
||
|
|
7: "TV"
|
||
|
|
}
|
||
|
|
self.configurations = {
|
||
|
|
'Dataset012_OAR_TV': {
|
||
|
|
'2d': {
|
||
|
|
'model': 'nnUNetTrainer__nnUNetPlans__2d',
|
||
|
|
'path': '/mnt/1248/onlylian/T1C_OAR_TV/output_predictions/2d',
|
||
|
|
'folds': range(5)
|
||
|
|
},
|
||
|
|
'3d_fullres': {
|
||
|
|
'model': 'nnUNetTrainer__nnUNetPlans__3d_fullres',
|
||
|
|
'path': '/mnt/1248/onlylian/T1C_OAR_TV/output_predictions/3d_fullres',
|
||
|
|
'folds': range(5)
|
||
|
|
},
|
||
|
|
'3d_lowres': {
|
||
|
|
'model': 'nnUNetTrainer__nnUNetPlans__3d_lowres',
|
||
|
|
'path': '/mnt/1248/onlylian/T1C_OAR_TV/output_predictions/3d_lowres',
|
||
|
|
'folds': range(5)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
# 設定並行處理的核心數
|
||
|
|
self.n_jobs = max(1, multiprocessing.cpu_count() - 1)
|
||
|
|
self.setup_paths()
|
||
|
|
self.create_output_folders()
|
||
|
|
|
||
|
|
def setup_paths(self):
|
||
|
|
"""設置所有需要的路徑"""
|
||
|
|
self.dataset_paths = {}
|
||
|
|
for dataset, model_configs in self.configurations.items():
|
||
|
|
self.dataset_paths[dataset] = {
|
||
|
|
'ground_truth': os.path.join(
|
||
|
|
"/mnt/1248/onlylian/nnUNet/nnUNet_raw",
|
||
|
|
dataset,
|
||
|
|
"labelsTs"
|
||
|
|
),
|
||
|
|
'predictions': {}
|
||
|
|
}
|
||
|
|
for model_type, config in model_configs.items():
|
||
|
|
self.dataset_paths[dataset]['predictions'][model_type] = {
|
||
|
|
fold: os.path.join(config['path'], f"fold_{fold}")
|
||
|
|
for fold in config['folds']
|
||
|
|
}
|
||
|
|
|
||
|
|
self.output_path = os.path.join(
|
||
|
|
self.base_path,
|
||
|
|
"analysis_results",
|
||
|
|
datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
|
|
)
|
||
|
|
|
||
|
|
def create_output_folders(self):
|
||
|
|
"""建立輸出資料夾"""
|
||
|
|
os.makedirs(self.output_path, exist_ok=True)
|
||
|
|
|
||
|
|
def setup_font_configuration(self):
|
||
|
|
"""設置基本字型配置"""
|
||
|
|
plt.rcParams.update(plt.rcParamsDefault)
|
||
|
|
plt.rcParams['font.family'] = 'DejaVu Sans'
|
||
|
|
plt.rcParams['axes.unicode_minus'] = False
|
||
|
|
plt.rcParams['font.size'] = 12
|
||
|
|
plt.rcParams['axes.titlesize'] = 14
|
||
|
|
plt.rcParams['axes.labelsize'] = 12
|
||
|
|
plt.rcParams['xtick.labelsize'] = 10
|
||
|
|
plt.rcParams['ytick.labelsize'] = 10
|
||
|
|
return True
|
||
|
|
|
||
|
|
def load_nifti(self, file_path):
|
||
|
|
"""載入 NIfTI 檔案"""
|
||
|
|
return sitk.GetArrayFromImage(sitk.ReadImage(file_path))
|
||
|
|
|
||
|
|
def separate_structures(self, combined_mask, structure_name=""):
|
||
|
|
"""分離左右結構，改進的版本"""
|
||
|
|
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: # 背景+至少兩個區域
|
||
|
|
# 只在實際需要時輸出警告
|
||
|
|
if np.sum(combined_mask) > 0: # 只有當mask中確實有內容時才輸出警告
|
||
|
|
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:
|
||
|
|
if np.sum(combined_mask) > 0: # 只有當mask中確實有內容時才輸出警告
|
||
|
|
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 compute_dice(self, 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 evaluate_case(self, gt_path, pred_path):
|
||
|
|
"""評估單一案例，包含腫瘤距離分析"""
|
||
|
|
gt_img = self.load_nifti(gt_path)
|
||
|
|
pred_img = self.load_nifti(pred_path)
|
||
|
|
|
||
|
|
# 計算腫瘤距離圖
|
||
|
|
tumor_mask = (gt_img == 7) # 7 是腫瘤的標籤
|
||
|
|
tumor_distance = distance_transform_edt(~tumor_mask)
|
||
|
|
|
||
|
|
results = {}
|
||
|
|
distance_results = {}
|
||
|
|
|
||
|
|
# Process eyes
|
||
|
|
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))
|
||
|
|
|
||
|
|
results['combined_eye'] = self.compute_dice(processed_pred_eye, combined_gt_eye)
|
||
|
|
eye_distances = tumor_distance[combined_gt_eye > 0]
|
||
|
|
if len(eye_distances) > 0:
|
||
|
|
distance_results['combined_eye'] = np.mean(eye_distances)
|
||
|
|
|
||
|
|
right_pred_eye, left_pred_eye = self.separate_structures(processed_pred_eye, "Eyes")
|
||
|
|
right_gt_eye, left_gt_eye = self.separate_structures(combined_gt_eye, "Eyes")
|
||
|
|
|
||
|
|
# Process optic nerves
|
||
|
|
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))
|
||
|
|
|
||
|
|
results['combined_nerve'] = self.compute_dice(processed_pred_nerve, combined_gt_nerve)
|
||
|
|
nerve_distances = tumor_distance[combined_gt_nerve > 0]
|
||
|
|
if len(nerve_distances) > 0:
|
||
|
|
distance_results['combined_nerve'] = np.mean(nerve_distances)
|
||
|
|
|
||
|
|
right_pred_nerve, left_pred_nerve = self.separate_structures(processed_pred_nerve, "Optic Nerves")
|
||
|
|
right_gt_nerve, left_gt_nerve = self.separate_structures(combined_gt_nerve, "Optic Nerves")
|
||
|
|
|
||
|
|
# Process individual structures
|
||
|
|
for label in self.organ_names.keys():
|
||
|
|
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
|
||
|
|
elif label == 7: # TV (tumor)
|
||
|
|
gt_binary = (gt_img == label).astype(np.uint8)
|
||
|
|
pred_binary = (pred_img == label).astype(np.uint8)
|
||
|
|
else:
|
||
|
|
continue
|
||
|
|
|
||
|
|
try:
|
||
|
|
dice = self.compute_dice(pred_binary, gt_binary)
|
||
|
|
results[label] = dice
|
||
|
|
|
||
|
|
# 計算每個器官到腫瘤的距離
|
||
|
|
if label != 7: # 不計算腫瘤自身的距離
|
||
|
|
organ_distances = tumor_distance[gt_binary > 0]
|
||
|
|
if len(organ_distances) > 0:
|
||
|
|
distance_results[label] = np.mean(organ_distances)
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Error processing label {label}: {str(e)}")
|
||
|
|
results[label] = np.nan
|
||
|
|
distance_results[label] = np.nan
|
||
|
|
|
||
|
|
return results, distance_results
|
||
|
|
|
||
|
|
def process_single_case(self, args):
|
||
|
|
"""處理單一案例的包裝函數"""
|
||
|
|
gt_path, pred_path = args
|
||
|
|
return self.evaluate_case(gt_path, pred_path)
|
||
|
|
|
||
|
|
def process_fold(self, model_type, fold, model_config, dataset):
|
||
|
|
"""處理單一fold的函數"""
|
||
|
|
results = []
|
||
|
|
pred_folder = os.path.join(model_config['path'], f"fold_{fold}")
|
||
|
|
gt_folder = self.dataset_paths[dataset]['ground_truth']
|
||
|
|
|
||
|
|
if not os.path.exists(pred_folder):
|
||
|
|
print(f"Warning: Prediction folder does not exist: {pred_folder}")
|
||
|
|
return results
|
||
|
|
|
||
|
|
pred_files = [f for f in os.listdir(pred_folder) if f.endswith('.nii.gz')]
|
||
|
|
|
||
|
|
# 準備參數
|
||
|
|
case_args = [(os.path.join(gt_folder, pred_file),
|
||
|
|
os.path.join(pred_folder, pred_file))
|
||
|
|
for pred_file in pred_files
|
||
|
|
if os.path.exists(os.path.join(gt_folder, pred_file))]
|
||
|
|
|
||
|
|
# 使用 joblib 進行並行處理
|
||
|
|
processed_results = Parallel(n_jobs=self.n_jobs)(
|
||
|
|
delayed(self.process_single_case)(args) for args in tqdm(case_args,
|
||
|
|
desc=f"Processing fold {fold}")
|
||
|
|
)
|
||
|
|
|
||
|
|
# 整理結果
|
||
|
|
for (dice_results, distance_results), (gt_path, pred_path) in zip(processed_results, case_args):
|
||
|
|
pred_file = os.path.basename(pred_path)
|
||
|
|
for key, value in dice_results.items():
|
||
|
|
if isinstance(key, int):
|
||
|
|
result_entry = {
|
||
|
|
'model_type': model_type,
|
||
|
|
'fold': fold,
|
||
|
|
'case_id': pred_file,
|
||
|
|
'structure': self.organ_names[key],
|
||
|
|
'dice_score': value,
|
||
|
|
'distance_to_tumor': distance_results.get(key, np.nan)
|
||
|
|
}
|
||
|
|
results.append(result_entry)
|
||
|
|
|
||
|
|
return results
|
||
|
|
|
||
|
|
def analyze_and_visualize(self):
|
||
|
|
"""分析並視覺化結果，包含距離分析"""
|
||
|
|
self.setup_font_configuration()
|
||
|
|
plots_path = os.path.join(self.output_path, "plots")
|
||
|
|
os.makedirs(plots_path, exist_ok=True)
|
||
|
|
|
||
|
|
for dataset in self.configurations.keys():
|
||
|
|
all_results = []
|
||
|
|
|
||
|
|
for model_type, model_config in self.configurations[dataset].items():
|
||
|
|
print(f"\nProcessing {dataset} - {model_type}")
|
||
|
|
|
||
|
|
# 並行處理每個 fold
|
||
|
|
fold_results = Parallel(n_jobs=self.n_jobs)(
|
||
|
|
delayed(self.process_fold)(model_type, fold, model_config, dataset)
|
||
|
|
for fold in model_config['folds']
|
||
|
|
)
|
||
|
|
|
||
|
|
# 整合所有結果
|
||
|
|
for results in fold_results:
|
||
|
|
all_results.extend(results)
|
||
|
|
|
||
|
|
# 創建DataFrame並視覺化
|
||
|
|
if all_results:
|
||
|
|
df = pd.DataFrame(all_results)
|
||
|
|
|
||
|
|
# 基本的模型比較圖
|
||
|
|
plt.figure(figsize=(12, 6))
|
||
|
|
sns.boxplot(data=df, x='model_type', y='dice_score')
|
||
|
|
plt.title(f'{dataset} - Model Comparison')
|
||
|
|
plt.xlabel('Model Type')
|
||
|
|
plt.ylabel('Dice Score')
|
||
|
|
plt.tight_layout()
|
||
|
|
plt.savefig(os.path.join(plots_path, f'{dataset}_model_comparison.png'))
|
||
|
|
plt.close()
|
||
|
|
|
||
|
|
# 結構wise比較圖
|
||
|
|
plt.figure(figsize=(15, 8))
|