CBT_project/imaging/segmentation.py
Xiao Furen 0274e954be refactor(core): improve cylinder parameter snapping and optimization logic
Refactor the cylinder parameter selection and optimization process to
improve accuracy and robustness.

- Implement `snap_to_discrete_values_xfr` using a KDTree for efficient
  mapping of continuous diameter and length values to a predefined
  set of discrete points.
- Update `objective_function_xfr` to use the new snapping mechanism
  and introduce a weighted loss component for diameter and length.
- Adjust PSO optimization bounds and search ranges in `run_pso_torch_xfr`
  to better align with image dimensions and anatomical constraints.
- Refine scoring logic in `cl_score_torch_xfr` with updated penalty
  weights for overlaps and out-of-bone voxels.
- Update `config/constant.py` with new allowed diameter and length
  ranges.
- Improve `imaging/preprocessing.py` by making `PROGRESS_FILE` a
  parameter to allow per-output-directory progress tracking.
- Update `xfr_debug.py` with improved error handling and directory
  paths for batch processing.
2026-08-26 22:36:07 +08:00

108 lines
No EOL
3.7 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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, True) #fullyConnected
# 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
...
"""