151 lines
No EOL
5.5 KiB
Python
151 lines
No EOL
5.5 KiB
Python
import nibabel as nib
|
||
import matplotlib.pyplot as plt
|
||
import numpy as np
|
||
import os
|
||
|
||
def load_nifti(file_path):
|
||
"""載入NIfTI檔案"""
|
||
return nib.load(file_path).get_fdata()
|
||
|
||
def create_colored_mask(data):
|
||
"""創建彩色遮罩"""
|
||
color_dict = {
|
||
1: [144/255, 238/255, 144/255], # Brainstem - 綠色
|
||
2: [255/255, 218/255, 150/255], # Right_Eye - 淺黃色
|
||
3: [205/255, 170/255, 125/255], # Left_Eye - 棕色
|
||
4: [135/255, 206/255, 235/255], # Optic_Chiasm - 藍色
|
||
5: [255/255, 99/255, 71/255], # Right_Optic_Nerve - 亮珊瑚紅
|
||
6: [255/255, 160/255, 122/255], # Left_Optic_Nerve - 淺珊瑚紅
|
||
7: [255/255, 0/255, 0/255] # TV - 亮紅色
|
||
}
|
||
|
||
mask = np.zeros((*data.shape, 4))
|
||
for label, color in color_dict.items():
|
||
organ_mask = (data == label)
|
||
if np.any(organ_mask):
|
||
mask[organ_mask] = [*color, 1.0]
|
||
return mask
|
||
|
||
def get_display_bounds(image):
|
||
"""直接返回較大的顯示區域"""
|
||
h, w = image.shape
|
||
# 使用圖像中心點
|
||
center_h = h // 2
|
||
center_w = w // 2
|
||
# 取圖像大小的 85% 作為顯示範圍
|
||
radius_h = int(h * 0.425)
|
||
radius_w = int(w * 0.425)
|
||
|
||
return (
|
||
max(center_h - radius_h, 0),
|
||
min(center_h + radius_h, h),
|
||
max(center_w - radius_w, 0),
|
||
min(center_w + radius_w, w)
|
||
)
|
||
|
||
def check_tv_in_slice(slice_data):
|
||
"""檢查切片中是否包含TV（label 7）"""
|
||
return 7 in np.unique(slice_data)
|
||
|
||
def find_tv_slices(ground_truth):
|
||
"""找出所有包含TV的切片"""
|
||
tv_slices = []
|
||
tv_pixels = []
|
||
total_slices = ground_truth.shape[2]
|
||
|
||
for slice_num in range(total_slices):
|
||
slice_data = ground_truth[:, :, slice_num]
|
||
if check_tv_in_slice(slice_data):
|
||
tv_pixel_count = np.sum(slice_data == 7)
|
||
tv_slices.append(slice_num)
|
||
tv_pixels.append(tv_pixel_count)
|
||
|
||
sorted_indices = np.argsort(tv_pixels)[::-1]
|
||
sorted_slices = [tv_slices[i] for i in sorted_indices]
|
||
sorted_pixels = [tv_pixels[i] for i in sorted_indices]
|
||
|
||
return sorted_slices, sorted_pixels
|
||
|
||
def save_slice(image, ground_truth, prediction, slice_num, save_path):
|
||
"""保存指定切片的三列比較圖（原始影像、Ground Truth和預測結果）"""
|
||
plt.clf()
|
||
fig, axes = plt.subplots(1, 3, figsize=(36, 11))
|
||
plt.subplots_adjust(wspace=0.01)
|
||
|
||
img_slice = np.rot90(image[:, :, slice_num])
|
||
rmin, rmax, cmin, cmax = get_display_bounds(img_slice)
|
||
|
||
# 調整圖像對比度
|
||
p2, p98 = np.percentile(img_slice, (2, 98))
|
||
img_slice = np.clip(img_slice, p2, p98)
|
||
img_slice = (img_slice - p2) / (p98 - p2)
|
||
|
||
# 1. Original Image
|
||
axes[0].imshow(img_slice[rmin:rmax, cmin:cmax], cmap='gray', zorder=1)
|
||
axes[0].set_title('Original Image', fontsize=32, pad=20, weight='bold')
|
||
axes[0].axis('off')
|
||
|
||
# 2. Ground Truth
|
||
gt_slice = np.rot90(ground_truth[:, :, slice_num])
|
||
gt_mask = create_colored_mask(gt_slice)
|
||
|
||
axes[1].imshow(img_slice[rmin:rmax, cmin:cmax], cmap='gray', zorder=1)
|
||
axes[1].imshow(gt_mask[rmin:rmax, cmin:cmax], zorder=2)
|
||
axes[1].set_title('Ground Truth', fontsize=32, pad=20, weight='bold')
|
||
axes[1].axis('off')
|
||
|
||
# 3. Prediction
|
||
pred_slice = np.rot90(prediction[:, :, slice_num])
|
||
pred_mask = create_colored_mask(pred_slice)
|
||
|
||
axes[2].imshow(img_slice[rmin:rmax, cmin:cmax], cmap='gray', zorder=1)
|
||
axes[2].imshow(pred_mask[rmin:rmax, cmin:cmax], zorder=2)
|
||
axes[2].set_title('Prediction', fontsize=32, pad=20, weight='bold')
|
||
axes[2].axis('off')
|
||
|
||
plt.suptitle('')
|
||
plt.savefig(save_path, bbox_inches='tight', dpi=300, pad_inches=0.05)
|
||
plt.close()
|
||
print(f"已保存切片 {slice_num} 至: {save_path}")
|
||
|
||
if __name__ == "__main__":
|
||
base_dir = "/mnt/1248/onlylian"
|
||
image_path = os.path.join(base_dir, "nnUNet/nnUNet_raw/Dataset012_OAR_TV/imagesTs/OAR_TV_092_0000.nii.gz")
|
||
gt_path = os.path.join(base_dir, "nnUNet/nnUNet_raw/Dataset012_OAR_TV/labelsTs/OAR_TV_092.nii.gz")
|
||
pred_path = os.path.join(base_dir, "T1C_OAR_TV/output_predictions/2d/fold_2/OAR_TV_092.nii.gz") # 更新預測結果路徑
|
||
save_dir = os.path.join(base_dir, "T1C_OAR_TV")
|
||
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
|
||
print("正在載入影像...")
|
||
image = load_nifti(image_path)
|
||
ground_truth = load_nifti(gt_path)
|
||
prediction = load_nifti(pred_path) # 載入預測結果
|
||
print("影像載入完成")
|
||
|
||
print("\n尋找包含TV的切片...")
|
||
tv_slices, tv_pixels = find_tv_slices(ground_truth)
|
||
|
||
if tv_slices:
|
||
print(f"\n找到 {len(tv_slices)} 個包含TV的切片:")
|
||
for slice_num, pixel_count in zip(tv_slices, tv_pixels):
|
||
print(f"切片 {slice_num}: {pixel_count} 像素")
|
||
else:
|
||
print("未找到包含TV的切片")
|
||
|
||
print("\n開始切片選擇和保存過程...")
|
||
while True:
|
||
try:
|
||
slice_num = int(input("\n請輸入要保存的切片編號 (-1 退出): "))
|
||
if slice_num == -1:
|
||
print("程序結束")
|
||
break
|
||
if 0 <= slice_num < ground_truth.shape[2]:
|
||
save_path = os.path.join(save_dir, f'comparison_OAR_TV_092_3_{slice_num:03d}.png')
|
||
save_slice(image, ground_truth, prediction, slice_num, save_path)
|
||
else:
|
||
print(f"切片編號必須在 0 到 {ground_truth.shape[2]-1} 之間")
|
||
except ValueError:
|
||
print("請輸入有效的數字")
|
||
except Exception as e:
|
||
print(f"發生錯誤: {str(e)}") |