CBT_project/imaging/segmentation.py

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)

    # 1. 提取標籤 n 的二值遮罩 (將標籤 n 設為 1，其餘為 0)
    binary_mask = sitk.BinaryThreshold(resampled_sitk_lbl,  n, n, 1, 0)
    
    # 2. 獲取所有連通區域
    # 連通區域濾波器會將 binary_mask 中的不同物體標記為 1, 2, 3...
    cc_image = sitk.ConnectedComponent(binary_mask)
    
    # 3. 根據區域大小（像素/體積）重新標記
    # RelabelComponent 會按大小排序，最大的物體標籤會被設為 1
    relabeled_cc = sitk.RelabelComponent(cc_image, sortByObjectSize=True)
    largest_mask = relabeled_cc == 1
    
    # 4. 計算形狀統計信息
    shape_stats = sitk.LabelShapeStatisticsImageFilter()
    shape_stats.Execute(relabeled_cc)
    
    # 檢查是否有找到任何組件
    if shape_stats.GetNumberOfLabels() < 1:
        return None
        
    # 5. 獲取最大組件（標籤為 1）的邊界框
    # 格式通常為 [x_start, y_start, z_start, x_size, y_size, z_size]
    bbox2 = shape_stats.GetBoundingBox(1)

    roi    = sitk.RegionOfInterest(resampled_sitk_img, bbox2[3:], bbox2[:3])
    roi_path = os.path.join(output_base, f"{label_name}_roi.nii.gz")
    sitk.WriteImage(roi, roi_path)

    # label2 = sitk.RegionOfInterest(resampled_sitk_lbl, bbox2[3:], bbox2[:3])
    # binary = sitk.BinaryThreshold(label2, lowerThreshold=n, upperThreshold=n, outsideValue=0, insideValue=1)
    binary = sitk.RegionOfInterest(largest_mask, bbox2[3:], bbox2[:3])
    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 = sitk.Mask(roi, binary)
    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])

    stats = sitk.LabelStatisticsImageFilter()
    stats.UseHistogramsOn() # Required for median calculation
    stats.Execute(roi, binary)

    # Get the median for the region where mask == 1
    threshold = stats.GetMedian(1)

    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
        ...
"""