Xiao Furen
2ae08ac2cd
Refactor the preprocessing and segmentation pipeline to handle AP orientation variations and improve anatomical boundary detection. Key changes include: - Implement automated AP orientation detection in `process_single_image` to handle prone scans by flipping CT and labels when necessary. - Enhance `segment_spinous_process` using a gap-based approach to identify the spinal canal, providing more stable thresholds for spinous process and vertebral body segmentation. - Improve optimization search space by using the vertebral body (VBODY) projection for x/z bounding box calculation instead of the whole bone. - Refactor `render_bone_figure` to unify 2D/3D visualization and support detailed anatomical coloring (VBODY, spinous process). - Update `cl_score_torch_xfr` with more robust penalty handling for out-of-bone and null-voxel regions. - Add `retry_robust` utility to handle transient NFS file system errors. - Update `xfr_preprocess.py` to include anatomical segmentation coloring in rotated level visualizations.
841 lines
No EOL
37 KiB
Python
841 lines
No EOL
37 KiB
Python
import csv
|
||
import os
|
||
from datetime import datetime
|
||
|
||
import matplotlib
|
||
matplotlib.use("Agg")
|
||
import matplotlib.pyplot as plt
|
||
from matplotlib.colors import to_rgba
|
||
from matplotlib.lines import Line2D
|
||
|
||
import numpy as np
|
||
import SimpleITK as sitk
|
||
from scipy.ndimage import map_coordinates
|
||
|
||
from imaging.orientation import (azimuth_rotation, analyze_vertebral_tilt_contour,
|
||
best_symmetry_plane, best_upper_endplate_plane,
|
||
segment_spinous_process, segment_vertebral_body)
|
||
from utils.helpers import get_unique_filepath, retry_robust, save_with_unique_name
|
||
|
||
|
||
# 體積吸收渲染（Beer-Lambert），與 res_plot_3d 相同：
|
||
# 每 voxel 不透明度 = 1 - exp(-mu * voxel_width)
|
||
BONE_MU_CORTICAL = 0.02 # 1/mm
|
||
BONE_MU_TRABECULAR = 0.005 # 1/mm
|
||
BONE_MARKER_SIZE = 3.0 # 骨散點點面積 (pt^2)
|
||
BONE_SUBSAMPLE = 1 # 抽稀（1 = 全畫）
|
||
|
||
|
||
def set_axes_equal_3d(ax):
|
||
"""讓 3D 座標軸等比例，球/立方不變形（同 res_plot_3d）。"""
|
||
x_limits = ax.get_xlim3d()
|
||
y_limits = ax.get_ylim3d()
|
||
z_limits = ax.get_zlim3d()
|
||
|
||
x_range = abs(x_limits[1] - x_limits[0])
|
||
x_middle = np.mean(x_limits)
|
||
y_range = abs(y_limits[1] - y_limits[0])
|
||
y_middle = np.mean(y_limits)
|
||
z_range = abs(z_limits[1] - z_limits[0])
|
||
z_middle = np.mean(z_limits)
|
||
|
||
plot_radius = 0.5 * max([x_range, y_range, z_range])
|
||
|
||
ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius])
|
||
ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius])
|
||
ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius])
|
||
try:
|
||
ax.set_box_aspect([1, 1, 1])
|
||
except AttributeError:
|
||
pass
|
||
|
||
|
||
def _load_mask(path):
|
||
"""讀取 nifti 二值遮罩 -> (z,y,x) bool；缺檔/讀取失敗/空 -> None。"""
|
||
if path is None or not os.path.exists(path):
|
||
return None
|
||
try:
|
||
arr = sitk.GetArrayFromImage(sitk.ReadImage(path, sitk.sitkUInt8))
|
||
except Exception as e:
|
||
print(f"[warn] 讀取失敗 {path}: {e}")
|
||
return None
|
||
m = arr > 0
|
||
if int(m.sum()) == 0:
|
||
return None
|
||
return m
|
||
|
||
|
||
def _rgba_block(n, color, a):
|
||
arr = np.empty((n, 4))
|
||
arr[:] = to_rgba(color)
|
||
arr[:, 3] = a
|
||
return arr
|
||
|
||
|
||
# ======================================================================
|
||
# 旋轉對齊（index space）：
|
||
# 1) 最佳鏡稱面 normal -> +X（使鏡稱面平行 x=0）
|
||
# 2) 再繞 X 旋轉，使上終板面 normal 的 y 分量为 0（與 y 無關）
|
||
# 所有向量採 (x, y, z) 順序；array 為 (z, y, x)。
|
||
# ======================================================================
|
||
def _unit(v):
|
||
v = np.asarray(v, dtype=float)
|
||
n = np.linalg.norm(v)
|
||
return v / n if n > 0 else v
|
||
|
||
|
||
def _skew(ax):
|
||
ax = np.asarray(ax, dtype=float)
|
||
return np.array([[0.0, -ax[2], ax[1]],
|
||
[ax[2], 0.0, -ax[0]],
|
||
[-ax[1], ax[0], 0.0]])
|
||
|
||
|
||
def rotation_from_to(u, v):
|
||
"""3x3 旋轉矩陣，把單位向量 u 映射到單位向量 v（最小角旋转）。"""
|
||
u = _unit(u)
|
||
v = _unit(v)
|
||
c = float(np.clip(np.dot(u, v), -1.0, 1.0))
|
||
if c > 1.0 - 1e-12:
|
||
return np.eye(3)
|
||
if c < -1.0 + 1e-12: # 反平行：繞任一垂直軸轉 180 度
|
||
if abs(u[0]) < 0.9:
|
||
ax = np.cross([1.0, 0.0, 0.0], u)
|
||
else:
|
||
ax = np.cross([0.0, 1.0, 0.0], u)
|
||
ax = ax / np.linalg.norm(ax)
|
||
K = _skew(ax)
|
||
return np.eye(3) + 2.0 * (K @ K)
|
||
ax = np.cross(u, v)
|
||
s = np.linalg.norm(ax)
|
||
ax = ax / s
|
||
K = _skew(ax)
|
||
ang = np.arccos(c)
|
||
return np.eye(3) + np.sin(ang) * K + (1.0 - np.cos(ang)) * (K @ K)
|
||
|
||
|
||
def rotation_about_x(theta):
|
||
ct, st = np.cos(theta), np.sin(theta)
|
||
return np.array([[1.0, 0.0, 0.0],
|
||
[0.0, ct, -st],
|
||
[0.0, st, ct]])
|
||
|
||
|
||
def compute_normalizing_rotation(mask, sym, symp):
|
||
"""回傳 (R, c_xyz)。
|
||
R：使鏡稱面 normal -> +X、上終板面 normal 的 y 分量 -> 0（z 分量保持 >0）。
|
||
若 symp 為 None，只做鏡稱面對齊。
|
||
c_xyz：旋轉中心（array 中心），(x, y, z) 順序。
|
||
"""
|
||
nz, ny, nx = mask.shape
|
||
c_xyz = np.array([(nx - 1) / 2.0, (ny - 1) / 2.0, (nz - 1) / 2.0])
|
||
n_sym = _unit(sym["normal"])
|
||
R1 = rotation_from_to(n_sym, np.array([1.0, 0.0, 0.0]))
|
||
if symp is None:
|
||
return R1, c_xyz
|
||
n_end = _unit(symp["normal"])
|
||
w = R1 @ n_end
|
||
if abs(w[1]) + abs(w[2]) < 1e-12:
|
||
return R1, c_xyz
|
||
theta = float(np.arctan2(w[1], w[2]))
|
||
R2 = rotation_about_x(theta)
|
||
return R2 @ R1, c_xyz
|
||
|
||
|
||
def rotate_volume(arr, R, c_xyz, order=1, cval=0.0):
|
||
"""以 R（作用於 (x,y,z) 向量）在 index space 旋轉 (z,y,x) 體積，
|
||
繞 c_xyz 中心。用逆向映射 + 多項式插值（order=0 最近邻 / 1 三线性）。"""
|
||
nz, ny, nx = arr.shape
|
||
idx = np.indices((nz, ny, nx), dtype=np.float64)
|
||
xx, yy, zz = idx[2], idx[1], idx[0]
|
||
V = np.stack([xx, yy, zz], axis=0) # (3, nz,ny,nx) = (x,y,z)
|
||
c = c_xyz[:, None, None, None]
|
||
W = np.tensordot(R.T, V - c, axes=([1], [0])) + c # (3, nz,ny,nx) = (in_x,in_y,in_z)
|
||
return map_coordinates(arr, [W[2], W[1], W[0]], order=order, cval=cval, mode="constant")
|
||
|
||
|
||
def rotated_sitk_image(src_image, arr):
|
||
"""以相同幾何（origin/spacing/direction）包裝旋轉後的 array。"""
|
||
out = sitk.GetImageFromArray(arr)
|
||
out.CopyInformation(src_image)
|
||
return out
|
||
|
||
|
||
def rotated_grid(box_size_zyx, R, c_xyz, margin=4):
|
||
"""未旋轉（緊密裁切）grid 的 8 角點經 (R, c_xyz) 剛性旋轉（forward
|
||
dest = R(src - c) + c）後取最小包圍軸對齊盒，外擴 margin。
|
||
回傳 (start, size)，皆為 (x, y, z) 序、模板 index 系；size 含旋轉
|
||
擴張，旋轉後整顆骨頭保留、不被裁切（同尺寸旋轉會切掉角落）。"""
|
||
nz, ny, nx = box_size_zyx
|
||
c = np.asarray(c_xyz, dtype=np.float64)
|
||
Rf = np.asarray(R, dtype=np.float64)
|
||
corners = np.array([[x, y, z]
|
||
for x in (0.0, float(nx))
|
||
for y in (0.0, float(ny))
|
||
for z in (0.0, float(nz))], dtype=np.float64)
|
||
pts = (corners - c) @ Rf.T + c
|
||
lo = np.floor(pts.min(axis=0)).astype(int) - int(margin)
|
||
hi = np.ceil(pts.max(axis=0)).astype(int) + int(margin)
|
||
return lo, hi - lo + 1
|
||
|
||
|
||
def rotate_volume_to(arr, R, c_xyz, start, size, order=1, cval=0.0):
|
||
"""以 R（作用於 (x,y,z)）逆向映射旋轉 (z,y,x) 體積，繞 c_xyz（模板
|
||
index 系），輸出到指定 grid (start, size)（模板 index 系，(x, y, z)
|
||
序）；arr 外區域以 cval 填充。回傳 (size_z, size_y, size_x) 陣列。"""
|
||
sx, sy, sz = (int(v) for v in start)
|
||
nx, ny, nz = (int(v) for v in size)
|
||
idx = np.indices((nz, ny, nx), dtype=np.float64)
|
||
V = np.stack([idx[2] + sx, idx[1] + sy, idx[0] + sz], axis=0)
|
||
c = np.asarray(c_xyz, dtype=np.float64)[:, None, None, None]
|
||
Rf = np.asarray(R, dtype=np.float64)
|
||
W = np.tensordot(Rf.T, V - c, axes=([1], [0])) + c
|
||
return map_coordinates(arr, [W[2], W[1], W[0]], order=order, cval=cval,
|
||
mode="constant")
|
||
|
||
|
||
def rotated_sitk_image_at(src_image, arr, start):
|
||
"""包裝 arr（size 可不同於 src_image）：spacing/direction 同 src_image；
|
||
origin = start（(x, y, z)，src_image index 系）對應的實體座標。"""
|
||
out = sitk.GetImageFromArray(arr)
|
||
out.SetSpacing(src_image.GetSpacing())
|
||
d = src_image.GetDirection()
|
||
out.SetDirection([d[i * 3 + j] for i in range(3) for j in range(3)])
|
||
out.SetOrigin(src_image.TransformIndexToPhysicalPoint(
|
||
(int(start[0]), int(start[1]), int(start[2]))))
|
||
return out
|
||
|
||
|
||
def _shift_plane_params(plane, start):
|
||
"""平面（與模板同原點系，n·x = d）換到陣列原點 = 模板原點 + start 的
|
||
局部系：x = v + start -> n·v = d - n·start（方向 u, v 不變）。"""
|
||
n = np.asarray(plane["plane"][:3], dtype=float)
|
||
d = float(plane["plane"][3]) - float(n @ np.asarray(start, dtype=float))
|
||
out = dict(plane)
|
||
out["plane"] = (float(n[0]), float(n[1]), float(n[2]), d)
|
||
return out
|
||
|
||
|
||
def _rotate_plane_params(plane, R, c_xyz):
|
||
"""把平面 (a,b,c,d) 與 in-plane 方向 u,v 一併旋轉；回傳同結構 dict。"""
|
||
n = np.array(plane["plane"][:3], dtype=float)
|
||
d = float(plane["plane"][3])
|
||
n_rot = R @ n
|
||
d_rot = float(n_rot @ c_xyz) + d - float(n @ c_xyz)
|
||
u_rot = R @ np.array(plane["u"], dtype=float)
|
||
v_rot = R @ np.array(plane["v"], dtype=float)
|
||
out = {
|
||
"plane": (n_rot[0], n_rot[1], n_rot[2], d_rot),
|
||
"normal": (n_rot[0], n_rot[1], n_rot[2]),
|
||
"offset": d_rot,
|
||
"u": (u_rot[0], u_rot[1], u_rot[2]),
|
||
"v": (v_rot[0], v_rot[1], v_rot[2]),
|
||
}
|
||
for k in ("ratio", "inlier_ratio", "tilt_deg", "theta_deg", "phi_deg"):
|
||
if k in plane:
|
||
out[k] = plane[k]
|
||
return out
|
||
|
||
|
||
def _mask_to_array(m):
|
||
"""骨頭 / 皮質遮罩：path (nifti) / (z,y,x) ndarray / torch tensor
|
||
-> (z,y,x) bool；None / 缺檔 / 空 -> None。"""
|
||
if m is None:
|
||
return None
|
||
if isinstance(m, str):
|
||
return _load_mask(m)
|
||
arr = m.cpu().numpy() if hasattr(m, 'cpu') else m
|
||
arr = np.asarray(arr)
|
||
if arr.size == 0:
|
||
return None
|
||
b = arr > 0
|
||
return b if int(b.sum()) > 0 else None
|
||
|
||
|
||
def render_bone_figure(volume_id, level, binary_path, cortical_path,
|
||
base_folder="/mnt/1248/open2/cyrou/Output",
|
||
spacing=(0.5, 0.5, 0.5), way="CBT",
|
||
planes_only=False, output_path=None, rotation=None,
|
||
best_position_l=None, best_position_r=None,
|
||
diameter_l=None, length_l=None,
|
||
diameter_r=None, length_r=None,
|
||
image2_path=None, device=None, grid=None,
|
||
swarm_size=None, max_iter=None, total_time=None,
|
||
write_csv=True):
|
||
"""統一骨頭 X-ray 四視角圖（合併原 render_bone_figure + res_plt_2_torch）。
|
||
|
||
四視角：預設 / axial（俯視 XY）/ 冠狀（後視）/ 矢狀。
|
||
內容：皮質 vs 鬆質吸收骨、椎體(gold)、棘突(purple，獨立層覆蓋椎體)、
|
||
中矢狀鏡稱面(orange)、上終板面(green)；螺絲模式另畫中心線(紅)
|
||
+ 圓柱（L darkcyan / R blue，o 層粉）。
|
||
繪製採固定分層（不依深度排序）：基底骨 < VBODY < 棘突 < 終板 < 鏡稱面 < 螺絲。
|
||
|
||
骨骼輸入：
|
||
binary_path 骨頭遮罩：path (nifti) 或 (z,y,x) ndarray / torch tensor
|
||
cortical_path 皮質遮罩：同型（path / ndarray / tensor）；None / 缺檔
|
||
時全部視為鬆質骨
|
||
volume_id/level 可為 None：由路徑反推（…/{vol}/rotated/{level}_*.nii.gz，
|
||
rotated 的上一層 = vol）
|
||
image2_path 未給定時 = binary_path（path 情形）；TPS 模式用其算 2D
|
||
參考 az/alt（Azimuth/Altitude 用相對角）
|
||
|
||
平面 / 分割：
|
||
planes_only=True 只畫平面，不做棘突 / 椎體（VBODY）分割
|
||
rotation=(R, c_xyz) 點雲與平面同依 R 旋轉（繞 c_xyz），用於「對齊後
|
||
（rotated）」的平面圖；R 作用於 (x,y,z) 向量
|
||
（螺絲模式一律做分割：R 側 loss 有 VBODY 獎勵、gold
|
||
顯示需椎體，planes_only 被忽略）
|
||
|
||
螺絲（可選，給定 best_position_l/r 時啟用，optimizer 的輸出）：
|
||
best_position = (z, y, x, az, alt[, d, L])，須搭配該側 diameter/length；
|
||
device/grid 供圓柱生成；swarm_size/max_iter/total_time 供圖面註記；
|
||
write_csv 時在 base_folder/{date}/{vol}/output.csv 依欄位名稱 append
|
||
L/R 兩行（舊 schema 自動重映射）。
|
||
|
||
輸出：
|
||
output_path 給定 -> 存該路徑（含自動加 _1/_2 防覆蓋）；None 時
|
||
螺絲模式: base_folder/{date}/{vol}/{level}_{way}_L{d}_{l}_R{d}_{l}_{swarm}_{iter}.png
|
||
否則: base_folder/{date}/{vol}/{vol} {level}_{way}.png
|
||
回傳存檔路徑；無有效骨頭遮罩時回傳 None。
|
||
"""
|
||
# ---- 路徑反推：volume_id / level / image2_path ----
|
||
if image2_path is None and isinstance(binary_path, str):
|
||
image2_path = binary_path
|
||
if (volume_id is None or level is None) and image2_path:
|
||
parent = os.path.dirname(image2_path)
|
||
vol_p = (os.path.basename(os.path.dirname(parent))
|
||
if os.path.basename(parent) == 'rotated'
|
||
else os.path.basename(parent))
|
||
volume_id = volume_id or vol_p
|
||
level = level or os.path.basename(image2_path).split('_')[0]
|
||
if volume_id is None or level is None:
|
||
raise ValueError('volume_id / level unknown: 需提供 image2_path，'
|
||
'或顯式傳 volume_id 與 level')
|
||
|
||
# ---- 螺絲：決定畫哪些側 ----
|
||
side_cfg = {'L': (best_position_l, diameter_l, length_l, 'Left'),
|
||
'R': (best_position_r, diameter_r, length_r, 'Right')}
|
||
screw_sides = []
|
||
for s in ('L', 'R'):
|
||
pos, d, L, _cn = side_cfg[s]
|
||
if pos is None:
|
||
continue
|
||
if d is None or L is None:
|
||
raise ValueError(f'{s} 側：best_position 須搭配 diameter / length')
|
||
screw_sides.append(s)
|
||
screw_mode = bool(screw_sides)
|
||
CBT = str(way).upper() == 'CBT'
|
||
|
||
# ---- 骨骼 / 皮質 ----
|
||
spine = _mask_to_array(binary_path)
|
||
if spine is None:
|
||
print(f"[skip] {volume_id} {level}: 無/空骨頭遮罩 {binary_path}")
|
||
return None
|
||
cortical = _mask_to_array(cortical_path)
|
||
if cortical is None:
|
||
cortical = np.zeros_like(spine)
|
||
|
||
image_shape = spine.shape
|
||
voxel_mm = float(spacing[0])
|
||
alpha_cortical = 1.0 - np.exp(-BONE_MU_CORTICAL * voxel_mm)
|
||
alpha_trabecular = 1.0 - np.exp(-BONE_MU_TRABECULAR * voxel_mm)
|
||
|
||
# 骨頭 voxel 拆成皮質 / 鬆質兩組（同 res_plt_2_torch）
|
||
z_corti, y_corti, x_corti = np.where((spine) & (cortical))
|
||
z_trab, y_trab, x_trab = np.where((spine) & (~cortical))
|
||
|
||
# ---- 方向 / 解剖分割 ----
|
||
sym = best_symmetry_plane(spine)
|
||
symp = best_upper_endplate_plane(spine)
|
||
|
||
if planes_only and not screw_mode:
|
||
# 只做方向平面，不做棘突 / 椎體分割
|
||
vb_mask = None
|
||
sp_corti = sp_trab = None
|
||
vb_corti = vb_trab = None
|
||
sp_info = {}
|
||
vb_info = {}
|
||
else:
|
||
# 棘突：鏡稱面中線帶（|s|<=w）且在 AP 谷底之後側；棘突缺如
|
||
#（先前 laminectomy / 棘突切除）時 sp_mask=None，不標示
|
||
sp_mask, sp_th, sp_info = segment_spinous_process(spine, sym)
|
||
if sp_info.get('mode') == 'no_spinous':
|
||
if screw_mode:
|
||
top_off = f"{sp_info['top_off']:.1f}" \
|
||
if sp_info.get('top_off') is not None else 'n/a'
|
||
print(f"[NO-SP] 中線後側缺如（先前 laminectomy / 棘突切除）: "
|
||
f"deficit={sp_info.get('deficit', float('nan')):.1f} voxel "
|
||
f"({sp_info.get('deficit', 0.0) * 0.5:.1f} mm), "
|
||
f"rear3={sp_info.get('rear3')} voxel, top_off={top_off} "
|
||
f"-> 不標示棘突；椎體用放寬後側谷底切分")
|
||
sp_corti = sp_trab = None
|
||
elif sp_mask is not None and sp_mask.any():
|
||
sp_corti = sp_mask[z_corti, y_corti, x_corti]
|
||
sp_trab = sp_mask[z_trab, y_trab, x_trab]
|
||
if screw_mode:
|
||
sp_n_bone = max(int(spine.sum()), 1)
|
||
print(f"[SPINOUS] n={sp_info['n_sp']} "
|
||
f"({100.0 * sp_info['n_sp'] / sp_n_bone:.1f}% of bone) "
|
||
f"band=+/-{sp_info['band_w']:.1f} voxel "
|
||
f"AP>={sp_info['ap_thresh']:.1f} mode={sp_info['mode']}")
|
||
else:
|
||
sp_corti = sp_trab = None
|
||
|
||
# 椎體：上終板之下 + 中線 AP 谷底之前側（對齊基準系下與 loss 的
|
||
# VBODY 獎勵完全同 input / 同參數，顯示的椎體＝計分的椎體）
|
||
vb_mask, vb_th, vb_info = segment_vertebral_body(spine, sym, symp,
|
||
sp_th, sp_info["mode"])
|
||
if vb_mask is not None and vb_mask.any():
|
||
vb_corti = vb_mask[z_corti, y_corti, x_corti]
|
||
vb_trab = vb_mask[z_trab, y_trab, x_trab]
|
||
if screw_mode:
|
||
print(f"[VBODY] n={vb_info['n_vb']} "
|
||
f"({100.0 * vb_info['n_vb'] / max(int(spine.sum()), 1):.1f}% of bone) "
|
||
f"AP<{vb_info['ap_thresh']:.1f} mode={vb_info['mode']}")
|
||
if vb_info['mode'] == 'quantile':
|
||
print(f"[VBODY] WARNING: 未找到體/弓後側谷底，閾值退回 55 百分位 "
|
||
f"（可能切進椎體內），建議人工核對該 level 的椎體邊界")
|
||
else:
|
||
vb_corti = vb_trab = None
|
||
if screw_mode:
|
||
print(f"[VBODY] skipped: {vb_info['mode']}")
|
||
|
||
# ---- 骨頭點雲（體積吸收）----
|
||
x_bone = np.concatenate([x_corti, x_trab])
|
||
y_bone = np.concatenate([y_corti, y_trab])
|
||
z_bone = np.concatenate([z_corti, z_trab])
|
||
bone_rgba = np.concatenate([
|
||
_rgba_block(len(x_corti), "lightblue", float(alpha_cortical)),
|
||
_rgba_block(len(x_trab), "lightblue", float(alpha_trabecular)),
|
||
])
|
||
bone_size = np.full(len(x_bone), BONE_MARKER_SIZE)
|
||
|
||
if BONE_SUBSAMPLE > 1:
|
||
x_bone = x_bone[::BONE_SUBSAMPLE]
|
||
y_bone = y_bone[::BONE_SUBSAMPLE]
|
||
z_bone = z_bone[::BONE_SUBSAMPLE]
|
||
bone_rgba = bone_rgba[::BONE_SUBSAMPLE]
|
||
bone_size = bone_size[::BONE_SUBSAMPLE]
|
||
|
||
# VBODY / 棘突的 voxel flags（corti+trab 接合陣列上）；兩 mask 重合時
|
||
# 歸棘突（與 label map 2 覆蓋 1 一致）
|
||
vb_flag = None
|
||
sp_flag = None
|
||
if vb_corti is not None:
|
||
vb_flag = np.concatenate([vb_corti, vb_trab]).astype(bool)
|
||
if BONE_SUBSAMPLE > 1:
|
||
vb_flag = vb_flag[::BONE_SUBSAMPLE]
|
||
if sp_corti is not None:
|
||
sp_flag = np.concatenate([sp_corti, sp_trab]).astype(bool)
|
||
if BONE_SUBSAMPLE > 1:
|
||
sp_flag = sp_flag[::BONE_SUBSAMPLE]
|
||
if vb_flag is None:
|
||
vb_flag = np.zeros(len(x_bone), dtype=bool)
|
||
if sp_flag is None:
|
||
sp_flag = np.zeros(len(x_bone), dtype=bool)
|
||
vb_flag &= ~sp_flag
|
||
|
||
# ---- 旋轉對齊（若給定）：骨頭點雲與平面同旋轉 ----
|
||
if rotation is not None:
|
||
R, c_xyz = rotation
|
||
P = np.stack([x_bone - c_xyz[0], y_bone - c_xyz[1], z_bone - c_xyz[2]], axis=0)
|
||
P = R @ P
|
||
x_bone = P[0] + c_xyz[0]
|
||
y_bone = P[1] + c_xyz[1]
|
||
z_bone = P[2] + c_xyz[2]
|
||
sym = _rotate_plane_params(sym, R, c_xyz)
|
||
if symp is not None:
|
||
symp = _rotate_plane_params(symp, R, c_xyz)
|
||
|
||
# ---- 拆三層：基底骨 / VBODY (gold) / 棘突 (purple) ----
|
||
# x_bone 保留完整點雲（含 VBODY / SP）供下方平面 patch 算範圍
|
||
#（VBODY 前側是整顆骨最前緣，剔除後綠色終板 patch 會縮小）；
|
||
# mpl 3D scatter 在同一 collection 內依深度排序 markers，
|
||
# 「棘突覆蓋椎體、螺絲覆蓋骨頭」改以固定 zorder 分層達成（見 _fill_ax）
|
||
base_idx = ~(vb_flag | sp_flag)
|
||
x_base, y_base, z_base = x_bone[base_idx], y_bone[base_idx], z_bone[base_idx]
|
||
rgba_base, size_base = bone_rgba[base_idx], bone_size[base_idx]
|
||
vb_idx = vb_flag & ~sp_flag
|
||
x_vb, y_vb, z_vb = x_bone[vb_idx], y_bone[vb_idx], z_bone[vb_idx]
|
||
x_sp, y_sp, z_sp = x_bone[sp_flag], y_bone[sp_flag], z_bone[sp_flag]
|
||
|
||
# ---- 螺絲：圓柱 + 中心線 + loss（lazy import torch / core.*；
|
||
# 無螺絲的 preprocess 路徑不會載入 torch）----
|
||
side_info = {}
|
||
side_azlat, side_acep = {}, {}
|
||
theta_v = tau_y = tau_x = float('nan')
|
||
azi = alt = float('nan')
|
||
x_screw = y_screw = z_screw = None
|
||
screw_rgba = screw_size = None
|
||
if screw_mode:
|
||
if device is None:
|
||
raise ValueError('螺絲模式需要 device（torch device）')
|
||
spacing = list(spacing) # core.cylinder 以 list 比對 spacing
|
||
import torch
|
||
from core.cylinder import (generate_cylinder_n_torch,
|
||
generate_cylinder_o_torch)
|
||
from core.intersection import center_line_intersections_torch
|
||
from core.scoring import cl_score_torch, cl_score_torch_xfr
|
||
|
||
# TPS 用 2D 參考角（CBT 無 2D 參考，az/alt 為 nan）
|
||
if not CBT:
|
||
if image2_path:
|
||
azi = float(azimuth_rotation(image2_path))
|
||
alt = float(analyze_vertebral_tilt_contour(
|
||
image2_path, edge_type='superior',
|
||
show_plot=False, debug=False)['superior']['tilt_angle_deg'])
|
||
else:
|
||
print('[warn] TPS 模式缺 image2_path；'
|
||
'Azimuth/Altitude 退回原始角')
|
||
|
||
spine_tensor = torch.from_numpy(spine.astype(np.uint8)).to(device)
|
||
cortical_tensor = torch.from_numpy(cortical.astype(np.uint8)).to(device)
|
||
vbody_tensor = None
|
||
if vb_mask is not None and vb_mask.any():
|
||
# R 側 loss 使用與 PSO 目標函數相同的 VBODY 獎勵
|
||
#（回報分數與優化一致）
|
||
vbody_tensor = torch.from_numpy(vb_mask.astype(np.uint8)).to(device=device)
|
||
|
||
# 角度註記（與 CSV 同參數）：
|
||
# Azimuth_Lateral = 螺絲在鏡稱面內相對 AP 軸的發散角（+ = L 側往外，− = R 側）
|
||
# Altitude_Endplate = 螺絲相對上終板面的仰角
|
||
sym_n = np.asarray(sym['normal'], dtype=float)
|
||
theta_v = float(np.degrees(np.arctan2(sym_n[1], sym_n[0])))
|
||
if symp is not None:
|
||
e_n = np.asarray(symp['normal'], dtype=float)
|
||
e_n = e_n / np.linalg.norm(e_n)
|
||
tau_y = float(np.degrees(np.arctan2(e_n[1], e_n[2])))
|
||
tau_x = float(np.degrees(np.arctan2(e_n[0], e_n[2])))
|
||
else:
|
||
e_n = None
|
||
|
||
def _rel_angles(az_deg, alt_deg):
|
||
az_r = np.radians(az_deg)
|
||
alt_r = np.radians(alt_deg)
|
||
d_v = np.array([np.cos(az_r) * np.sin(alt_r),
|
||
np.sin(az_r) * np.sin(alt_r),
|
||
np.cos(alt_r)])
|
||
az_lateral = az_deg - theta_v - 90.0
|
||
alt_cep = (90.0 - float(np.degrees(np.arccos(
|
||
np.clip(d_v @ e_n, -1.0, 1.0))))
|
||
if e_n is not None else float('nan'))
|
||
return az_lateral, alt_cep
|
||
|
||
for s in screw_sides:
|
||
pos, d, L, _cn = side_cfg[s]
|
||
# 螺絲方向向量（與 generate_cylinder_n_torch 同慣例）：
|
||
# d = (cos(az)·sin(alt), sin(az)·sin(alt), cos(alt))，alt = 相對 +z 的極角
|
||
cyl_n = generate_cylinder_n_torch(d, L, pos[0], pos[1], pos[2],
|
||
pos[3], pos[4],
|
||
image_shape, spacing, device, grid)
|
||
cyl_o = generate_cylinder_o_torch(d, L, pos[0], pos[1], pos[2],
|
||
pos[3], pos[4],
|
||
image_shape, spacing, device, grid)
|
||
inter, line_mask = center_line_intersections_torch(
|
||
pos[0], pos[1], pos[2], pos[3], pos[4],
|
||
int(L), spine_tensor, spacing, device)
|
||
if s == 'L':
|
||
loss = cl_score_torch(cortical_tensor, spine_tensor,
|
||
cyl_n, cyl_o, inter)
|
||
else:
|
||
loss = cl_score_torch_xfr(cortical_tensor, spine_tensor,
|
||
cyl_n, cyl_o, inter,
|
||
vbody_tensor=vbody_tensor)
|
||
cyl_points = int(torch.sum(cyl_n).item())
|
||
ovc = (100.0 * int(((cortical_tensor == 1) & (cyl_n == 1)).sum().item())
|
||
/ cyl_points) if cyl_points else 0.0
|
||
ovb = (100.0 * int(((spine_tensor == 1) & (cyl_n == 1)).sum().item())
|
||
/ cyl_points) if cyl_points else 0.0
|
||
side_info[s] = {
|
||
'line': np.where(line_mask.cpu().numpy() == 1),
|
||
'cyl_n': np.where(cyl_n.cpu().numpy() == 1),
|
||
'cyl_o': np.where(cyl_o.cpu().numpy() == 1),
|
||
'loss': float(loss),
|
||
'inter': inter,
|
||
'cyl_points': cyl_points,
|
||
'ovc': ovc,
|
||
'ovb': ovb,
|
||
}
|
||
side_azlat[s], side_acep[s] = _rel_angles(float(pos[3]), float(pos[4]))
|
||
|
||
# 螺絲點雲（存在的側接合）：中心線紅、圓柱 n = L darkcyan / R blue、o = 粉。
|
||
# 點序固定為舊 res_plt_2_torch 的接合序 [L線, R線, L_n, L_o, R_n, R_o]：
|
||
# mpl 3D 的 per-point 深度排序對近同深 markers 以輸入序決 tie，換序會讓
|
||
# 同一點集產生亞像素級抗鋸齒邊界差（A/B 實測 ~0.07% 邊緣像素），
|
||
# 固定點序保持與舊圖位元級一致。
|
||
_CYL_COLOR = {'L': 'darkcyan', 'R': 'blue'}
|
||
parts_x, parts_y, parts_z, parts_c, parts_s = [], [], [], [], []
|
||
for s in ('L', 'R'): # 中心線
|
||
if s not in side_info:
|
||
continue
|
||
z_p, y_p, x_p = side_info[s]['line']
|
||
parts_x.append(x_p); parts_y.append(y_p); parts_z.append(z_p)
|
||
parts_c.append(_rgba_block(len(x_p), 'r', 1.0))
|
||
parts_s.append(np.full(len(x_p), 3))
|
||
for s in ('L', 'R'): # 圓柱：各側 n 後 o（L_n, L_o, R_n, R_o）
|
||
if s not in side_info:
|
||
continue
|
||
for key, color, size in (('cyl_n', None, 36), ('cyl_o', 'pink', 36)):
|
||
z_p, y_p, x_p = side_info[s][key]
|
||
parts_x.append(x_p); parts_y.append(y_p); parts_z.append(z_p)
|
||
parts_c.append(_rgba_block(len(x_p), _CYL_COLOR[s] if color is None else color, 1.0))
|
||
parts_s.append(np.full(len(x_p), size))
|
||
if parts_x:
|
||
x_screw = np.concatenate(parts_x)
|
||
y_screw = np.concatenate(parts_y)
|
||
z_screw = np.concatenate(parts_z)
|
||
screw_rgba = np.concatenate(parts_c)
|
||
screw_size = np.concatenate(parts_s)
|
||
|
||
# ---- 中矢狀（鏡稱）平面 ----
|
||
_a, _b, _c, _d = sym["plane"]
|
||
_n = np.array([_a, _b, _c])
|
||
_u = np.array(sym["u"])
|
||
_v = np.array(sym["v"])
|
||
_p0 = _d * _n
|
||
xyz_bone = np.stack([x_bone - _p0[0], y_bone - _p0[1], z_bone - _p0[2]], axis=1)
|
||
_pu = xyz_bone @ _u
|
||
_pv = xyz_bone @ _v
|
||
_U_, _V_ = np.meshgrid(np.linspace(_pu.min(), _pu.max(), 8),
|
||
np.linspace(_pv.min(), _pv.max(), 8))
|
||
_Xp = _p0[0] + _U_ * _u[0] + _V_ * _v[0]
|
||
_Yp = _p0[1] + _U_ * _u[1] + _V_ * _v[1]
|
||
_Zp = _p0[2] + _U_ * _u[2] + _V_ * _v[2]
|
||
|
||
# ---- 上終板平面 ----
|
||
_EX = _EY = _EZ = None
|
||
if symp is not None:
|
||
_ea, _eb, _ec, _ed = symp["plane"]
|
||
_en = np.array([_ea, _eb, _ec])
|
||
_eu = np.array(symp["u"])
|
||
_ev = np.array(symp["v"])
|
||
_ep0 = _ed * _en
|
||
xz_ep = np.stack([x_bone - _ep0[0], y_bone - _ep0[1], z_bone - _ep0[2]], axis=1)
|
||
_pu_ep = xz_ep @ _eu
|
||
_pv_ep = xz_ep @ _ev
|
||
_EU, _EV = np.meshgrid(np.linspace(_pu_ep.min(), _pu_ep.max(), 8),
|
||
np.linspace(_pv_ep.min(), _pv_ep.max(), 8))
|
||
_EX = _ep0[0] + _EU * _eu[0] + _EV * _ev[0]
|
||
_EY = _ep0[1] + _EU * _eu[1] + _EV * _ev[1]
|
||
_EZ = _ep0[2] + _EU * _eu[2] + _EV * _ev[2]
|
||
|
||
# ---- 圖 ----
|
||
fig = plt.figure(figsize=(12, 12))
|
||
|
||
legend_handles = []
|
||
for s in ("L", "R"):
|
||
if s in side_info:
|
||
legend_handles.append(Line2D([], [], marker="o", ls="", ms=6,
|
||
color="darkcyan" if s == "L" else "blue",
|
||
label=f"Cylinder({s})"))
|
||
if vb_corti is not None:
|
||
legend_handles.append(Line2D([], [], marker="o", ls="", ms=6, color="gold", label="VertebralBody"))
|
||
if sp_corti is not None:
|
||
legend_handles.append(Line2D([], [], marker="o", ls="", ms=6, color="purple", label="SpinousProcess"))
|
||
|
||
def _fill_ax(ax):
|
||
# 固定分層（關 depth zorder，否則半透明骨頭會被重繪到螺絲上方）：
|
||
# 基底骨(5) < VBODY gold(6) < 棘突 purple(6.5) < 終板(7) < 鏡稱面(8) < 螺絲(10)
|
||
ax.computed_zorder = False
|
||
sc_bone = ax.scatter(x_base, y_base, z_base, c=rgba_base, s=size_base, marker="o")
|
||
sc_bone.set_zorder(5)
|
||
if x_vb.size:
|
||
sc_vb = ax.scatter(x_vb, y_vb, z_vb, c=to_rgba("gold", 0.95),
|
||
s=BONE_MARKER_SIZE, marker="o")
|
||
sc_vb.set_zorder(6)
|
||
if x_sp.size:
|
||
sc_sp = ax.scatter(x_sp, y_sp, z_sp, c=to_rgba("purple", 0.95),
|
||
s=BONE_MARKER_SIZE, marker="o")
|
||
sc_sp.set_zorder(6.5)
|
||
if _EX is not None:
|
||
ep = ax.plot_surface(_EX, _EY, _EZ, color="green", alpha=0.35,
|
||
linewidth=1.0, edgecolor="green", rstride=1, cstride=1)
|
||
ep.set_zorder(7)
|
||
plane = ax.plot_surface(_Xp, _Yp, _Zp, color="orange", alpha=0.30,
|
||
linewidth=1.0, edgecolor="orange", rstride=1, cstride=1)
|
||
plane.set_zorder(8)
|
||
if x_screw is not None:
|
||
sc_screw = ax.scatter(x_screw, y_screw, z_screw,
|
||
c=screw_rgba, s=screw_size, marker="o")
|
||
sc_screw.set_zorder(10)
|
||
|
||
ax1 = fig.add_subplot(221, projection="3d")
|
||
_fill_ax(ax1)
|
||
ax1.set_xlabel("X-axis"); ax1.set_ylabel("Y-axis"); ax1.set_zlabel("Z-axis")
|
||
set_axes_equal_3d(ax1)
|
||
|
||
ax2 = fig.add_subplot(222, projection="3d")
|
||
ax2.view_init(elev=90, azim=-90, roll=0)
|
||
_fill_ax(ax2)
|
||
ax2.set_xlabel("X-axis"); ax2.set_ylabel("Y-axis"); ax2.set_zlabel("Z-axis")
|
||
set_axes_equal_3d(ax2)
|
||
if legend_handles:
|
||
ax2.legend(handles=legend_handles)
|
||
|
||
ax3 = fig.add_subplot(223, projection="3d")
|
||
# 後視圖：相機在 −y 後側，x 軸畫面左小右大
|
||
ax3.view_init(elev=0, azim=-90, roll=0)
|
||
_fill_ax(ax3)
|
||
ax3.set_xlabel("X-axis"); ax3.set_ylabel("Y-axis"); ax3.set_zlabel("Z-axis")
|
||
set_axes_equal_3d(ax3)
|
||
|
||
ax4 = fig.add_subplot(224, projection="3d")
|
||
ax4.view_init(elev=0, azim=0, roll=0)
|
||
_fill_ax(ax4)
|
||
ax4.set_xlabel("X-axis"); ax4.set_ylabel("Y-axis"); ax4.set_zlabel("Z-axis")
|
||
set_axes_equal_3d(ax4)
|
||
|
||
label_str = f"{volume_id} {level}"
|
||
if screw_mode:
|
||
fig.text(0.5, 0.98, f"{label_str} Best Position", ha="center", fontsize=15)
|
||
|
||
d_l = f"{diameter_l} mm, {length_l} mm" if diameter_l is not None else ''
|
||
d_r = f"{diameter_r} mm, {length_r} mm" if diameter_r is not None else ''
|
||
t_s = f"Total time = {total_time:.2f} s" if total_time is not None else ''
|
||
fig.text(
|
||
0.5, 0.44,
|
||
f"L: Diameter = {d_l}, "
|
||
f"R: Diameter = {d_r}, "
|
||
f"Swarm size = {swarm_size}, Iteration = {max_iter}, {t_s}",
|
||
ha="center", fontsize=12
|
||
)
|
||
|
||
def _ang_segs(az_v, alt_v, azlat_v, acep_v):
|
||
# CBT 沒有 2D 參考面，Azimuth/Altitude 直接用最佳化出的原始角；
|
||
# TPS 沿用 2D 參考之相對角。終板面擬合失敗（nan）時該段自動略過。
|
||
segs = [f"Azimuth = {az_v:.2f}", f"Altitude = {alt_v:.2f}"]
|
||
if np.isfinite(azlat_v):
|
||
segs.append(f"Azimuth_Lateral = {azlat_v:.2f}")
|
||
if np.isfinite(acep_v):
|
||
segs.append(f"Altitude_Endplate = {acep_v:.2f}")
|
||
return ', '.join(segs)
|
||
|
||
def _side_footer(s, y):
|
||
pos, _d, _L, cn_name = side_cfg[s]
|
||
if pos is None:
|
||
return
|
||
if CBT:
|
||
segs = _ang_segs(float(pos[3]), float(pos[4]),
|
||
side_azlat[s], side_acep[s])
|
||
else:
|
||
segs = _ang_segs(90.0 - float(pos[3]) - azi,
|
||
90.0 - float(pos[4]) - alt,
|
||
side_azlat[s], side_acep[s])
|
||
di = side_info[s]
|
||
cb_ratio = di['ovc'] / di['ovb'] if di['ovb'] else 0.0
|
||
fig.text(
|
||
0.5, y,
|
||
f"{cn_name} : Position = ({pos[2]:.2f}, {pos[1]:.2f}, {pos[0]:.2f}), "
|
||
f"{segs}, "
|
||
f"Intersection = {di['inter']}, "
|
||
f"Score = {di['ovc']:.2f} / {di['ovb']:.2f} / {cb_ratio:.2f}",
|
||
ha="center", fontsize=8
|
||
)
|
||
|
||
_side_footer('L', 0.03)
|
||
_side_footer('R', 0.01)
|
||
else:
|
||
base_tag = "planes only" if planes_only else "no screws"
|
||
if rotation is not None:
|
||
base_tag += ", rotated"
|
||
fig.text(0.5, 0.98, f"{label_str} ({base_tag})", ha="center", fontsize=15)
|
||
|
||
ratio = float(sym.get("ratio", float("nan")))
|
||
if planes_only:
|
||
ep_ratio = float(symp.get("inlier_ratio", float("nan"))) if symp is not None else float("nan")
|
||
info = (f"sym_ratio={ratio:.3f} "
|
||
f"endplane_ratio={ep_ratio:.3f} "
|
||
f"endplate={'yes' if symp is not None else 'no'}")
|
||
else:
|
||
info = (f"sym_ratio={ratio:.3f} "
|
||
f"spinous={sp_info.get('n_sp', 0)} ({sp_info.get('mode', '?')}) "
|
||
f"vertebral_body={vb_info.get('n_vb', 0)} ({vb_info.get('mode', '?')})")
|
||
fig.text(0.5, 0.03, info, ha="center", fontsize=8)
|
||
|
||
fig.tight_layout()
|
||
|
||
date_str = datetime.now().strftime("%Y%m%d")
|
||
output_folder = os.path.join(base_folder, date_str, volume_id)
|
||
retry_robust(os.makedirs, output_folder, exist_ok=True)
|
||
|
||
# ---- CSV（螺絲模式）：L/R 兩行；舊 schema 檔案依欄位名稱重映射後改寫，
|
||
# 避免 append 欄位錯位 ----
|
||
if screw_mode and write_csv:
|
||
csv_path = os.path.join(output_folder, 'output.csv')
|
||
# CBT 模式下恆為 nan 的 2D 參考欄不寫入 CSV
|
||
headers = [
|
||
'Label', 'Side', 'Diameter', 'Length', 'Swarm_Size', 'Max_Iter',
|
||
'Position_XYZ', 'Raw_Azimuth', 'Raw_Altitude',
|
||
'Intersections', 'Best_Loss', 'cyl_points', 'Overlap_Cortical', 'Overlap_Bone',
|
||
'Cortical_Bone_Ratio',
|
||
'Sym_Theta_v_deg', 'Endplate_Tau_y_deg', 'Endplate_Tau_x_deg',
|
||
'Azimuth_Lateral_deg',
|
||
'Altitude_Cephalad_Endplate_deg',
|
||
'Total_Time'
|
||
]
|
||
|
||
def _fmt(v):
|
||
return '' if not np.isfinite(v) else f"{float(v):.2f}"
|
||
|
||
file_exists = os.path.isfile(csv_path)
|
||
if file_exists:
|
||
with retry_robust(open, csv_path, newline='') as f:
|
||
old_rows = [row for row in csv.reader(f) if any(c.strip() for c in row)]
|
||
if not old_rows or old_rows[0] != headers:
|
||
old_h = old_rows[0] if old_rows else None
|
||
with retry_robust(open, csv_path, 'w', newline='') as f:
|
||
w = csv.writer(f)
|
||
w.writerow(headers)
|
||
for r in (old_rows[1:] if old_rows else []):
|
||
if old_h:
|
||
d = dict(zip(old_h, r))
|
||
w.writerow([d.get(h, '') for h in headers])
|
||
else:
|
||
w.writerow(r + [''] * max(0, len(headers) - len(r)))
|
||
|
||
try:
|
||
with retry_robust(open, csv_path, 'a', newline='') as csvfile:
|
||
writer = csv.writer(csvfile)
|
||
if not file_exists:
|
||
writer.writerow(headers)
|
||
for s in ('L', 'R'):
|
||
pos, d, L, _cn = side_cfg[s]
|
||
if pos is None:
|
||
continue
|
||
di = side_info[s]
|
||
writer.writerow([
|
||
label_str, s, d, L, swarm_size, max_iter,
|
||
f"({pos[2]:.2f}, {pos[1]:.2f}, {pos[0]:.2f})",
|
||
f"{pos[3]:.2f}", f"{pos[4]:.2f}",
|
||
di['inter'], f"{di['loss']:.2f}", di['cyl_points'],
|
||
f"{di['ovc']:.2f}", f"{di['ovb']:.2f}",
|
||
f"{(di['ovc'] / di['ovb'] if di['ovb'] else 0):.2f}",
|
||
_fmt(theta_v), _fmt(tau_y), _fmt(tau_x),
|
||
_fmt(side_azlat[s]), _fmt(side_acep[s]),
|
||
f"{total_time:.2f}" if total_time is not None else ''
|
||
])
|
||
print(f"[CSV Saved] {csv_path}")
|
||
except Exception as e:
|
||
print(f"[Error] Failed to write CSV: {e}")
|
||
|
||
if output_path is not None:
|
||
path = get_unique_filepath(output_path)
|
||
elif screw_mode:
|
||
path = save_with_unique_name(
|
||
output_folder, level, way,
|
||
diameter_l if diameter_l is not None else '',
|
||
length_l if length_l is not None else '',
|
||
diameter_r if diameter_r is not None else '',
|
||
length_r if length_r is not None else '',
|
||
swarm_size if swarm_size is not None else '',
|
||
max_iter if max_iter is not None else '',
|
||
)
|
||
else:
|
||
path = get_unique_filepath(
|
||
os.path.join(output_folder, f"{volume_id} {level}_{way}.png"))
|
||
|
||
retry_robust(os.makedirs, os.path.dirname(path) or ".", exist_ok=True)
|
||
retry_robust(fig.savefig, path, dpi=200, bbox_inches="tight")
|
||
print("[Saved figure]", path)
|
||
plt.close(fig)
|
||
return path |