Xiao Furen
523ec7ee16
Introduces a new scoring component for vertebral body (VBODY) rewards to improve optimization accuracy. The update includes: - Added `vbody_tensor` to `OptimizationContext` and scoring functions to reward screw placement within the vertebral body. - Enhanced `segment_spinous_process` with diagnostic capabilities to detect spinous process absence (e.g., post-laminectomy). - Improved `resample_img` to prevent physical boundary clipping and handle interpolation more robustly for CT and label data. - Implemented a metadata cache using TinyDB in the preprocessing pipeline to skip low-resolution or insufficient scans efficiently. - Added robust error handling for NFS-based file operations and directory creation. - Added new visualization tools for bone figures and level plotting. refactor(imaging): improve segmentation and resampling precision - Refactored `seg_bone` to support original resolution binary masks and Signed Maurer Distance Maps (SMD) for more accurate boundary handling. - Updated `resample_img` to use `ceil` for output size calculation to ensure full physical coverage. - Optimized `process_single_image` to utilize metadata for skipping processing of invalid or low-quality scans.
161 lines
No EOL
8.3 KiB
Python
161 lines
No EOL
8.3 KiB
Python
import os
|
||
import SimpleITK as sitk
|
||
from config.constant import LABEL_MAP
|
||
from imaging.resample import resample_img
|
||
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 _largest_cc_bbox(mask_img):
|
||
"""26-連通的最大连通區域 + 其 bbox（RelabelComponent 依大小排序，最大者=1）。
|
||
回傳 (largest_mask, bbox2)；沒有任何組件時回傳 None。
|
||
bbox2 格式：[x_start, y_start, z_start, x_size, y_size, z_size]。"""
|
||
cc_image = sitk.ConnectedComponent(mask_img, True) # fullyConnected
|
||
relabeled_cc = sitk.RelabelComponent(cc_image, sortByObjectSize=True)
|
||
shape_stats = sitk.LabelShapeStatisticsImageFilter()
|
||
shape_stats.Execute(relabeled_cc)
|
||
if shape_stats.GetNumberOfLabels() < 1:
|
||
return None
|
||
return (relabeled_cc == 1), shape_stats.GetBoundingBox(1)
|
||
|
||
def _bbox_roi(img, bbox, margin=0):
|
||
"""裁 bbox（對稱外擴 margin 個 voxel，clamp 到影像邊界）。
|
||
bbox 格式：[x_start, y_start, z_start, x_size, y_size, z_size]。"""
|
||
n = img.GetSize() # (x, y, z)
|
||
index = [max(0, int(bbox[i]) - margin) for i in range(3)]
|
||
size = [min(n[i] - index[i], int(bbox[i + 3]) + 2 * margin) for i in range(3)]
|
||
return sitk.RegionOfInterest(img, size, index)
|
||
|
||
def seg_bone(n, name, resampled_sitk_img, resampled_sitk_lbl, output_base=None, label_map=LABEL_MAP,
|
||
original_label=None):
|
||
|
||
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]
|
||
|
||
# ============ 原解析度（未插值）chain ============
|
||
# 1. 提取標籤 n 的二值遮罩 (將標籤 n 設為 1，其餘為 0)，最大連通區域
|
||
smd_path = resampled_path = binary_sdf_path = binary_erode_path = None
|
||
binary_linear_path = binary_nn_path = None
|
||
if original_label is not None:
|
||
bin_orig = sitk.BinaryThreshold(original_label, n, n, 1, 0)
|
||
cc_orig = _largest_cc_bbox(bin_orig)
|
||
if cc_orig is None:
|
||
return None
|
||
largest_orig, bbox_orig = cc_orig
|
||
|
||
# _binary.nii.gz：原解析度【未插值】遮罩（最大连通區域、裁到物件 bbox，
|
||
# 不重取樣、不插值 —— 原始 label 的忠實二值版本）
|
||
binary_path = os.path.join(output_base, f"{label_name}_binary.nii.gz")
|
||
sitk.WriteImage(_bbox_roi(largest_orig, bbox_orig), binary_path)
|
||
|
||
# _smd.nii.gz：SignedMaurerDistanceMap（ITK 慣例：物件內負 / 外正；
|
||
# 距離以原始 index（pixel）單位、不隨各向异性 spacing 縮放——
|
||
# _binary_sdf 的 0.5 閾值（mid-gap）正是依賴這個 index 單位慣例）。
|
||
# 裁 bbox_orig 外擴 4 voxel 的背景輪：
|
||
# 沒有背景輪時填充值直接貼着物件邊緣，重取樣會在裁切邊界產生
|
||
# 假的閾值穿越。
|
||
# 這個 SimpleITK build 的 3D SignedMaurerDistanceMap 只支援整數輸入，
|
||
# 先 Cast 到 uint8 再回傳 float32 輸出
|
||
smd_full = sitk.SignedMaurerDistanceMap(sitk.Cast(largest_orig, sitk.sitkUInt8))
|
||
smd_full = sitk.Cast(smd_full, sitk.sitkFloat32)
|
||
smd_margined = _bbox_roi(smd_full, bbox_orig, margin=4)
|
||
smd_path = os.path.join(output_base, f"{label_name}_smd.nii.gz")
|
||
sitk.WriteImage(smd_margined, smd_path)
|
||
|
||
# SMD 體積【线性插值】重取樣到 0.5mm（reference = 0.5mm CT，與
|
||
# resampled_sitk_img 同 grid，之後可直接用 0.5mm bbox 裁切）。
|
||
# SMD 在每個 input voxel 內分段線性，线性重取樣近似精確、無漣波，
|
||
# 各等值面（物件邊界等）不變。填充值 = 裁切角落（背景側，正值）；
|
||
# 若為負（物件貼影像邊界的病態情況）用 0。
|
||
corner = float(sitk.GetArrayViewFromImage(smd_margined).flat[0])
|
||
rs = sitk.ResampleImageFilter()
|
||
rs.SetReferenceImage(resampled_sitk_img)
|
||
rs.SetInterpolator(sitk.sitkLinear)
|
||
rs.SetDefaultPixelValue(corner if corner > 0 else 0.0)
|
||
smd_res_full = rs.Execute(smd_margined)
|
||
|
||
# 0.5mm linear 二值化 mask（full extent，與 resampled_sitk_img 同 grid）：
|
||
# 原數據 5mm 切片上採樣 10x 到 0.5mm，最近邻會在邊界產生 10 體素厚的
|
||
# 階梯鋸齒；线性插值使邊界落在次體素位置（rotation 前的邊界更平滑）。
|
||
bin_lin = resample_img(sitk.Cast(bin_orig, sitk.sitkFloat32))
|
||
arr = (sitk.GetArrayFromImage(bin_lin) > 0.5).astype(np.uint8)
|
||
if arr.shape != resampled_sitk_img.GetSize()[::-1]:
|
||
raise RuntimeError(
|
||
f"linear binary resample shape mismatch: {arr.shape} vs {resampled_sitk_img.GetSize()}")
|
||
binary_mask = sitk.GetImageFromArray(arr)
|
||
binary_mask.CopyInformation(resampled_sitk_img)
|
||
else:
|
||
# 無原始 label：舊版路徑（0.5mm label 閾值），不產生 SMD/SDF chain
|
||
binary_mask = sitk.BinaryThreshold(resampled_sitk_lbl, n, n, 1, 0)
|
||
binary_path = None
|
||
|
||
# 2. 0.5mm 最大連通區域（26-連通）+ 邊界框；所有 0.5mm 輸出裁到同一 bbox
|
||
cc_res = _largest_cc_bbox(binary_mask)
|
||
if cc_res is None:
|
||
return None
|
||
largest_mask, bbox2 = cc_res
|
||
|
||
if binary_path is None:
|
||
binary_path = os.path.join(output_base, f"{label_name}_binary.nii.gz")
|
||
sitk.WriteImage(sitk.RegionOfInterest(largest_mask, bbox2[3:], bbox2[:3]), binary_path)
|
||
|
||
if smd_path is not None:
|
||
# _smd_resampled.nii.gz：线性重取樣到 0.5mm 的 SMD（浮點，裁 bbox2）
|
||
resampled_path = os.path.join(output_base, f"{label_name}_smd_resampled.nii.gz")
|
||
sitk.WriteImage(sitk.RegionOfInterest(smd_res_full, bbox2[3:], bbox2[:3]), resampled_path)
|
||
|
||
# _binary_sdf.nii.gz：_smd_resampled 於 0.5 閾值 -> 0.5mm 平滑 mask
|
||
# （SMD 內負/外正；0.5 介於內殼 ≈0 與外殼 ≈+1 之間，即原解析度
|
||
# label 邊界的 mid-gap 位置，physical volume 與 _binary/_binary_nn
|
||
# 一致，sub-voxel 表面、無 NN 階梯）
|
||
bin_sdf_full = sitk.GetImageFromArray(
|
||
(sitk.GetArrayFromImage(smd_res_full) < 0.5).astype(np.uint8))
|
||
bin_sdf_full.CopyInformation(resampled_sitk_img)
|
||
binary_sdf_path = os.path.join(output_base, f"{label_name}_binary_sdf.nii.gz")
|
||
sitk.WriteImage(sitk.RegionOfInterest(bin_sdf_full, bbox2[3:], bbox2[:3]), binary_sdf_path)
|
||
|
||
# _binary_nn.nii.gz：0.5mm label 最近邻（舊版），裁自己的 bbox，僅供對比
|
||
nn_mask = sitk.BinaryThreshold(resampled_sitk_lbl, n, n, 1, 0)
|
||
nn_res = _largest_cc_bbox(nn_mask)
|
||
if nn_res is not None:
|
||
nn_largest, nn_bbox = nn_res
|
||
binary_nn_path = os.path.join(output_base, f"{label_name}_binary_nn.nii.gz")
|
||
sitk.WriteImage(sitk.RegionOfInterest(nn_largest, nn_bbox[3:], nn_bbox[:3]),
|
||
binary_nn_path)
|
||
|
||
# 3. roi（0.5mm）
|
||
# _roi2 不再存檔；_cortical 改由 xfr_preprocess 的旋轉後處理產出
|
||
# （rotated/{level}_cortical.nii.gz，定義不變：門檻 = 骨頭 mask 內 median HU）
|
||
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)
|
||
|
||
return roi_path, binary_path, None, None, binary_nn_path, \
|
||
binary_linear_path, smd_path, resampled_path, binary_sdf_path, binary_erode_path
|
||
|
||
"""
|
||
Dataset/
|
||
└── standardized/
|
||
└── subject001/
|
||
├── L1_binary.nii.gz # 原解析度【未插值】遮罩（最大连通區域、裁物件 bbox）
|
||
├── L1_smd.nii.gz # SignedMaurerDistanceMap（內負/外正，原始 index 單位；
|
||
│ # bbox 外扩 4 voxel 背景輪，供重取樣插值用）
|
||
├── L1_smd_resampled.nii.gz # _smd 經线性插值重取樣到 0.5mm（浮點，裁 0.5mm bbox）
|
||
├── L1_binary_sdf.nii.gz # _smd_resampled 於 0.5 閾值 -> 0.5mm 平滑 mask（0/1，裁同 bbox）
|
||
├── L1_binary_nn.nii.gz # 最近邻版 0/1（對比用，各自 bbox）
|
||
├── L1_roi.nii.gz
|
||
├── L2_binary.nii.gz
|
||
...
|
||
""" |