107 lines
No EOL
4.6 KiB
Python
107 lines
No EOL
4.6 KiB
Python
import os
|
|
import glob
|
|
import SimpleITK as sitk
|
|
from difflib import get_close_matches
|
|
|
|
def process_images(input_base_path, output_base_path, output_prefix):
|
|
# 定義組織標籤對應關係
|
|
labels = {
|
|
"background": "0",
|
|
"Struct_Brain_Stem": "1",
|
|
"Struct_Right_Eye": "2",
|
|
"Struct_Left_Eye": "3",
|
|
"Struct_Optic_Chiasm": "4",
|
|
"Struct_Right_Optic_Nerve": "5",
|
|
"Struct_Left_Optic_Nerve": "6",
|
|
"TV": "7" # Tumor
|
|
}
|
|
|
|
# 確保輸出資料夾存在
|
|
os.makedirs(output_base_path, exist_ok=True)
|
|
|
|
# 定義重新取樣函數
|
|
def resample_image(image, reference_image):
|
|
resample = sitk.ResampleImageFilter()
|
|
resample.SetReferenceImage(reference_image)
|
|
resample.SetInterpolator(sitk.sitkNearestNeighbor) # 使用最近鄰插值避免標籤變化
|
|
resample.SetDefaultPixelValue(0)
|
|
resampled_image = resample.Execute(image)
|
|
return resampled_image
|
|
|
|
# 遍歷所有病人的資料夾
|
|
for patient_id in os.listdir(input_base_path):
|
|
patient_folder_path = os.path.join(input_base_path, patient_id)
|
|
if not os.path.isdir(patient_folder_path):
|
|
continue # 跳過非資料夾的項目
|
|
|
|
# 確定病人資料夾下的日期資料夾
|
|
date_folder = os.listdir(patient_folder_path)[0]
|
|
ORGAN_folder_path = os.path.join(patient_folder_path, date_folder)
|
|
|
|
# 創建一個空的影像來合併所有標籤
|
|
combined_image = None
|
|
|
|
# 列出病人資料夾內的所有檔案
|
|
all_files = os.listdir(ORGAN_folder_path)
|
|
|
|
# 遍歷每個標籤並合併
|
|
for organ_name, label in labels.items():
|
|
if organ_name == "TV":
|
|
# 動態搜尋 TV 檔案
|
|
tv_files = glob.glob(os.path.join(ORGAN_folder_path, "TV-*.nii.gz"))
|
|
if len(tv_files) == 0:
|
|
print(f"TV file not found for patient {patient_id}, skipping TV...")
|
|
continue
|
|
organ_path = tv_files[0] # 取第一個找到的檔案
|
|
else:
|
|
# 使用模糊比對尋找最接近的檔案名
|
|
matches = get_close_matches(organ_name, all_files, n=1, cutoff=0.6)
|
|
if len(matches) == 0:
|
|
print(f"No close match found for {organ_name} for patient {patient_id}, skipping...")
|
|
continue
|
|
organ_filename = matches[0]
|
|
organ_path = os.path.join(ORGAN_folder_path, organ_filename)
|
|
|
|
if not os.path.exists(organ_path):
|
|
print(f"File {organ_path} not found for patient {patient_id}, skipping...")
|
|
continue
|
|
|
|
organ_image = sitk.ReadImage(organ_path)
|
|
|
|
if combined_image is None:
|
|
# 初始化合併影像的大小和參考
|
|
combined_image = sitk.Image(organ_image.GetSize(), sitk.sitkUInt8)
|
|
combined_image.CopyInformation(organ_image)
|
|
else:
|
|
# 如果器官影像的尺寸和合併影像不同，進行重新取樣
|
|
if organ_image.GetSize() != combined_image.GetSize():
|
|
print(f"Resampling organ {organ_name} for patient {patient_id} due to size mismatch...")
|
|
organ_image = resample_image(organ_image, combined_image)
|
|
|
|
# 將器官影像添加到合併影像中
|
|
organ_array = sitk.GetArrayFromImage(organ_image)
|
|
combined_array = sitk.GetArrayFromImage(combined_image)
|
|
|
|
# 更新合併影像的陣列，根據標籤值來設定非背景區域
|
|
combined_array[organ_array > 0] = int(label)
|
|
|
|
# 將更新的陣列轉回 SimpleITK 影像
|
|
combined_image = sitk.GetImageFromArray(combined_array)
|
|
combined_image.CopyInformation(organ_image)
|
|
|
|
# 保存合併後的影像
|
|
output_filename = f'{output_prefix}_label_{patient_id}.nii.gz'
|
|
output_path = os.path.join(output_base_path, output_filename)
|
|
sitk.WriteImage(combined_image, output_path)
|
|
print(f'Saved combined image for patient {patient_id} to {output_path}')
|
|
|
|
if __name__ == "__main__":
|
|
# 處理 T1C 影像
|
|
T1C_input_path = '/home/onlylian/sorted_OAR_Target_files/T1C'
|
|
T1C_output_path = '/home/onlylian/T1C_OAR_Target_label_folder'
|
|
process_images(T1C_input_path, T1C_output_path, "T1C")
|
|
|
|
# 處理 CT+CTC 影像
|
|
CT_CTC_input_path = '/home/onlylian/sorted_OAR_Target_files/T1C' # 注意：這裡使用相同的輸入路徑，如果需要不同的路徑請修改
|
|
CT_CTC_output_path = '/home/onlylian/CT+CTC_OAR_Target_label_folder'
|
|
process_images(CT_CTC_input_path, CT_CTC_output_path, "CT+CTC") |