CBT_project/visualization/res_cyl_nifti.py
2026-04-10 13:25:27 +08:00

83 lines
No EOL
3.1 KiB
Python

import SimpleITK as sitk
import os
from datetime import datetime
from core.cylinder import generate_cylinder_numpy
def save_cyl(best_position_l, best_position_r, spacing, IMAGE3, output_base, CBT=True):
image3 = sitk.ReadImage(IMAGE3) #Reading image2 or image3 can get the same shape and origin actually
image3_array = sitk.GetArrayFromImage(image3)
image3_shape = image3_array.shape
origin = image3.GetOrigin()
diameter_l = best_position_l[5]
diameter_r = best_position_r[5]
length_l = best_position_l[6]
length_r = best_position_r[6]
cyl_l = generate_cylinder_numpy(diameter_l,
length_l,
best_position_l[0],
best_position_l[1],
best_position_l[2],
best_position_l[3],
best_position_l[4],
image3_shape,
spacing)
cyl_r = generate_cylinder_numpy(diameter_r,
length_r,
best_position_r[0],
best_position_r[1],
best_position_r[2],
best_position_r[3],
best_position_r[4],
image3_shape,
spacing)
cyl_L = sitk.GetImageFromArray(cyl_l)
cyl_R = sitk.GetImageFromArray(cyl_r)
cyl_L.SetOrigin(origin)
cyl_R.SetOrigin(origin)
image_path = IMAGE3
date_str = datetime.now().strftime("%Y%m%d")
patient_id = os.path.basename(os.path.dirname(image_path))
file_name = os.path.basename(image_path)
level = file_name.split('_')[0]
output_folder = os.path.join(output_base, date_str, patient_id)
os.makedirs(output_folder, exist_ok=True)
def get_unique_filename(path):
if not os.path.exists(path):
return path
# 特別處理 .nii.gz 檔案
if path.endswith(".nii.gz"):
base = path[:-7] # 去除 .nii.gz
ext = ".nii.gz"
else:
base, ext = os.path.splitext(path)
i = 1
new_path = f"{base}_{i}{ext}"
while os.path.exists(new_path):
i += 1
new_path = f"{base}_{i}{ext}"
return new_path
if CBT:
output_path_l = get_unique_filename(f'{output_folder}/{level}_{diameter_l}_{length_l}_CBT_L.nii.gz')
output_path_r = get_unique_filename(f'{output_folder}/{level}_{diameter_r}_{length_r}_CBT_R.nii.gz')
sitk.WriteImage(cyl_L, output_path_l)
sitk.WriteImage(cyl_R, output_path_r)
else:
output_path_l = get_unique_filename(f'{output_folder}/{level}_{diameter_l}_{length_l}_TPS_L.nii.gz')
output_path_r = get_unique_filename(f'{output_folder}/{level}_{diameter_r}_{length_r}_TPS_R.nii.gz')
sitk.WriteImage(cyl_L, output_path_l)
sitk.WriteImage(cyl_R, output_path_r)
return output_path_l, output_path_r