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))