CBT_project/imaging/segmentation.py

75 lines
2.3 KiB
Python
Raw Normal View History

2026-04-10 05:25:27 +00:00
import os
import SimpleITK as sitk
from config.constant import LABEL_MAP
import numpy as np
"""
# 沿用原本 LABEL_MAP
seg_bone(n, name, img, lbl)
# user 自定義
my_map = {1: "L1", 2: "L2", 3: "L3"}
seg_bone(n, name, img, lbl, label_map=my_map)
"""
def seg_bone(n, name, resampled_sitk_img, resampled_sitk_lbl, output_base=None, label_map=LABEL_MAP):
if output_base==None:
output_base=='Dataset'
if n not in label_map:
raise ValueError(f"Label {n} not found in label_map")
label_name = label_map[n]
lssif = sitk.LabelShapeStatisticsImageFilter()
lssif.Execute(resampled_sitk_lbl)
if not lssif.HasLabel(n):
raise RuntimeError(f"Label {n} not found")
bbox2 = lssif.GetBoundingBox(n)
roi = sitk.RegionOfInterest(resampled_sitk_img, bbox2[3:], bbox2[:3])
label2 = sitk.RegionOfInterest(resampled_sitk_lbl, bbox2[3:], bbox2[:3])
roi_path = os.path.join(output_base, f"{label_name}_roi.nii.gz")
sitk.WriteImage(roi, roi_path)
binary = sitk.BinaryThreshold(label2, lowerThreshold=n, upperThreshold=n, outsideValue=0, insideValue=1)
binary_path = os.path.join(output_base, f"{label_name}_binary.nii.gz")
sitk.WriteImage(binary, binary_path)
roi_pixel_type = roi.GetPixelID()
binary_cast = sitk.Cast(binary, roi_pixel_type)
roi2 = roi * binary_cast
roi2_path = os.path.join(output_base, f"{label_name}_roi2.nii.gz")
sitk.WriteImage(roi2, roi2_path)
lsif = sitk.LabelStatisticsImageFilter()
label2_int = sitk.Cast(label2, sitk.sitkUInt16)
lsif.Execute(roi2, label2_int)
labels_in_roi = lsif.GetLabels()
if n in labels_in_roi:
roi_hu = sitk.GetArrayFromImage(roi2)
threshold = np.percentile(roi_hu, 60)
else:
threshold = lsif.GetMedian(labels_in_roi[0])
cortical = sitk.BinaryThreshold(roi2, lowerThreshold=threshold, upperThreshold=10000, outsideValue=0, insideValue=1)
cortical_path = os.path.join(output_base, f"{label_name}_cortical.nii.gz")
sitk.WriteImage(cortical, cortical_path)
return roi_path, binary_path, roi2_path, cortical_path
"""
Dataset/
standardized/
subject001/
L1_roi.nii.gz
L1_binary.nii.gz
L1_roi2.nii.gz
L1_cortical.nii.gz
L2_roi.nii.gz
L2_binary.nii.gz
...
"""