CBT_project/visualization/res_bone_figure.py
Xiao Furen 523ec7ee16 feat(core): implement vertebral body rewards and enhanced segmentation logic
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.
2026-09-05 04:30:10 +08:00

451 lines
No EOL
18 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
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 (best_symmetry_plane, best_upper_endplate_plane,
segment_spinous_process, segment_vertebral_body)
from utils.helpers import get_unique_filepath
# 體積吸收渲染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 分量 -> 0z 分量保持 >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):
"""包裝 arrsize 可不同於 src_imagespacing/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 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):
"""繪製單一 (volume, level) 骨頭 X-ray 圖(不畫螺絲)。
四視角:預設/axial/coronal/sagittal。
內容:皮質 vs 鬆質吸收骨、椎體(gold)、棘突(purple)、
中矢狀鏡稱面(orange)、上終板面(green)。無圓柱/中心線。
planes_only=True只畫骨頭 + 中矢狀鏡稱面 + 上終板面,
不做棘突 / 椎體VBODY分割。
output_path若給定直接存到該路徑含自動加 _1/_2 防覆蓋);
否則存到 base_folder/{date}/{volume_id}/{volume_id} {level}_{way}.png。
rotation(R, c_xyz)。給定時,把骨頭點雲與兩個平面都依 R 旋轉(繞 c_xyz
用於畫「對齊後rotated」的平面圖R 作用於 (x,y,z) 向量。
cortical_path皮質遮罩路徑或 (z,y,x) 0/1 陣列(須與 binary_path 同 grid
None / 缺檔時全部視為鬆質骨。
回傳存檔路徑;無有效骨頭遮罩時回傳 None。
"""
spine = _load_mask(binary_path)
if spine is None:
print(f"[skip] {volume_id} {level}: 無/空骨頭遮罩 {binary_path}")
return None
if isinstance(cortical_path, np.ndarray):
cortical = cortical_path > 0
else:
cortical = _load_mask(cortical_path)
if cortical is None:
cortical = np.zeros_like(spine)
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:
# 只做方向平面,不做棘突 / 椎體分割
sp_corti = sp_trab = None
vb_corti = vb_trab = None
sp_info = {}
vb_info = {}
else:
sp_mask, sp_th, sp_info = segment_spinous_process(spine, sym)
if 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]
else:
sp_corti = sp_trab = None
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]
else:
vb_corti = vb_trab = None
# ---- 骨頭點雲(體積吸收)----
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]
if vb_corti is not None:
vb_flag = np.concatenate([vb_corti, vb_trab])
if BONE_SUBSAMPLE > 1:
vb_flag = vb_flag[::BONE_SUBSAMPLE]
bone_rgba[vb_flag] = to_rgba("gold", 0.95)
if sp_corti is not None:
sp_flag = np.concatenate([sp_corti, sp_trab])
if BONE_SUBSAMPLE > 1:
sp_flag = sp_flag[::BONE_SUBSAMPLE]
bone_rgba[sp_flag] = to_rgba("purple", 0.95)
# ---- 旋轉對齊(若給定):骨頭點雲與平面同旋轉 ----
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)
# ---- 中矢狀(鏡稱)平面 ----
_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 = []
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="Spinous"))
def _fill_ax(ax):
ax.computed_zorder = False
sc_bone = ax.scatter(x_bone, y_bone, z_bone, c=bone_rgba, s=bone_size, marker="o")
sc_bone.set_zorder(5)
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 _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)
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")
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}"
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()
if output_path is not None:
path = get_unique_filepath(output_path)
else:
date_str = datetime.now().strftime("%Y%m%d")
output_folder = os.path.join(base_folder, date_str, volume_id)
output_file = os.path.join(output_folder, f"{volume_id} {level}_{way}.png")
path = get_unique_filepath(output_file)
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
fig.savefig(path, dpi=200, bbox_inches="tight")
plt.close(fig)
return path