From 523ec7ee1647cbc70ba1c51fe5298f7e4b3762d7 Mon Sep 17 00:00:00 2001 From: Xiao Furen Date: Sat, 5 Sep 2026 04:30:10 +0800 Subject: [PATCH] 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. --- .gitignore | 2 + core/objective.py | 6 +- core/optimizer.py | 519 ++++++++++++++++++------------- core/scoring.py | 16 +- imaging/orientation.py | 447 +++++++++++++++++++++++--- imaging/preprocessing.py | 165 ++++++++-- imaging/resample.py | 49 ++- imaging/segmentation.py | 187 +++++++---- visualization/res_bone_figure.py | 451 +++++++++++++++++++++++++++ visualization/res_plot_3d.py | 293 +++++++++++++---- xfr_check_spinous.py | 310 ++++++++++++++++++ xfr_debug.py | 178 +++++++---- xfr_plot_level.py | 97 ++++++ xfr_preprocess.py | 421 ++++++++++++++++++++++++- 14 files changed, 2655 insertions(+), 486 deletions(-) create mode 100644 visualization/res_bone_figure.py create mode 100644 xfr_check_spinous.py create mode 100644 xfr_plot_level.py diff --git a/.gitignore b/.gitignore index af1ac9e..c0a009b 100644 --- a/.gitignore +++ b/.gitignore @@ -217,4 +217,6 @@ __marimo__/ .kilo/ logs/ +results/ progress.json +xfr_image_metadata* diff --git a/core/objective.py b/core/objective.py index 5851514..0b3da0f 100644 --- a/core/objective.py +++ b/core/objective.py @@ -34,6 +34,8 @@ class OptimizationContext: Optional: grid: precomputed coordinate grid (z_t, y_t, x_t) use_tip_penalty: add the tip-cylinder penalty to the loss + vbody_tensor: VBODY (vertebral body) mask 0/1 — adds the VBODY + voxel rewards in cl_score_torch_xfr The *_array / spine_roi_tensor / diameter / length fields are kept for compatibility and debugging; the loss itself does not use them. """ @@ -44,6 +46,7 @@ class OptimizationContext: device: torch.device grid: Optional[tuple] = None use_tip_penalty: bool = False + vbody_tensor: Optional[torch.Tensor] = None spine_roi_tensor: Optional[torch.Tensor] = None image1_array: Optional[np.ndarray] = None image2_array: Optional[np.ndarray] = None @@ -121,7 +124,8 @@ def cylinder_circle_line_intersection_loss_deductions_torch( loss_value = cl_score_torch_xfr( ctx.cortical_tensor, ctx.spine_tensor, cyl_fwd, cyl_opp, intersections, - cylinder_tip_torch=cyl_tip + cylinder_tip_torch=cyl_tip, + vbody_tensor=ctx.vbody_tensor ) return loss_value diff --git a/core/optimizer.py b/core/optimizer.py index 1d08cd1..9ffe872 100644 --- a/core/optimizer.py +++ b/core/optimizer.py @@ -1,10 +1,12 @@ +import json +import os import time from datetime import datetime import SimpleITK as sitk import torch from imaging.orientation import (azimuth_rotation, analyze_vertebral_tilt_contour, best_symmetry_plane, best_upper_endplate_plane, - segment_spinous_process) + segment_spinous_process, segment_vertebral_body) from config.constant import ALLOWED_DIAMETERS, ALLOWED_LENGTHS from core.objective import OptimizationContext, make_objective_function, make_objective_function_xfr from pyswarm import pso @@ -16,12 +18,30 @@ from visualization.res_plot_3d import res_plt_2_torch LATERAL_REFINE_MIN_IN_BONE = 0.97 +# 入口柱向前（+y，椎體方向）幾個 voxel 內進入 VBODY 即視為「入口在椎體上」： +# 涵蓋椎體後側皮質邊緣（mask 外 1~2 voxel 的分割界線帶）； +# 真正的後側要素與椎體間隔（椎間孔）大於此值，不受影響。 +VBODY_ENTRY_EDGE = 4 + + +def _makedirs_retry(path, retries=5, delay=0.5): + """NFS 上建目錄重試（同 res_plot_3d._retry_robust 處理的瞬時錯誤）""" + for i in range(retries): + try: + os.makedirs(path, exist_ok=True) + return + except OSError: + if i == retries - 1: + raise + time.sleep(delay) + def refine_lateral_longer( z, x, azimuth, altitude, diameter_raw, length_raw, side, az_bounds, x_bounds, y_indices, image_shape, spacing, device, grid, cortical_tensor, spine_tensor, + vbody_tensor=None, ): """ Deterministic local refinement after PSO: try aiming more laterally and @@ -43,8 +63,9 @@ def refine_lateral_longer( if cyl.sum().item() == 0: return None inter, _ = center_line_intersections_torch(z_c, y_c, x_c, az_c, alt_c, - L_c, spine_tensor, spacing, device) - loss = cl_score_torch_xfr(cortical_tensor, spine_tensor, cyl, cyl_o, inter) + L_c, spine_tensor, spacing, device) + loss = cl_score_torch_xfr(cortical_tensor, spine_tensor, cyl, cyl_o, inter, + vbody_tensor=vbody_tensor) in_bone = ((spine_tensor == 1) & (cyl == 1)).sum().item() / cyl.sum().item() return {'pos': cand, 'loss': loss, 'in_bone': in_bone} @@ -86,7 +107,20 @@ def _validate_bounds(lb, ub, name): raise ValueError(f'PSO bounds invalid for {name}: dim {i} lower {lo} >= upper {hi}') -def get_first_nonzero_y(arr): +def get_first_nonzero_y(arr, endplate_plane=None, vbody_mask=None): + """每個 (z, x) 柱的第一個 bone voxel 的 y（voxel index）。 + + endplate_plane 為 {'plane': (a, b, c, d), ...}（法線朝上，c > 0）時， + 入口點 (x, y, z) 落在終板面之上的柱也設 OUTSIDE_VALUE， + 避免從終板上方的柱選入射點。 + + vbody_mask（(z, y, x) 0/1，VBODY 椎體 mask）提供時， + 入口點落在椎體上的柱設 OUTSIDE_VALUE，避免螺絲入口點放在椎體上： + 入口 voxel 本身在 VBODY 內、或往前（+y，椎體方向） VBODY_ENTRY_EDGE + 個 voxel 內進入 VBODY（椎體後側皮質邊緣，mask 外 0.5~2mm 的交界帶） + 都算。真正的後側要素（椎弓根/椎板/關節突）厚度遠大於該邊緣， + 且與椎體之間有椎間孔間隔，不會被誤排。 + """ OUTSIDE_VALUE = -100 # 1. Create a boolean mask where elements are non-zero @@ -104,7 +138,35 @@ def get_first_nonzero_y(arr): y_indices = np.where(y_indices < arr.shape[1] * .1, OUTSIDE_VALUE, y_indices) y_indices = np.where(y_indices > arr.shape[1] * .4, OUTSIDE_VALUE, y_indices) - + + # 5. 入口點在終板面上方（a*x + b*y + c*z - d > 0）的柱設 OUTSIDE_VALUE + if endplate_plane is not None and y_indices.max() >= 0: + a, b, c, d = endplate_plane['plane'] + z_grid, x_grid = np.meshgrid(np.arange(arr.shape[0]), np.arange(arr.shape[2]), + indexing='ij') + above = (y_indices >= 0) & (a * x_grid + b * y_indices + c * z_grid > d) + y_indices = np.where(above, OUTSIDE_VALUE, y_indices) + + # 6. 入口點落在椎體上的柱設 OUTSIDE_VALUE（入口不放椎體）： + # 入口 voxel 本身在 VBODY 內，或往前 VBODY_ENTRY_EDGE 個 voxel 內進入 + # VBODY（椎體後側皮質邊緣，mask 外 0.5~2mm 的分割界線帶）都排除。 + if vbody_mask is not None: + vb = np.asarray(vbody_mask) > 0 + valid = y_indices >= 0 + y_i = np.where(valid, y_indices, 0).astype(np.int64) + z_grid, x_grid = np.meshgrid(np.arange(arr.shape[0]), np.arange(arr.shape[2]), + indexing='ij') + on_vbody = np.zeros(y_indices.shape, dtype=bool) + for k in range(VBODY_ENTRY_EDGE + 1): + y_k = np.minimum(y_i + k, arr.shape[1] - 1) + on_vbody |= valid & vb[z_grid, y_k, x_grid] + n_vb = int(on_vbody.sum()) + if n_vb: + y_indices = np.where(on_vbody, OUTSIDE_VALUE, y_indices) + print(f"[Y-INDEX] VBODY entry excluded: {n_vb} columns whose entry " + f"is on/within {VBODY_ENTRY_EDGE} vox of VBODY removed from " + f"entry surface") + return y_indices.astype(np.float32) def constraint_y(x, y_indices): @@ -125,6 +187,11 @@ def run_pso_torch_xfr( grid=None, debug=False, omega = 0.9, + side: str = 'both', + level: str = None, + patient_id: str = None, + side_dir: str = None, + run_id: str = None, ): """ Main function to run PSO. @@ -169,19 +236,19 @@ def run_pso_torch_xfr( length=length if not optimize_size else None, ) - azi = azimuth_rotation(image2_path) - res = analyze_vertebral_tilt_contour(image2_path, edge_type='superior', show_plot=False, debug=False) - alt = res['superior']['tilt_angle_deg'] + if not CBT: + azi = azimuth_rotation(image2_path) + res = analyze_vertebral_tilt_contour(image2_path, edge_type='superior', show_plot=False, debug=False) + alt = res['superior']['tilt_angle_deg'] - # ===== az/alt 搜尋範圍錨點：以椎體自身座標系取代 2D 近似 ===== - # 鏡稱對稱面（最佳 mirror plane）：法線 (s_nx, s_ny, s_nz)，s_nx>0 固定符號 - # theta_v = atan2(s_ny, s_nx) 是椎體真實左右軸相對 +x 的旋轉角， - # 取代舊 2D 輪廓前點角 azi（舊式 az 中心 90-azi == 新式 90+theta_v）。 - # 上終板面（normal 朝上, e_nz>0）： - # tau_y = atan2(e_ny, e_nz)：終板 AP 面傾斜；負 = 面向 +y（終板側）升高， - # 軌跡需爬得更陡 → altitude 中心減小（舊式 65-alt 的 3D 版）。 - # tau_x = atan2(e_nx, e_nz)：終板 LR 面傾斜；正 = 面向小 x（L 側）升高， - # L 側需更陡（-tau_x）、R 側較平（+tau_x）。 + # ===== 平面：影像已由 xfr_preprocess 對齊旋轉到椎體基準系 ===== + # （鏡稱面法線 -> +x、上終板 normal y=0 / z>0），az/alt 範圍改用固定 + # 約束（見下方 CBT bounds），不再逐椎以 theta_v / tau 重新錨定。 + # 此處平面仅供： + # - 棘突移除（segment_spinous_process，sym） + # - 入口面終板上方剪除（get_first_nonzero_y，endplate） + # - VBODY 椎體分割（segment_vertebral_body，sym + endplate） + # - 對齊健全性檢查：基準系下 theta_v ≈ 0、tau_y ≈ 0（偏大 = 預處理失效） sym_plane = best_symmetry_plane(image2_array) # 入口面 y_indices 取「每個 (z,x) 柱第一個 bone voxel（最後側）」， @@ -189,19 +256,51 @@ def run_pso_torch_xfr( # 見 segment_spinous_process）從 image2_array 移除再取 surface， # 讓 y_indices 永不落在棘突上（loss 用 spine_tensor，不受影響）。 sp_mask, sp_th, sp_info = segment_spinous_process(image2_array, sym_plane) - if sp_mask is not None and sp_mask.any(): + + # 上終板面與椎體都在「完整 mask（SP 移除前）」上計算： + # - 終板面由前側頂面擬合，SP 移除不改變平面； + # - 椎體與 res_plt_2_torch 的 gold 顯示完全同 input / 同參數， + # 確保顯示出來的椎體就是 loss 裡 VBODY 獎勵的區域。 + # 棘突缺如（laminectomy，mode='no_spinous'）時不該把殘留後側要素 + # 當「棘突」移除（會鏟進椎體後側），入口面維持完整 mask。 + endplate_plane = best_upper_endplate_plane(image2_array) + + # VBODY 評分獎勵：兩平面（鏡稱面 + 上終板）切出的椎體 mask + vb_mask_np, vb_th, vb_info = segment_vertebral_body(image2_array, sym_plane, + endplate_plane, sp_th, sp_info['mode']) + vbody_tensor = None + if vb_mask_np is not None and vb_mask_np.any(): + vbody_tensor = torch.from_numpy(vb_mask_np.astype(np.uint8)).to(device=device) + print(f"[VBODY-SCORE] n={vb_info['n_vb']} " + f"({100.0 * vb_info['n_vb'] / max(int(image2_array.sum()), 1):.1f}% of bone) " + f"AP<{vb_info['ap_thresh']:.1f} mode={vb_info['mode']} -> added to loss") + if vb_info['mode'] == 'quantile': + print(f"[VBODY-SCORE] WARNING: 未找到體/弓後側谷底，閾值退回 55 百分位 " + f"（可能切進椎體內），建議人工核對該 level 的椎體邊界") + else: + print(f"[VBODY-SCORE] skipped: {vb_info['mode']}") + ctx.vbody_tensor = vbody_tensor + + if sp_info['mode'] == 'no_spinous': + 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['deficit']:.1f} vox ({sp_info['deficit'] * 0.5:.1f} mm), " + f"rear3={sp_info['rear3']} vox, top_off={top_off} vox " + f"-> 不移除 SP（完整 mask 取入口面），椎體用放寬後側谷底切分") + elif sp_mask is not None and sp_mask.any(): n_sp = int(sp_mask.sum()) image2_array[sp_mask] = 0 print(f"[Y-INDEX] spinous process removed from entry surface: {n_sp} vox " f"(band=+/-{sp_info['band_w']:.1f} voxel, AP>={sp_info['ap_thresh']:.1f}, " f"mode={sp_info['mode']})") - y_indices = get_first_nonzero_y(image2_array) + y_indices = get_first_nonzero_y(image2_array, endplate_plane, vb_mask_np) + objective_fn = make_objective_function_xfr(ctx, y_indices) + # 對齊健全性檢查（基準系下 theta_v ≈ 0、tau_y ≈ 0；不參與 bounds 計算） s_nx, s_ny, s_nz = sym_plane['normal'] theta_v = float(np.degrees(np.arctan2(s_ny, s_nx))) - endplate_plane = best_upper_endplate_plane(image2_array) if endplate_plane is not None: e_nx, e_ny, e_nz = endplate_plane['normal'] tau_y = float(np.degrees(np.arctan2(e_ny, e_nz))) @@ -217,7 +316,7 @@ def run_pso_torch_xfr( f" (tau_y={tau_y:+.2f} deg, tau_x={tau_x:+.2f} deg, " f"inlier={endplate_plane['inlier_ratio']:.2f})") else: - print("[PLANE] endplate: 資料不足，altitude 錨點退回 tau_y=tau_x=0") + print("[PLANE] endplate: 資料不足，終板入口剪除停用") # flat_min_index = np.argmin(y_indices) # z_border, x_border = np.unravel_index(flat_min_index, y_indices.shape) @@ -282,18 +381,21 @@ def run_pso_torch_xfr( if x_bounds_right[0] >= x_bounds_right[1]: x_bounds_right = (x_bounds_right[1] - min_band, x_bounds_right[1]) + # 固定約束（影像已由 xfr_preprocess 對齊旋轉到椎體基準系：鏡稱面法線 +x、 + # 終板 normal y=0 / z>0；基準系下 theta_v≈0、tau_y≈0，舊式逐椎錨定 + # (98+theta_v, 105+theta_v) / (60+tau_y∓tau_x*sin_mid, 70+tau_y∓tau_x*sin_mid) + # 不再需要）： + # azimuth (Lateral)：az=90° 為 AP 直向，每側向外發散 8~20° + # L = 90 + (8~20) = 98~110、R = 90 - (8~20) = 70~82 + # altitude (Cephalad)：相對終板面 25~30°（+z 極角 90 - 25~30）= 60~65 # 舊 2D 版本（以輪廓角 azi / 矢狀面傾斜 alt 平移固定範圍），保留供對照： # azimuth_bounds_l = ((98-azi), (120-azi)) # azimuth_bounds_r = ((60-azi), (82-azi)) # altitude_bounds = ((60-alt), (70-alt)) - # 新：範圍以「椎體自身座標系」為中心 —— 椎體系中 az=90° 是 AP 直向、 - # L 帶 = AP 後退 8~30°（偏 -x）、R 帶 = AP 前進 8~30°（偏 +x）、 - # altitude 60~70。再換算回影像系：az 整體加 theta_v（鏡稱面）， - # altitude 加 tau_y（終板 AP 傾斜）並逐側加 -/+tau_x（終板 LR 傾斜）。 - azimuth_bounds_l = ((98+theta_v), (120+theta_v)) - azimuth_bounds_r = ((60+theta_v), (82+theta_v)) - altitude_bounds_l = ((60+tau_y-tau_x), (70+tau_y-tau_x)) - altitude_bounds_r = ((60+tau_y+tau_x), (70+tau_y+tau_x)) + azimuth_bounds_l = (98, 110) + azimuth_bounds_r = (70, 82) + altitude_bounds_l = (60, 65) + altitude_bounds_r = (60, 65) else: z_bounds = (0, image_shape[0] - 1) @@ -340,7 +442,8 @@ def run_pso_torch_xfr( diameter_bounds = (min(ALLOWED_DIAMETERS), max(ALLOWED_DIAMETERS)*1.01) length_bounds = (min(ALLOWED_LENGTHS), max(ALLOWED_LENGTHS)*1.01) - # bounds 現在有 7 個參數（altitude 分 L/R 兩側，由終板面 tau_x 決定） + # bounds 現在有 6 個參數 [z, x, az, alt, d, L]（y 由 y_indices 取）； + # az/alt 兩側用同一組固定約束，x 帶依 L/R 分開 lb_l = [z_bounds[0], x_bounds_left[0], azimuth_bounds_l[0], altitude_bounds_l[0], diameter_bounds[0], length_bounds[0]] ub_l = [z_bounds[1], x_bounds_left[1], azimuth_bounds_l[1], @@ -365,220 +468,200 @@ def run_pso_torch_xfr( ub_r = [z_bounds[1], y_bounds[1], x_bounds_right[1], azimuth_bounds_r[1], altitude_bounds_r[1]] if True or debug: - print(lb_l) - print(ub_l) - print(lb_r) - print(ub_r) + for b in (lb_l, ub_l, lb_r, ub_r): + print('[' + ', '.join(f'{v:10.2f}' for v in b) + ']') best_loss_l = float('inf') best_loss_r = float('inf') best_position_l = None best_position_r = None - # Left side optimization - print(f"\n=== {label_str} 左側 ===") - _validate_bounds(lb_l, ub_l, f'{label_str} L') - position_l, loss_l = pso(objective_fn, lb_l, ub_l, - # ieqcons=[constraint_y], - swarmsize=swarm_size, - omega = omega, - maxiter=max_iter, debug=debug) + # L/R 是兩次獨立 PSO（bounds 不同、目標函數共用），互不依賴： + # side='both' 維持原行為（同一次呼叫先 L 後 R）； + # side='L'/'R' 只跑該側，讓兩側可排到不同 GPU worker。 + # 單側模式的合併輸出（3D 圖 + CSV）由「較晚完成」的一側在 + # 下方「單側收尾」段觸發。 + def _side_bounds(s_i): + if s_i == 'L': + return lb_l, ub_l, azimuth_bounds_l, x_bounds_left + return lb_r, ub_r, azimuth_bounds_r, x_bounds_right - z, x, azimuth, altitude, diameter, length = position_l - az_pso, x_pso, L_pso = azimuth, x, length - y = y_indices[round(z), round(x)] + def _run_one_side(s_i): + lb_s, ub_s, az_bounds_s, x_bounds_s = _side_bounds(s_i) + tag_cn = '左側' if s_i == 'L' else '右側' + tag = 'LEFT' if s_i == 'L' else 'RIGHT' + print(f"\n=== {label_str} {tag_cn} ===") + _validate_bounds(lb_s, ub_s, f'{label_str} {s_i}') + position, loss = pso(objective_fn, lb_s, ub_s, + # ieqcons=[constraint_y], + swarmsize=swarm_size, + omega=omega, + maxiter=max_iter, debug=debug) - z, x, azimuth, altitude, diameter, length, adopted_l, ref_l = refine_lateral_longer( - z, x, azimuth, altitude, diameter, length, - "L", azimuth_bounds_l, x_bounds_left, y_indices, - image_shape, spacing, device, grid, - cortical_tensor, spine_tensor, - ) - y = y_indices[round(z), round(x)] - if adopted_l: - loss_l = ref_l['loss'] - print(f"[LEFT] lateral-refine: az {az_pso:.2f} -> {azimuth:.2f}, x {x_pso:.2f} -> {x:.2f}, " - f"L {L_pso:.2f} -> {length:.2f}, in-bone {ref_l['in_bone']*100:.1f}%") - else: - print("[LEFT] lateral-refine: no improvement") + # 如果需要 retry（loss > 0）：重跑 PSO 取較好者 + # （原 L/R 各一份註解版 retry，此處合併為一式） + # max_retries = 0 + # retries = 0 + # while loss > 0 and retries < max_retries: + # position, loss = pso(objective_fn, lb_s, ub_s, + # swarmsize=swarm_size, maxiter=max_iter) + # retries += 1 - position_l = z, y, x, azimuth, altitude, diameter, length + z, x, azimuth, altitude, diameter, length = position + az_pso, x_pso, L_pso = azimuth, x, length + y = y_indices[round(z), round(x)] - overlap_l, diameter_l, length_l = eval_overlap_from_position( - position_l, "L", optimize_size, spine_tensor, image_shape, spacing - ) - print(f"[LEFT] overlap: {overlap_l*100:.1f}%") + z, x, azimuth, altitude, diameter, length, adopted, ref = refine_lateral_longer( + z, x, azimuth, altitude, diameter, length, + s_i, az_bounds_s, x_bounds_s, y_indices, + image_shape, spacing, device, grid, + cortical_tensor, spine_tensor, + vbody_tensor, + ) + y = y_indices[round(z), round(x)] + if adopted: + loss = ref['loss'] + print(f"[{tag}] lateral-refine: az {az_pso:.2f} -> {azimuth:.2f}, x {x_pso:.2f} -> {x:.2f}, " + f"L {L_pso:.2f} -> {length:.2f}, in-bone {ref['in_bone']*100:.1f}%") + else: + print(f"[{tag}] lateral-refine: no improvement") - if optimize_size: - print(f"[LEFT] Position: {position_l[:5]}") - print(f"[LEFT] Diameter: {diameter_l} mm (raw: {position_l[5]:.2f})") - print(f"[LEFT] Length: {length_l} mm (raw: {position_l[6]:.2f})") - print(f"[LEFT] Loss: {loss_l}\n") - best_position_l = list(position_l[:5]) + [diameter_l, length_l] - else: - print(f"[LEFT] Position: {position_l}") - best_position_l = position_l + position = z, y, x, azimuth, altitude, diameter, length - best_loss_l = loss_l - best_overlap_l = overlap_l # 新增 - - # max_retries = 0 - # retries = 0 + overlap_s, d_snap, L_snap = eval_overlap_from_position( + position, s_i, optimize_size, spine_tensor, image_shape, spacing + ) + print(f"[{tag}] overlap: {overlap_s*100:.1f}%") - # 左側 retry：loss 要 <=0 且 overlap >= 0.5 才算過關 - # while (best_loss_l > 0 or best_overlap_l < OVERLAP_THRESH) and retries < max_retries: - # position_l, loss_l = pso(objective_fn, lb_l, ub_l, swarmsize=swarm_size, maxiter=max_iter) - # overlap_l, diameter_l, length_l = eval_overlap_from_position( - # position_l, "L", optimize_size, spine_tensor, image_shape, spacing - # ) + if optimize_size: + print(f"[{tag}] Position: {position[:5]}") + print(f"[{tag}] Diameter: {d_snap} mm (raw: {position[5]:.2f})") + print(f"[{tag}] Length: {L_snap} mm (raw: {position[6]:.2f})") + print(f"[{tag}] Loss: {loss}\n") + best_pos = list(position[:5]) + [d_snap, L_snap] + else: + print(f"[{tag}] Position: {position}") + best_pos = position + return best_pos, loss, overlap_s - # 只要找到更好的 loss（或你想用 loss+overlap 綜合排序也行）就更新 best - # 安全版本：優先選「合格解」；沒有合格解時才用 loss 最小的當備案 - # candidate_pos = (list(position_l[:5]) + [diameter_l, length_l]) if optimize_size else position_l + sides = ('L', 'R') if side == 'both' else (side,) + side_time = {} + for s_i in sides: + t0 = time.time() + pos_s, loss_s, _ = _run_one_side(s_i) + side_time[s_i] = time.time() - t0 + if s_i == 'L': + best_position_l, best_loss_l = pos_s, loss_s + else: + best_position_r, best_loss_r = pos_s, loss_s - # candidate_ok = is_solution_ok(loss_l, overlap_l, OVERLAP_THRESH) - # best_ok = is_solution_ok(best_loss_l, best_overlap_l, OVERLAP_THRESH) - - # if candidate_ok and (not best_ok or loss_l < best_loss_l): - # best_position_l = candidate_pos - # best_loss_l = loss_l - # best_overlap_l = overlap_l - # print(f"[LEFT][retry {retries+1}] ✅ ok | loss={loss_l:.4f}, overlap={overlap_l*100:.1f}%") - # elif (not best_ok) and (loss_l < best_loss_l): - # best 還不合格時，先用更小 loss 的當暫存（至少越來越好） - # best_position_l = candidate_pos - # best_loss_l = loss_l - # best_overlap_l = overlap_l - # print(f"[LEFT][retry {retries+1}] ⚠️ not ok | loss improved={loss_l:.4f}, overlap={overlap_l*100:.1f}%") - # else: - # print(f"[LEFT][retry {retries+1}] ❌ no improve | loss={loss_l:.4f}, overlap={overlap_l*100:.1f}%") - - # retries += 1 - - # Right side optimization - print(f"\n=== {label_str} 右側 ===") - _validate_bounds(lb_r, ub_r, f'{label_str} R') - position_r, loss_r = pso(objective_fn, lb_r, ub_r, - # ieqcons=[constraint_y], - swarmsize=swarm_size, - omega = omega, - maxiter=max_iter, debug=debug) - - z, x, azimuth, altitude, diameter, length = position_r - az_pso, x_pso, L_pso = azimuth, x, length - y = y_indices[round(z), round(x)] - - z, x, azimuth, altitude, diameter, length, adopted_r, ref_r = refine_lateral_longer( - z, x, azimuth, altitude, diameter, length, - "R", azimuth_bounds_r, x_bounds_right, y_indices, - image_shape, spacing, device, grid, - cortical_tensor, spine_tensor, - ) - y = y_indices[round(z), round(x)] - if adopted_r: - loss_r = ref_r['loss'] - print(f"[RIGHT] lateral-refine: az {az_pso:.2f} -> {azimuth:.2f}, x {x_pso:.2f} -> {x:.2f}, " - f"L {L_pso:.2f} -> {length:.2f}, in-bone {ref_r['in_bone']*100:.1f}%") - else: - print("[RIGHT] lateral-refine: no improvement") - - position_r = z, y, x, azimuth, altitude, diameter, length - - overlap_r, diameter_r, length_r = eval_overlap_from_position( - position_r, "R", optimize_size, spine_tensor, image_shape, spacing - ) - print(f"[RIGHT] overlap: {overlap_r*100:.1f}%") - - if optimize_size: - # diameter_r, length_r = snap_to_discrete_values(position_r[5], position_r[6]) - diameter_r, length_r = snap_to_discrete_values_xfr(position_r[5], position_r[6]) - print(f"[RIGHT] Position: {position_r[:5]}") - print(f"[RIGHT] Diameter: {diameter_r} mm (raw: {position_r[5]:.2f})") - print(f"[RIGHT] Length: {length_r} mm (raw: {position_r[6]:.2f})") - print(f"[RIGHT] Loss: {loss_r}\n") - - best_position_r = list(position_r[:5]) + [diameter_r, length_r] - else: - print(f"[RIGHT] Position: {position_r}") - print(f"[RIGHT] Loss: {loss_r}\n") - best_position_r = position_r - - best_loss_r = loss_r - best_overlap_r = overlap_r - - # 如果需要 retry（loss > 0） - # max_retries = 10 - # retries = 0 - - # while (best_loss_r > 0 or best_overlap_r < OVERLAP_THRESH) and retries < max_retries: - # position_r, loss_r = pso(objective_fn, lb_r, ub_r, swarmsize=swarm_size, maxiter=max_iter) - # overlap_r, diameter_r, length_r = eval_overlap_from_position( - # position_r, "R", optimize_size, spine_tensor, image_shape, spacing - # ) - - # 只要找到更好的 loss（或你想用 loss+overlap 綜合排序也行）就更新 best - # 這裡給你一個更安全的版本：優先選「合格解」；沒有合格解時才用 loss 最小的當備案 - # candidate_pos = (list(position_r[:5]) + [diameter_r, length_r]) if optimize_size else position_r - - # candidate_ok = is_solution_ok(loss_r, overlap_r, OVERLAP_THRESH) - # best_ok = is_solution_ok(best_loss_r, best_overlap_r, OVERLAP_THRESH) - - # if candidate_ok and (not best_ok or loss_r < best_loss_r): - # best_position_r = candidate_pos - # best_loss_r = loss_r - # best_overlap_r = overlap_r - # print(f"[RIGHT][retry {retries+1}] ✅ ok | loss={loss_r:.4f}, overlap={overlap_r*100:.1f}%") - # elif (not best_ok) and (loss_r < best_loss_r): - # best 還不合格時，先用更小 loss 的當暫存（至少越來越好） - # best_position_r = candidate_pos - # best_loss_r = loss_r - # best_overlap_r = overlap_r - # print(f"[RIGHT][retry {retries+1}] ⚠️ not ok | loss improved={loss_r:.4f}, overlap={overlap_r*100:.1f}%") - # else: - # print(f"[RIGHT][retry {retries+1}] ❌ no improve | loss={loss_r:.4f}, overlap={overlap_r*100:.1f}%") - - # retries += 1 - end_time = time.time() total_time = end_time - start_time # 提取最終的 diameter 和 length if optimize_size: - final_diameter_l = best_position_l[5] - final_length_l = best_position_l[6] - final_diameter_r = best_position_r[5] - final_length_r = best_position_r[6] - - print(f"\n=== {label_str} 最終結果 ===") - print(f"Left - Diameter: {final_diameter_l} mm, Length: {final_length_l} mm") - print(f"Right - Diameter: {final_diameter_r} mm, Length: {final_length_r} mm") + final_diameter_l = best_position_l[5] if best_position_l is not None else float('nan') + final_length_l = best_position_l[6] if best_position_l is not None else float('nan') + final_diameter_r = best_position_r[5] if best_position_r is not None else float('nan') + final_length_r = best_position_r[6] if best_position_r is not None else float('nan') else: final_diameter_l = diameter final_length_l = length final_diameter_r = diameter final_length_r = length - res_plt_2_torch( - spine_tensor, - cortical_tensor, - image_shape, - image2_path, - folder, - label_str, - final_diameter_l, - final_length_l, - final_diameter_r, - final_length_r, - best_position_l, - best_position_r, - swarm_size, - max_iter, - total_time, - spacing, - CBT, - device, - grid) - + def _plot_combined(d_l, l_l, d_r, l_r, pos_l, pos_r, total_t): + res_plt_2_torch( + spine_tensor, + cortical_tensor, + image_shape, + image2_path, + folder, + label_str, + d_l, + l_l, + d_r, + l_r, + pos_l, + pos_r, + swarm_size, + max_iter, + total_t, + spacing, + CBT, + device, + grid, + ) + + if side == 'both': + print(f"\n=== {label_str} 最終結果 ===") + print(f"Left - Diameter: {final_diameter_l} mm, Length: {final_length_l} mm") + print(f"Right - Diameter: {final_diameter_r} mm, Length: {final_length_r} mm") + _plot_combined(final_diameter_l, final_length_l, + final_diameter_r, final_length_r, + best_position_l, best_position_r, total_time) + return best_position_l, best_loss_l, best_position_r, best_loss_r, total_time + + # ---- 單側收尾：寫該側結果檔；兩側都完成時由較晚的一側觸發合併輸出 ---- + # 兩側可能在不同 GPU worker：各自把結果寫 _.json + # （先寫 tmp 再 os.replace，對端讀到的一定是完整檔）。先完成者看不到 + # 對端檔就跳過；後完成者看到兩側齊了、搶到 plot lock（O_EXCL， + # 確保合併輸出只跑一次）才載入對端結果跑 res_plt_2_torch + # （3D 圖 + CSV 兩行，與 'both' 模式相同）。json / lock 保留供事後 + # 檢查；若該側流程死在 plotting 中段，該 (volume, level) 重跑即可 + # （run_id 是新的一次，不會互相干擾）。 + other = 'R' if side == 'L' else 'L' + own_pos = best_position_l if side == 'L' else best_position_r + own_d = final_diameter_l if side == 'L' else final_diameter_r + own_L = final_length_l if side == 'L' else final_length_r + own_cn = 'Left' if side == 'L' else 'Right' + side_cn = '左' if side == 'L' else '右' + print(f"\n=== {label_str} 最終結果（{side_cn}側） ===") + print(f"{own_cn} - Diameter: {own_d} mm, Length: {own_L} mm") + + missing_pair = [n for n, v in (('level', level), ('patient_id', patient_id), + ('side_dir', side_dir), ('run_id', run_id)) if not v] + if missing_pair: + print(f"[SIDE] 合併輸出跳過（未提供 {', '.join(missing_pair)}）") + else: + patient_dir = os.path.join(side_dir, run_id, patient_id) + own_path = os.path.join(patient_dir, f'{level}_{side}.json') + other_path = os.path.join(patient_dir, f'{level}_{other}.json') + _makedirs_retry(patient_dir) + own = {'position': [float(v) for v in own_pos], + 'diameter': float(own_d), + 'length': float(own_L), + 'time': float(side_time[side])} + tmp_path = f'{own_path}.{os.getpid()}.tmp' + with open(tmp_path, 'w') as f: + json.dump(own, f) + os.replace(tmp_path, own_path) + + if not os.path.isfile(other_path): + print(f"[SIDE] {other} 側尚未完成，合併輸出（圖 + CSV）待該側完成時觸發") + else: + lock_path = os.path.join(patient_dir, f'{level}.plot.lock') + try: + fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + os.close(fd) + except FileExistsError: + print(f'[SIDE] 合併輸出已由 {other} 側觸發，跳過') + else: + with open(other_path) as f: + other_res = json.load(f) + pos_other = [float(v) for v in other_res['position']] + if side == 'L': + d_l, l_l, pos_l = own_d, own_L, list(own_pos) + d_r, l_r, pos_r = other_res['diameter'], other_res['length'], pos_other + else: + d_r, l_r, pos_r = own_d, own_L, list(own_pos) + d_l, l_l, pos_l = other_res['diameter'], other_res['length'], pos_other + print(f"[SIDE] {side} + {other} 兩側完成 -> 合併輸出（圖 + CSV）") + # total_time 用兩側各自耗時相加（各含一次影像載入/平面計算， + # 比原同流程 wall time 略大，僅影響 CSV 的時間欄） + _plot_combined(d_l, l_l, d_r, l_r, pos_l, pos_r, + side_time[side] + other_res['time']) + return best_position_l, best_loss_l, best_position_r, best_loss_r, total_time diff --git a/core/scoring.py b/core/scoring.py index a4679fd..06bd593 100644 --- a/core/scoring.py +++ b/core/scoring.py @@ -10,13 +10,23 @@ def cl_score_torch_xfr( intersections: int, diameter: float = None, length: float = None, - cylinder_tip_torch: torch.Tensor = None # 新增：尖端 mask + cylinder_tip_torch: torch.Tensor = None, # 新增：尖端 mask + vbody_tensor: torch.Tensor = None # VBODY（椎體）mask (z,y,x) 0/1，None = 不計 VBODY 獎勵 ) -> float: """ 漸進式評分：優先確保找到骨頭，再改善細節 """ cyl_total = cylinder_torch.sum().item() overlap = ((cortical_tensor == 1) & (cylinder_torch == 1)).sum().item() + # VBODY 獎勵：螺絲落在 (cortical + VBODY) 內的 voxel 每個 100 分 + # （100 分項由純 cortical 擴展到 cortical∪VBODY；cortical voxel 分數不變）， + # 其中落在 VBODY 的 voxel 每個再加 10 分 + if vbody_tensor is not None: + in_vbody = ((vbody_tensor == 1) & (cylinder_torch == 1)).sum().item() + in_corti_vb = (((cortical_tensor == 1) | (vbody_tensor == 1)) & (cylinder_torch == 1)).sum().item() + else: + in_vbody = 0 + in_corti_vb = overlap null_vox = ((cortical_tensor == 0) & (cylinder_torch == 1)).sum().item() null_vox2 = ((spine_tensor == 1) & (cylinder_o_torch == 1)).sum().item() @@ -38,10 +48,12 @@ def cl_score_torch_xfr( score += 20 * in_bone # 10 實在太低 score += 100 * overlap + score += 100 * in_corti_vb # (cortical + VBODY) 每 voxel 100 分 + score += 50 * in_vbody # VBODY 每 voxel 再加 10 分 # score -= 2000 * max(0, not_in_bone-10) # score -= 1000 * max(0, null_vox2-10) score -= 1000 * not_in_bone - score -= 1000 * null_vox2 + score -= 2000 * null_vox2 return float(-score) diff --git a/imaging/orientation.py b/imaging/orientation.py index eae41df..11ad440 100644 --- a/imaging/orientation.py +++ b/imaging/orientation.py @@ -1,8 +1,15 @@ import numpy as np import SimpleITK as sitk -from scipy.ndimage import center_of_mass, rotate +from scipy.ndimage import center_of_mass, rotate, distance_transform_edt, gaussian_filter import matplotlib.pyplot as plt +# 鏡稱面法線的左右（LR，=x 軸）分量大下。正常左右鏡稱面法線幾乎沿 x； +# 若 |a| < 0.7，代表搜尋掉進前後（coronal）面局部極大（椎體是厚實塊狀， +# 易自鏡射拿高分，例：0019 L3/L4），以該平面定義的「中線帶」變成斜切 +# 板層，其後側中線缺如判定不可信 → diagnose 回 bad_mirror_plane、 +# 不判 no_spinous（避免假 [NO-SP] 誤改椎體切分）。 +MIRROR_MIN_LR = 0.7 + def _best_vertical_split(mask2d): """ binary 2D 陣列（最後一軸 = x）：找讓 mask 與其鏡射重疊最大的垂直線 x = t。 @@ -158,8 +165,83 @@ def best_symmetry_plane(mask_zyx, phi_max=45.0, subsample=7, 'v': (float(v[0]), float(v[1]), float(v[2])), } +def diagnose_spinous_process(mask_zyx, sym, band_frac=0.06, min_band=6.0, + deficit_min=4.0, rear_margin=3.0, rear3_max=20, + narrow_frac=0.4): + """ + 診斷該椎體的棘突（中線後側構造）是否缺如 + （先前手術如 laminectomy / 棘突切除所造成）。 + + 棘突完整時，其尖端是全椎體最後側的骨構造且位於中線，因此： + deficit = 全骨最後側 AP - 中線帶最後側 AP ~= 0， + 且中線窄帶 |s| <= narrow_frac*w 在最後側骨 rear_margin 內有充足骨 voxel。 + 切除後，中線帶的後側止於殘留後側要素，最後側骨偏到側方的 + 殘餘結構（椎板/關節突）→ deficit 大 + 後側中線空洞。 + 判定：no_spinous = deficit >= deficit_min 且 rear3 <= rear3_max。 + + 回傳 dict： + n_bone : 骨 voxel 總數 + band_w : 中線帶半寬 w（voxel） + deficit : 後側中線缺如量（voxel，0.5 mm/vox） + rear3 : 最後側 3 voxel 內、中線窄帶的骨 voxel 數 + rear6 : 最後側 6 voxel 內、中線窄帶的骨 voxel 數 + top_off : axial 投影（y,x）最後側一行相對鏡稱中線的偏移（voxel） + no_spinous : True = 後側中線缺如（棘突已切除） + reason : 'ok' / 'no_spinous' / 'too_few_bone' / 'too_few_midline' + / 'bad_mirror_plane' + """ + m = np.asarray(mask_zyx) > 0 + zz, yy, xx = np.nonzero(m) + out = {'n_bone': int(zz.size), 'band_w': None, 'deficit': 0.0, + 'rear3': None, 'rear6': None, 'top_off': None, + 'no_spinous': False, 'reason': 'ok'} + if zz.size < 100: + out['reason'] = 'too_few_bone' + return out + a, b, c, d = sym['plane'] + if abs(a) < MIRROR_MIN_LR: + # 鏡稱面非左右為主（前後 coronal 局部極大）→ 缺如判定不可信， + # 按 ok 回傳（寧可漏報，不可假切除） + out['reason'] = 'bad_mirror_plane' + return out + X = xx.astype(np.float64) + Y = yy.astype(np.float64) + Z = zz.astype(np.float64) + s = X * a + Y * b + Z * c - d + u_ap = _ap_axis(sym) # +AP = 後側（y 小側） + ap = X * u_ap[0] + Y * u_ap[1] + Z * u_ap[2] + w = max(float(min_band), float(band_frac) * float(s.max() - s.min())) + mid = np.abs(s) <= w + aps = ap[mid] + out['band_w'] = float(w) + if aps.size < 50: + out['reason'] = 'too_few_midline' + return out + narrow = np.abs(s) <= w * float(narrow_frac) + ap_max_all = float(ap.max()) + deficit = ap_max_all - float(aps.max()) + rear3 = int(((ap >= ap_max_all - rear_margin) & narrow).sum()) + rear6 = int(((ap >= ap_max_all - 2.0 * rear_margin) & narrow).sum()) + out['deficit'] = float(deficit) + out['rear3'] = rear3 + out['rear6'] = rear6 + proj = m.max(axis=0) # (y, x) + ys_p, xs_p = np.where(proj > 0) + # 只在鏡稱面 x 分量主導時才算 top_off（最後側一行對中線的 x 偏移）； + # 脊椎在面內大角度旋轉時 |a| 小，(d-b*y-c*z)/a 會放大成無意義的巨值。 + if ys_p.size and abs(a) >= 0.5 * max(abs(b), abs(c)): + y_min = int(ys_p.min()) + zc = (m.shape[0] - 1) / 2.0 + x_mid = (d - b * y_min - c * zc) / a + out['top_off'] = float(abs(xs_p[ys_p == y_min].mean() - x_mid)) + if deficit >= float(deficit_min) and rear3 <= int(rear3_max): + out['no_spinous'] = True + out['reason'] = 'no_spinous' + return out + + def segment_spinous_process(mask_zyx, sym, band_frac=0.06, min_band=6.0, - min_mass_frac=0.05): + min_mass_frac=0.05, expand_cap=5): """ 以 best_symmetry_plane 的結果 sym 從 3D bone mask (z, y, x) 切出棘突。 棘突是中線後側構造，利用鏡稱面 a·x+b·y+c·z=d 定義： @@ -170,8 +252,14 @@ def segment_spinous_process(mask_zyx, sym, band_frac=0.06, min_band=6.0, 3) 中線帶的 AP 分佈呈兩大叢（椎體在前、椎弓/棘突在後）， 以兩叢間的 AP 谷底為界，AP >= 谷底 的中線帶 voxel = 棘突（含中線椎弓）； 無明顯谷底（如骨橋）fallback 取中線帶後側 15%。 + 4) 中線帶側緣補回（_expand_spinous_runs，expand_cap）：帶由鏡稱面定義， + 棘突楔若略偏中線，側緣薄條會留在帶外成為 other bone；每條 (y, z) + 線把棘突 run 向左右各補至多 expand_cap 個 bone voxel。 + 先做後側中線缺如判定（diagnose_spinous_process，laminectomy/棘突切除）； + 缺如、或谷底後側叢只剩 <5% 殘片時 mode='no_spinous'、sp_mask=None， + 由呼叫端（入口面 / 椎體）改用對應策略。 回傳 (sp_mask (z,y,x) bool, ap_thresh, info dict)； - 資料過少時 sp_mask = None（info['mode'] 說明原因）。 + 資料過少或棘突缺如時 sp_mask = None（info['mode'] 說明原因）。 """ m = np.asarray(mask_zyx) > 0 zz, yy, xx = np.nonzero(m) @@ -197,6 +285,19 @@ def segment_spinous_process(mask_zyx, sym, band_frac=0.06, min_band=6.0, if aps.size < 50: info['mode'] = 'too_few_midline' return None, None, info + + # 後側中線缺如（laminectomy / 棘突切除）：最後側骨偏到側方殘餘、 + # 中線帶後側無質量，谷底搜尋沒有意義 → no_spinous， + # 呼叫端（y_indices 入口面 / segment_vertebral_body）改用對應策略。 + diag = diagnose_spinous_process(mask_zyx, sym, band_frac=band_frac, min_band=min_band) + info['deficit'] = diag['deficit'] + info['rear3'] = diag['rear3'] + info['rear6'] = diag['rear6'] + if diag['no_spinous']: + info['mode'] = 'no_spinous' + info['top_off'] = diag['top_off'] + return None, None, info + lo = int(np.floor(aps.min())) hi = int(np.ceil(aps.max())) th = None @@ -216,6 +317,14 @@ def segment_spinous_process(mask_zyx, sym, band_frac=0.06, min_band=6.0, if score > best_score: best_score, best_i = score, i if best_i is not None: + post_frac = (total - csum[best_i + 1]) / total + if post_frac < 0.05 and abs(a) >= MIRROR_MIN_LR: + # 谷底後側叢只剩小殘片（<5% 中線帶質量）＝棘突幾乎全除 + # （部分切除殘餘）：當作缺如，椎體切分不採用此閾值。 + # 僅在鏡稱面左右為主時才算數（斜板層時 post_frac 不可信） + info['mode'] = 'no_spinous' + info['post_frac'] = float(post_frac) + return None, None, info th = float(0.5 * (edges[best_i] + edges[best_i + 1])) mode = 'gap' if th is None: @@ -223,19 +332,247 @@ def segment_spinous_process(mask_zyx, sym, band_frac=0.06, min_band=6.0, sp_mask = np.zeros(m.shape, dtype=bool) sel = mid & (ap >= th) sp_mask[zz[sel], yy[sel], xx[sel]] = True - info.update(n_sp=int(sel.sum()), ap_thresh=th, mode=mode) + sp_mask = _expand_spinous_runs(sp_mask, m, cap=expand_cap) + info.update(n_sp=int(sp_mask.sum()), ap_thresh=th, mode=mode) return sp_mask, th, info -def best_upper_endplate_plane(mask_zyx, angle_max=45.0, thresh=3.0, + +def _expand_spinous_runs(sp_mask, bone, cap=5): + """對每條 (y, z) 水平線（x = 左右方向），把既有棘突 run 向左右兩側各延伸 + 至多 cap 個 bone voxel：補回落在中線帶 |s|<=w 之外的棘突楔側緣薄條 + （帶由鏡稱面定義，楔略偏中線時側緣會被漏成 other bone）。 + 遇到非 bone、既有棘突 voxel 或達 cap 即停；只延伸既有 run，不新建 + run（該線無棘突時絕不標記）。椎弓厚處最多少數 cap 個「棘突基部」 + 相鄰 bone 一併補入，屬棘突連續結構。""" + out = sp_mask.copy() + if not out.any(): + return out + nz, ny, nx = out.shape + for y in range(ny): + sp_y = out[:, y, :] + if not sp_y.any(): + continue + bone_y = bone[:, y, :] + for z in np.flatnonzero(sp_y.any(axis=1)): + idx = np.flatnonzero(sp_y[z]) + runs = np.split(idx, np.flatnonzero(np.diff(idx) > 1) + 1) + for r in runs: + lo, hi = int(r[0]), int(r[-1]) + x = lo - 1 + added = 0 + while x >= 0 and added < cap and bone_y[z, x] and not out[z, y, x]: + out[z, y, x] = True + x -= 1 + added += 1 + x = hi + 1 + added = 0 + while x < nx and added < cap and bone_y[z, x] and not out[z, y, x]: + out[z, y, x] = True + x += 1 + added += 1 + return out + +def smooth_mask_sdf(mask_zyx, sigma=1.5): + """對 3D boolean mask 做有號距離場（SDF）高斯平滑： + 把表面沿法線做近似平均曲率流（mean curvature flow），圓化凸出的 + 鋸齒/尖刺、填進凹陷的缺口，但不會均勻收縮整體形状（與直接對 mask + 做高斯Blur再 thresh 不同——後者會把薄結構一起磨細、體積明顯縮小）。 + 這是消除「三线性旋轉 > 0.5 之後仍剩餘的 voxel 級鋸齒」最有效的方法。 + + 參數： + mask_zyx : (z, y, x) boolean / 0-1 陣列 + sigma : 平滑強度（voxel）。1.5 ≈ 表面積 -4%、體積 -0.6%； + 越大越圓，但薄結構（棘突尖）會被磨損越多。 + 回傳同 shape 的 boolean mask。 + """ + m = np.asarray(mask_zyx) > 0 + if not m.any() or m.all(): + return m.copy() + d_in = distance_transform_edt(m) + d_out = distance_transform_edt(~m) + sdf = d_in - d_out + return gaussian_filter(sdf, sigma=float(sigma)) > 0.0 + +def shape_based_smooth(vol01, sigma=1.5): + """Shape-based interpolation（SBI）平滑：以 0..1 體積的 0.5 等值面為 + 物件形狀（shape），計算有號距離場（SDF，內正外負），在距離域做高斯 + 插值（interpolation），再於 0 重新閾值，得到圓滑的 smooth mask。 + 與 smooth_mask_sdf 的差異：邊界層（0 < v < 1）內三线性場值滿足 + v ≈ 0.5 + d（d = 以 voxel 為單位的有號距離），故直接以 v - 0.5 替換 + EDT 量值，得到次體素精確的 SDF（層外仍由 distance_transform_edt 給出 + 精確距離）；表面梯度連續，不再是被 EDT 量化成 ±1 的硬階梯，平滑後 + 的邊界更貼合 soft boundary 的真實位置。 + 參數： + vol01 : (z, y, x) 0..1 浮點體積（例如旋轉後的 _binary_linear） + sigma : 高斯平滑強度（voxel），意義與 smooth_mask_sdf 相同。 + 回傳：同 shape 的 boolean mask。 + """ + v = np.clip(np.asarray(vol01, dtype=np.float32), 0.0, 1.0) + b = v > 0.5 + if not b.any() or b.all(): + return b.copy() + d_in = distance_transform_edt(b) + d_out = distance_transform_edt(~b) + sdf = d_in - d_out + band = (v > 0) & (v < 1) + sdf[band] = v[band] - 0.5 + return gaussian_filter(sdf, sigma=float(sigma)) > 0.0 + +def _ap_axis(sym): + """鏡稱面內的前後（AP）軸：(u,v) 中 |y| 分量大者，正規化成 +AP = 後側（y 小側） + （本資料系 y 往前遞增，與 segment_spinous_process 同慣例）。""" + u = np.array(sym['u']) + v = np.array(sym['v']) + u_ap = u if abs(u[1]) >= abs(v[1]) else v + if u_ap[1] > 0: + u_ap = -u_ap + return u_ap + + +def _full_ap_valley(aps, min_side_frac=0.15, max_ratio=0.85, smooth=3): + """終板下骨體 AP 分佈的平滑谷底：最深相對谷底（sm[i] 對鄰近峰的最小比值）， + 要求兩側各有 >= min_side_frac 的質量（拒絕對小尾巴的偽谷底）。 + 回傳 (th 或 None, ratio 或 None)。""" + lo = int(np.floor(aps.min())) + hi = int(np.ceil(aps.max())) + if hi - lo < 10: + return None, None + hist, edges = np.histogram(aps, bins=range(lo, hi + 1)) + sm = np.convolve(hist, np.ones(smooth) / float(smooth), mode='same') + csum = np.concatenate([[0], np.cumsum(hist)]) + total = csum[-1] + best_i, best_ratio = None, 1.0 + for i in range(1, len(hist) - 1): + if not (sm[i] <= sm[i - 1] and sm[i] <= sm[i + 1]): + continue + if csum[i] < min_side_frac * total or (total - csum[i + 1]) < min_side_frac * total: + continue + flank = min(float(sm[:i].max()), float(sm[i + 1:].max())) + if flank <= 0: + continue + ratio = float(sm[i]) / flank + if ratio <= max_ratio and ratio < best_ratio: + best_ratio, best_i = ratio, i + if best_i is None: + return None, None + return float(0.5 * (edges[best_i] + edges[best_i + 1])), best_ratio + + +def _posterior_min_threshold(aps, rear_frac=0.40, min_side_frac=0.08, smooth=3): + """AP 分佈後側 `rear_frac` 區間內、平滑直方圖的最小值位置（該處後側質量 + 比例需 >= min_side_frac）：棘突缺如椎體中 _full_ap_valley 失敗時的 + 「椎體後側末端 = 體/弓最薄處」備援。 + 回傳 th 或 None。""" + lo = int(np.floor(aps.min())) + hi = int(np.ceil(aps.max())) + if hi - lo < 10: + return None + hist, edges = np.histogram(aps, bins=range(lo, hi + 1)) + sm = np.convolve(hist, np.ones(smooth) / float(smooth), mode='same') + csum = np.concatenate([[0], np.cumsum(hist)]) + total = csum[-1] + i0 = int(len(hist) * (1.0 - rear_frac)) + best_i, best_val = None, np.inf + for i in range(i0, len(hist)): + if (total - csum[i + 1]) < min_side_frac * total: + continue + if sm[i] < best_val: + best_val, best_i = float(sm[i]), i + if best_i is None: + return None + return float(0.5 * (edges[best_i] + edges[best_i + 1])) + + +def segment_vertebral_body(mask_zyx, sym, endplate, ap_thresh, sp_mode, margin=2.0): + """ + 以兩平面從 3D bone mask (z, y, x) 切出椎體（前側中央主體塊）： + 1) best_upper_endplate_plane 的上終板面（法線朝上 a·x+b·y+c·z=d）： + 只保留終板下側（身體側，e·p-d <= margin）的 bone voxel， + 排除跨在終板上方的後側構造（椎板/棘突）。 + 2) best_symmetry_plane 的鏡稱面：提供平面內 AP 軸（_ap_axis， + +AP = 後側，與 segment_spinous_process 同慣例）；AP 切點取三層優先： + a) sp_mode == 'gap'：中線帶（椎管）的體/弓谷底 ap_thresh + （segment_spinous_process 回傳值，最可靠）； + b) sp_mode == 'no_spinous'（棘突已切除）：中線帶沒有空的體/弓谷底 + 可用；椎體後側末端 = 終板下整體 AP 分佈的局部最小（體/弓最薄處）。 + 切除後殘留後側要素會部分填滿椎管、谷底比正常椎體淺， + 因此放寬 _full_ap_valley 條件（min_side_frac 0.15->0.10、 + max_ratio 0.85->0.95）；再失敗則取後側 40% 區間的平滑谷底 + （_posterior_min_threshold）。 + c) 否則：終板下整體 AP 分佈的平滑谷底（_full_ap_valley， + 處理中線骨橋等中線搜尋 fallback 的情形）； + d) 最後 fallback：AP 分佈 55 百分位（可能切進椎體內，會打 WARNING）。 + 椎體 = AP < 切點（切點之前側）且終板下側的 bone。 + 回傳 (vb_mask (z,y,x) bool, ap_thresh, info dict)； + 資料不足或上終板面缺位時 vb_mask = None（info['mode'] 說明原因）。 + """ + m = np.asarray(mask_zyx) > 0 + zz, yy, xx = np.nonzero(m) + info = {'n_bone': int(zz.size), 'n_vb': 0, 'ap_thresh': None, 'mode': 'empty'} + if zz.size < 50: + return None, None, info + if endplate is None: + info['mode'] = 'no_endplate' + return None, None, info + a, b, c, d = endplate['plane'] + X = xx.astype(np.float64) + Y = yy.astype(np.float64) + Z = zz.astype(np.float64) + below = (X * a + Y * b + Z * c - d) <= float(margin) + u_ap = _ap_axis(sym) + ap = X * u_ap[0] + Y * u_ap[1] + Z * u_ap[2] + aps = ap[below] + if aps.size < 50: + info['mode'] = 'too_few_below' + return None, None, info + th, mode = None, 'quantile' + if sp_mode == 'gap' and ap_thresh is not None: + th, mode = float(ap_thresh), 'midline_gap' + elif sp_mode == 'no_spinous': + # 棘突已切除：用放寬條件的整體谷底（見 docstring b） + th, _ = _full_ap_valley(aps, min_side_frac=0.10, max_ratio=0.95) + if th is not None: + mode = 'nosp_gap' + else: + th = _posterior_min_threshold(aps) + if th is not None: + mode = 'nosp_post_min' + else: + th, _ = _full_ap_valley(aps) + if th is not None: + mode = 'full_gap' + if th is None: + th = float(np.quantile(aps, 0.55)) + sel = below & (ap < th) + if not sel.any(): + info['mode'] = 'no_body_voxels' + return None, None, info + vb_mask = np.zeros(m.shape, dtype=bool) + vb_mask[zz[sel], yy[sel], xx[sel]] = True + info.update(n_vb=int(sel.sum()), ap_thresh=th, mode=mode) + return vb_mask, th, info + +def best_upper_endplate_plane(mask_zyx, angle_max=45.0, thresh=4.0, n_iter=500, seed=42): """ 3D bone mask (z, y, x) 的最佳「上終板」近似平面 a·x + b·y + c·z = d （voxel index 座標；(a,b,c) 為朝上的單位法線）。 - 1) 每個 (y, x) 欄位取最上方 bone voxel 作為頂面點（僅前側半邊， - y >= COM_y，避開後方元素，與 2D superior endplate 定義一致） - 2) RANSAC 三點擬平面：法線限制在與 +z 軸 ≤ angle_max° 內， - 計數 ±thresh voxel 內的頂面點為 inlier，取 inlier 最多者 - 3) SVD 最小二乘微調 + + 區域：只用該 level 的前側半邊（y 為 AP 方向，anterior = y 較大側）： + 切點 = bone 的 y 範圍中點（y >= (y_min+y_max)//2），避開後方 + 棘突 / 弓根 / 椎管結構，與 2D superior endplate 定義一致。 + + 兩階段穩健擬合（舊式單次 3D RANSAC 會被側旁結構 / 上層椎體侵入 + 的頂面污染擬出過陡平面，例：L1 擬到 21°）： + 1) Stage A（矢狀線）：每個 (y, x) 欄位取最上方 bone voxel 作為頂面點， + 再收斂成 per-y 最上 z，RANSAC 擬合矢狀 (y-z) 線 z = s·y + i， + 取得穩健的矢狀斜率 s； + 2) Stage B（側向線）：先把頂面點筛到矢狀線 ±max(2·thresh, 5) 帶內 + （移除離帶的高/低污染點），再 RANSAC 擬合 (z − s·y − i) = c·x + k， + 取得側向斜率 c； + 3) 平面 z = s·y + c·x + (i+k) → 單位法線 (−c,−s,1)/√(1+s²+c²)、 + d = (i+k)/√(1+s²+c²)、tilt_deg = 法線與 +z 軸夾角； + inlier_ratio = 全部前側半頂面點落在平面 ±thresh 的比例。 回傳 dict（資料不足時回傳 None）： plane : (a, b, c, d) normal / offset @@ -244,51 +581,67 @@ def best_upper_endplate_plane(mask_zyx, angle_max=45.0, thresh=3.0, n_points / n_inliers u, v : 平面內正交方向（供繪製用） """ + from sklearn.linear_model import RANSACRegressor + m = np.asarray(mask_zyx) > 0 nz, ny, nx = m.shape + if not m.any(): + return None idx = np.where(m, np.arange(nz)[:, None, None], -1) ztop = idx.max(axis=0) # (y, x) 每欄最上 z - y_split = int(round(center_of_mass(m)[1])) if m.sum() else 0 + # 前側半邊：bone y 範圍中點為切點（前側 = y >= 切點） + y_present = np.where(np.any(m, axis=(0, 2)))[0] + y_split = int((y_present[0] + y_present[-1]) // 2) # ztop 是 (y, x)：條件作用在 y 軸（axis 0） sel = (ztop >= 0) & (np.arange(ny)[:, None] >= y_split) yy, xx = np.where(sel) - if xx.size < 8: + if xx.size < 20: return None - P = np.stack((xx, yy, ztop[yy, xx]), axis=1).astype(np.float64) - n = P.shape[0] + xa = xx.astype(np.float64) + ya = yy.astype(np.float64) + za = ztop[yy, xx].astype(np.float64) + n = za.size - rng = np.random.default_rng(seed) - cos_min = np.cos(np.deg2rad(angle_max)) - best_cnt, best_nv, best_d = -1, None, 0.0 - for _ in range(n_iter): - i, j, k = rng.choice(n, 3, replace=False) - cr = np.cross(P[j] - P[i], P[k] - P[i]) - ln = np.linalg.norm(cr) - if ln < 1e-6: - continue - nv = cr / ln - if nv[2] < 0: - nv = -nv - if nv[2] < cos_min: # 法線必須朝上 - continue - d = float(nv @ P[i]) - cnt = int(np.count_nonzero(np.abs(P @ nv - d) <= thresh)) - if cnt > best_cnt: - best_cnt, best_nv, best_d = cnt, nv, d - if best_nv is None: + def _line1d(x1, y1, residual): + """RANSAC 擬線 y = k·x + b；失敗時退回普通最小二乘。""" + try: + fit = RANSACRegressor(residual_threshold=float(residual), + max_trials=int(n_iter), + random_state=int(seed)).fit(x1.reshape(-1, 1), y1) + return float(fit.estimator_.coef_[0]), float(fit.estimator_.intercept_) + except Exception: + k, b = np.polyfit(x1, y1, 1) + return float(k), float(b) + + # ---- Stage A：矢狀 (y-z) 線 → 穩健矢狀斜率 s ---- + zt = np.full(ny, -1.0) + np.maximum.at(zt, yy, za) + yv = np.where(zt >= 0)[0] + if yv.size < 10: return None - # SVD 微調 - inl = P[np.abs(P @ best_nv - best_d) <= thresh] - if inl.shape[0] < 3: + s, i = _line1d(yv.astype(np.float64), zt[yv], 5.0) + if abs(float(np.degrees(np.arctan(s)))) > angle_max: return None - mean = inl.mean(axis=0) - _, _, Vt = np.linalg.svd(inl - mean, full_matrices=False) - nv = Vt[2] - if nv[2] < 0: - nv = -nv - d = float(nv @ mean) - dist = np.abs(P @ nv - d) - inl = P[dist <= thresh] + + # ---- Stage B：固定 s，擬側向 (x) 線 (z − s·y − i) = c·x + k ---- + # 合起來 z = s·y + c·x + (i + k) → n0·p = d0，n0 = (−c, −s, 1), d0 = i + k + res_sag = za - s * ya - i + band = np.abs(res_sag) <= max(2.0 * float(thresh), 5.0) + if int(band.sum()) < 20: + band = np.ones(n, dtype=bool) + c, k = _line1d(xa[band], res_sag[band], float(thresh)) + d0 = i + k + + n0 = np.array([-c, -s, 1.0]) + norm0 = float(np.linalg.norm(n0)) + nv = n0 / norm0 + d = float(d0 / norm0) + tilt_deg = float(np.degrees(np.arccos(np.clip(nv[2], -1.0, 1.0)))) + if tilt_deg > angle_max: + return None + + dist = np.abs(-c * xa - s * ya + za - d0) / norm0 + n_in = int(np.count_nonzero(dist <= float(thresh))) u = np.cross(nv, [1.0, 0.0, 0.0]) u = u / np.linalg.norm(u) v = np.cross(nv, u) @@ -296,10 +649,10 @@ def best_upper_endplate_plane(mask_zyx, angle_max=45.0, thresh=3.0, 'plane': (float(nv[0]), float(nv[1]), float(nv[2]), d), 'normal': (float(nv[0]), float(nv[1]), float(nv[2])), 'offset': d, - 'tilt_deg': float(np.degrees(np.arccos(np.clip(nv[2], -1.0, 1.0)))), - 'inlier_ratio': float(inl.shape[0] / n), + 'tilt_deg': tilt_deg, + 'inlier_ratio': float(n_in / n), 'n_points': int(n), - 'n_inliers': int(inl.shape[0]), + 'n_inliers': n_in, 'u': (float(u[0]), float(u[1]), float(u[2])), 'v': (float(v[0]), float(v[1]), float(v[2])), } diff --git a/imaging/preprocessing.py b/imaging/preprocessing.py index 0a5e69e..2bc297d 100644 --- a/imaging/preprocessing.py +++ b/imaging/preprocessing.py @@ -20,13 +20,89 @@ def save_progress(progress, PROGRESS_FILE): with open(PROGRESS_FILE, "w") as f: json.dump(progress, f, indent=2) -def process_single_image(image_path, label_path, output_dir_base=None): +def process_single_image(image_path, label_path, output_dir_base=None, max_z_spacing=None, allowed_levels=None, min_levels=None, metadata_cache=None): + """metadata_cache：可選，需提供 get(name) -> dict|None 與 + put(name, dict)。dict 可含 spacing=[x,y,z]、labels=[label id]。 + 兩者都在 db 裡時整支跳過判定不需讀影像 / label 檔。""" - image = sitk.ReadImage(image_path) - label = sitk.ReadImage(label_path) file_name = os.path.basename(image_path) name = file_name.replace(".nii.gz", "") + # pixel spacing / labels 優先用 metadata db，避免每次 run 都讀檔 + meta = metadata_cache.get(name) if metadata_cache is not None else None + image = None + label = None + spacing = meta.get("spacing") if meta is not None else None + existing_labels = meta.get("labels") if meta is not None else None + + if spacing is None: + image = sitk.ReadImage(image_path) + spacing = [float(v) for v in image.GetSpacing()] + if metadata_cache is not None: + metadata_cache.put(name, {"spacing": spacing}) + + # z spacing 過大（低解析度掃描）跳過整支 pipeline + if max_z_spacing is not None and spacing[2] > max_z_spacing: + z_spacing = spacing[2] + print(f"z spacing {z_spacing:.2f} mm > {max_z_spacing} mm, " + f"skipping pipeline for {name}") + return { + "processed_labels": [], + "missing_labels": [], + "skipped": True, + "skip_reason": f"z spacing {z_spacing:.2f} mm > {max_z_spacing} mm" + } + + if existing_labels is None: + label = sitk.ReadImage(label_path) + # 取得現有 label + lssif = sitk.LabelShapeStatisticsImageFilter() + lssif.Execute(label) + existing_labels = [int(v) for v in lssif.GetLabels()] # 例如 [1,2,3,20,21] + if metadata_cache is not None: + metadata_cache.put(name, {"labels": existing_labels}) + else: + print(f"Metadata db hit for {name} (spacing={spacing}, labels={existing_labels})") + + print(f"Existing labels in {os.path.basename(label_path)}: {existing_labels}") + + allowed_label_list = [n for n in existing_labels + if n in LABEL_MAP and (allowed_levels is None + or LABEL_MAP[n] in allowed_levels)] + + # 沒有任何符合 allowed levels（如 lumbar）的 label：不建立輸出資料夾、 + # 不做重取樣，整檔跳過（不計入 max_images） + if not allowed_label_list: + print(f"No label matches allowed levels {allowed_levels} in {name}, " + f"skipping (no output folder)") + return { + "processed_labels": [], + "missing_labels": [], + "skipped": True, + "skip_reason": f"no label in allowed levels {allowed_levels}" + } + + # 符合 allowed levels 的 label 少於 min_levels（如 lumbar < 2 層）： + # 整支 pipeline 跳過，不建立輸出資料夾（不計入 max_images） + if min_levels is not None and len(allowed_label_list) < min_levels: + have = ", ".join(f"{n} ({LABEL_MAP[n]})" for n in allowed_label_list) + print(f"Only {len(allowed_label_list)} allowed-level label(s) [{have}] in {name} " + f"< {min_levels}, skipping pipeline (no output folder)") + return { + "processed_labels": [], + "missing_labels": [], + "skipped": True, + "skip_reason": (f"only {len(allowed_label_list)} label(s) in allowed " + f"levels < {min_levels}") + } + + # 進入 pipeline：若上面的 spacing / labels 來自 metadata db， + # 影像與 label 檔還沒讀，這裡補讀 + if image is None: + image = sitk.ReadImage(image_path) + if label is None: + label = sitk.ReadImage(label_path) + # LabelStatisticsImageFilter computes statistics (e.g., mean, minimum, maximum, median) of pixel values in an image, segmented by labels in a corresponding label image. lsif = sitk.LabelStatisticsImageFilter() lsif.Execute(image, label) @@ -34,14 +110,7 @@ def process_single_image(image_path, label_path, output_dir_base=None): # Assume to have some sitk image (itk_image) and label (itk_label) resampled_sitk_img = resample_img(image, out_spacing=[0.5, 0.5, 0.5], is_label=False) resampled_sitk_lbl = resample_img(label, out_spacing=[0.5, 0.5, 0.5], is_label=True) - - # 取得現有 label - lssif = sitk.LabelShapeStatisticsImageFilter() - lssif.Execute(label) - existing_labels = lssif.GetLabels() # 這會回傳 list，例如 [1,2,3,20,21] - - print(f"Existing labels in {os.path.basename(label_path)}: {existing_labels}") - + # 建立每個檔案的輸出資料夾 file_name = os.path.basename(image_path) name = file_name.replace(".nii.gz", "") @@ -54,21 +123,47 @@ def process_single_image(image_path, label_path, output_dir_base=None): for lab in existing_labels: f.write(f"{lab}\t{LABEL_MAP.get(lab, 'Unknown')}\n") - # 遍歷現有 label 做分割 + # 遍歷現有 label 做分割；個別 label 失敗（label_map 缺該 label、 + # 或 seg_bone 報錯）只跳過該 label，不中斷整檔處理 + processed = [] + skipped = [] for n in existing_labels: + if n not in LABEL_MAP: + print(f"Label {n} not found in label_map, skipping this label (file continues).") + skipped.append(n) + continue + if allowed_levels is not None and LABEL_MAP[n] not in allowed_levels: + print(f"Label {n} ({LABEL_MAP[n]}) not in allowed levels, " + f"skipping this label (file continues).") + skipped.append(n) + continue try: - roi_path, binary_path, roi2_path, cortical_path = seg_bone(n, name, resampled_sitk_img, resampled_sitk_lbl, output_dir, label_map=LABEL_MAP) - for path in [roi_path, binary_path, roi2_path, cortical_path]: - standardize_affine(path, output_dir) + res = seg_bone(n, name, resampled_sitk_img, resampled_sitk_lbl, output_dir, + label_map=LABEL_MAP, original_label=label) + if res is None: + print(f"Label {n}: empty after largest-CC extraction, skipping this label.") + skipped.append(n) + continue + (roi_path, binary_path, roi2_path, cortical_path, binary_nn_path, + binary_linear_path, smd_path, resampled_path, binary_sdf_path, + binary_erode_path) = res + for path in [roi_path, binary_path, roi2_path, cortical_path, + binary_nn_path, binary_linear_path, + smd_path, resampled_path, binary_sdf_path, + binary_erode_path]: + if path is not None: + standardize_affine(path, output_dir) + processed.append(n) except RuntimeError as e: print(f"Label {n} could not be processed, skipping. Error: {e}") + skipped.append(n) return { - "processed_labels": [lab for lab in existing_labels], - "missing_labels": [] + "processed_labels": processed, + "missing_labels": skipped } -def process_dataset(image_dir, label_dir, output_dir, labels_to_process=None): +def process_dataset(image_dir, label_dir, output_dir, labels_to_process=None, max_images=None, post_process=None, max_z_spacing=None, allowed_levels=None, min_levels=None, metadata_cache=None): image_files = sorted(glob.glob(os.path.join(image_dir, "*.nii.gz"))) total_files = len(image_files) print(f"Total files: {total_files}") @@ -77,8 +172,15 @@ def process_dataset(image_dir, label_dir, output_dir, labels_to_process=None): progress = load_progress(PROGRESS_FILE) all_file_summary = [] + # max_images 統計「實際進入 pipeline 的檔數」： + # z spacing 超標被跳過的檔不計入配額，繼續掃描後續檔案 + processed_count = 0 for idx, image_path in enumerate(image_files, 1): + if max_images is not None and processed_count >= max_images: + print(f"Reached max_images={max_images}, stopping early.") + break + file_name = os.path.basename(image_path) name = file_name.replace(".nii.gz", "") label_path = os.path.join(label_dir, file_name.replace(".nii.gz", "_seg.nii.gz")) @@ -93,25 +195,40 @@ def process_dataset(image_dir, label_dir, output_dir, labels_to_process=None): print(f"[{idx}/{total_files}] Already finished: {file_name}") file_summary["current_labels"] = progress[name].get("processed_labels", []) file_summary["missing_labels"] = progress[name].get("missing_labels", []) + processed_count += 1 all_file_summary.append(file_summary) continue if not os.path.exists(label_path): print(f"[{idx}/{total_files}] Warning: label not found for {file_name}") file_summary["note"] = "Label file not found" + processed_count += 1 all_file_summary.append(file_summary) continue try: - result = process_single_image(image_path, label_path, output_dir_base=output_dir) + result = process_single_image(image_path, label_path, output_dir_base=output_dir, + max_z_spacing=max_z_spacing, + allowed_levels=allowed_levels, + min_levels=min_levels, + metadata_cache=metadata_cache) # print(result) # exit() except Exception as e: print(f"[{idx}/{total_files}] Error processing {file_name}: {e}") file_summary["note"] = f"Error: {e}" + processed_count += 1 all_file_summary.append(file_summary) continue + if result.get("skipped"): + # z spacing 超標：不計入 max_images 配額，繼續掃描後續檔案 + print(f"[{idx}/{total_files}] Skipped ({result['skip_reason']}): {file_name}") + file_summary["note"] = result["skip_reason"] + all_file_summary.append(file_summary) + continue + + processed_count += 1 file_summary["current_labels"] = result["processed_labels"] file_summary["missing_labels"] = result["missing_labels"] all_file_summary.append(file_summary) @@ -123,7 +240,15 @@ def process_dataset(image_dir, label_dir, output_dir, labels_to_process=None): } save_progress(progress, PROGRESS_FILE) - print(f"[{idx}/{total_files}] Finished: {file_name} | Missing labels: {result['missing_labels'] or 'None'}") + count_msg = f" | processed {processed_count}/{max_images}" if max_images is not None else "" + print(f"[{idx}/{total_files}] Finished: {file_name} | " + f"Missing labels: {result['missing_labels'] or 'None'}{count_msg}") + + if post_process is not None: + try: + post_process(os.path.join(output_dir, name), result["processed_labels"]) + except Exception as e: + print(f"[{idx}/{total_files}] post_process error for {name}: {e}") # --- Summary --- summary_path = os.path.join(output_dir, "all_files_label_summary.txt") diff --git a/imaging/resample.py b/imaging/resample.py index 4051cf2..239398b 100644 --- a/imaging/resample.py +++ b/imaging/resample.py @@ -2,28 +2,49 @@ import numpy as np import SimpleITK as sitk from config.constant import LABEL_MAP -def resample_img(sitk_image, out_spacing=[0.5, 0.5, 0.5], is_label=False): - - # Resample images to 2mm spacing with SimpleITK - original_spacing = sitk_image.GetSpacing() - original_size = sitk_image.GetSize() - - out_size = [ - int(np.round(original_size[0] * (original_spacing[0] / out_spacing[0]))), - int(np.round(original_size[1] * (original_spacing[1] / out_spacing[1]))), - int(np.round(original_size[2] * (original_spacing[2] / out_spacing[2])))] +def resample_img(sitk_image, out_spacing=[0.5, 0.5, 0.5], is_label=False, + interpolator=None, cval=None): + """重採樣到 out_spacing，完整保留物理範圍（邊界 margin 不裁剪）： + - CT（is_label=False）：线性插值。本流程是上採樣（~0.75-1mm -> 0.5mm）， + 线性即無漣波的理想重建；B-spline 會在身體 margin 等強邊界產生 + 非物理的 undershoot/overshoot（量測：~0.5% 體素低於 air floor + -1024 HU，最低 -1274，沿體輪廓形成黑暈）。 + - 預設填充值 = 影像最小值（air），而非 GetPixelIDValue() + （對 int16 CT 回傳佔位值 2 ≈ 軟組織，會把 margin 填成軟組織）。 + - 輸出尺寸用 ceil 覆蓋原物理尺寸（round 會讓遠端端點被裁 ≤0.25mm）。 + label（is_label=True）：最近邻、填充 0。 + interpolator：明確指定插值器（None = 依 is_label 取 Linear / NearestNeighbor； + 例如 SMD 等 piecewise-linear 場用 sitk.sitkBSpline，三阶 B 样条对分段 + 线性场为精确重建、無漣波）。 + cval：明確指定填充值（None = 上述預設）。 + """ + original_spacing = np.array(sitk_image.GetSpacing(), dtype=float) + original_size = np.array(sitk_image.GetSize()) + out_spacing = np.array(out_spacing, dtype=float) + physical = original_size * original_spacing + out_size = [max(1, int(np.ceil(physical[i] / out_spacing[i] - 1e-6))) + for i in range(3)] resample = sitk.ResampleImageFilter() - resample.SetOutputSpacing(out_spacing) + resample.SetOutputSpacing(out_spacing.tolist()) resample.SetSize(out_size) resample.SetOutputDirection(sitk_image.GetDirection()) resample.SetOutputOrigin(sitk_image.GetOrigin()) resample.SetTransform(sitk.Transform()) - resample.SetDefaultPixelValue(sitk_image.GetPixelIDValue()) if is_label: - resample.SetInterpolator(sitk.sitkNearestNeighbor) + resample.SetInterpolator(interpolator or sitk.sitkNearestNeighbor) + resample.SetDefaultPixelValue(cval if cval is not None else 0) + elif interpolator is not None: + resample.SetInterpolator(interpolator) + # SMD 等 signed 場的填充：預設 0 = 表面層值（比影像 min/max 安全， + # 不會製造假的零穿越環）；可用品值可用 cval 覆蓋 + resample.SetDefaultPixelValue(float(cval) if cval is not None else 0.0) else: - resample.SetInterpolator(sitk.sitkBSpline) + resample.SetInterpolator(sitk.sitkLinear) + # air 值 = 影像最小 HU（statistics 濾波器 streaming 計算，不載入整張 array） + stats = sitk.StatisticsImageFilter() + stats.Execute(sitk_image) + resample.SetDefaultPixelValue(float(stats.GetMinimum())) return resample.Execute(sitk_image) \ No newline at end of file diff --git a/imaging/segmentation.py b/imaging/segmentation.py index 0d68778..936ca54 100644 --- a/imaging/segmentation.py +++ b/imaging/segmentation.py @@ -1,6 +1,7 @@ import os import SimpleITK as sitk from config.constant import LABEL_MAP +from imaging.resample import resample_img import numpy as np """ @@ -12,7 +13,28 @@ 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): +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' @@ -22,87 +44,118 @@ def seg_bone(n, name, resampled_sitk_img, resampled_sitk_lbl, output_base=None, label_name = label_map[n] - # lssif = sitk.LabelShapeStatisticsImageFilter() - # lssif.Execute(resampled_sitk_lbl) + # ============ 原解析度（未插值）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 - # if not lssif.HasLabel(n): - # raise RuntimeError(f"Label {n} not found") + # _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) - # bbox2 = lssif.GetBoundingBox(n) + # _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) - # 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: + # 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 - - # 5. 獲取最大組件（標籤為 1）的邊界框 - # 格式通常為 [x_start, y_start, z_start, x_size, y_size, z_size] - bbox2 = shape_stats.GetBoundingBox(1) + 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) - # 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 + 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 - ├── L1_binary.nii.gz - ├── L1_roi2.nii.gz - ├── L1_cortical.nii.gz - ├── L2_roi.nii.gz ├── L2_binary.nii.gz ... """ \ No newline at end of file diff --git a/visualization/res_bone_figure.py b/visualization/res_bone_figure.py new file mode 100644 index 0000000..9914f54 --- /dev/null +++ b/visualization/res_bone_figure.py @@ -0,0 +1,451 @@ +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 分量 -> 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 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 \ No newline at end of file diff --git a/visualization/res_plot_3d.py b/visualization/res_plot_3d.py index 3ab072d..8c9cb94 100644 --- a/visualization/res_plot_3d.py +++ b/visualization/res_plot_3d.py @@ -5,6 +5,8 @@ from matplotlib.colors import to_rgba from matplotlib.lines import Line2D from mpl_toolkits.mplot3d.art3d import Poly3DCollection import os +import errno +import time from datetime import datetime import csv @@ -13,7 +15,7 @@ from core.intersection import center_line_intersections_torch from core.scoring import cl_score_torch, compute_overlap_ratio_from_cylinder_mask, cl_score_torch_xfr from imaging.orientation import (azimuth_rotation, analyze_vertebral_tilt_contour, best_symmetry_plane, best_upper_endplate_plane, - segment_spinous_process) + segment_spinous_process, segment_vertebral_body) from utils.helpers import save_with_unique_name # Volume absorption 渲染（Beer-Lambert）：每 voxel 不透明度 = 1 - exp(-mu * voxel_width) @@ -25,6 +27,17 @@ BONE_MU_TRABECULAR = 0.005 # 1/mm → 每 voxel = 1-exp(-0.005*0.5) ~ 0.0025 BONE_MARKER_SIZE = 3.0 # 骨骼散點點面積 (pt^2)；略大以補償抽稀後的顆粒感 BONE_SUBSAMPLE = 1 # 每 10 個骨 voxel 畫 1 個，降低堆疊不透明度（1=全畫） +def _retry_robust(fn, *args, retries=20, delay=0.5, **kwargs): + """對 ENOENT/EEXIST 重試：NFS 上輸出樹被外部刪除（或多 worker 併發建同一 + output 目錄）會有短暫的 ENOENT 窗口，重試可恢復；其他錯誤直接丟出。""" + for i in range(retries): + try: + return fn(*args, **kwargs) + except OSError as e: + if e.errno not in (errno.ENOENT, errno.EEXIST) or i == retries - 1: + raise + time.sleep(delay) + def set_axes_equal_3d(ax): """ Make axes of 3D plot have equal scale so that spheres appear as spheres, @@ -156,11 +169,15 @@ def res_plt_2_torch( device ) # loss_r = cl_score_torch(cortical_tensor, spine_tensor, cyl_r, cyl_ro, intersections_r) - loss_r = cl_score_torch_xfr(cortical_tensor, spine_tensor, cyl_r, cyl_ro, intersections_r) + # loss_r 放在下方 VBODY mask 計算之後：計入與 PSO 目標函數相同的 VBODY voxel 獎勵 - azi = azimuth_rotation(image2_path) - res = analyze_vertebral_tilt_contour(image2_path, edge_type='superior', show_plot=False, debug=False) - alt = res['superior']['tilt_angle_deg'] + if CBT: + azi = float('nan') + alt = float('nan') + else: + azi = azimuth_rotation(image2_path) + res = analyze_vertebral_tilt_contour(image2_path, edge_type='superior', show_plot=False, debug=False) + alt = res['superior']['tilt_angle_deg'] # Move data to CPU for plotting line_mask_l_cpu = line_mask_l.cpu().numpy() @@ -191,8 +208,16 @@ def res_plt_2_torch( sym = best_symmetry_plane(spine_cpu) # 棘突：鏡稱面中線帶（|s|<=w）且在 AP 谷底之後側的骨 voxel，換不同顏色標示 + # 棘突缺如（先前 laminectomy / 棘突切除）時 sp_mask=None，不標示。 sp_mask, sp_th, sp_info = segment_spinous_process(spine_cpu, sym) - if sp_mask is not None and sp_mask.any(): + if sp_info['mode'] == 'no_spinous': + 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['deficit']:.1f} voxel ({sp_info['deficit'] * 0.5:.1f} mm), " + f"rear3={sp_info['rear3']} voxel, top_off={top_off} voxel " + 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] sp_n_bone = max(int(spine_cpu.sum()), 1) @@ -203,6 +228,34 @@ def res_plt_2_torch( else: sp_corti = sp_trab = None + # 上終板平面：RANSAC 擬合骨頭頂面（前側）的最佳 a·x + b·y + c·z = d + symp = best_upper_endplate_plane(spine_cpu) + + # 椎體：上終板之下（排除跨終板的後側構造）且中線 AP 谷底之前側的骨 voxel， + # 換不同顏色標示（見 segment_vertebral_body；谷底優先取中線帶， + # 中線搜尋 fallback 時退回終板下整體 AP 分佈谷底） + vb_mask, vb_th, vb_info = segment_vertebral_body(spine_cpu, 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] + print(f"[VBODY] n={vb_info['n_vb']} " + f"({100.0 * vb_info['n_vb'] / max(int(spine_cpu.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 + print(f"[VBODY] skipped: {vb_info['mode']}") + + # loss_r 使用與 PSO 目標函數相同的 VBODY 獎勵（回報分數與優化一致） + vbody_tensor = None + if vb_mask is not None and vb_mask.any(): + vbody_tensor = torch.from_numpy(vb_mask.astype(np.uint8)).to(device=device) + loss_r = cl_score_torch_xfr(cortical_tensor, spine_tensor, cyl_r, cyl_ro, intersections_r, + vbody_tensor=vbody_tensor) + # X-ray 外觀：骨頭合成一個半透明體積吸收點雲（下方）； # 螺絲（中心線 + 圓柱 + 入口軌跡延长）合成一個全不透明點雲，永遠畫在骨頭之上 def _rgba_block(n, color, a): @@ -227,11 +280,41 @@ def res_plt_2_torch( z_bone = z_bone[::BONE_SUBSAMPLE] bone_rgba = bone_rgba[::BONE_SUBSAMPLE] bone_size = bone_size[::BONE_SUBSAMPLE] - if sp_corti is not None: - sp_flag = np.concatenate([sp_corti, sp_trab]) + # 平面 patch 的範圍用「完整骨頭」（含 VBODY / SP）算： + # VBODY 前側是整顆骨最前緣，剔除後綠色終板 patch 會縮小 + x_bone_all, y_bone_all, z_bone_all = x_bone, y_bone, z_bone + # VBODY / 棘突拆成獨立上層（在 _fill_ax 內畫）： + # 繪製順序 基底骨(5) < VBODY gold(6) < SP purple(6.5) < 終板(7) < 鏡稱面(8) < 螺絲(10) + # axial 視角（ax3）相機在 +y 前側，椎體 physically 擋在棘突與相機之間， + # 棘突最後畫 -> 紫色不被 gold 遮住 + vb_flag = np.zeros(x_bone.shape, dtype=bool) + sp_flag = np.zeros(x_bone.shape, dtype=bool) + if vb_corti is not None: + f = np.concatenate([vb_corti, vb_trab]).astype(bool) if BONE_SUBSAMPLE > 1: - sp_flag = sp_flag[::BONE_SUBSAMPLE] - bone_rgba[sp_flag] = to_rgba('purple', 0.95) + f = f[::BONE_SUBSAMPLE] + vb_flag |= f + if sp_corti is not None: + f = np.concatenate([sp_corti, sp_trab]).astype(bool) + if BONE_SUBSAMPLE > 1: + f = f[::BONE_SUBSAMPLE] + sp_flag |= f + + vbody_pts = None + vbody_mask = vb_flag & ~sp_flag + if vbody_mask.any(): + vbody_pts = (x_bone[vbody_mask], y_bone[vbody_mask], z_bone[vbody_mask]) + sp_pts = None + if sp_flag.any(): + sp_pts = (x_bone[sp_flag], y_bone[sp_flag], z_bone[sp_flag]) + + # 基底骨層去掉 VBODY / SP voxel（由上面的專屬圖層畫） + base_mask = ~(vb_flag | sp_flag) + x_bone = x_bone[base_mask] + y_bone = y_bone[base_mask] + z_bone = z_bone[base_mask] + bone_rgba = bone_rgba[base_mask] + bone_size = bone_size[base_mask] x_screw = np.concatenate([x_lin1, x_lin2, x_cyl_l1, x_cyl_l2, x_cyl_r1, x_cyl_r2]) y_screw = np.concatenate([y_lin1, y_lin2, y_cyl_l1, y_cyl_l2, y_cyl_r1, y_cyl_r2]) @@ -242,7 +325,7 @@ def res_plt_2_torch( _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) + xyz_bone = np.stack([x_bone_all - _p0[0], y_bone_all - _p0[1], z_bone_all - _p0[2]], axis=1) _pu = xyz_bone @ _u _pv = xyz_bone @ _v _U_, _V_ = np.meshgrid(np.linspace(_pu.min(), _pu.max(), 8), @@ -251,8 +334,6 @@ def res_plt_2_torch( _Yp = _p0[1] + _U_ * _u[1] + _V_ * _v[1] _Zp = _p0[2] + _U_ * _u[2] + _V_ * _v[2] - # 上終板平面：RANSAC 擬合骨頭頂面（前側）的最佳 a·x + b·y + c·z = d - symp = best_upper_endplate_plane(spine_cpu) _EX = _EY = _EZ = None if symp is not None: _ea, _eb, _ec, _ed = symp['plane'] @@ -260,7 +341,7 @@ def res_plt_2_torch( _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) + xz_ep = np.stack([x_bone_all - _ep0[0], y_bone_all - _ep0[1], z_bone_all - _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), @@ -284,32 +365,32 @@ def res_plt_2_torch( fig = plt.figure(figsize=(12, 12)) - # 圖例色塊提高到可讀不透明度（實際渲染仍用真實吸收 alpha） - _leg_alpha_c = max(float(alpha_cortical), 0.35) - _leg_alpha_t = max(float(alpha_trabecular), 0.2) legend_handles = [ - Line2D([], [], marker='o', ls='', ms=5, color=to_rgba('lightblue', _leg_alpha_c), label='Spine (cortical)'), - Line2D([], [], marker='o', ls='', ms=5, color=to_rgba('lightblue', _leg_alpha_t), label='Spine (trabecular)'), - Line2D([], [], marker='o', ls='', ms=2, color='r', label='Centerline'), Line2D([], [], marker='o', ls='', ms=6, color='darkcyan', label='Cylinder(L)'), Line2D([], [], marker='o', ls='', ms=6, color='blue', label='Cylinder(R)'), - Line2D([], [], marker='o', ls='', ms=6, color='pink', label='Entry track (outer)'), -Line2D([], [], color='orange', lw=2, alpha=0.6, - label=f"Mirror plane {sym['plane'][0]:+.2f}x {sym['plane'][1]:+.2f}y {sym['plane'][2]:+.2f}z = {sym['plane'][3]:.1f}"), ] + 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=f"Spinous process (mirror-plane midline, {sp_info['n_sp']} vox)")) - if symp is not None: - legend_handles.append(Line2D([], [], color='green', lw=2, alpha=0.7, - label=f"Upper endplate plane tilt {symp['tilt_deg']:.1f} deg")) + legend_handles.append( + Line2D([], [], marker='o', ls='', ms=6, color='purple', label='SpinousProcess')) def _fill_ax(ax): # X-ray 外觀：關閉 mplot3d 依深度自動排序 zorder（否則半透明骨頭會被重繪到 - # 螺絲上方）；改為固定分層：吸收骨頭 zorder=5，全不透明螺絲 zorder=10 + # 螺絲上方）；改為固定分層： + # 基底骨 zorder=5 < VBODY 6 < SP 6.5 < 終板 7 < 鏡稱面 8 < 螺絲 10 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) + if vbody_pts is not None: + sc_vb = ax.scatter(vbody_pts[0], vbody_pts[1], vbody_pts[2], + c=to_rgba('gold', 0.95), s=BONE_MARKER_SIZE, marker='o') + sc_vb.set_zorder(6) + if sp_pts is not None: + sc_sp = ax.scatter(sp_pts[0], sp_pts[1], sp_pts[2], + c=to_rgba('purple', 0.95), s=BONE_MARKER_SIZE, marker='o') + sc_sp.set_zorder(6.5) sc_screw = ax.scatter(x_screw, y_screw, z_screw, c=screw_rgba, s=screw_size, marker='o') sc_screw.set_zorder(10) # 中矢狀面（理論左右對稱切分面）：半透明橘色平面 x = x_mid @@ -366,28 +447,90 @@ Line2D([], [], color='orange', lw=2, alpha=0.6, user_azimuth_l = 90 - best_position_l[3] - azi user_azimuth_r = 90 - best_position_r[3] - azi + # 螺絲方向向量（與 generate_cylinder_n_torch 同慣例）： + # d = (cos(az)·sin(alt), sin(az)·sin(alt), cos(alt))，alt = 相对 +z 的極角 + # Azimuth 相对鏡稱面（法線 s，theta_v = atan2(sy, sx)）： + # Azimuth_Lateral = az - theta_v - 90 （面內 AP 軸起的帶號發散角，+ = L 側往外，− = R 側） + # Altitude 相對上終板面（法線 e，朝上）： + # Altitude_Cephalad_Endplate = 90 - (d 與 e 的夾角) + 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 + tau_y = float('nan') + tau_x = float('nan') + + 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 + if e_n is not None: + ang_norm = float(np.degrees(np.arccos(np.clip(d_v @ e_n, -1.0, 1.0)))) + alt_cep = 90.0 - ang_norm + else: + alt_cep = float('nan') + return az_lateral, alt_cep + + azlat_l, acep_l = _rel_angles(best_position_l[3], best_position_l[4]) + azlat_r, acep_r = _rel_angles(best_position_r[3], best_position_r[4]) + date_str = datetime.now().strftime("%Y%m%d") - patient_id = os.path.basename(os.path.dirname(image2_path)) + # 旋轉後影像存在 /rotated/ 下：parent 是 'rotated'， + # 再上一層才是 volume id（未旋轉路徑不受影響） + img_parent = os.path.dirname(image2_path) + patient_id = os.path.basename(os.path.dirname(img_parent)) \ + if os.path.basename(img_parent) == 'rotated' else os.path.basename(img_parent) output_folder = os.path.join(base_folder, date_str, patient_id) - os.makedirs(output_folder, exist_ok=True) + _retry_robust(os.makedirs, output_folder, exist_ok=True) csv_path = os.path.join(output_folder, 'output.csv') - # 檢查檔案是否存在 (決定是否寫入標題) - file_exists = os.path.isfile(csv_path) - - # 欄位標題 (Header) + # 欄位標題 (Header)。CBT 模式下恆為 nan 的 2D 參考欄 + # (Azimuth_Diff / Altitude_Diff / User_Azimuth / User_Altitude) 不寫入 CSV。 headers = [ - 'Label', 'Side', 'Diameter', 'Length', 'Swarm_Size', 'Max_Iter', - 'Position_XYZ', 'Raw_Azimuth', 'Azimuth_Diff', 'Raw_Altitude', 'Altitude_Diff', - 'Intersections', 'Best_Loss', 'cyl_points', 'Overlap_Cortical', 'Overlap_Bone', - 'Cortical_Bone_Ratio', 'User_Azimuth', 'User_Altitude', 'Total_Time' + '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}" + + # 檢查檔案是否存在 (決定是否寫入標題)；舊 schema 的檔案按欄位名稱重映射後 + # 改以新 Header 重寫（舊檔多出的欄位捨去、缺的欄位補空白），避免 append 欄位錯位 + 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 open(csv_path, 'a', newline='') as csvfile: + with _retry_robust(open, csv_path, 'a', newline='') as csvfile: writer = csv.writer(csvfile) - - # 如果是新檔案，寫入 Header + + # 新檔案寫入 Header if not file_exists: writer.writerow(headers) @@ -402,17 +545,18 @@ Line2D([], [], color='orange', lw=2, alpha=0.6, # f"({best_position_l[0]:.2f}, {best_position_l[1]:.2f}, {best_position_l[2]:.2f})", f"({best_position_l[2]:.2f}, {best_position_l[1]:.2f}, {best_position_l[0]:.2f})", f"{best_position_l[3]:.2f}", - f"{best_position_l[3]-azi:.2f}", f"{best_position_l[4]:.2f}", - f"{best_position_l[4]-alt:.2f}", intersections_l, f"{loss_l:.2f}", cyl_points_l, f"{overlap_cortical_l:.2f}", f"{overlap_vertebral_l:.2f}", f"{(overlap_cortical_l/overlap_vertebral_l if overlap_vertebral_l!=0 else 0):.2f}", - f"{user_azimuth_l:.2f}", - f"{user_altitude_l:.2f}", + _fmt(theta_v), + _fmt(tau_y), + _fmt(tau_x), + _fmt(azlat_l), + _fmt(acep_l), f"{total_time:.2f}" ]) @@ -427,17 +571,18 @@ Line2D([], [], color='orange', lw=2, alpha=0.6, # f"({best_position_r[0]:.2f}, {best_position_r[1]:.2f}, {best_position_r[2]:.2f})", f"({best_position_r[2]:.2f}, {best_position_r[1]:.2f}, {best_position_r[0]:.2f})", f"{best_position_r[3]:.2f}", - f"{best_position_r[3]-azi:.2f}", f"{best_position_r[4]:.2f}", - f"{best_position_r[4]-alt:.2f}", intersections_r, f"{loss_r:.2f}", cyl_points_r, f"{overlap_cortical_r:.2f}", f"{overlap_vertebral_r:.2f}", f"{(overlap_cortical_r/overlap_vertebral_r if overlap_vertebral_r!=0 else 0):.2f}", - f"{user_azimuth_r:.2f}", - f"{user_altitude_r:.2f}", + _fmt(theta_v), + _fmt(tau_y), + _fmt(tau_x), + _fmt(azlat_r), + _fmt(acep_r), f"{total_time:.2f}" ]) print(f"[CSV Saved] {csv_path}") @@ -453,21 +598,41 @@ Line2D([], [], color='orange', lw=2, alpha=0.6, f'Swarm size = {swarm_size}, Iteration = {max_iter}, Total time = {total_time:.2f} s', ha='center', fontsize=12 ) + + # 角度註記：CBT 沒有 2D 參考面（user az/alt 為 nan），Azimuth/Altitude 直接顯示 + # 最佳化出的原始角（＝CSV 的 Raw_Azimuth / Raw_Altitude）；TPS 沿用 2D 參考之相對角。 + # 另補上相對骨骼的角度（與 CSV 同參數）： + # Azimuth_Lateral = 螺絲在鏡稱面內相對 AP 軸的發散角（+ = L 側往外，− = R 側） + # Altitude_Endplate = 螺絲相對上終板面的仰角 + # 終板面擬合失敗（nan）時該段自動略過。 + def _fig_angle_segs(az, alt, azlat, acep): + segs = [f'Azimuth = {az:.2f}', f'Altitude = {alt:.2f}'] + if np.isfinite(azlat): + segs.append(f'Azimuth_Lateral = {azlat:.2f}') + if np.isfinite(acep): + segs.append(f'Altitude_Endplate = {acep:.2f}') + return ', '.join(segs) + + if CBT: + _ang_l = _fig_angle_segs(float(best_position_l[3]), float(best_position_l[4]), azlat_l, acep_l) + _ang_r = _fig_angle_segs(float(best_position_r[3]), float(best_position_r[4]), azlat_r, acep_r) + else: + _ang_l = _fig_angle_segs(user_azimuth_l, user_altitude_l, azlat_l, acep_l) + _ang_r = _fig_angle_segs(user_azimuth_r, user_altitude_r, azlat_r, acep_r) + fig.text( 0.5, 0.03, - # f'Left : Position = ({best_position_l[0]:.2f}, {best_position_l[1]:.2f}, {best_position_l[2]:.2f}), ' f'Left : Position = ({best_position_l[2]:.2f}, {best_position_l[1]:.2f}, {best_position_l[0]:.2f}), ' - f'Azimuth = {user_azimuth_l:.2f}, Altitude = {user_altitude_l:.2f}, ' + f'{_ang_l}, ' f'Intersection = {intersections_l}, Score = {overlap_cortical_l:.2f} / {overlap_vertebral_l:.2f} / {cb_ratio_l:.2f}', - ha='center', fontsize=9 + ha='center', fontsize=8 ) fig.text( 0.5, 0.01, - # f'Right : Position = ({best_position_r[0]:.2f}, {best_position_r[1]:.2f}, {best_position_r[2]:.2f}), ' f'Right : Position = ({best_position_r[2]:.2f}, {best_position_r[1]:.2f}, {best_position_r[0]:.2f}), ' - f'Azimuth = {user_azimuth_r:.2f}, Altitude = {user_altitude_r:.2f}, ' + f'{_ang_r}, ' f'Intersection = {intersections_r}, Score = {overlap_cortical_r:.2f} / {overlap_vertebral_r:.2f} / {cb_ratio_r:.2f}', - ha='center', fontsize=9 + ha='center', fontsize=8 ) fig.tight_layout() @@ -476,7 +641,6 @@ Line2D([], [], color='orange', lw=2, alpha=0.6, file_name = os.path.basename(image2_path) level = file_name.split('_')[0] output_folder = os.path.join(base_folder, date_str, patient_id) - os.makedirs(output_folder, exist_ok=True) if CBT == True: way = 'CBT' @@ -484,11 +648,18 @@ Line2D([], [], color='orange', lw=2, alpha=0.6, else: way = 'TPS' - path = save_with_unique_name(output_folder, label_str, way, - diameter_l, length_l, diameter_r, length_r, - swarm_size, max_iter) + # 建目錄 + 存檔一起重試：輸出樹被外部刪除（NFS 刪除競態）時，重試會重建目錄 + path = None + def _save_fig_once(): + nonlocal path + _retry_robust(os.makedirs, output_folder, exist_ok=True) + # 檔名只用 level（volume id 已在資料夾名裡，不重複） + path = save_with_unique_name(output_folder, level, way, + diameter_l, length_l, diameter_r, length_r, + swarm_size, max_iter) + fig.savefig(path, dpi=200, bbox_inches="tight") - fig.savefig(path, dpi=200, bbox_inches="tight") + _retry_robust(_save_fig_once) print("[Saved figure]", path) plt.close(fig) diff --git a/xfr_check_spinous.py b/xfr_check_spinous.py new file mode 100644 index 0000000..02b8775 --- /dev/null +++ b/xfr_check_spinous.py @@ -0,0 +1,310 @@ +#!/home/xfr/.conda/envs/cbt/bin/python +""" +檢查各 (volume, level) 的棘突（midline 後側構造）是否缺如 +（先前手術如 laminectomy / 棘突切除，例：0005 L5、0770 L5）。 + +判定原理（見 imaging.orientation.diagnose_spinous_process）： + 棘突完整時，尖端是全椎體最後側的骨且位於中線 → + deficit（全骨最後側 - 中線帶最後側）≈ 0，後側中線窄帶有骨。 + 切除後，最後側骨偏到側方殘餘、後側中線空洞 → deficit 大 + rear3 ≈ 0。 + no_spinous = deficit >= 4 vox 且 rear3 <= 20 vox（正常 level deficit 恒為 0）。 + +兩種模式： + quick（預設，掃全資料集用）→ 兩段式： + stage 1：axial 投影快速篩檢（~0.5s/level）。用「寬行（椎體行）」擬合 + 中線軌跡（對傾斜/旋轉穩健，不受後側殘片污染），找出「缺棘突候選」。 + 高召回（幾乎不漏）、但因純投影在傾斜脊椎上仍會有少數假候補。 + stage 2：僅對候選跑 full（鏡稱面）驗證（~3s/候選），以 full 為準給出 + 最終判定（confirmed / false alarm）。全資料集 ~25min + 候數×3s。 + full（指定單一 volume+level 時預設）：跑鏡稱面 + 棘突/椎體完整分割， + 另回報 VBODY 分割 mode（nosp_gap / nosp_post_min / quantile...）， + 確認該 level 的椎體切分落在「椎體後側末端（體/弓最薄處）」。 + +識別出的 no_spinous level 在 xfr_debug.py 執行時會打 [NO-SP] 標記， +並自動改用放寬後側谷底的椎體切分（不再把殘留後側要素當棘突移除）。 + +Usage: + python xfr_check_spinous.py # quick，全部 volume L1~L5 + python xfr_check_spinous.py 0005 # quick，該 volume L1~L5 + python xfr_check_spinous.py 0005 L5 # full（精確）單一 level + python xfr_check_spinous.py --full 0005 L5 # 強制 full + python xfr_check_spinous.py --quick # 強制 quick（全資料集掃） +""" + +import csv +import os +import sys +import time + +import numpy as np +import SimpleITK as sitk + +from imaging.orientation import (best_symmetry_plane, segment_spinous_process, + best_upper_endplate_plane, segment_vertebral_body, + diagnose_spinous_process) + +standardized_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr/' + +LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5') + +LOG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'results') + +# 與 diagnose_spinous_process 相同的判定閾值 +DEFICIT_MIN = 4.0 +REAR3_MAX = 20 +# quick：中線帶內一行的最小 bone voxel 數（< 此值視為殘片游離 voxel，不算「尖端在」） +MIN_MID_VOX = 3 + + +# ===================================================================== +# quick screen：不依賴鏡稱面（~0.3s/level） +# ===================================================================== +def quick_screen(m): + """axial 投影 (y, x) 快速篩檢（不跑鏡稱面，對傾斜/旋轉穩健）： + - 「寬行」= 行內骨 voxel 數 >= 最大行的 50%（即椎體行；椎體左右 + 大致對稱 → 寬行 x 重心 ≈ 中線上的點）。 + - 中線軌跡 x(y) = p0 + p1·y 只用寬行（椎體行）擬合： + * 用全部行擬合 → laminectomy 後側殘片行污染中線，deficit 被吃掉 + （0005 L5 假正常）； + * 用固定垂直中線 → 傾斜棘突尖端被誤判為側方偏移 + （0555/0609/0637 假切除）。椎體行兩者都不沾，故最穩。 + - deficit = 最後側骨行 (y_min) 到「中線軌跡 ± 25% 椎體寬度 band + 內有骨的最後側行」的行距（尖端在中線上 → 0；切除後 → ≈切除量）。 + - rear3/rear6 = 最後側 3/6 行內、中線軌跡 ± 10% 寬度（窄中線帶）的 + 骨 voxel 數。 + - top_off = 最後側一行 x 重心到中線軌跡的距離。 + 回傳 dict 或 None（資料不足）。""" + proj = m.max(axis=0) + ys_p, xs_p = np.where(proj > 0) + if ys_p.size < 50: + return None + ny = proj.shape[0] + cnt = np.bincount(ys_p, minlength=ny) + max_cnt = int(cnt.max()) + if max_cnt < 4: + return None + xsum = np.bincount(ys_p, weights=xs_p.astype(np.float64), minlength=ny) + centroid = xsum / cnt.astype(float) + wide = cnt >= (0.5 * max_cnt) # 椎體行 + if int(wide.sum()) < 10: + return None + yv = np.where(wide)[0] + p1, p0 = np.polyfit(yv, centroid[wide], 1) # 中線軌跡（椎體行擬合） + xmax = np.full(ny, -1.0) + xmin = np.full(ny, np.inf) + np.maximum.at(xmax, ys_p, xs_p) + np.minimum.at(xmin, ys_p, xs_p) + span_c = float(np.median(xmax[yv] - xmin[yv])) + if span_c < 10: + return None + line = p0 + p1 * ys_p + band = 0.25 * span_c # central band（椎體中央區） + narrow = 0.10 * span_c # 窄中線帶 + y_min = int(ys_p.min()) + # 「中線有骨」需 >= MIN_MID_VOX 個 voxel 才算數：棘突尖端在帶內有 + # 多個 voxel，切除後的側方殘片在帶內通常只有 1-2 個游離 voxel + # （否則 0005 L5 的單 voxel 殘端會被當成尖端 → 假正常）。 + bcnt = np.bincount(ys_p[np.abs(xs_p - line) <= band], minlength=ny) + idx = np.where(bcnt >= MIN_MID_VOX)[0] + y_mid = int(idx.min()) if idx.size else ny - 1 + deficit = float(y_mid - y_min) + nn = np.abs(xs_p - line) <= narrow + rear3 = int(((ys_p <= y_min + 3.0) & nn).sum()) + rear6 = int(((ys_p <= y_min + 6.0) & nn).sum()) + top_off = float(abs(xs_p[ys_p == y_min].mean() - (p0 + p1 * y_min))) + return {'deficit': deficit, 'rear3': rear3, 'rear6': rear6, + 'top_off': top_off, 'no_spinous': bool(deficit >= DEFICIT_MIN + and rear3 <= REAR3_MAX)} + + +# ===================================================================== +# full：鏡稱面 + 完整棘突/椎體分割（精確，較慢） +# ===================================================================== +def full_check(m): + sym = best_symmetry_plane(m) + diag = diagnose_spinous_process(m, sym) + sp_mask, sp_th, sp_info = segment_spinous_process(m, sym) + ep = best_upper_endplate_plane(m) + vb, vb_th, vb_info = segment_vertebral_body(m, sym, ep, sp_th, sp_info['mode']) + row = {'deficit': float(diag['deficit']), 'rear3': diag['rear3'], + 'rear6': diag['rear6'], 'top_off': diag['top_off'], + 'no_spinous': bool(diag['no_spinous']), + 'sp_mode': sp_info['mode'], + 'sp_pct': 100.0 * sp_info['n_sp'] / max(int(m.sum()), 1), + 'vb_mode': vb_info['mode'], + 'vb_pct': (100.0 * vb_info['n_vb'] / max(int(m.sum()), 1)) + if vb is not None else None, + 'vb_ap_th': vb_th} + return row + + +# ===================================================================== +def load_binary(volume_dir, level): + # _binary.nii.gz 現為原解析度（未插值）；0.5mm 分析優 _binary_sdf（SDF 平滑遮罩） + p = os.path.join(volume_dir, f'{level}_binary_sdf.nii.gz') + if not os.path.exists(p): + p = os.path.join(volume_dir, f'{level}_binary.nii.gz') + if not os.path.exists(p): + return None + m = sitk.GetArrayFromImage(sitk.ReadImage(p, sitk.sitkUInt8)) > 0 + if int(m.sum()) < 1000 or min(m.shape) < 8: + return None + return m + + +def check_level(vid, level, use_full): + m = load_binary(os.path.join(standardized_dir, vid), level) + if m is None: + return {'vid': vid, 'level': level, 'note': 'missing/degenerate mask'} + t0 = time.time() + row = quick_screen(m) if not use_full else full_check(m) + row.update(vid=vid.rsplit('.', 1)[-1], level=level, + n_bone=int(m.sum()), sec=round(time.time() - t0, 2)) + return row + + +def parse_args(argv): + use_full = None + pos = [] + for a in argv: + if a == '--full': + use_full = True + elif a == '--quick': + use_full = False + else: + pos.append(a) + if len(pos) > 2: + sys.exit('Usage: python xfr_check_spinous.py [--full|--quick] [volume_id] [level]') + vid_arg = pos[0] if len(pos) >= 1 else None + level_arg = pos[1] if len(pos) >= 2 else None + if level_arg and level_arg.upper() not in LEVELS: + sys.exit(f'Invalid level: {level_arg} (choose from {"/".join(LEVELS)})') + if level_arg and not vid_arg: + sys.exit('level requires volume_id') + if use_full is None: + use_full = bool(vid_arg and level_arg) # 指定單 level → 精確模式 + return vid_arg, (level_arg.upper() if level_arg else None), use_full + + +def main(): + os.makedirs(LOG_DIR, exist_ok=True) + vid_arg, level_arg, use_full = parse_args(sys.argv[1:]) + + volumes = [d for d in sorted(os.listdir(standardized_dir)) + if os.path.isdir(os.path.join(standardized_dir, d))] + if vid_arg is not None: + key = vid_arg.lower() + vols = [v for v in volumes + if v.lower() == key or v.rsplit('.', 1)[-1] == key] + if not vols: + sys.exit(f'Volume not found: {vid_arg}') + volumes = vols + levels = (level_arg,) if level_arg else LEVELS + tasks = [(v, lvl) for v in volumes for lvl in levels] + mode = ('FULL (mirror plane)' if use_full + else 'QUICK screen -> FULL-verify candidates (2-stage)') + print(f'Mode: {mode} | {len(volumes)} volumes x {len(levels)} levels = {len(tasks)} tasks', + flush=True) + + t0 = time.time() + rows = [] + candidates = [] # (row 索引, volume 目錄名) + for i, (vid, lvl) in enumerate(tasks, 1): + try: + row = check_level(vid, lvl, use_full) + except Exception as e: + row = {'vid': vid.rsplit('.', 1)[-1], 'level': lvl, 'note': f'error: {e}'} + rows.append(row) + if 'note' in row: + print(f"[{i}/{len(tasks)}] {row['vid']} {row['level']}: {row['note']}", flush=True) + elif use_full: + top_str = 'n/a' if row['top_off'] is None else f"{row['top_off']:.1f}" + print(f"[{i}/{len(tasks)}] {row['vid']} {row['level']}: " + f"deficit={row['deficit']:6.1f} rear3={row['rear3']} " + f"top_off={top_str} {'<<< NO-SPINUS' if row['no_spinous'] else 'ok'} " + f"sp={row['sp_mode']} {row['sp_pct']:.1f}% " + f"vb={row['vb_mode']} {row['vb_pct']:.1f}%", flush=True) + else: + mark = 'candidate' if row['no_spinous'] else 'ok' + print(f"[{i}/{len(tasks)}] {row['vid']} {row['level']}: " + f"deficit={row['deficit']:6.1f} rear3={row['rear3']:4d} " + f"rear6={row['rear6']:5d} top_off={row['top_off']:5.1f} {mark}", + flush=True) + if not use_full and row.get('no_spinous'): + candidates.append((len(rows) - 1, vid)) + + # ---- Stage 2（quick 模式）：候選用 full（鏡稱面）驗證，以 full 為準 ---- + flagged = [] + cleared = [] + if use_full: + for r in rows: + if r.get('no_spinous'): + flagged.append(f"{r['vid']} {r['level']}") + elif candidates: + print(f'\n=== stage 2: {len(candidates)} quick candidate(s) -> ' + f'full (mirror plane) verification ===', flush=True) + for j, (k, vid_dir) in enumerate(candidates, 1): + row = rows[k] + try: + m = load_binary(os.path.join(standardized_dir, vid_dir), row['level']) + frow = full_check(m) if m is not None else {} + except Exception as e: + frow = {'note': f'verify error: {e}'} + for kk in ('deficit', 'rear3', 'rear6', 'top_off', 'no_spinous', + 'sp_mode', 'sp_pct', 'vb_mode', 'vb_pct', 'vb_ap_th'): + if kk in frow: + row[kk] = frow[kk] + if not frow: + flag_str = '-> UNVERIFIED (mask missing/full failed)' + cleared.append(f"{row['vid']} {row['level']} (UNVERIFIED)") + elif row.get('no_spinous'): + flag_str = '<<< NO-SPINUS (confirmed)' + flagged.append(f"{row['vid']} {row['level']}") + else: + flag_str = '-> normal (quick false alarm)' + cleared.append(f"{row['vid']} {row['level']} " + f"(full deficit={row.get('deficit', float('nan')):.1f})") + to = row.get('top_off') + print(f"[verify {j}/{len(candidates)}] {row['vid']} {row['level']}: " + f"deficit={row.get('deficit', float('nan')):6.1f} " + f"rear3={row.get('rear3', '?')} " + f"top_off={'n/a' if to is None else format(to, '5.1f')} " + f"sp={row.get('sp_mode', '?')} " + f"vb={row.get('vb_mode', '?')} {row.get('vb_pct') or 0:.1f}% {flag_str}", + flush=True) + + # ---- CSV ---- + ts = time.strftime('%Y%m%d_%H%M%S') + fields = ['vid', 'level', 'n_bone', 'deficit', 'rear3', 'rear6', 'top_off', + 'no_spinous', 'sp_mode', 'sp_pct', 'vb_mode', 'vb_pct', 'vb_ap_th', + 'sec', 'note'] + csv_path = os.path.join(LOG_DIR, f'spinous_check_{ts}.csv') + with open(csv_path, 'w', newline='') as f: + wtr = csv.DictWriter(f, fieldnames=fields, extrasaction='ignore') + wtr.writeheader() + for r in rows: + r['vid'] = r.get('vid', '') + wtr.writerow(r) + + dt = (time.time() - t0) / 60.0 + print('\n' + '=' * 64) + print(f'Done in {dt:.1f} min | {len(rows)} levels checked (mode: {mode})') + if flagged: + print(f'\n>>> {len(flagged)} level(s) with NO SPINOUS PROCESS (prior laminectomy/resection):') + for f_ in flagged: + print(f' {f_}') + print(' (xfr_debug.py 執行這些 level 時會打 [NO-SP] 並自動改用放寬後側谷底的椎體切分)') + else: + print('\nNo missing spinous process detected.') + if cleared: + print(f' (quick 候選、經 full 驗證為正常: {", ".join(cleared)})') + if use_full: + quant = [f"{r['vid']} {r['level']}" for r in rows + if r.get('vb_mode') == 'quantile'] + if quant: + print(f'\nWARNING: VBODY 退回 55 百分位切分（可能切進椎體）: {", ".join(quant)}') + print(f'CSV: {csv_path}') + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/xfr_debug.py b/xfr_debug.py index 581dfe4..9e04e4f 100644 --- a/xfr_debug.py +++ b/xfr_debug.py @@ -1,3 +1,6 @@ +#!/home/xfr/.conda/envs/cbt/bin/python + +import logging import os import re import sys @@ -15,7 +18,7 @@ from core.objective import set_global_context from core.optimizer import run_pso_torch, run_de_torch, run_nm_torch, run_pso_torch_xfr from imaging.orientation import azimuth_rotation, analyze_vertebral_tilt_contour -standardized_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr/' +standardized_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3/' azimuth_rotation_dir = '/mnt/1248/open2/cyrou/azimuth_rotation' tilt_contour_dir = '/mnt/1248/open2/cyrou/tilt_contour' @@ -25,6 +28,12 @@ LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5') LOG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logs') +# 單側任務的中繼結果（//_.json + plot lock）： +# 同一 level 的 L/R 可跑在不同 GPU，較晚完成的一側讀到兩側結果後跑合併輸出 +SIDE_RESULT_DIR = os.path.join(LOG_DIR, 'side_results') + +logger = logging.getLogger('xfr_debug') + # 目前任務標記（每個 worker 流程各自一份），讓交錯的 log 可以歸屬到 (volume, level, side) _TASK_TAG = {'vid': None, 'level': None, 'side': None} @@ -32,10 +41,10 @@ _TASK_TAG = {'vid': None, 'level': None, 'side': None} _NP_WRAP_RE = re.compile(r'\bnp\.[A-Za-z_][A-Za-z0-9_]*\(([^()]*)\)') -def set_task_tag(volume_id, level): +def set_task_tag(volume_id, level, side=None): _TASK_TAG['vid'] = volume_id _TASK_TAG['level'] = level - _TASK_TAG['side'] = None + _TASK_TAG['side'] = side class _Tee: @@ -123,6 +132,15 @@ def setup_tee(log_path): sys.stderr = _Tee(sys.stderr, log_fh) +def _setup_logging(): + """须在 setup_tee 之後呼叫，讓 handler 寫入 Tee（同時進 console 與 log 檔）""" + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + + def get_device(gpu_id=None): if torch.cuda.is_available(): if gpu_id is None: @@ -135,10 +153,10 @@ def get_device(gpu_id=None): max_free = free_mem gpu_id = i device = torch.device(f"cuda:{gpu_id}") - print(f"Using GPU {gpu_id}: {torch.cuda.get_device_name(gpu_id)}") + logger.info(f"Using GPU {gpu_id}: {torch.cuda.get_device_name(gpu_id)}") else: device = torch.device("cpu") - print("CUDA not available, using CPU") + logger.info("CUDA not available, using CPU") return device def debug_orientation(volume_id, level): @@ -146,9 +164,10 @@ def debug_orientation(volume_id, level): volume_dir = os.path.join(standardized_dir, volume_id) - cortical_path = os.path.join(volume_dir, f'{level}_cortical.nii.gz') - binary_path = os.path.join(volume_dir, f'{level}_binary.nii.gz') - roi_path = os.path.join(volume_dir, f'{level}_roi2.nii.gz') + # _binary.nii.gz 現為原解析度；0.5mm 用 _binary_sdf（SDF 平滑遮罩） + sdf = os.path.join(volume_dir, f'{level}_binary_sdf.nii.gz') + binary_path = sdf if os.path.exists(sdf) \ + else os.path.join(volume_dir, f'{level}_binary.nii.gz') # azi = azimuth_rotation(binary_path) # res = analyze_vertebral_tilt_contour(binary_path, edge_type='superior', show_plot=False, debug=False) @@ -156,11 +175,11 @@ def debug_orientation(volume_id, level): res = analyze_vertebral_tilt_contour(binary_path, edge_type='superior', show_plot=True, debug=False, save_plt=True, output_path=f'{tilt_contour_dir}/{level}_{volume_id}.png') alt = res['superior']['tilt_angle_deg'] - print(binary_path) + logger.info(binary_path) # print(f'Azimuth: {azi}, Alt: {alt}') - print(f'Alt: {alt}') + logger.info(f'Alt: {alt}') -def debug_pso(volume_id, level, device=None): +def debug_pso(volume_id, level, device=None, side='both', run_id=None): # ====== PSO ====== swarm_size = 100 max_iter = 100 @@ -174,10 +193,16 @@ def debug_pso(volume_id, level, device=None): CBT = True volume_dir = os.path.join(standardized_dir, volume_id) + # _cortical 現只產出旋轉版（rotated/ 子資料夾）；_roi2 不再存檔， + # 改用旋轉後 _roi（CT）。 + rotated_dir = os.path.join(volume_dir, 'rotated') + cortical_path = os.path.join(rotated_dir, f'{level}_cortical.nii.gz') + binary_path = os.path.join(rotated_dir, f'{level}_binary_sdf.nii.gz') + roi_path = os.path.join(rotated_dir, f'{level}_roi.nii.gz') - cortical_path = os.path.join(volume_dir, f'{level}_cortical.nii.gz') - binary_path = os.path.join(volume_dir, f'{level}_binary.nii.gz') - roi_path = os.path.join(volume_dir, f'{level}_roi2.nii.gz') + missing = [p for p in (cortical_path, binary_path, roi_path) if not os.path.exists(p)] + if missing: + raise ValueError(f'{volume_id} {level}: missing {missing}') cortical_image = sitk.ReadImage(cortical_path) binary_image = sitk.ReadImage(binary_path) @@ -222,6 +247,11 @@ def debug_pso(volume_id, level, device=None): device=device, optimize_size=True, grid=grid, + side=side, + level=level, + patient_id=volume_id, + side_dir=SIDE_RESULT_DIR, + run_id=run_id, ) # exit() @@ -239,43 +269,62 @@ def list_gpu_ids(): return [] -def gpu_worker(gpu_id, log_path, task_queue, result_queue): +def gpu_worker(gpu_id, log_path, run_id, task_queue, result_queue): """每張 GPU 一個工作流程：先鎖死該 GPU，再從共享隊列領 (volume, level) 任務""" # worker 是獨立流程，要自己把輸出 tee 到 log 檔 setup_tee(log_path) + _setup_logging() # 必須在任何 torch.cuda 呼叫前設定 os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id) torch.cuda.set_device(0) device = torch.device('cuda:0') - print(f'=== [GPU {gpu_id}] worker started ===', flush=True) + logger.info(f'=== [GPU {gpu_id}] worker started ===') while True: # 注意：multiprocessing.Queue 沒有 task_done()（只有 queue.Queue 有），別加回來 item = task_queue.get() if item is None: break - volume_id, level = item - set_task_tag(volume_id, level) + volume_id, level, side = item + # tag 沿用 _Tee 的 LEFT/RIGHT 寫法（'L'/'R' 是 optimizer 的 side 值） + set_task_tag(volume_id, level, 'LEFT' if side == 'L' else 'RIGHT') try: - debug_pso(volume_id, level, device) - result_queue.put(('task', gpu_id, volume_id, level, True, '')) + debug_pso(volume_id, level, device, side=side, run_id=run_id) + result_queue.put(('task', gpu_id, volume_id, level, side, True, '')) except Exception as e: - print(f'[GPU {gpu_id}] Error in {volume_id} {level}: {e}', flush=True) - result_queue.put(('task', gpu_id, volume_id, level, False, str(e))) + logger.error(f'[GPU {gpu_id}] Error in {volume_id} {level} {side}: {e}') + result_queue.put(('task', gpu_id, volume_id, level, side, False, str(e))) result_queue.put(('done', gpu_id)) - print(f'=== [GPU {gpu_id}] worker finished ===', flush=True) + logger.info(f'=== [GPU {gpu_id}] worker finished ===') -def _run_sequential(tasks): +def _run_sequential(tasks, run_id): """沒有（或只有一張）GPU 時的回退：單流程串行""" device = get_device() - for volume_id, level in tasks: - set_task_tag(volume_id, level) + for volume_id, level, side in tasks: + set_task_tag(volume_id, level, 'LEFT' if side == 'L' else 'RIGHT') try: - debug_pso(volume_id, level, device) + debug_pso(volume_id, level, device, side=side, run_id=run_id) except Exception as e: - print(f'Error in {volume_id} {level}: {e}') + logger.error(f'Error in {volume_id} {level} {side}: {e}') + + +USAGE = 'Usage: python xfr_debug.py [volume_id] [level]' + + +def parse_args(argv): + """volume_id 可用完整 ID 或末段（如 0005）；level 為 LEVELS 之一（L1~L5）。 + 兩者可省略（＝全部）；level 必須搭配 volume_id 使用。""" + vid_arg = argv[0] if len(argv) >= 1 else None + level_arg = argv[1] if len(argv) >= 2 else None + if len(argv) > 2: + sys.exit(f'{USAGE}\nToo many arguments') + if level_arg and level_arg.upper() not in LEVELS: + sys.exit(f'{USAGE}\nInvalid level: {level_arg} (choose from {"/".join(LEVELS)})') + if level_arg and not vid_arg: + sys.exit(f'{USAGE}\nlevel requires volume_id') + return vid_arg, (level_arg.upper() if level_arg else None) def main(): @@ -289,31 +338,48 @@ def main(): # 要處理的 volume 數（並行模式下是「最多嘗試的 volume 數」） MAX_SUCCESSFUL_VOLUMES = 100 - MAX_SUCCESSFUL_VOLUMES = 1 + MAX_SUCCESSFUL_VOLUMES = 10 # log 檔（console 與檔案同時輸出；各 GPU worker 也會 append 進同一個檔） os.makedirs(LOG_DIR, exist_ok=True) - log_path = os.path.join(LOG_DIR, f'xfr_debug_{time.strftime("%Y%m%d_%H%M%S")}.log') + run_id = time.strftime("%Y%m%d_%H%M%S") + log_path = os.path.join(LOG_DIR, f'xfr_debug_{run_id}.log') setup_tee(log_path) - print(f'Log file: {log_path}', flush=True) - print(f'Command: {sys.executable} {" ".join(sys.argv)}', flush=True) - print(f'Working directory: {os.getcwd()}', flush=True) + _setup_logging() + logger.info(f'Log file: {log_path}') + logger.info(f'Command: {sys.executable} {" ".join(sys.argv)}') + logger.info(f'Working directory: {os.getcwd()}') volumes = [d for d in sorted(os.listdir(standardized_dir)) if os.path.isdir(os.path.join(standardized_dir, d))] + + vid_arg, level_arg = parse_args(sys.argv[1:]) + if vid_arg is not None: + key = vid_arg.lower() + vols = [v for v in volumes + if v.lower() == key or v.rsplit('.', 1)[-1] == key] + if not vols: + sys.exit(f'Volume not found: {vid_arg}') + volumes = vols volumes = volumes[:MAX_SUCCESSFUL_VOLUMES] - tasks = [(vid, level) for vid in volumes for level in LEVELS] - print(f'Total {len(volumes)} volumes / {len(tasks)} (volume, level) tasks', flush=True) + levels = (level_arg,) if level_arg else LEVELS + # 任務粒度 = (volume, level, side)：同一 level 的 L/R 是兩個獨立任務， + # 可被不同 GPU 的 worker 領走並行執行；順序 L1 L -> L1 R -> L2 L -> L2 R -> ... + tasks = [(vid, level, s) for vid in volumes for level in levels for s in ('L', 'R')] + logger.info(f'Total {len(volumes)} volumes / {len(tasks)} (volume, level, side) tasks ' + f'(levels: {", ".join(levels)})') + if vid_arg or level_arg: + logger.info(f'Filter: volume_id={vid_arg!r} level={level_arg!r}') gpu_ids = list_gpu_ids() if len(gpu_ids) <= 1: - print(f'Only {len(gpu_ids)} GPU(s) available, running sequentially', flush=True) - _run_sequential(tasks) + logger.info(f'Only {len(gpu_ids)} GPU(s) available, running sequentially') + _run_sequential(tasks, run_id) return - print(f'Found {len(gpu_ids)} GPUs: {gpu_ids}, starting {len(gpu_ids)} workers (one per GPU)', flush=True) + logger.info(f'Found {len(gpu_ids)} GPUs: {gpu_ids}, starting {len(gpu_ids)} workers (one per GPU)') ctx = mp.get_context('spawn') task_queue = ctx.Queue() @@ -323,8 +389,8 @@ def main(): for _ in gpu_ids: task_queue.put(None) # 每個 worker 一個結束哨兵 - procs = [ctx.Process(target=gpu_worker, args=(g, log_path, task_queue, result_queue), name=f'cbt-gpu-{g}') - for g in gpu_ids] + procs = [ctx.Process(target=gpu_worker, args=(g, log_path, run_id, task_queue, result_queue), name=f'cbt-gpu-{g}') + for g in gpu_ids] for p in procs: p.start() @@ -333,7 +399,7 @@ def main(): finished = 0 while finished < len(gpu_ids): if not any(p.is_alive() for p in procs): - print('Warning: a worker exited early; reaping remaining tasks...', flush=True) + logger.warning('A worker exited early; reaping remaining tasks...') break try: msg = result_queue.get(timeout=5) @@ -357,27 +423,31 @@ def main(): p.join(timeout=60) total_time = time.time() - start_time - ok = [r for r in results if r[4]] - fail = [r for r in results if not r[4]] + ok = [r for r in results if r[5]] + fail = [r for r in results if not r[5]] per_volume = {} - for _, _, vid, level, success, _ in results: + per_volume_n = {} + for _, _, vid, level, side, success, _ in results: per_volume.setdefault(vid, set()).add(success) + per_volume_n[vid] = per_volume_n.get(vid, 0) + 1 missing = len(tasks) - len(results) - # 一個 volume 算「成功」必須它的(level)全部執行過且全部成功 - n_success_volumes = sum(1 for vid in volumes - if per_volume.get(vid, set()) == {True}) + # 一個 volume 算「成功」必須它的所有 (level, side) 任務都執行過且全部成功 + n_success_volumes = sum( + 1 for vid in volumes + if per_volume.get(vid, set()) == {True} + and per_volume_n.get(vid, 0) == len(levels) * 2) print('=' * 60) - print(f'Finished in {total_time / 60:.1f} min | ' - f'tasks {len(results)}/{len(tasks)} (ok {len(ok)} / failed {len(fail)} / not-run {missing})') + logger.info(f'Finished in {total_time / 60:.1f} min | ' + f'tasks {len(results)}/{len(tasks)} (ok {len(ok)} / failed {len(fail)} / not-run {missing})') if missing: - print(f'Warning: {missing} task(s) were never executed (worker crash?)') - print(f'Successful volumes (all levels OK): {n_success_volumes}/{len(volumes)}') + logger.warning(f'{missing} task(s) were never executed (worker crash?)') + logger.info(f'Successful volumes (all levels OK): {n_success_volumes}/{len(volumes)}') if fail: - print('Failed tasks:') - for _, g, vid, level, _, err in fail: - print(f' [GPU {g}] {vid} {level}: {err}') + logger.info('Failed tasks:') + for _, g, vid, level, side, _, err in fail: + logger.error(f'[GPU {g}] {vid} {level} {side}: {err}') if __name__ == '__main__': diff --git a/xfr_plot_level.py b/xfr_plot_level.py new file mode 100644 index 0000000..588268f --- /dev/null +++ b/xfr_plot_level.py @@ -0,0 +1,97 @@ +#!/home/xfr/.conda/envs/cbt/bin/python +""" +為每個 (volume, level) 的骨頭遮罩（{level}_binary_sdf.nii.gz，缺則 _binary）繪製 +X-ray 四視角圖（不畫螺絲），存到 Output/{date}/{volume}/。 + +輸出檔名：{volume} {level}_CBT.png +（例：1.3.6.1.4.1.9328.50.4.0005 L1_CBT.png） + +Usage: + python xfr_plot_level.py # 全部 volume 的 L1~L5 + python xfr_plot_level.py 0005 # 該 volume 的 L1~L5 + python xfr_plot_level.py 0005 L1 # 單一 (volume, level) + python xfr_plot_level.py --dir --output +""" + +import argparse +import os +import sys + +from visualization.res_bone_figure import render_bone_figure + +standardized_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr/' +output_base = '/mnt/1248/open2/cyrou/Output' + +LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5') + +USAGE = 'Usage: python xfr_plot_level.py [volume_id] [level] [--dir D] [--output O]' + + +def parse_args(argv): + parser = argparse.ArgumentParser(description='繪製各 lumbar level 的骨頭 X-ray 圖（不畫螺絲）') + parser.add_argument('volume_id', nargs='?', default=None, + help='volume ID（完整 UID 或末段，如 0005）；省略＝全部') + parser.add_argument('level', nargs='?', default=None, + help=f'level（{" / ".join(LEVELS)}）；省略＝全部') + parser.add_argument('--dir', default=standardized_dir, + help=f'standardized 資料夾（預設 {standardized_dir}）') + parser.add_argument('--output', default=output_base, + help=f'輸出根目錄（預設 {output_base}）') + return parser.parse_args(argv) + + +def main(): + args = parse_args(sys.argv[1:]) + + if args.level is not None: + level_key = args.level.upper() + if level_key not in LEVELS: + print(f'Invalid level: {args.level} (choose from {"/".join(LEVELS)})') + sys.exit(1) + if args.level is not None and args.volume_id is None: + print(f'{USAGE}\nlevel requires volume_id') + sys.exit(1) + + volumes = [d for d in sorted(os.listdir(args.dir)) + if os.path.isdir(os.path.join(args.dir, d))] + if args.volume_id is not None: + key = args.volume_id.lower() + vols = [v for v in volumes + if v.lower() == key or v.rsplit('.', 1)[-1] == key] + if not vols: + print(f'Volume not found: {args.volume_id}') + sys.exit(1) + volumes = vols + + levels = (args.level.upper(),) if args.level else LEVELS + tasks = [(vid, lvl) for vid in volumes for lvl in levels] + print(f'{len(volumes)} volume(s) x {len(levels)} level(s) = {len(tasks)} figure(s)', + flush=True) + + ok, skip = [], [] + for i, (vid, lvl) in enumerate(tasks, 1): + vol_dir = os.path.join(args.dir, vid) + # _binary.nii.gz 現為原解析度；0.5mm 用 _binary_sdf（SDF 平滑遮罩） + sdf_path = os.path.join(vol_dir, f'{lvl}_binary_sdf.nii.gz') + binary_path = sdf_path if os.path.exists(sdf_path) \ + else os.path.join(vol_dir, f'{lvl}_binary.nii.gz') + cortical_path = os.path.join(vol_dir, f'{lvl}_cortical.nii.gz') + path = render_bone_figure(vid, lvl, binary_path, cortical_path, + base_folder=args.output) + if path is None: + skip.append((vid, lvl)) + print(f'[{i}/{len(tasks)}] {vid} {lvl}: skipped', flush=True) + else: + ok.append((vid, lvl)) + print(f'[{i}/{len(tasks)}] {vid} {lvl}: saved {path}', flush=True) + + print('=' * 60) + print(f'Done. saved={len(ok)} skipped={len(skip)}') + if skip: + print('Skipped (missing/empty mask):') + for vid, lvl in skip: + print(f' {vid} {lvl}') + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/xfr_preprocess.py b/xfr_preprocess.py index 32c6b5c..73e50b6 100644 --- a/xfr_preprocess.py +++ b/xfr_preprocess.py @@ -1,10 +1,32 @@ +#!/home/xfr/.conda/envs/cbt/bin/python + +import argparse +import json +import logging import os +import sys +import time + +import numpy as np +import SimpleITK as sitk +from tinydb import TinyDB, Query from imaging.preprocessing import process_dataset +from config.constant import LABEL_MAP +from imaging.orientation import (best_symmetry_plane, best_upper_endplate_plane, + segment_spinous_process, segment_vertebral_body, + smooth_mask_sdf) +from visualization.res_bone_figure import (render_bone_figure, + compute_normalizing_rotation, + rotate_volume_to, rotated_grid, + rotated_sitk_image_at, + _rotate_plane_params, + _shift_plane_params) data_root = '/mnt/1220/Public/dataset/Spine/CTSpine1K/data/' label_root = '/mnt/1220/Public/dataset/Spine/CTSpine1K/label/' output_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-2/' +output_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3/' label_map = { 'colon': 'conlon', @@ -13,12 +35,407 @@ label_map = { 'liver': 'Liver', } -def main(): +LUMBAR_LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5', 'L6') +MAX_Z_SPACING_MM = 4.0 # z spacing (mm) 大於此值者跳過整支 pipeline +MIN_LUMBAR_LEVELS = 2 # 符合 lumbar 的 label 少於此值者跳過整支 pipeline + +LOG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logs') +METADATA_DB = os.path.join(os.path.dirname(os.path.abspath(__file__)), + 'xfr_image_metadata.json') + +logger = logging.getLogger('xfr_preprocess') + + +def _upsert_by_name(table, name, meta): + """以 'name' 欄位為鍵 upsert（TinyDB 4.9 insert 不支援自訂 doc_id）。""" + q = Query().name == name + entry = dict(table.get(q) or {}) + entry.update(meta) + entry['name'] = name + if table.get(q) is None: + table.insert(entry) + else: + table.update(entry, q) + + +def _migrate_legacy_metadata(path): + """舊版純 JSON 檔（{name: {...}}、無 TinyDB 的 _default 結構）： + 改名成 .legacy- 備份，內容匯入新 TinyDB。""" + if not os.path.exists(path): + return + try: + with open(path) as f: + data = json.load(f) + except (json.JSONDecodeError, OSError): + return + legacy = {k: v for k, v in data.items() + if isinstance(v, dict)} if isinstance(data, dict) else {} + if not (isinstance(data, dict) and '_default' not in data and legacy): + return + bak = f'{path}.legacy-{time.strftime("%Y%m%d_%H%M%S")}' + os.replace(path, bak) + print(f'Metadata db {path} is old pure-JSON format; renamed to {bak} ' + f'and migrating {len(legacy)} entrie(s)') + t = TinyDB(path).table('images') + for k, v in legacy.items(): + _upsert_by_name(t, k, dict(v)) + + +class ImageMetadataDB: + """影像 metadata 快取（TinyDB，table 'images'）： + 以影像檔名前綴（= 輸出資料夾名，如 1.3.6.1.4.1.9328.50.4.0001、 + liver_100、volume-covid19-A-0011_ct）為 key（存於文件 'name' 欄位）， + 記錄 {'spacing': [x, y, z], 'labels': [label id]}。 + + 供 process_single_image 做跳過判定（z spacing / lumbar 層數）， + 命中時不必讀影像 / label 檔。TinyDB 每次操作自動落檔， + 中途重跑可繼承。 + """ + + def __init__(self, path=METADATA_DB): + self.path = path + _migrate_legacy_metadata(path) + self._table = TinyDB(path).table('images') + + def get(self, name): + return self._table.get(Query().name == name) + + def put(self, name, meta): + _upsert_by_name(self._table, name, meta) + + def __len__(self): + return len(self._table) + + +class _Tee: + """同時輸出到 console 與 log 檔，逐行寫入。""" + + def __init__(self, console, log_fh): + self.console = console + self.log_fh = log_fh + self.buf = '' + + def _emit(self, line): + self.console.write(line + '\n') + self.log_fh.write(line + '\n') + + def write(self, data): + if not data: + return + self.buf += data + while True: + idx_n = self.buf.find('\n') + idx_r = self.buf.find('\r') + candidates = [i for i in (idx_n, idx_r) if i != -1] + if not candidates: + break + idx = min(candidates) + self._emit(self.buf[:idx]) + self.buf = self.buf[idx + 1:] + + def flush(self): + if self.buf: + self._emit(self.buf) + self.buf = '' + self.console.flush() + self.log_fh.flush() + + def isatty(self): + return False + + +def setup_tee(log_path): + """把這個流程的 stdout/stderr 同時寫到 log_path（append、line-buffered）""" + log_fh = open(log_path, 'a', buffering=1) + sys.stdout = _Tee(sys.stdout, log_fh) + sys.stderr = _Tee(sys.stderr, log_fh) + + +def _cortical_from_roi(roi_arr, bin_arr): + """皮質遮罩（與舊未旋轉 _cortical.nii.gz 同定義）： + 門檻 = 骨頭 mask（bin_arr）內 CT 的 median HU； + cortical = mask 內且 門檻 <= HU <= 10000。 + roi_arr / bin_arr 須同 grid（(z, y, x)）；回傳 uint8 0/1， + 輸入缺失或 mask 為空時回傳 None。""" + if roi_arr is None or bin_arr is None: + return None + inside = bin_arr > 0 + if int(inside.sum()) == 0: + return None + threshold = float(np.median(roi_arr[inside])) + return (inside & (roi_arr >= threshold) & (roi_arr <= 10000)).astype(np.uint8) + + +def _write_rotated_level(vol_dir, level, smd_path, mask_path, roi_path): + """對單一 level：計算對齊旋轉（鏡稱面 -> +X、上終板 normal y=0）， + 存到 /rotated/（先旋轉 _binary_nn，再以它的非零 bbox（±4 + 體素）作為所有旋轉輸出的裁切 grid，尺寸不求同未旋轉版、整顆骨頭 + 保留，origin 隨裁切平移）： + - _smd_resampled.nii.gz 三线性旋轉後 SMD（內負/外正） + - _binary_sdf.nii.gz 旋轉 SMD 於 0.5 閾值（與未旋轉版一致） + - _binary_nn.nii.gz NN 旋轉後 0/1 遮罩（對比用，裁切基準） + - _roi.nii.gz + - _cortical.nii.gz 旋轉 CT 以骨頭 mask 內 median HU 為門檻 + （取代舊的未旋轉 _cortical，定義相同） + 再畫「rotated」平面圖，並在旋轉體上做 VBODY / 棘突分割存 label map + （1=VBODY、2=棘突、3=other bone、0=background）。 + 平面參數經 R 剛性旋轉並換元到輸出 grid 的局部座標。 + + 旋轉來源 + 幾何模板：patient 未旋轉的 0.5mm _smd_resampled.nii.gz + （float、裁物件 bbox，SMD < 0.5 = 內部，與 _binary_sdf 同幾何）； + 缺失或格式不符時回退 _binary_nn.nii.gz（0.5mm 最近邻 0/1）， + 此情形只產出 _binary_nn（+ roi / cortical / 圖 / label）。 + mask_path（未旋轉 0.5mm 0/1 遮罩，如 _binary_sdf）只供畫圖 + （render_bone_figure 以 uint8 讀取；float 檔會被讀成全 0）， + 缺失時自動回退 _binary_nn。 + """ + volume_id = os.path.basename(vol_dir) + + template = None + smd_arr = None + if smd_path is not None and os.path.exists(smd_path): + sa = sitk.GetArrayFromImage(sitk.ReadImage(smd_path)).astype(np.float32) + if sa.size > 0 and np.isfinite(sa).all() and sa.min() < 0.0 and sa.max() > 0.0: + template = sitk.ReadImage(smd_path) # 0.5mm 幾何模板 + smd_arr = sa + else: + logger.warning(f'[rotated] {volume_id} {level}: _smd_resampled 不可用 ' + f'(value)，改用 _binary_nn') + if template is None: + nn_in_path = os.path.join(vol_dir, f'{level}_binary_nn.nii.gz') + if not os.path.exists(nn_in_path): + logger.warning(f'[rotated] {volume_id} {level}: _smd_resampled / _binary_nn 皆不可用, skip') + return + template = sitk.ReadImage(nn_in_path) + + # 平面估計輸入：SMD 時取 0.5 閾值內部，與 _binary_sdf 同定義 + if smd_arr is not None: + bin_arr = (smd_arr < 0.5).astype(np.uint8) + else: + bin_arr = (sitk.GetArrayFromImage(template) > 0).astype(np.uint8) + if int(bin_arr.sum()) == 0: + logger.warning(f'[rotated] {volume_id} {level}: empty binary, skip') + return + + # 圖用遮罩須與模板同幾何（SDF 與 SMD 皆裁 bbox2；NN fallback 只有 + # NN 遮罩一致）；0.5mm 0/1，float 檔會被 uint8 讀取成全 0 + nn_in_path = os.path.join(vol_dir, f'{level}_binary_nn.nii.gz') + if smd_arr is None: + mask_path = nn_in_path if os.path.exists(nn_in_path) else mask_path + if mask_path is None or not os.path.exists(mask_path): + mask_path = nn_in_path if os.path.exists(nn_in_path) else None + + sym = best_symmetry_plane(bin_arr) + symp = best_upper_endplate_plane(bin_arr) + R, c_xyz = compute_normalizing_rotation(bin_arr, sym, symp) + + # 輸出 grid：未旋轉緊密 bbox 的 8 角點旋轉後的最小包圍盒 + margin—— + # 尺寸不受未旋轉版約束，保證旋轉後整顆骨頭保留（同尺寸旋轉切角落） + start, size = rotated_grid(bin_arr.shape, R, c_xyz, margin=4) + + rotated_dir = os.path.join(vol_dir, 'rotated') + os.makedirs(rotated_dir, exist_ok=True) + + # 先旋轉 _binary_nn（order=0，未旋轉 _binary_nn 先重採樣到模板 grid + # 再旋轉），以其非零 bbox（外扩 4 體素）定所有旋轉輸出的裁切 grid + # （fstart/fsize，模板 index 系）；NN 缺失 / 為空時回退緊密包圍盒 + rot_nn = None + fstart = (int(start[0]), int(start[1]), int(start[2])) + fsize = (int(size[0]), int(size[1]), int(size[2])) + if os.path.exists(nn_in_path): + nn_img = sitk.ReadImage(nn_in_path) + nn_arr = (sitk.GetArrayFromImage( + sitk.Resample(nn_img, template, interpolator=sitk.sitkNearestNeighbor, + defaultPixelValue=0)) > 0).astype(np.uint8) + rot_nn_full = (rotate_volume_to(nn_arr, R, c_xyz, start, size, + order=0, cval=0.0) > 0.5).astype(np.uint8) + p_nn = os.path.join(rotated_dir, f'{level}_binary_nn.nii.gz') + if int(rot_nn_full.sum()) > 0: + zz, yy, xx = np.where(rot_nn_full > 0) + m = 4 + z0 = max(0, int(zz.min()) - m) + z1 = min(rot_nn_full.shape[0] - 1, int(zz.max()) + m) + y0 = max(0, int(yy.min()) - m) + y1 = min(rot_nn_full.shape[1] - 1, int(yy.max()) + m) + x0 = max(0, int(xx.min()) - m) + x1 = min(rot_nn_full.shape[2] - 1, int(xx.max()) + m) + fstart = (fstart[0] + x0, fstart[1] + y0, fstart[2] + z0) + fsize = (x1 - x0 + 1, y1 - y0 + 1, z1 - z0 + 1) + rot_nn = rot_nn_full[z0:z1 + 1, y0:y1 + 1, x0:x1 + 1] + else: + rot_nn = rot_nn_full + sitk.WriteImage(rotated_sitk_image_at(template, rot_nn, fstart), p_nn) + logger.info(f'[rotated] saved {p_nn} (grid={list(fsize)}, start={fstart})') + + # 旋轉後 SMD：未旋轉 _smd_resampled 的 SMD 場經三线性旋轉得以保留 + # （float32 來源，map_coordinates 輸出維持浮點，<0.5 即為真閾值）； + # 體外填充值取正 SMD（背景側），避免裁切邊界出現假的閾值穿越 + rot_sdf_bin = None + if smd_arr is not None: + rot_smd = rotate_volume_to(smd_arr, R, c_xyz, fstart, fsize, + order=1, cval=float(smd_arr.max())) + p_smd = os.path.join(rotated_dir, f'{level}_smd_resampled.nii.gz') + sitk.WriteImage(rotated_sitk_image_at(template, rot_smd.astype(np.float32), + fstart), p_smd) + logger.info(f'[rotated] saved {p_smd}') + + # _binary_sdf.nii.gz：旋轉 SMD 於 0.5 閾值（與未旋轉版一致） + rot_sdf_bin = (rot_smd < 0.5).astype(np.uint8) + p_sdf = os.path.join(rotated_dir, f'{level}_binary_sdf.nii.gz') + sitk.WriteImage(rotated_sitk_image_at(template, rot_sdf_bin, fstart), p_sdf) + logger.info(f'[rotated] saved {p_sdf}') + + # 旋轉後的 roi（三线性）：先重採樣到模板 grid，再旋轉到裁切 grid + roi_arr = None + rot_roi_arr = None + if roi_path is not None and os.path.exists(roi_path): + roi_img = sitk.ReadImage(roi_path) + roi_arr = sitk.GetArrayFromImage( + sitk.Resample(roi_img, template, defaultPixelValue=0.0)) + rot_roi = rotate_volume_to(roi_arr, R, c_xyz, fstart, fsize, + order=1, cval=0.0) + p_roi = os.path.join(rotated_dir, f'{level}_roi.nii.gz') + sitk.WriteImage(rotated_sitk_image_at(template, rot_roi, fstart), p_roi) + logger.info(f'[rotated] saved {p_roi}') + rot_roi_arr = rot_roi + + # 分割 / cortical 輸入：SMD 時直接用旋轉 _binary_sdf（0.5 閾值、 + # 次體素平滑）；NN fallback 時以 smooth_mask_sdf 記憶體內平滑（不存檔） + if rot_sdf_bin is not None: + rot_bin_seg = rot_sdf_bin + else: + rot_bin_seg = smooth_mask_sdf(rot_nn, sigma=1.5).astype(np.uint8) + + # 旋轉後的 cortical：旋轉 CT 以骨頭 mask 內 median HU 為門檻 + #（取代舊的未旋轉 _cortical.nii.gz，定義相同） + rot_cort = _cortical_from_roi(rot_roi_arr, rot_bin_seg) + if rot_cort is not None: + p_cort = os.path.join(rotated_dir, f'{level}_cortical.nii.gz') + sitk.WriteImage(rotated_sitk_image_at(template, rot_cort, fstart), p_cort) + logger.info(f'[rotated] saved {p_cort}') + else: + logger.warning(f'[rotated] {volume_id} {level}: 無旋轉 CT / mask，跳過 _cortical') + + # 用旋轉後的平面畫圖（rotated 版 planes）；皮質著色由未旋轉 CT + mask + # 現算（未旋轉 _cortical 不再存檔） + p_fig = os.path.join(rotated_dir, f'{level}_planes.png') + if mask_path is not None: + fig_cortical = _cortical_from_roi(roi_arr, bin_arr) + fig = render_bone_figure(volume_id, level, mask_path, fig_cortical, + planes_only=True, rotation=(R, c_xyz), output_path=p_fig) + if fig is not None: + logger.info(f'[rotated] saved {fig}') + else: + logger.warning(f'[rotated] {volume_id} {level}: 無可用 0/1 遮罩，跳過平面圖') + + # 旋轉體上的 VBODY / 棘突分割：平面參數隨 R 剛性旋轉（rotated 體與 + # 原始體只差一個剛性變換，平面隨動即可，不需重新偵測）， + # label map 存同資料夾：1 = VBODY（椎體），2 = 棘突（spinous process）， + # 3 = other bone（其餘骨頭），0 = background。 + # 換元到輸出 grid 的局部座標（origin 在模板 index fstart 處） + sym_rot = _shift_plane_params(_rotate_plane_params(sym, R, c_xyz), fstart) + symp_rot = (_shift_plane_params(_rotate_plane_params(symp, R, c_xyz), fstart) + if symp is not None else None) + sp_mask, sp_th, sp_info = segment_spinous_process(rot_bin_seg, sym_rot) + vb_mask, vb_th, vb_info = segment_vertebral_body(rot_bin_seg, sym_rot, symp_rot, + sp_th, sp_info['mode']) + label_arr = np.zeros(rot_bin_seg.shape, dtype=np.uint8) + label_arr[rot_bin_seg > 0] = 3 # other bone + if vb_mask is not None: + label_arr[vb_mask] = 1 # VBODY + if sp_mask is not None: + label_arr[sp_mask] = 2 # spinous process + p_lbl = os.path.join(rotated_dir, f'{level}_label.nii.gz') + sitk.WriteImage(rotated_sitk_image_at(template, label_arr, fstart), p_lbl) + logger.info(f'[rotated] saved {p_lbl} ' + f'(vbody={int((label_arr == 1).sum())}, spinous={int((label_arr == 2).sum())}, ' + f'other={int((label_arr == 3).sum())}, ' + f'sp_mode={sp_info["mode"]}, vb_mode={vb_info["mode"]})') + + +def make_lumbar_post_process(): + """每個 volume 處理完後，對其 lumbar level： + 1) 畫「骨頭 + 方向平面」圖（不畫螺絲、不做棘突 / 椎體分割）-> /lumbar/ + 2) 計算對齊旋轉，存旋轉後的 smd_resampled / binary_sdf / binary_nn / roi + + cortical + 旋轉平面圖 + label map -> /rotated/ + + post_process 由 process_dataset 呼叫：(volume_dir, processed_labels)。 + processed_labels 為該 volume 實際存在的 label id（int），對照 LABEL_MAP。 + """ + + def _post_process(vol_dir, processed_labels): + volume_id = os.path.basename(vol_dir) + lumbar_dir = os.path.join(vol_dir, 'lumbar') + for n in processed_labels: + level = LABEL_MAP.get(int(n)) + if level not in LUMBAR_LEVELS: + continue + smd_res_path = os.path.join(vol_dir, f'{level}_smd_resampled.nii.gz') + nn_path = os.path.join(vol_dir, f'{level}_binary_nn.nii.gz') + if not os.path.exists(smd_res_path) and not os.path.exists(nn_path): + continue + # 畫圖用 0/1 遮罩：優 _binary_sdf（0.5mm 平滑），缺則 _binary_nn + sdf_path = os.path.join(vol_dir, f'{level}_binary_sdf.nii.gz') + mask_path = sdf_path if os.path.exists(sdf_path) else nn_path + roi_path = os.path.join(vol_dir, f'{level}_roi.nii.gz') + + # 1) 原始（未旋轉）planes 圖；皮質著色由未旋轉 CT + mask 現算 + #（未旋轉 _cortical 不再存檔） + output_path = os.path.join(lumbar_dir, f'{level}_planes.png') + if os.path.exists(mask_path): + fig_cortical = None + if os.path.exists(roi_path): + mask_img = sitk.ReadImage(mask_path) + roi_img = sitk.ReadImage(roi_path) + roi_arr = sitk.GetArrayFromImage( + sitk.Resample(roi_img, mask_img, defaultPixelValue=0.0)) + fig_cortical = _cortical_from_roi( + roi_arr, sitk.GetArrayFromImage(mask_img)) + path = render_bone_figure(volume_id, level, mask_path, fig_cortical, + planes_only=True, output_path=output_path) + if path is not None: + logger.info(f'[lumbar] saved {path}') + + # 2) 旋轉對齊：rotated/ 的 smd_resampled + binary_sdf + binary_nn + # + roi + cortical + planes 圖 + label + _write_rotated_level(vol_dir, level, smd_res_path, mask_path, roi_path) + + return _post_process + + +def main(): + parser = argparse.ArgumentParser(description='Preprocess CT spine dataset.') + parser.add_argument('--max-images', type=int, default=None, dest='max_images', + help='Process at most this number of images per dataset (default: all).') + args = parser.parse_args() + + # log 檔（console 與檔案同時輸出） + os.makedirs(LOG_DIR, exist_ok=True) + log_path = os.path.join(LOG_DIR, f'xfr_preprocess_{time.strftime("%Y%m%d_%H%M%S")}.log') + setup_tee(log_path) + # basicConfig 在 setup_tee 之後，讓 handler 寫入 Tee（同時進 console 與 log 檔） + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + logger.info(f'Log file: {log_path}') + logger.info(f'Command: {sys.executable} {" ".join(sys.argv)}') + logger.info(f'Working directory: {os.getcwd()}') + + # metadata db：跳過判定（z spacing / lumbar 層數）命中時免讀影像 / label 檔 + metadata_db = ImageMetadataDB() + logger.info(f'Metadata db: {metadata_db.path} ({len(metadata_db)} entries)') + + post_process = make_lumbar_post_process() for key, value in label_map.items(): data_dir = os.path.join(data_root, key) label_dir = os.path.join(label_root, value) - process_dataset(data_dir, label_dir, output_dir) + process_dataset(data_dir, label_dir, output_dir, max_images=args.max_images, + post_process=post_process, max_z_spacing=MAX_Z_SPACING_MM, + allowed_levels=LUMBAR_LEVELS, min_levels=MIN_LUMBAR_LEVELS, + metadata_cache=metadata_db) if __name__ == '__main__': main() \ No newline at end of file