LIENWEIYING/preprocessing/final/nnunet_Dataset021_OAR_preparation.py

128 lines
4.6 KiB
Python
Raw Normal View History

2026-07-22 06:50:09 +00:00
import json
import os
import shutil
import glob
import re
import pandas as pd
# 創建 dataset.json 內容
dataset_json = {
"channel_names": {
"0": "CT+CTC",
"1": "T1C"
},
"labels": {
"background": "0",
"Brainstem": "1",
"Right_Eye": "2",
"Left_Eye": "3",
"Optic_Chiasm": "4",
"Right_Optic_Nerve": "5",
"Left_Optic_Nerve": "6"
},
"numTraining": 0,
"file_ending": ".nii.gz"
}
# 定義目錄
source_image1_dir = "/home/onlylian/resampled_CT+CTC_image"
source_image2_dir = "/home/onlylian/resampled_T1C_image"
source_label_dir = "/home/onlylian/resampled_label"
target_base_dir = "/home/onlylian/nnUNet/nnUNet_raw/Dataset021_OAR"
imagesTr_dir = os.path.join(target_base_dir, "imagesTr")
labelsTr_dir = os.path.join(target_base_dir, "labelsTr")
# 創建目錄
os.makedirs(imagesTr_dir, exist_ok=True)
os.makedirs(labelsTr_dir, exist_ok=True)
def extract_patient_id(filename):
match = re.search(r'(CT\+CTC|T1C|label)_([^.]+)\.nii\.gz', os.path.basename(filename))
return match.group(2) if match else None
# 獲取並排序原始圖像文件
source_image1_files = sorted(glob.glob(os.path.join(source_image1_dir, "*.nii.gz")))
source_image2_files = sorted(glob.glob(os.path.join(source_image2_dir, "*.nii.gz")))
processed_files = 0
patient_mapping = []
# 創建病歷號到文件名的映射
label_files = glob.glob(os.path.join(source_label_dir, "*.nii.gz"))
label_map = {extract_patient_id(f): f for f in label_files}
print(f"找到 {len(label_map)} 個標籤文件")
print(f"找到 {len(source_image1_files)} 個 CT+CTC 圖像文件")
print(f"找到 {len(source_image2_files)} 個 T1C 圖像文件")
# 處理所有資料
for index, (source_image1_file, source_image2_file) in enumerate(zip(source_image1_files, source_image2_files), start=1):
if index % 10 == 0: # 每處理10個文件顯示一次進度
print(f"\n正在處理檔案 {index}/{len(source_image1_files)}:")
patient_id1 = extract_patient_id(source_image1_file)
patient_id2 = extract_patient_id(source_image2_file)
if patient_id1 != patient_id2 or patient_id1 is None or patient_id2 is None:
print(f"警告: 病歷號不匹配或無法提取,跳過這對文件: {source_image1_file}, {source_image2_file}")
continue
# 使用新的文件名格式
new_filename = f"OAR_{index:03d}"
# 複製並重命名圖像文件
target_image1_file = os.path.join(imagesTr_dir, f"{new_filename}_0000.nii.gz")
target_image2_file = os.path.join(imagesTr_dir, f"{new_filename}_0001.nii.gz")
shutil.copyfile(source_image1_file, target_image1_file)
shutil.copyfile(source_image2_file, target_image2_file)
# 處理標籤文件
if patient_id1 in label_map:
source_label_file = label_map[patient_id1]
target_label_file = os.path.join(labelsTr_dir, f"{new_filename}.nii.gz")
shutil.copyfile(source_label_file, target_label_file)
processed_files += 1
# 添加到映射列表
patient_mapping.append({
"編號": f"{index:03d}",
"病歷號": patient_id1
})
# 輸出處理細節
print(f"CT+CTC: {source_image1_file}")
print(f"T1C: {source_image2_file}")
print(f"複製圖像檔案: {source_image1_file}{target_image1_file}")
print(f"複製圖像檔案: {source_image2_file}{target_image2_file}")
print(f"複製標籤檔案: {source_label_file}{target_label_file}")
print() # 空行分隔每筆記錄
else:
print(f"警告: 未找到 {patient_id1} 的標籤檔案。")
print("\n所有資料處理完畢。")
# 更新 dataset.json
dataset_json["numTraining"] = processed_files
# 寫入 dataset.json 檔案
dataset_json_path = os.path.join(target_base_dir, "dataset.json")
with open(dataset_json_path, 'w') as f:
json.dump(dataset_json, f, indent=4)
print(f"dataset.json 已寫入到 {dataset_json_path}")
print(f"總共處理了 {processed_files} 個完整的資料集(圖像對加標籤)")
# 創建 Excel 文件
df = pd.DataFrame(patient_mapping)
excel_path = os.path.join(target_base_dir, "patient_mapping.xlsx")
df.to_excel(excel_path, index=False)
print(f"病人映射表已保存到 {excel_path}")
# 輸出處理摘要
print("\n處理摘要:")
print(f"總共找到的 CT+CTC 圖像: {len(source_image1_files)}")
print(f"總共找到的 T1C 圖像: {len(source_image2_files)}")
print(f"總共找到的標籤文件: {len(label_map)}")
print(f"成功處理的完整資料集: {processed_files}")
print(f"跳過的文件對數量: {len(source_image1_files) - processed_files}")