diff --git a/config/constant.py b/config/constant.py index db2016e..91f859d 100644 --- a/config/constant.py +++ b/config/constant.py @@ -29,7 +29,7 @@ ALLOWED_LENGTHS = [ 35, 40, 45, - 50, + # 50, # 60, # 70, # 80, diff --git a/core/cylinder.py b/core/cylinder.py index ee0db8b..9a3cfaf 100644 --- a/core/cylinder.py +++ b/core/cylinder.py @@ -189,11 +189,16 @@ def generate_cylinder_o_torch( shape: tuple[int, int, int], spacing: list[float], device: torch.device, - grid: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None + grid: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + butt_mm: float = 5.0 ) -> torch.Tensor: """ Generate an "opposite" (negative z) cylinder mask in 3D space using PyTorch tensors. Returns a binary mask (torch.uint8) on the specified device. + + butt_mm: 入口端（z_rot=0，螺絲 butt 面）附近前 butt_mm 不計入此 mask， + 讓 bone surface 附近允許最多 butt_mm 的 miss（入口端略陷骨面內 + 該帶 voxel 不再被 scoring 的 cyl_o∩bone 懲罰）。 """ if grid is None: @@ -228,14 +233,14 @@ def generate_cylinder_o_torch( radius = diameter / 2.0 mask = ( (x_rot**2 + y_rot**2 <= radius**2) - & (z_rot <= 0) - & (z_rot <= length) + & (z_rot <= -butt_mm) + & (z_rot >= -length) ) elif spacing == [0.5, 0.5, 0.5]: radius = (diameter / 2.0) * 2 mask = ( (x_rot**2 + y_rot**2 <= radius**2) - & (z_rot <= 0) + & (z_rot <= -butt_mm * 2) & (z_rot >= -length * 2) ) else: @@ -280,12 +285,14 @@ def generate_cylinder_butt_torch( position_z, position_y, position_x, azimuth, altitude, shape, spacing, device, grid=None, - butt_mm=0.25 # 入口端最後 butt_mm（mm） + butt_mm=5.0 # 入口端最後 butt_mm（mm）：but 豁免帶寬度 ) -> torch.Tensor: """生成圓柱「最後端」mask：入口端（z_rot=0，遠離 VBODY 的一端） 最靠近的 butt_mm 圓柱短柱（與 generate_cylinder_n_torch 同慣例、同直徑）。 尖端在 z_rot=length（靠 VBODY 端），不在此 mask。 - 此 mask 是完整圓柱的子集，可直接用於 not_in_bone 豁免。""" + 此 mask 是完整圓柱的子集，scoring 中 mask 內 not_in_bone 的 voxel + 豁免 10000/voxel 扣分與 HARD_BREACH_PER_VOX 硬約束（螺絲入口端允許 + 伸出骨外 butt_mm）。""" if grid is None: z_t, y_t, x_t = create_coordinate_grid(shape, device) @@ -330,6 +337,68 @@ def generate_cylinder_butt_torch( return mask.to(torch.uint8) +def generate_cone_torch( + diameter: float, + length: float, + position_z: float, + position_y: float, + position_x: float, + azimuth: float, + altitude: float, + shape: tuple[int, int, int], + spacing: list[float], + device: torch.device, + grid: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, +) -> torch.Tensor: + """螺絲圖用的錐體 mask（與 generate_cylinder_n_torch 同慣例/同座標系）： + 入口端（z_rot=0）半徑 = diameter/2，沿 +z 線性收細到尖端 + （z_rot=length，半徑=0）。純顯示用（render_bone_figure 的螺絲體）， + 評分維持圓柱、不經此 mask。""" + + if grid is None: + z_t, y_t, x_t = create_coordinate_grid(shape, device) + else: + z_t, y_t, x_t = grid + + azimuth_rad_t = torch.deg2rad(torch.tensor(azimuth, device=device, dtype=torch.float32)) + altitude_rad_t = torch.deg2rad(torch.tensor(altitude, device=device, dtype=torch.float32)) + + z_t = z_t - position_z + y_t = y_t - position_y + x_t = x_t - position_x + + x_rot = ( + x_t * torch.cos(azimuth_rad_t) * torch.cos(altitude_rad_t) + + y_t * torch.sin(azimuth_rad_t) * torch.cos(altitude_rad_t) + - z_t * torch.sin(altitude_rad_t) + ) + y_rot = -x_t * torch.sin(azimuth_rad_t) + y_t * torch.cos(azimuth_rad_t) + z_rot = ( + x_t * torch.cos(azimuth_rad_t) * torch.sin(altitude_rad_t) + + y_t * torch.sin(azimuth_rad_t) * torch.sin(altitude_rad_t) + + z_t * torch.cos(altitude_rad_t) + ) + + if spacing == [1, 1, 1]: + radius = diameter / 2.0 + total_len = length + elif spacing == [0.5, 0.5, 0.5]: + radius = (diameter / 2.0) * 2 + total_len = length * 2 + else: + raise ValueError(f"Unsupported spacing: {spacing}") + + taper = torch.clamp(z_rot / total_len, min=0.0, max=1.0) + r = radius * (1.0 - taper) + + mask = ( + (x_rot**2 + y_rot**2 <= r**2) + & (z_rot >= 0) + & (z_rot <= total_len) + ) + return mask.to(torch.uint8) + + def generate_cylinder_tip_torch( diameter, length, position_z, position_y, position_x, diff --git a/core/objective.py b/core/objective.py index b22ffed..f78cf78 100644 --- a/core/objective.py +++ b/core/objective.py @@ -120,7 +120,7 @@ def cylinder_circle_line_intersection_loss_deductions_torch( ctx.image2_shape, ctx.spacing, ctx.device, ctx.grid ) - # 最後端（入口端，遠離 VBODY 的一端）0.25mm 豁免 mask + # 最後端（入口端，遠離 VBODY 的一端）0.5mm 豁免 mask cyl_butt = generate_cylinder_butt_torch( diameter, position_z, position_y, position_x, diff --git a/core/optimizer.py b/core/optimizer.py index 66b4dea..d3dbe61 100644 --- a/core/optimizer.py +++ b/core/optimizer.py @@ -5,19 +5,17 @@ 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_vertebral_body) +from imaging.transforms import load_rotated_planes_from_path, load_rotated_pedicle_from_path from config.constant import ALLOWED_DIAMETERS, ALLOWED_LENGTHS from core.objective import OptimizationContext, make_objective_function, make_objective_function_xfr from pyswarm import pso -from core.cylinder import generate_cylinder_n_torch, generate_cylinder_o_torch, snap_to_discrete_values, create_coordinate_grid, snap_to_discrete_values_xfr, generate_cylinder_butt_torch -from core.intersection import center_line_intersections_torch -from core.scoring import compute_overlap_ratio_from_cylinder_mask, is_solution_ok, cl_score_torch_xfr +from scipy.ndimage import minimum_filter, generate_binary_structure +from core.cylinder import generate_cylinder_n_torch, generate_cylinder_butt_torch, snap_to_discrete_values, create_coordinate_grid, snap_to_discrete_values_xfr +from core.scoring import compute_overlap_ratio_from_cylinder_mask, is_solution_ok from config.constant import OVERLAP_THRESH from visualization.res_bone_figure import render_bone_figure -LATERAL_REFINE_MIN_IN_BONE = 0.97 - # 入口柱向前（+y，椎體方向）幾個 voxel 內進入 VBODY 即視為「入口在椎體上」： # 涵蓋椎體後側皮質邊緣（mask 外 1~2 voxel 的分割界線帶）； # 真正的後側要素與椎體間隔（椎間孔）大於此值，不受影響。 @@ -36,81 +34,27 @@ def _makedirs_retry(path, retries=5, delay=0.5): 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 - using a longer screw. Only accepts a candidate if it stays - >= LATERAL_REFINE_MIN_IN_BONE inside bone AND improves the score. - - left (x-lower half): more lateral = larger azimuth, entry shifted toward -x - right (x-upper half): more lateral = smaller azimuth, entry shifted toward +x - """ - def _score_candidate(cand): - z_c, x_c, az_c, alt_c, d_c, L_c = cand - y_c = y_indices[round(z_c), round(x_c)] - if y_c < 0: - return None - cyl = generate_cylinder_n_torch(d_c, L_c, z_c, y_c, x_c, az_c, alt_c, - image_shape, spacing, device, grid) - cyl_o = generate_cylinder_o_torch(d_c, L_c, z_c, y_c, x_c, az_c, alt_c, - image_shape, spacing, device, grid) - 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) - cyl_butt = generate_cylinder_butt_torch(d_c, z_c, y_c, x_c, az_c, alt_c, - image_shape, spacing, device, grid) - loss = cl_score_torch_xfr(cortical_tensor, spine_tensor, cyl, cyl_o, inter, - diameter=d_c, - vbody_tensor=vbody_tensor, - cylinder_butt_torch=cyl_butt) - in_bone = ((spine_tensor == 1) & (cyl == 1)).sum().item() / cyl.sum().item() - return {'pos': cand, 'loss': loss, 'in_bone': in_bone} - - d_snap, L_snap = snap_to_discrete_values_xfr(diameter_raw, length_raw) - az_lo, az_hi = az_bounds - x_lo, x_hi = x_bounds - - az_steps = [0.0, 5.0, 10.0, 15.0] if side == "L" else [0.0, -5.0, -10.0, -15.0] - x_shifts = [0.0, -4.0, -8.0] if side == "L" else [0.0, 4.0, 8.0] - - best = _score_candidate((z, x, azimuth, altitude, d_snap, L_snap)) - if best is None: - return z, x, azimuth, altitude, diameter_raw, length_raw, False, None - - for ds in az_steps: - az_c = min(az_hi - 0.01, max(az_lo + 0.01, azimuth + ds)) - if abs(az_c - azimuth) < 0.5 and ds != 0: - continue - for dx in x_shifts: - x_c = min(x_hi, max(x_lo, x + dx)) - if abs(x_c - x) < 0.5 and dx != 0: - continue - for L_c in sorted({L_snap} | {l for l in ALLOWED_LENGTHS if l > L_snap}): - cand = _score_candidate((z, x_c, az_c, altitude, d_snap, L_c)) - if cand is None: - continue - if cand['in_bone'] >= LATERAL_REFINE_MIN_IN_BONE and cand['loss'] < best['loss']: - best = cand - - z_r, x_r, az_r, alt_r, d_r, L_r = best['pos'] - adopted = not (abs(z_r - z) < 1e-6 and abs(x_r - x) < 1e-6 - and abs(az_r - azimuth) < 1e-6 and abs(L_r - L_snap) < 1e-6) - return z_r, x_r, az_r, alt_r, d_r, L_r, adopted, best - - def _validate_bounds(lb, ub, name): for i, (lo, hi) in enumerate(zip(lb, ub)): if lo >= hi: raise ValueError(f'PSO bounds invalid for {name}: dim {i} lower {lo} >= upper {hi}') +# 螺絲直徑搜索範圍（xfr 管線）：上限 = 該側椎弓根寬度（measure_pedicle_width， +# xfr_preprocess 記錄於 transform.json，raw 值已 clip 到 MAX_PEDICLE_WIDTH_MM）、 +# 下限 = 該值減 DIA_PEDICLE_MARGIN_MM；該側量測不到時退回 ALLOWED 範圍（fallback）。 +DIA_PEDICLE_MARGIN_MM = 1.5 + + +def _diameter_bounds_from_pedicle(ped_mm, fallback): + if ped_mm is not None and ped_mm > DIA_PEDICLE_MARGIN_MM: + # return (ped_mm - DIA_PEDICLE_MARGIN_MM, ped_mm) + ub = min(ped_mm, max(ALLOWED_DIAMETERS)+0.4) + lb = max(ub - DIA_PEDICLE_MARGIN_MM, min(ALLOWED_DIAMETERS)) + return (lb, max(lb+0.1, ub)) + return fallback + + def get_first_nonzero_y(arr, endplate_plane=None, vbody_mask=None): """每個 (z, x) 柱的第一個 bone voxel 的 y（voxel index）。 @@ -196,9 +140,11 @@ def run_pso_torch_xfr( patient_id: str = None, side_dir: str = None, run_id: str = None, + screw_shape: str = 'cone', ): """ - Main function to run PSO. + Main function to run PSO. + screw_shape: 最終合併圖的螺絲體外形（'cone' 錐體預設 / 'cylinder' 圓柱，只影響繪製） 如果 optimize_size=True，diameter 和 length 也會被最佳化 如果 optimize_size=False，使用預設值（向後兼容） """ @@ -245,15 +191,26 @@ def run_pso_torch_xfr( res = analyze_vertebral_tilt_contour(image2_path, edge_type='superior', show_plot=False, debug=False) alt = res['superior']['tilt_angle_deg'] - # ===== 平面：影像已由 xfr_preprocess 對齊旋轉到椎體基準系 ===== - # （鏡稱面法線 -> +x、上終板 normal y=0 / z>0），az/alt 範圍改用固定 - # 約束（見下方 CBT bounds），不再逐椎以 theta_v / tau 重新錨定。 + # ===== 平面：讀 transform.json rotated 段（xfr_preprocess 記錄，經 R + # 剛性旋轉 + start 平移到 rotated/ 局部系）===== + # 不在影像上重新估算（舊式逐椎 best_symmetry_plane / + # best_upper_endplate_plane 擬合已移除）；影像已由 xfr_preprocess 對齊 + # 旋轉到椎體基準系（鏡稱面法線 -> +x、上終板 normal y=0 / z>0）， + # az/alt 範圍為固定約束（見下方 CBT bounds）。記錄平面與棘突 / 椎體 + # 分割皆以完整 mask（SP 移除前）為基準（xfr_preprocess 記錄時同式）。 # 此處平面仅供： # - 棘突移除（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) + sym_plane, endplate_plane = load_rotated_planes_from_path(image2_path, level) + + # 雙側椎弓根寬度（measure_pedicle_width，預處理旋轉時記錄於 transform.json）： + # 決定 diameter bounds（上限 = 椎弓根寬度、下限 = 該值減 DIA_PEDICLE_MARGIN_MM）； + # 舊世代未記錄 / 該側量測不到時退回 ALLOWED 範圍（見 _diameter_bounds_from_pedicle） + pedicle_rec = load_rotated_pedicle_from_path(image2_path, level) + ped_mm_l = pedicle_rec.get('L_mm') if pedicle_rec else None + ped_mm_r = pedicle_rec.get('R_mm') if pedicle_rec else None # 入口面 y_indices 取「每個 (z,x) 柱第一個 bone voxel（最後側）」， # 中線柱會落在棘突上（入口太靠內後）。先把棘突（鏡稱面中線後側， @@ -261,13 +218,8 @@ def run_pso_torch_xfr( # 讓 y_indices 永不落在棘突上（loss 用 spine_tensor，不受影響）。 sp_mask, sp_th, sp_info = segment_spinous_process(image2_array, sym_plane) - # 上終板面與椎體都在「完整 mask（SP 移除前）」上計算： - # - 終板面由前側頂面擬合，SP 移除不改變平面； - # - 椎體與 render_bone_figure 的 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, @@ -285,6 +237,12 @@ def run_pso_torch_xfr( print(f"[VBODY-SCORE] skipped: {vb_info['mode']}") ctx.vbody_tensor = vbody_tensor + # 完整 mask（SP 移除前）備份，供合併圖使用：圖的骨頭點雲 + 棘突(紫) / + # 椎體(金) 著色由圖端以記錄平面對完整 mask 切出，與 optimizer 的 + # segmentation 完全同 input / 同參數（移除版 image2_array 只供 + # y_indices / loss——spine_tensor 與其共享記憶體，顯示 tensor 不受影響） + image2_array_full = image2_array.copy() + 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 / 棘突切除）: " @@ -300,8 +258,23 @@ def run_pso_torch_xfr( y_indices = get_first_nonzero_y(image2_array, endplate_plane, vb_mask_np) + # 評分 spine mask：向內侵蝕 1 voxel（各方向 1 pixel，6 鄰接）， + # 讓螺絲與骨頭邊界保持 1 voxel 安全邊距（容錯分割界線誤差）； + # 繪圖（spine_display_tensor）仍用完整未侵蝕 mask。 + spine_score_tensor = torch.from_numpy( + (minimum_filter(image2_array, + footprint=generate_binary_structure(3, 1), + mode='constant', cval=0) >= 1) + .astype(np.uint8) + ).to(device=device) + ctx.spine_tensor = spine_score_tensor + objective_fn = make_objective_function_xfr(ctx, y_indices) + # 合併圖用完整 mask 的 tensor（未移除 SP；與 image2_array_full 同 grid） + spine_display_tensor = torch.from_numpy(image2_array_full).to( + device=device, dtype=torch.uint8) + # 對齊健全性檢查（基準系下 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))) @@ -313,12 +286,13 @@ def run_pso_torch_xfr( tau_y, tau_x = 0.0, 0.0 print(f"[PLANE] mirror : {sym_plane['plane'][0]:+.3f}x {sym_plane['plane'][1]:+.3f}y " f"{sym_plane['plane'][2]:+.3f}z = {sym_plane['plane'][3]:.1f}" - f" (theta_v={theta_v:+.2f} deg, mirror ratio={sym_plane['ratio']:.3f})") + f" (theta_v={theta_v:+.2f} deg, mirror ratio=" + f"{sym_plane.get('ratio', float('nan')):.3f})") if endplate_plane is not None: p = endplate_plane['plane'] print(f"[PLANE] endplate: {p[0]:+.3f}x {p[1]:+.3f}y {p[2]:+.3f}z = {p[3]:.1f}" f" (tau_y={tau_y:+.2f} deg, tau_x={tau_x:+.2f} deg, " - f"inlier={endplate_plane['inlier_ratio']:.2f})") + f"inlier={endplate_plane.get('inlier_ratio', float('nan')):.2f})") else: print("[PLANE] endplate: 資料不足，終板入口剪除停用") @@ -410,16 +384,18 @@ def run_pso_torch_xfr( azimuth_bounds_l = (60-azi, 90-azi) azimuth_bounds_r = (90-azi, 120-azi) altitude_bounds = (65-alt, 80-alt) - # 非 CBT 分支維持單一式，兩側同用 +# 非 CBT 分支維持單一式，兩側同用 altitude_bounds_l = altitude_bounds altitude_bounds_r = altitude_bounds def eval_overlap_from_position(pos, side: str, optimize_size: bool, - spine_tensor: torch.Tensor, - image_shape, spacing): + spine_tensor: torch.Tensor, + image_shape, spacing): """ 根據 PSO 給的 position 生成 cylinder mask，再算 overlap ratio side: "L" or "R" 只是方便 debug + 回傳 (overlap, d, L, not_in_bone)：最終 log 用完整未侵蝕骨頭 mask + （spine_display_tensor）與去掉 butt 豁免帶（入口端 5mm）的螺絲圓柱計算 """ if optimize_size: # d, L = snap_to_discrete_values(pos[5], pos[6]) @@ -435,29 +411,48 @@ def run_pso_torch_xfr( params_5[3], params_5[4], image_shape, spacing, device, grid ) + butt_mask = generate_cylinder_butt_torch( + d, + params_5[0], params_5[1], params_5[2], + params_5[3], params_5[4], + image_shape, spacing, device, grid + ) + # 螺絲圓柱去掉入口端 5mm butt 帶 + cyl_eval = ((cyl_mask == 1) & (butt_mask == 0)).to(torch.uint8) - overlap = compute_overlap_ratio_from_cylinder_mask(cyl_mask, spine_tensor) - return overlap, d, L + overlap = compute_overlap_ratio_from_cylinder_mask(cyl_eval, spine_tensor) + outside = (spine_tensor == 0) & (cyl_eval == 1) + not_in_bone = outside.sum().item() + return overlap, d, L, not_in_bone if optimize_size: # 模式 1：優化 diameter 和 length print("=== 最佳化模式：最佳化位置、角度、直徑和長度 ===") # 設定 diameter 和 length 的 bounds（連續範圍） - diameter_bounds = (min(ALLOWED_DIAMETERS), max(ALLOWED_DIAMETERS)*1.01) + # diameter 上限 = 該側椎弓根寬度、下限 = 該值減 DIA_PEDICLE_MARGIN_MM， + # 依 L/R 分開；該側量測不到時退回 ALLOWED 範圍 + dia_fb = (min(ALLOWED_DIAMETERS), max(ALLOWED_DIAMETERS)*1.01) + diameter_bounds_l = _diameter_bounds_from_pedicle(ped_mm_l, dia_fb) + diameter_bounds_r = _diameter_bounds_from_pedicle(ped_mm_r, dia_fb) length_bounds = (min(ALLOWED_LENGTHS), max(ALLOWED_LENGTHS)*1.01) + _ped_l_str = 'n/a' if ped_mm_l is None else f'{ped_mm_l:.2f}' + _ped_r_str = 'n/a' if ped_mm_r is None else f'{ped_mm_r:.2f}' + print(f"[BOUNDS] pedicle width L={_ped_l_str} mm, R={_ped_r_str} mm -> " + f"diameter bounds L=[{diameter_bounds_l[0]:.2f}, {diameter_bounds_l[1]:.2f}], " + f"R=[{diameter_bounds_r[0]:.2f}, {diameter_bounds_r[1]:.2f}]") # bounds 現在有 6 個參數 [z, x, az, alt, d, L]（y 由 y_indices 取）； - # az/alt 兩側用同一組固定約束，x 帶依 L/R 分開 + # az/alt 兩側用同一組固定約束，x 帶 / diameter 依 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]] + altitude_bounds_l[0], diameter_bounds_l[0], length_bounds[0]] ub_l = [z_bounds[1], x_bounds_left[1], azimuth_bounds_l[1], - altitude_bounds_l[1], diameter_bounds[1], length_bounds[1]] + altitude_bounds_l[1], diameter_bounds_l[1], length_bounds[1]] lb_r = [z_bounds[0], x_bounds_right[0], azimuth_bounds_r[0], - altitude_bounds_r[0], diameter_bounds[0], length_bounds[0]] + altitude_bounds_r[0], diameter_bounds_r[0], length_bounds[0]] ub_r = [z_bounds[1], x_bounds_right[1], azimuth_bounds_r[1], - altitude_bounds_r[1], diameter_bounds[1], length_bounds[1]] + altitude_bounds_r[1], diameter_bounds_r[1], length_bounds[1]] else: # 模式 2：固定 diameter 和 length（向後兼容） @@ -513,30 +508,15 @@ def run_pso_torch_xfr( # retries += 1 z, x, azimuth, altitude, diameter, length = position - az_pso, x_pso, L_pso = azimuth, x, length y = y_indices[round(z), round(x)] - 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") - position = z, y, x, azimuth, altitude, diameter, length - overlap_s, d_snap, L_snap = eval_overlap_from_position( - position, s_i, optimize_size, spine_tensor, image_shape, spacing + overlap_s, d_snap, L_snap, nib_s = eval_overlap_from_position( + position, s_i, optimize_size, spine_display_tensor, image_shape, spacing ) print(f"[{tag}] overlap: {overlap_s*100:.1f}%") + print(f"[{tag}] not_in_bone: {nib_s} vox") if optimize_size: print(f"[{tag}] Position: {position[:5]}") @@ -576,20 +556,26 @@ def run_pso_torch_xfr( final_length_r = length def _plot_combined(d_l, l_l, d_r, l_r, pos_l, pos_r, total_t): - # volume / level 由 image2_path 反推（…/{vol}/rotated/{level}_*.nii.gz） + # volume / level 由 image2_path 反推（…/{vol}/rotated/{level}_*.nii.gz）。 + # 平面傳 transform.json 記錄值（上方 load_rotated_planes_from_path）， + # 骨頭傳完整 mask（SP 移除前）的 spine_display_tensor：圖上的 + # SP（紫）/ VBODY（金）由記錄平面對完整 mask 切出，與 optimizer + # 分割完全一致；橙 / 綠平面同用上方 transform.json 記錄值。 render_bone_figure( None, None, - spine_tensor, + spine_display_tensor, cortical_tensor, folder, spacing=spacing, way='CBT' if CBT else 'TPS', + sym=sym_plane, symp=endplate_plane, best_position_l=pos_l, best_position_r=pos_r, diameter_l=d_l, length_l=l_l, diameter_r=d_r, length_r=l_r, image2_path=image2_path, device=device, grid=grid, swarm_size=swarm_size, max_iter=max_iter, total_time=total_t, + screw_shape=screw_shape, ) if side == 'both': @@ -1163,6 +1149,499 @@ def run_de_torch( return best_position_l, best_loss_l, best_position_r, best_loss_r, total_time + +def run_de_torch_xfr( + label_str: str, + image1_path: str, + image2_path: str, + image3_path: str, + folder: str, + swarm_size: int, + max_iter: int, + spacing: list, + CBT: bool, + device: torch.device, + optimize_size: bool = True, + grid=None, + debug=False, + side: str = 'both', + level: str = None, + patient_id: str = None, + side_dir: str = None, + run_id: str = None, + screw_shape: str = 'cone', +): + """ + Main function to run DE（xfr 版，與 run_pso_torch_xfr 同管線）。 + screw_shape: 最終合併圖的螺絲體外形（'cone' 錐體預設 / 'cylinder' 圓柱，只影響繪製） + 前處理、bounds、單側收尾完全沿用 run_pso_torch_xfr， + 只把優化器換成 differential_evolution（無 omega；popsize 換算成 + swarm_size // 維度，實際族群 ≈ swarm_size，與 PSO 可比較）。 + L/R 是兩次獨立 DE（bounds 不同、目標函數共用），互不依賴： + side='both' 維持原行為（同一次呼叫先 L 後 R）； + side='L'/'R' 只跑該側，讓兩側可排到不同 GPU worker。 + 單側模式的合併輸出（3D 圖 + CSV）由「較晚完成」的一側觸發。 + 如果 optimize_size=True，diameter 和 length 也會被最佳化 + 如果 optimize_size=False，使用預設值（向後兼容） + """ + start_time = time.time() + + # Use global references + global image1_array, image2_array, image2_shape, image3_array + global diameter, length # 這些現在只用於非最佳化模式 + global spine_tensor, cortical_tensor, spine_roi_tensor + + # Load images + image1 = sitk.ReadImage(image1_path) + image2 = sitk.ReadImage(image2_path) + image3 = sitk.ReadImage(image3_path) + image1_array = sitk.GetArrayFromImage(image1) + image2_array = sitk.GetArrayFromImage(image2) + image3_array = sitk.GetArrayFromImage(image3) + image2_shape = image2_array.shape + image_shape = image2_shape + + # Move arrays to torch + cortical_tensor = torch.from_numpy(image1_array).to(device=device, dtype=torch.uint8) + spine_tensor = torch.from_numpy(image2_array).to(device=device, dtype=torch.uint8) + spine_roi_tensor = torch.from_numpy(image3_array).to(device=device, dtype=torch.uint8) + + # 建立明確的優化上下文（取代舊的跨檔案 global 注入；狀態全部顯式傳入目標函數） + ctx = OptimizationContext( + cortical_tensor=cortical_tensor, + spine_tensor=spine_tensor, + spine_roi_tensor=spine_roi_tensor, + image1_array=image1_array, + image2_array=image2_array, + image3_array=image3_array, + image2_shape=image2_shape, + spacing=spacing, + device=device, + grid=grid, + diameter=diameter if not optimize_size else None, + length=length if not optimize_size else None, + ) + + 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'] + + # ===== 平面：讀 transform.json rotated 段（xfr_preprocess 記錄，經 R + # 剛性旋轉 + start 平移到 rotated/ 局部系）===== + # 不在影像上重新估算（舊式逐椎 best_symmetry_plane / + # best_upper_endplate_plane 擬合已移除）；影像已由 xfr_preprocess 對齊 + # 旋轉到椎體基準系（鏡稱面法線 -> +x、上終板 normal y=0 / z>0）， + # az/alt 範圍為固定約束（見下方 CBT bounds）。記錄平面與棘突 / 椎體 + # 分割皆以完整 mask（SP 移除前）為基準（xfr_preprocess 記錄時同式）。 + # 此處平面仅供： + # - 棘突移除（segment_spinous_process，sym） + # - 入口面終板上方剪除（get_first_nonzero_y，endplate） + # - VBODY 椎體分割（segment_vertebral_body，sym + endplate） + # - 對齊健全性檢查：基準系下 theta_v ≈ 0、tau_y ≈ 0（偏大 = 預處理失效） + sym_plane, endplate_plane = load_rotated_planes_from_path(image2_path, level) + + # 雙側椎弓根寬度（measure_pedicle_width，預處理旋轉時記錄於 transform.json）： + # 決定 diameter bounds（上限 = 椎弓根寬度、下限 = 該值減 DIA_PEDICLE_MARGIN_MM）； + # 舊世代未記錄 / 該側量測不到時退回 ALLOWED 範圍（見 _diameter_bounds_from_pedicle） + pedicle_rec = load_rotated_pedicle_from_path(image2_path, level) + ped_mm_l = pedicle_rec.get('L_mm') if pedicle_rec else None + ped_mm_r = pedicle_rec.get('R_mm') if pedicle_rec else None + + # 入口面 y_indices 取「每個 (z,x) 柱第一個 bone voxel（最後側）」， + # 中線柱會落在棘突上（入口太靠內後）。先把棘突（鏡稱面中線後側， + # 見 segment_spinous_process）從 image2_array 移除再取 surface， + # 讓 y_indices 永不落在棘突上（loss 用 spine_tensor，不受影響）。 + sp_mask, sp_th, sp_info = segment_spinous_process(image2_array, sym_plane) + + # 棘突缺如（laminectomy，mode='no_spinous'）時不該把殘留後側要素 + # 當「棘突」移除（會鏟進椎體後側），入口面維持完整 mask。 + + # 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 + + # 完整 mask（SP 移除前）備份，供合併圖使用：圖的骨頭點雲 + 棘突(紫) / + # 椎體(金) 著色由圖端以記錄平面對完整 mask 切出，與 optimizer 的 + # segmentation 完全同 input / 同參數（移除版 image2_array 只供 + # y_indices / loss——spine_tensor 與其共享記憶體，顯示 tensor 不受影響） + image2_array_full = image2_array.copy() + + 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, endplate_plane, vb_mask_np) + + # 評分 spine mask：向內侵蝕 1 voxel（各方向 1 pixel，6 鄰接）， + # 讓螺絲與骨頭邊界保持 1 voxel 安全邊距（容錯分割界線誤差）； + # 繪圖（spine_display_tensor）仍用完整未侵蝕 mask。 + spine_score_tensor = torch.from_numpy( + (minimum_filter(image2_array, + footprint=generate_binary_structure(3, 1), + mode='constant', cval=0) >= 1) + .astype(np.uint8) + ).to(device=device) + ctx.spine_tensor = spine_score_tensor + + objective_fn = make_objective_function_xfr(ctx, y_indices) + + # 合併圖用完整 mask 的 tensor（未移除 SP；與 image2_array_full 同 grid） + spine_display_tensor = torch.from_numpy(image2_array_full).to( + device=device, dtype=torch.uint8) + + # 對齊健全性檢查（基準系下 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))) + 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))) + tau_x = float(np.degrees(np.arctan2(e_nx, e_nz))) + else: + tau_y, tau_x = 0.0, 0.0 + print(f"[PLANE] mirror : {sym_plane['plane'][0]:+.3f}x {sym_plane['plane'][1]:+.3f}y " + f"{sym_plane['plane'][2]:+.3f}z = {sym_plane['plane'][3]:.1f}" + f" (theta_v={theta_v:+.2f} deg, mirror ratio=" + f"{sym_plane.get('ratio', float('nan')):.3f})") + if endplate_plane is not None: + p = endplate_plane['plane'] + print(f"[PLANE] endplate: {p[0]:+.3f}x {p[1]:+.3f}y {p[2]:+.3f}z = {p[3]:.1f}" + f" (tau_y={tau_y:+.2f} deg, tau_x={tau_x:+.2f} deg, " + f"inlier={endplate_plane.get('inlier_ratio', float('nan')):.2f})") + else: + print("[PLANE] endplate: 資料不足，終板入口剪除停用") + + # x/z 搜索空間邊界：把 VBODY 椎體投影到 xz 平面，取該投影的 bounding box， + # 再在 x 向切 L/R 兩 band + 中央缺口、z 向留 10%~90%（見下方 CBT bounds）。 + # CBT 入口點應落在椎體上（左/右 band），而非跨整段骨頭：整段骨頭包含 + # 肋、後側要素、橫突等極端，x 範圍比椎體寬，會把 L/R band 往外推。 + # VBODY mask 為 (z,y,x)；np.any(..., axis=1) 折疊 y 得 xz 投影 (z,x)。 + # x 範圍 = 有 VBODY 的欄（投影 axis=0 是 z，沿 z 做 any） + # z 範圍 = 有 VBODY 的行（投影 axis=1 是 x，沿 x 做 any） + # VBODY 分割失敗（None / 全 0）時退回整段骨頭 x/z 範圍。 + if vb_mask_np is not None and vb_mask_np.any(): + vb_xz = np.any(vb_mask_np, axis=1) # (z, x) 投影 + x1 = int(np.where(vb_xz.any(axis=0))[0][0]) + x2 = int(np.where(vb_xz.any(axis=0))[0][-1]) + z1 = int(np.where(vb_xz.any(axis=1))[0][0]) + z2 = int(np.where(vb_xz.any(axis=1))[0][-1]) + print(f"[BOUNDS] VBODY xz projection: z[{z1},{z2}] x[{x1},{x2}] " + f"(z_height={z2 - z1}, x_width={x2 - x1})") + else: + print("[BOUNDS] VBODY unavailable, falling back to whole-bone xz range") + x_with_nonzero = np.where(np.any(image2_array, axis=(0, 1)))[0] + z_with_nonzero = np.where(np.any(image2_array, axis=(0, 2)))[0] + x1 = int(x_with_nonzero[0]) + x2 = int(x_with_nonzero[-1]) + z1 = int(z_with_nonzero[0]) + z2 = int(z_with_nonzero[-1]) + x_width = x2 - x1 + z_height = z2 - z1 + + # 設定基本的 bounds + if CBT == True: + z_bounds = (0, z1+z_height*.7) + x_bounds_left = (x1+x_width*.1, x1+x_width*.4) + x_bounds_right = (x1+x_width*.6, x1+x_width*.9) + + # 脊椎若被體積邊界切到（真正偏心、骨頭貼著左/右邊緣）， + # 對應那側的 x band 下限會 >= 上限，優化器會丟 bounds 錯誤。 + # 出錯時 clamp 成同側最小寬度（5% 寬度）的合法 band。 + min_band = .05 * image_shape[2] + if x_bounds_left[1] <= x_bounds_left[0]: + x_bounds_left = (x_bounds_left[0], x_bounds_left[0] + min_band) + 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)) + azimuth_bounds_l = (98, 105) + azimuth_bounds_r = (75, 82) + altitude_bounds_l = (60, 65) + altitude_bounds_r = (60, 65) + + else: + z_bounds = (0, image_shape[0] - 1) + y_bounds = (image_shape[1]/5, image_shape[1]/2 - 1) + x_bounds_left = (0, image_shape[2]/2 - image_shape[2]/10 - 1) + x_bounds_right = (image_shape[2]/2 + image_shape[2]/10, image_shape[2] - 1) + azimuth_bounds_l = (60-azi, 90-azi) + azimuth_bounds_r = (90-azi, 120-azi) + altitude_bounds = (65-alt, 80-alt) + # 非 CBT 分支維持單一式，兩側同用 + altitude_bounds_l = altitude_bounds + altitude_bounds_r = altitude_bounds + + def eval_overlap_from_position(pos, side: str, optimize_size: bool, + spine_tensor: torch.Tensor, + image_shape, spacing): + """ + 根據 DE 給的 position 生成 cylinder mask，再算 overlap ratio + side: "L" or "R" 只是方便 debug + """ + if optimize_size: + d, L = snap_to_discrete_values_xfr(pos[5], pos[6]) + params_5 = pos[:5] + else: + d, L = diameter, length + params_5 = pos + + cyl_mask = generate_cylinder_n_torch( + d, L, + params_5[0], params_5[1], params_5[2], + params_5[3], params_5[4], + image_shape, spacing, device, grid + ) + + overlap = compute_overlap_ratio_from_cylinder_mask(cyl_mask, spine_tensor) + return overlap, d, L + + if optimize_size: + # 模式 1：優化 diameter 和 length + print("=== DE 最佳化模式：最佳化位置、角度、直徑和長度 ===") + + # 設定 diameter 和 length 的 bounds（連續範圍） + # diameter 上限 = 該側椎弓根寬度、下限 = 該值減 DIA_PEDICLE_MARGIN_MM， + # 依 L/R 分開；該側量測不到時退回 ALLOWED 範圍 + dia_fb = (min(ALLOWED_DIAMETERS), max(ALLOWED_DIAMETERS)*1.01) + diameter_bounds_l = _diameter_bounds_from_pedicle(ped_mm_l, dia_fb) + diameter_bounds_r = _diameter_bounds_from_pedicle(ped_mm_r, dia_fb) + length_bounds = (min(ALLOWED_LENGTHS), max(ALLOWED_LENGTHS)*1.01) + _ped_l_str = 'n/a' if ped_mm_l is None else f'{ped_mm_l:.2f}' + _ped_r_str = 'n/a' if ped_mm_r is None else f'{ped_mm_r:.2f}' + print(f"[BOUNDS] pedicle width L={_ped_l_str} mm, R={_ped_r_str} mm -> " + f"diameter bounds L=[{diameter_bounds_l[0]:.2f}, {diameter_bounds_l[1]:.2f}], " + f"R=[{diameter_bounds_r[0]:.2f}, {diameter_bounds_r[1]:.2f}]") + + # bounds 現在有 6 個參數 [z, x, az, alt, d, L]（y 由 y_indices 取）； + # search 空間不含 y，與 pso 版相同；az/alt 兩側各用固定約束，x 帶 / diameter 依 L/R 分開 + lb_l = [z_bounds[0], x_bounds_left[0], azimuth_bounds_l[0], + altitude_bounds_l[0], diameter_bounds_l[0], length_bounds[0]] + ub_l = [z_bounds[1], x_bounds_left[1], azimuth_bounds_l[1], + altitude_bounds_l[1], diameter_bounds_l[1], length_bounds[1]] + + lb_r = [z_bounds[0], x_bounds_right[0], azimuth_bounds_r[0], + altitude_bounds_r[0], diameter_bounds_r[0], length_bounds[0]] + ub_r = [z_bounds[1], x_bounds_right[1], azimuth_bounds_r[1], + altitude_bounds_r[1], diameter_bounds_r[1], length_bounds[1]] + + else: + # 模式 2：固定 diameter 和 length（向後兼容） + print("=== DE 固定尺寸模式：最佳化位置和角度 ===") + # 使用預設值（需要在調用時提供） + diameter = 4.5 # 或從參數傳入 + length = 45 # 或從參數傳入 + + lb_l = [z_bounds[0], y_bounds[0], x_bounds_left[0], azimuth_bounds_l[0], altitude_bounds_l[0]] + ub_l = [z_bounds[1], y_bounds[1], x_bounds_left[1], azimuth_bounds_l[1], altitude_bounds_l[1]] + + lb_r = [z_bounds[0], y_bounds[0], x_bounds_right[0], azimuth_bounds_r[0], altitude_bounds_r[0]] + ub_r = [z_bounds[1], y_bounds[1], x_bounds_right[1], azimuth_bounds_r[1], altitude_bounds_r[1]] + + if True or debug: + 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 + + # L/R 是兩次獨立 DE（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 + + 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} (DE) ===") + _validate_bounds(lb_s, ub_s, f'{label_str} {s_i}') + # DE 的實際族群 = popsize * len(bounds)；換算成 ≈ swarm_size， + # 跟 PSO 版的 swarm_size 可公平比較 + de_popsize = max(1, swarm_size // len(lb_s)) + res_s = differential_evolution(objective_fn, list(zip(lb_s, ub_s)), + popsize=de_popsize, maxiter=max_iter) + position, loss = list(res_s.x), float(res_s.fun) + + z, x, azimuth, altitude, diameter, length = position + y = y_indices[round(z), round(x)] + + position = z, y, x, azimuth, altitude, diameter, length + + overlap_s, d_snap, L_snap = eval_overlap_from_position( + position, s_i, optimize_size, spine_score_tensor, image_shape, spacing + ) + print(f"[{tag}] overlap: {overlap_s*100:.1f}%") + + 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 + + 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 + + end_time = time.time() + total_time = end_time - start_time + + # 提取最終的 diameter 和 length + if optimize_size: + 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 + + def _plot_combined(d_l, l_l, d_r, l_r, pos_l, pos_r, total_t): + # volume / level 由 image2_path 反推（…/{vol}/rotated/{level}_*.nii.gz）。 + # 平面傳 transform.json 記錄值（上方 load_rotated_planes_from_path）， + # 骨頭傳完整 mask（SP 移除前）的 spine_display_tensor：圖上的 + # SP（紫）/ VBODY（金）由記錄平面對完整 mask 切出，與 optimizer + # 分割完全一致；橙 / 綠平面同用上方 transform.json 記錄值。 + render_bone_figure( + None, None, + spine_display_tensor, + cortical_tensor, + folder, + spacing=spacing, + way='CBT' if CBT else 'TPS', + sym=sym_plane, symp=endplate_plane, + best_position_l=pos_l, best_position_r=pos_r, + diameter_l=d_l, length_l=l_l, + diameter_r=d_r, length_r=l_r, + image2_path=image2_path, + device=device, grid=grid, + swarm_size=swarm_size, max_iter=max_iter, total_time=total_t, + screw_shape=screw_shape, + ) + + 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 + + # ---- 單側收尾：寫該側結果檔；兩側都完成時由較晚的一側觸發合併輸出 ---- + # （與 run_pso_torch_xfr 完全相同：兩側可能在不同 GPU worker：各自把結果寫 + # _.json（先寫 tmp 再 os.replace，對端讀到的一定是完整檔）。 + # 先完成者看不到對端檔就跳過；後完成者看到兩側齊了、搶到 plot lock + # （O_EXCL，確保合併輸出只跑一次）才載入對端結果跑 render_bone_figure + # （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 + + def run_nm_torch( label_str: str, image1_path: str, diff --git a/core/scoring.py b/core/scoring.py index e15100d..7cf3b6e 100644 --- a/core/scoring.py +++ b/core/scoring.py @@ -3,9 +3,20 @@ from core.cylinder import generate_cylinder_n_torch, generate_cylinder_tip_torch from config.constant import OVERLAP_THRESH # 直徑偏好 bonus（每 mm）：ALLOWED_DIAMETERS 相鄰等級間隔 0.5 mm， -# 20000/mm ⇒ 每級差 10000 分，約等於容許 1 個額外 out-of-bone voxel -# （10000 分/voxel 懲罰）換 0.5 mm 直徑，讓結果在安全前提下偏好較粗螺絲。 +# 20000/mm ⇒ 每級差 10000 分（＝1 個 voxel 的舊軟性 10000/voxel 懲罰， +# 現已改由 HARD_BREACH_PER_VOX 硬約束管制）；bonus 只在同 breach 數 +# （＝同安全等級）的可行解之間讓優化偏好較粗螺絲。 DIAMETER_BONUS_PER_MM = 20000.0 +# DIAMETER_BONUS_PER_MM = 5000000.0 +# DIAMETER_BONUS_PER_MM = 2000000.0 +# 硬約束的每 breach voxel 懲罰（不計 butt 豁免內）。 +# 取 1e9：大於其餘所有評分項可能造成的總分差 +# （d≤6.0 / L≤45 / 0.5mm voxel 時 |score| 展佈約 ≤ 2.1e8）， +# 使 optimizer 先「把 breach voxel 數降到最低」（lexicographic）， +# 再在同 breach 數下比品質。結果：完全在骨內（除 butt 豁免帶）的 +# 螺絲永遠優於任何有 breach 的螺絲；PSO 找不到可行解時才回 +# 最少 breach 的候選。 +HARD_BREACH_PER_VOX = 1e9 def cl_score_torch_xfr( cortical_tensor: torch.Tensor, @@ -17,12 +28,14 @@ def cl_score_torch_xfr( length: float = None, cylinder_tip_torch: torch.Tensor = None, # 新增：尖端 mask vbody_tensor: torch.Tensor = None, # VBODY（椎體）mask (z,y,x) 0/1，None = 不計 VBODY 獎勵 - cylinder_butt_torch: torch.Tensor = None # 最後端（入口端，遠離 VBODY 的一端）0.25mm mask + cylinder_butt_torch: torch.Tensor = None # 最後端（入口端，遠離 VBODY 的一端）5mm 豁免 mask ) -> float: """ 漸進式評分：優先確保找到骨頭，再改善細節。 diameter 提供時附加直徑偏好 bonus（偏好較粗螺絲）。 - cylinder_butt_torch 提供時，該 mask 內 not_in_bone 的 voxel 免除 10000/voxel 扣分。 + cylinder_butt_torch 提供時，該 mask 內 not_in_bone 的 voxel 免除扣分 + （butt 豁免帶）；豁免帶外的 breach voxel 以 HARD_BREACH_PER_VOX + 硬約束處理——breach 數少的解永遠優於 breach 數多的解，與品質分無關。 """ cyl_total = cylinder_torch.sum().item() overlap = ((cortical_tensor == 1) & (cylinder_torch == 1)).sum().item() # in cortical @@ -40,8 +53,8 @@ def cl_score_torch_xfr( in_bone= ((spine_tensor == 1) & (cylinder_torch == 1)).sum().item() not_in_bone= ((spine_tensor == 0) & (cylinder_torch == 1)).sum().item() - # 最後端（入口端，遠離 VBODY 的一端）0.25mm 豁免： - # 這些 voxel 即使 not_in_bone 也不計 10000/voxel 扣分 + # 最後端（入口端，遠離 VBODY 的一端）5mm 豁免： + # 這些 voxel 即使 not_in_bone 也不計 10000/voxel 扣分、也不觸發硬約束 if cylinder_butt_torch is not None: not_in_bone -= ((cylinder_butt_torch == 1) & (cylinder_torch == 1) & (spine_tensor == 0)).sum().item() @@ -57,7 +70,8 @@ def cl_score_torch_xfr( score = cyl_total # 直徑偏好 bonus：相同幾何下偏好較大直徑（diameter 單位 mm）。 - # bonus 遠小於 breaching 懲罰（10000 分/voxel），不會把螺絲推出骨頭。 + # breach 是硬約束（HARD_BREACH_PER_VOX ≫ 任何品質項）， + # bonus 永遠不會把螺絲推出骨頭（butt 豁免帶除外）。 if diameter is not None and diameter > 0: score += DIAMETER_BONUS_PER_MM * diameter @@ -69,10 +83,14 @@ def cl_score_torch_xfr( score += 10 * in_bone # 10 實在太低 score += 100 * overlap # in cortical score += 200 * in_corti_vb # (cortical + VBODY) 每 voxel 再加分 - score += 50 * in_vbody # VBODY 每 voxel 再加分 + score += 20 * in_vbody # VBODY 每 voxel 再加分 score -= 10000 * not_in_bone + # 硬約束：butt 豁免外的 breach 每 voxel 扣 HARD_BREACH_PER_VOX， + # 任何一個 breach voxel 的懲罰都超過其餘所有項的總分差 → + # 「breach=0」的解永遠優於「breach>0」的解（lexicographic）。 + # score -= HARD_BREACH_PER_VOX * not_in_bone # score -= 10000 * max(0, not_in_bone-allowed_error) - score -= 20000 * null_vox2 + score -= 10000 * null_vox2 # score -= 2000 * max(0, null_vox2-10) return float(-score) diff --git a/imaging/orientation.py b/imaging/orientation.py index b5bcf4c..dd62775 100644 --- a/imaging/orientation.py +++ b/imaging/orientation.py @@ -1,6 +1,7 @@ import numpy as np import SimpleITK as sitk -from scipy.ndimage import center_of_mass, rotate, distance_transform_edt, gaussian_filter +from scipy.ndimage import (center_of_mass, rotate, distance_transform_edt, + gaussian_filter, label) import matplotlib.pyplot as plt # 鏡稱面法線的左右（LR，=x 軸）分量大下。正常左右鏡稱面法線幾乎沿 x； @@ -708,6 +709,142 @@ d) 最後 fallback：AP 分佈 55 百分位（可能切進椎體內，會打 WAR info.update(n_vb=int(sel.sum()), ap_thresh=th, mode=mode, lat_clip=lat_clip) return vb_mask, th, info + +def measure_pedicle_width(mask_zyx, sym, sp_info, vb_info, + band_frac=0.06, min_band=6.0, min_slice_vox=30, + min_slices=4): + """ + 旋轉基準系（rotated/ 局部 index）內量測雙側椎弓根（pedicle）的 + 左右（介軸-外側，mediolateral）寬度。 + 輸入須與 segment_spinous_process / segment_vertebral_body 同 mask、 + 同 sym（如 _write_rotated_level 的 rot_bin_seg 與換系後 sym_rot）， + 本函式的 s / ap 定義與之分毫不差（閾值才可對照使用）： + s : 到鏡稱面（中線）的有號距離 a*x+b*y+c*z-d； + s<0 = 左側、s>0 = 右側（optimizer 的 x band 同慣例：x 大 = R）。 + ap : (x,y,z)·u_ap（u_ap = _ap_axis(sym)，+ap = 後側）， + 與兩分割函式相同的無常數定義。 + 椎弓根區（每側）＝ 椎管 AP 窗內、中線帶外（|s| >= w）的骨頭： + ap_front = 椎體後側切點（vb_info['ap_thresh']，gap 模式時與 + sp_info['ap_thresh'] 同值；棘突缺如時為放寬椎體切點）； + ap_back = 椎管空隙後側端（sp_info['ap_thresh_sp']）；缺失時 + （no_spinous / fallback）用全骨最後側 ap； + w = sp_info['band_w']（缺失時以兩分割同源公式重算）。 + 量的切片限定在窗中央 AP 帶（[ap_front+0.3·g, ap_front+0.7·g]， + g = min(窗深, 40 voxel)）：窗前緣是椎體角 fringe、窗後緣是關節突 / + 橫突根部，兩者與椎弓根連通會把分塊寬度撐大；椎弓根本身橫跨全窗， + 中央帶有乾淨的純弓根斷面（ap_back 不可信而窗過深時，cap 40 voxel + 讓帶仍落在弓根 AP 範圍內、避開後側構造）。 + 逐 axial 切片（z）取該側「最靠中線」的 2D 連通分塊 = 椎弓根； + 切片寬度 = 分塊內每條 y row 的 x 向跨度（含端點體素）之中位數 + （row 級取中位數：橫突 / 椎體角只在少數 row 融合撐寬，row 中位數 + 即回到純弓根斷面）；該側寬度 = 合格切片寬度沿 z 的中位數（voxel）。 + 合格切片數 < min_slices 的側不給寬度（薄殘片 level 單一切片的中位數 + 不可靠）。 + 回傳 dict： + 'ap_front' / 'ap_back' / 'band_w' : 所用 AP 窗與中線帶半寬（voxel 系） + 'L' / 'R' : {'width_vox': float|None, 'n_slices': int, + 'z': [z_min, z_max]|None} + 'reason' : 'ok' / 'too_few_bone' / 'no_front_cut' / + 'too_few_zone' / 'no_pedicle_slices' / 'few_slices' + """ + m = np.asarray(mask_zyx) > 0 + ny, nx = m.shape[1], m.shape[2] + out = {'ap_front': None, 'ap_back': None, 'band_w': None, 'reason': 'ok'} + for side in ('L', 'R'): + out[side] = {'width_vox': None, 'n_slices': 0, 'z': None} + if int(m.sum()) < 100: + out['reason'] = 'too_few_bone' + return out + a, b, c, d = sym['plane'] + zz, yy, xx = np.nonzero(m) + 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 = X * u_ap[0] + Y * u_ap[1] + Z * u_ap[2] + ap_front = (vb_info or {}).get('ap_thresh') + if ap_front is None: + ap_front = (sp_info or {}).get('ap_thresh') + if ap_front is None: + out['reason'] = 'no_front_cut' + return out + ap_back = (sp_info or {}).get('ap_thresh_sp') + if ap_back is None or ap_back <= ap_front: + ap_back = float(ap.max()) + w = (sp_info or {}).get('band_w') + if w is None: + w = max(float(min_band), float(band_frac) * float(s.max() - s.min())) + out['ap_front'] = float(ap_front) + out['ap_back'] = float(ap_back) + out['band_w'] = float(w) + in_window = (ap >= ap_front) & (ap <= ap_back) & (np.abs(s) >= w) + if int(in_window.sum()) < 100: + out['reason'] = 'too_few_zone' + return out + # 量測切片限窗中央 AP 帶（見 docstring） + gap_eff = min(float(ap_back - ap_front), 40.0) + b_lo = ap_front + 0.3 * gap_eff + b_hi = ap_front + 0.7 * gap_eff + if b_hi - b_lo < 8.0: + b_lo, b_hi = float(ap_front), float(ap_back) + sel = in_window & (ap >= b_lo) & (ap <= b_hi) + s_z = s[sel] + y_z = yy[sel] + x_z = xx[sel] + z_z = zz[sel] + any_side = False + any_widths = False + struct3 = np.ones((3, 3), dtype=int) + for side, sign in (('L', -1.0), ('R', 1.0)): + keep = (s_z * sign) > 0 + if int(keep.sum()) < min_slice_vox: + continue + y_s = y_z[keep] + x_s = x_z[keep] + z_s = z_z[keep] + widths = [] + zlist = [] + for z in np.unique(z_s): + r = (z_s == z) + plane = np.zeros((ny, nx), dtype=bool) + plane[y_s[r], x_s[r]] = True + lab, ncomp = label(plane, structure=struct3) + best_cc = 0 + best_medial = None + for cc in range(1, ncomp + 1): + cm = (lab == cc) + if int(cm.sum()) < min_slice_vox: + continue + cx = np.nonzero(cm)[1] + medial = int(cx.min()) if sign > 0 else int(cx.max()) + if (best_medial is None + or (medial < best_medial if sign > 0 + else medial > best_medial)): + best_medial = medial + best_cc = cc + if best_cc: + cm_best = (lab == best_cc) + row_w = [] + for yrow in np.flatnonzero(cm_best.any(axis=1)): + xs_row = np.flatnonzero(cm_best[yrow]) + row_w.append(xs_row.max() - xs_row.min() + 1) + if row_w: + widths.append(float(np.median(row_w))) + zlist.append(int(z)) + if widths: + any_widths = True + if len(widths) < int(min_slices): + continue # 合格切片太少：該側不給寬度 + out[side] = {'width_vox': float(np.median(widths)), + 'n_slices': len(widths), + 'z': [min(zlist), max(zlist)]} + any_side = True + if not any_side: + out['reason'] = 'few_slices' if any_widths else 'no_pedicle_slices' + return out + + def best_upper_endplate_plane(mask_zyx, angle_max=45.0, thresh=4.0, n_iter=500, seed=42): """ diff --git a/imaging/transforms.py b/imaging/transforms.py index 0fa1ef9..e44d349 100644 --- a/imaging/transforms.py +++ b/imaging/transforms.py @@ -38,7 +38,23 @@ transform.json 記錄每一步所需的參數（寫檔時即為 ground truth， "R": [[..]], // 作用於 (x,y,z)；forward dest = R(src-c)+c "center": [cx, cy, cz], // 未旋轉檔 disk（standardize 後）index 系 "start": [sx, sy, sz], // rotated 檔 origin（模板 index 系） - "size": [nx, ny, nz] + "size": [nx, ny, nz], + "sym": {...}, // 鏡稱面（best_symmetry_plane 結構）， + // 經 R 剛性旋轉 + start 平移到 rotated 檔局部系 + "endplate": {...} | null, // 上終板面（best_upper_endplate_plane 結構），同換系 + // （core/optimizer 的 xfr 管線直接讀此兩面， + // 不在影像上重新估算） + "pedicle": {...} // 雙側椎弓根寬度（measure_pedicle_width， + // 同 rot_bin_seg / sym_rot 量測）： + // {'L_mm','R_mm'（mediolateral 中位數寬度 mm, + // clip 到 MAX_PEDICLE_WIDTH_MM, + // None=該側量測不到）, + // 'L/R_raw_mm'（未 clip 原始值）, + // 'L/R_clipped'（真值超限=橫突/肋融合）, + // 'L/R_slices'（合格切片數）, + // 'L/R_z'（合格切片 z 範圍）, 'ap_window' + // （[前,後] ap 窗, rotated 局部系）, 'band_w', + // 'reason'} } } } @@ -114,6 +130,68 @@ def merge_rotated_into_transform(vol_dir, level, sec): save_transform(vol_dir, meta) +def load_rotated_planes(vol_dir, level): + """讀 transform.json 的 levels[level]['rotated'] 記錄平面： + (sym, endplate)。sym 為 best_symmetry_plane 同結構 dict（plane / normal / + u / v / ratio …），endplate 為 best_upper_endplate_plane 同結構 dict， + 可為 None（rotated/ 局部 index 系，與 _write_rotated_level 的 SP / VBODY + 分割同參數）。舊世代（未記錄）raise KeyError。""" + meta = load_transform(vol_dir) + lv = meta['levels'].get(level) + rot = lv.get('rotated') if isinstance(lv, dict) else None + if not rot or 'sym' not in rot: + raise KeyError(f'{vol_dir} {level}: transform.json rotated 段未記錄平面 ' + f"sym/endplate（舊世代資料；先跑 xfr_record_planes.py 補記，" + f"或重跑 xfr_preprocess 的 rotated 步驟）") + sym = rot['sym'] + for k in ('plane', 'normal', 'u', 'v'): + if k not in sym: + raise ValueError(f'{vol_dir} {level}: 記錄平面 sym 缺 {k!r}') + endplate = rot.get('endplate') + if endplate is not None and 'plane' not in endplate: + raise ValueError(f'{vol_dir} {level}: 記錄平面 endplate 缺 "plane"') + return sym, endplate + + +def load_rotated_planes_from_path(image_path, level=None): + """…/{vol}/rotated/{level}_*.nii.gz -> load_rotated_planes(vol_dir, level)； + level 未給定時由檔名前綴（split('_')[0]）取。不在 rotated/ 下時 raise。""" + parent = os.path.dirname(image_path) + if os.path.basename(parent) != 'rotated': + raise ValueError(f'{image_path} 不在 /rotated/ 下，無法讀記錄平面' + f'（xfr 管線須用 rotated/ 輸入）') + if level is None: + level = os.path.basename(image_path).split('_')[0] + return load_rotated_planes(os.path.dirname(parent), level) + + +def load_rotated_pedicle(vol_dir, level): + """讀 transform.json levels[level]['rotated']['pedicle'] 記錄的雙側椎弓根寬度 + dict（measure_pedicle_width，xfr_preprocess 旋轉時寫入）： + {'L_mm','R_mm'（mediolateral 中位數寬度 mm，clip 到 MAX_PEDICLE_WIDTH_MM， + None = 該側量測不到）, 'L/R_raw_mm', 'L/R_clipped', ...}。 + rotated 段或 pedicle 欄位缺失（舊世代資料只補記了平面）回傳 None。""" + meta = load_transform(vol_dir) + lv = meta['levels'].get(level) + rot = lv.get('rotated') if isinstance(lv, dict) else None + if not isinstance(rot, dict): + return None + ped = rot.get('pedicle') + return ped if isinstance(ped, dict) else None + + +def load_rotated_pedicle_from_path(image_path, level=None): + """…/{vol}/rotated/{level}_*.nii.gz -> load_rotated_pedicle(vol_dir, level)； + level 未給定時由檔名前綴（split('_')[0]）取。不在 rotated/ 下時 raise。""" + parent = os.path.dirname(image_path) + if os.path.basename(parent) != 'rotated': + raise ValueError(f'{image_path} 不在 /rotated/ 下，無法讀記錄椎弓根寬度' + f'（xfr 管線須用 rotated/ 輸入）') + if level is None: + level = os.path.basename(image_path).split('_')[0] + return load_rotated_pedicle(os.path.dirname(parent), level) + + def img_geom(img): """sitk img -> {"size"(x,y,z), "spacing", "direction"(row-major 9), "origin"}。""" return { diff --git a/utils/helpers.py b/utils/helpers.py index b89295d..0a52cc9 100644 --- a/utils/helpers.py +++ b/utils/helpers.py @@ -33,9 +33,16 @@ def get_unique_filepath(path: str) -> str: return current_path +def _fmt_length(v): + """螺絲長度進檔名時不帶小數點（35.0 -> 35）； + 直徑（width）維持原樣（4.5 / 5.0，仍帶一位小數）。 + '' （單側缺失時的佔位）原樣保留。""" + return str(int(round(float(v)))) if v != '' else v + + def save_with_unique_name(folder, label_str, way, diameter_l, length_l, diameter_r, length_r, swarm_size, max_iter): - - base_name = f'{label_str}_{way}_L{diameter_l}_{length_l}_R{diameter_r}_{length_r}_{swarm_size}_{max_iter}.png' + + base_name = f'{label_str}_{way}_L{diameter_l}_{_fmt_length(length_l)}_R{diameter_r}_{_fmt_length(length_r)}_{swarm_size}_{max_iter}.png' file_path = os.path.join(folder, base_name) count = 1 diff --git a/visualization/res_bone_figure.py b/visualization/res_bone_figure.py index 5baa1cd..4d96b39 100644 --- a/visualization/res_bone_figure.py +++ b/visualization/res_bone_figure.py @@ -258,18 +258,26 @@ 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, + sym=None, symp=None, best_position_l=None, best_position_r=None, diameter_l=None, length_l=None, diameter_r=None, length_r=None, image2_path=None, device=None, grid=None, swarm_size=None, max_iter=None, total_time=None, - write_csv=True): + write_csv=True, screw_shape='cone'): """統一骨頭 X-ray 四視角圖（合併原 render_bone_figure + res_plt_2_torch）。 四視角：預設 / axial（俯視 XY）/ 冠狀（後視）/ 矢狀。 內容：皮質 vs 鬆質吸收骨、椎體(gold)、棘突(purple 半透明，獨立層覆蓋椎體)、 - 中矢狀鏡稱面(orange)、上終板面(green)；螺絲模式另畫中心線(紅) - + 圓柱（L darkcyan / R blue，o 層粉）。 + 中矢狀鏡稱面(orange)、上終板面(green)。螺絲模式下，左上的預設 3D 視角 + = 完整骨頭點雲（含 VBODY / SP voxel、正常皮質 / 鬆質色，不做 + gold / purple 著色）+ 螺絲（半透明，嵌在骨內、前後骨頭透見）、 + 不畫平面；其余三個正視圖（右上 axial / 左下 coronal / 右下 + sagittal）全畫，平面填實不畫格線，螺絲不透明；planes-only 圖 + 四視角維持格線全畫。 + 螺絲模式另畫中心線(紅) + 螺絲體（screw_shape 預設 cone 錐體： + 入口端全直徑、線性收細至尖端；cylinder 舊式圓柱； + L darkcyan / R blue，o 層粉）。 繪製採固定分層（不依深度排序）： 基底骨 < VBODY < 螺絲(棘突後方) < 棘突(半透明) < 終板 < 鏡稱面 < 螺絲(棘突前方) 螺絲與棘突另依各視角相機深度拆分：比棘突中位深度深的螺絲畫在棘突之下， @@ -286,16 +294,35 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, 平面 / 分割： planes_only=True 只畫平面，不做棘突 / 椎體（VBODY）分割 - rotation=(R, c_xyz) 點雲與平面同依 R 旋轉（繞 c_xyz），用於「對齊後 - （rotated）」的平面圖；R 作用於 (x,y,z) 向量 - （螺絲模式一律做分割：R 側 loss 有 VBODY 獎勵、gold - 顯示需椎體，planes_only 被忽略） + rotation=(R, c_xyz[, shift]) 點雲與平面同依 R 旋轉（繞 c_xyz）， + 用於「對齊後（rotated）」的平面圖；R 作用於 (x,y,z) + 向量；給 shift（rotated/ 存檔 grid origin，fstart） + 時輸出換到存檔檔局部 index 系（與 CBT 螺絲圖同系） + （螺絲模式一律做分割：R 側 loss 有 VBODY 獎勵、gold + 顯示需椎體，planes_only 被忽略） + sym / symp 給定記錄平面 dict（best_symmetry_plane / + best_upper_endplate_plane 回傳同結構、與所給骨頭 + mask 同座標系）時，圖上平面 + 棘突 / 椎體著色 + + 角度註記全用記錄值、不在影像上重新估算（xfr 管線 + 的 CBT 螺絲圖傳 transform.json 記錄平面，圖與 + optimizer 完全一致）；未給定時維持重新估算 + （舊世代 / 未旋轉路徑） 螺絲（可選，給定 best_position_l/r 時啟用，optimizer 的輸出）： best_position = (z, y, x, az, alt[, d, L])，須搭配該側 diameter/length； device/grid 供圓柱生成；swarm_size/max_iter/total_time 供圖面註記； + screw_shape='cone'（預設）圖上螺絲體畫成錐體（入口端全直徑、 + 線性收細至尖端）；'cylinder' 維持舊式圓柱。 + 只影響繪製：評分（Best_Loss / Overlap / cyl_points） + 與 CSV、輸出檔名仍以圓柱為準； write_csv 時在 base_folder/{date}/{vol}/output.csv 依欄位名稱 append - L/R 兩行（舊 schema 自動重映射）。 + L/R 兩行（舊 schema 自動重映射）。CSV 的 Position_XYZ 維持 voxel + index（xfr_rerender_spinous 重新渲染 / optimizer 直接讀用）； + spacing 為各軸實體间距（mm 序 x, y, z）。 + + 座標單位：圖上所有座標（軸刻度、點雲 / 平面 / 螺絲、L/R Position + 註記）皆以實體 mm 顯示（index * spacing）；內部幾何仍以 voxel index + 計算、存檔座標系（rotated/*.nii.gz 局部 index、transform.json）不變。 輸出： output_path 給定 -> 存該路徑（含自動加 _1/_2 防覆蓋）；None 時 @@ -330,6 +357,10 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, screw_sides.append(s) screw_mode = bool(screw_sides) CBT = str(way).upper() == 'CBT' + screw_shape = str(screw_shape).lower() + if screw_shape not in ('cone', 'cylinder'): + raise ValueError(f"invalid screw_shape: {screw_shape!r} (use 'cone' or 'cylinder')") + _SCREW_BODY_NAME = 'Cone' if screw_shape == 'cone' else 'Cylinder' # ---- 骨骼 / 皮質 ---- spine = _mask_to_array(binary_path) @@ -351,8 +382,13 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, z_trab, y_trab, x_trab = np.where((spine) & (~cortical)) # ---- 方向 / 解剖分割 ---- - sym = best_symmetry_plane(spine) - symp = best_upper_endplate_plane(spine) + # sym / symp 給定（記錄平面、與 mask 同系）時直接用，確保圖上平面 / + # 著色 / 角度註記與 optimizer 所用一致；未給定才在 mask 上重新估算 + # （給定 mask 可能已被改過——如 CBT 圖的 SP 中線孔——重擬合會漂移）。 + if sym is None: + sym = best_symmetry_plane(spine) + if symp is None: + symp = best_upper_endplate_plane(spine) if planes_only and not screw_mode: # 只做方向平面，不做棘突 / 椎體分割 @@ -442,8 +478,13 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, vb_flag &= ~sp_flag # ---- 旋轉對齊（若給定）：骨頭點雲與平面同旋轉 ---- + # rotation = (R, c_xyz) 或 (R, c_xyz, shift)：shift 為旋轉後座標系的 + # 額外平移（rotated/ 存檔 grid origin fstart）時，點雲與平面落在存檔檔 + # 局部 index 系（V = R(W-c)+c，file idx = V - shift）——與 CBT 螺絲圖 + # 直接讀的 rotated/*.nii.gz 同座標系；shift 未給定時維持 V 系（未裁切）。 if rotation is not None: - R, c_xyz = rotation + R, c_xyz = rotation[0], rotation[1] + shift = rotation[2] if len(rotation) > 2 else None 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] @@ -452,6 +493,13 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, sym = _rotate_plane_params(sym, R, c_xyz) if symp is not None: symp = _rotate_plane_params(symp, R, c_xyz) + if shift is not None: + x_bone = x_bone - shift[0] + y_bone = y_bone - shift[1] + z_bone = z_bone - shift[2] + sym = _shift_plane_params(sym, shift) + if symp is not None: + symp = _shift_plane_params(symp, shift) # ---- 拆三層：基底骨 / VBODY (gold) / 棘突 (purple) ---- # x_bone 保留完整點雲（含 VBODY / SP）供下方平面 patch 算範圍 @@ -479,7 +527,8 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, raise ValueError('螺絲模式需要 device（torch device）') spacing = list(spacing) # core.cylinder 以 list 比對 spacing import torch - from core.cylinder import (generate_cylinder_n_torch, + from core.cylinder import (generate_cone_torch, + generate_cylinder_n_torch, generate_cylinder_o_torch, generate_cylinder_butt_torch) from core.intersection import center_line_intersections_torch @@ -559,9 +608,17 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, / cyl_points) if cyl_points else 0.0 ovb = (100.0 * int(((spine_tensor == 1) & (cyl_n == 1)).sum().item()) / cyl_points) if cyl_points else 0.0 + # 螺絲繪製體：cone（預設）入口端全直徑線性收細至尖端 / + # cylinder 舊式圓柱；評分（loss/ovc/ovb/CSV）仍用 cyl_n 圓柱 + body = cyl_n + if screw_shape == 'cone': + body = generate_cone_torch(d, L, pos[0], pos[1], pos[2], + pos[3], pos[4], + image_shape, spacing, device, grid) side_info[s] = { 'line': np.where(line_mask.cpu().numpy() == 1), 'cyl_n': np.where(cyl_n.cpu().numpy() == 1), + 'body': np.where(body.cpu().numpy() == 1), 'cyl_o': np.where(cyl_o.cpu().numpy() == 1), 'loss': float(loss), 'inter': inter, @@ -572,7 +629,8 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, side_azlat[s], side_acep[s] = _rel_angles(float(pos[3]), float(pos[4])) # 螺絲點雲（存在的側接合）：中心線紅、圓柱 n = L darkcyan / R blue、o = 粉。 - # 點序固定為舊 res_plt_2_torch 的接合序 [L線, R線, L_n, L_o, R_n, R_o]： + # 點序固定為舊 res_plt_2_torch 的接合序 [L線, R線, L體, L_o, R體, R_o] + #（body = cone 預設 / cylinder 舊式，佔 cyl_n 的位置）： # mpl 3D 的 per-point 深度排序對近同深 markers 以輸入序決 tie，換序會讓 # 同一點集產生亞像素級抗鋸齒邊界差（A/B 實測 ~0.07% 邊緣像素）， # 固定點序保持與舊圖位元級一致。 @@ -585,10 +643,10 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, parts_x.append(x_p); parts_y.append(y_p); parts_z.append(z_p) parts_c.append(_rgba_block(len(x_p), 'r', 1.0)) parts_s.append(np.full(len(x_p), 3)) - for s in ('L', 'R'): # 圓柱：各側 n 後 o（L_n, L_o, R_n, R_o） + for s in ('L', 'R'): # 螺絲體（cone 預設 / cylinder）：各側 body 後 o（L體, L_o, R體, R_o） if s not in side_info: continue - for key, color, size in (('cyl_n', None, 36), ('cyl_o', 'pink', 36)): + for key, color, size in (('body', None, 36), ('cyl_o', 'pink', 36)): z_p, y_p, x_p = side_info[s][key] parts_x.append(x_p); parts_y.append(y_p); parts_z.append(z_p) parts_c.append(_rgba_block(len(x_p), _CYL_COLOR[s] if color is None else color, 1.0)) @@ -632,6 +690,21 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, _EY = _ep0[1] + _EU * _eu[1] + _EV * _ev[1] _EZ = _ep0[2] + _EU * _eu[2] + _EV * _ev[2] + # ---- 顯示座標系：index (voxel) -> 實體 mm ---- + # 把所有要繪製的點雲 / 平面 / 螺絲座標陣列乘以各軸 spacing（mm）： + # 圖上軸刻度與 Position 註記自茲顯示實體 mm。內部幾何（旋轉 / 分割 / + # 深度拆分 / 平面範圍）皆已在上方以 index 系算畢，此處只換顯示單位。 + sx_mm, sy_mm, sz_mm = float(spacing[0]), float(spacing[1]), float(spacing[2]) + x_bone = x_bone * sx_mm; y_bone = y_bone * sy_mm; z_bone = z_bone * sz_mm + x_base = x_base * sx_mm; y_base = y_base * sy_mm; z_base = z_base * sz_mm + x_vb = x_vb * sx_mm; y_vb = y_vb * sy_mm; z_vb = z_vb * sz_mm + x_sp = x_sp * sx_mm; y_sp = y_sp * sy_mm; z_sp = z_sp * sz_mm + if x_screw is not None: + x_screw = x_screw * sx_mm; y_screw = y_screw * sy_mm; z_screw = z_screw * sz_mm + _Xp = _Xp * sx_mm; _Yp = _Yp * sy_mm; _Zp = _Zp * sz_mm + if _EX is not None: + _EX = _EX * sx_mm; _EY = _EY * sy_mm; _EZ = _EZ * sz_mm + # ---- 圖 ---- fig = plt.figure(figsize=(12, 12)) @@ -640,7 +713,7 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, if s in side_info: legend_handles.append(Line2D([], [], marker="o", ls="", ms=6, color="darkcyan" if s == "L" else "blue", - label=f"Cylinder({s})")) + label=f"{_SCREW_BODY_NAME}({s})")) if vb_corti is not None: legend_handles.append(Line2D([], [], marker="o", ls="", ms=6, color="gold", label="VertebralBody")) if sp_corti is not None: @@ -654,7 +727,10 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, np.cos(e) * np.sin(a), np.sin(e)]) - def _fill_ax(ax, elev, azim): + def _fill_ax(ax, elev, azim, deco=True): + # deco=False 的視角畫完整骨頭點雲（含 VBODY / SP voxel、正常皮質 / + # 鬆質色，不做 gold / purple 著色）+ 螺絲（無平面、無棘突遮蔽 + # 拆分）——螺絲圖的左上預設 3D 視角用。 # 固定分層（關 depth zorder，否則半透明骨頭會被重繪到螺絲上方）： # 基底骨(5) < VBODY gold(6) < 螺絲「棘突後方」(6.2) # < 棘突 purple(6.5, 半透明) < 終板(7) < 鏡稱面(8) < 螺絲「棘突前方」(10) @@ -662,20 +738,43 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, # 比棘突中位深度深（farther）的螺絲先畫、被半透明紫色擋住但仍透見 # （正確遮蔽螺絲路徑）；較淺的螺絲照舊畫在最上層。無棘突時不拆分。 ax.computed_zorder = False - sc_bone = ax.scatter(x_base, y_base, z_base, c=rgba_base, s=size_base, marker="o") - sc_bone.set_zorder(5) - if x_vb.size: - # (1,4) 單行 2D 陣列：整組點共用同一色，避免 *c* 被視為數值映射 - sc_vb = ax.scatter(x_vb, y_vb, z_vb, c=np.array([to_rgba("gold", 0.95)]), - s=BONE_MARKER_SIZE, marker="o") - sc_vb.set_zorder(6) + if deco: + sc_bone = ax.scatter(x_base, y_base, z_base, c=rgba_base, s=size_base, marker="o") + sc_bone.set_zorder(5) + if x_vb.size: + # (1,4) 單行 2D 陣列：整組點共用同一色，避免 *c* 被視為數值映射 + sc_vb = ax.scatter(x_vb, y_vb, z_vb, c=np.array([to_rgba("gold", 0.95)]), + s=BONE_MARKER_SIZE, marker="o") + sc_vb.set_zorder(6) + else: + # 完整骨頭（x_bone 含 VBODY / SP voxel），正常皮質 / 鬆質著色； + # per-point alpha 提昇 10 倍讓骨頭雲較不透明。 + # 螺絲（半透明 ×0.4）併入同一 collection：mpl 3D scatter 對 + # 同集合全部點依相機深度由遠到近绘制（alpha 混合），螺絲前方 + # 的骨頭自然畫在螺絲之上、後方的骨頭在其下——螺絲看起來嵌在 + # 骨內，而非浮在最上層（分開畫兩層 + zorder 做不到的效果）。 + c_bone = np.array(bone_rgba, copy=True) + c_bone[:, 3] = np.clip(c_bone[:, 3] * 10.0, 0.0, 1.0) + if x_screw is not None and len(x_screw): + c_screw = np.array(screw_rgba, copy=True) + c_screw[:, 3] *= 0.4 + x_all = np.concatenate([x_bone, x_screw]) + y_all = np.concatenate([y_bone, y_screw]) + z_all = np.concatenate([z_bone, z_screw]) + c_all = np.concatenate([c_bone, c_screw]) + s_all = np.concatenate([bone_size, screw_size]) + else: + x_all, y_all, z_all = x_bone, y_bone, z_bone + c_all, s_all = c_bone, bone_size + sc_bone = ax.scatter(x_all, y_all, z_all, c=c_all, s=s_all, marker="o") + sc_bone.set_zorder(5) # 螺絲拆「棘突前方 / 後方」（front = 本視角較淺或等深）：無螺絲時 - # None；無棘突時全 True（維持舊行為）；有棘突時依相機深度 vs 棘突 - # 中位深度拆分 + # None；無棘突 / 無著色視角時全 True（維持舊行為）；有棘突時依 + # 相機深度 vs 棘突中位深度拆分 if x_screw is None: front = None - elif x_sp.size: + elif deco and x_sp.size: cam = _view_dir(elev, azim) d_screw = x_screw * cam[0] + y_screw * cam[1] + z_screw * cam[2] d_sp = x_sp * cam[0] + y_sp * cam[1] + z_sp * cam[2] @@ -688,32 +787,46 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, sc_bh = ax.scatter(x_screw[m], y_screw[m], z_screw[m], c=screw_rgba[m], s=screw_size[m], marker="o") sc_bh.set_zorder(6.2) - if x_sp.size: + if deco and x_sp.size: sc_sp = ax.scatter(x_sp, y_sp, z_sp, c=np.array([to_rgba("purple", float(alpha_spinous))]), s=BONE_MARKER_SIZE, marker="o") sc_sp.set_zorder(6.5) - if _EX is not None: + # 螺絲圖（如 CBT 四視角圖）兩平面填實、不畫格線——格線會蓋在 + # 螺絲 / 骨頭點雲上；planes-only（preprocess）圖維持格線。 + # 左上的預設 3D 視角（deco=False）不畫平面與 SP / VBODY 著色； + # 上 / 冠狀 / 矢狀三個正視圖保留（平面在正視圖是完整直線）。 + if deco and _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) + linewidth=1.0, + edgecolor="none" if screw_mode else "green", + rstride=1, cstride=1) ep.set_zorder(7) - plane = ax.plot_surface(_Xp, _Yp, _Zp, color="orange", alpha=0.30, - linewidth=1.0, edgecolor="orange", rstride=1, cstride=1) - plane.set_zorder(8) - if front is not None and front.any(): + if deco: + plane = ax.plot_surface(_Xp, _Yp, _Zp, color="orange", alpha=0.30, + linewidth=1.0, + edgecolor="none" if screw_mode else "orange", + rstride=1, cstride=1) + plane.set_zorder(8) + # deco 視角：螺絲已併入骨頭 collection 依深度交織繪製（見上方 else 分支）。 + # 其餘三視角維持舊行為：不透明、單層最上（含棘突前後拆分）。 + if deco and front is not None and front.any(): sc_screw = ax.scatter(x_screw[front], y_screw[front], z_screw[front], - c=screw_rgba[front], s=screw_size[front], marker="o") + c=screw_rgba[front], s=screw_size[front], marker="o") sc_screw.set_zorder(10) ax1 = fig.add_subplot(221, projection="3d") - _fill_ax(ax1, 30, -60) # matplotlib 3D 預設視角 - ax1.set_xlabel("X-axis"); ax1.set_ylabel("Y-axis"); ax1.set_zlabel("Z-axis") + # 螺絲圖：左上預設視角 = 完整骨頭點雲（含 VBODY / SP、正常骨色，不 + # 做 gold / purple 著色）+ 螺絲，不畫平面（見 _fill_ax deco）； + # planes-only 圖四視角都全畫 + _fill_ax(ax1, 30, -60, deco=not screw_mode) + ax1.set_xlabel("X (mm)"); ax1.set_ylabel("Y (mm)"); ax1.set_zlabel("Z (mm)") set_axes_equal_3d(ax1) ax2 = fig.add_subplot(222, projection="3d") ax2.view_init(elev=90, azim=-90, roll=0) _fill_ax(ax2, 90, -90) - ax2.set_xlabel("X-axis"); ax2.set_ylabel("Y-axis"); ax2.set_zlabel("Z-axis") + ax2.set_xlabel("X (mm)"); ax2.set_ylabel("Y (mm)"); ax2.set_zlabel("Z (mm)") set_axes_equal_3d(ax2) if legend_handles: ax2.legend(handles=legend_handles) @@ -722,13 +835,13 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, # 後視圖：相機在 −y 後側，x 軸畫面左小右大 ax3.view_init(elev=0, azim=-90, roll=0) _fill_ax(ax3, 0, -90) - ax3.set_xlabel("X-axis"); ax3.set_ylabel("Y-axis"); ax3.set_zlabel("Z-axis") + ax3.set_xlabel("X (mm)"); ax3.set_ylabel("Y (mm)"); ax3.set_zlabel("Z (mm)") set_axes_equal_3d(ax3) ax4 = fig.add_subplot(224, projection="3d") ax4.view_init(elev=0, azim=0, roll=0) _fill_ax(ax4, 0, 0) - ax4.set_xlabel("X-axis"); ax4.set_ylabel("Y-axis"); ax4.set_zlabel("Z-axis") + ax4.set_xlabel("X (mm)"); ax4.set_ylabel("Y (mm)"); ax4.set_zlabel("Z (mm)") set_axes_equal_3d(ax4) label_str = f"{volume_id} {level}" @@ -771,7 +884,7 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, cb_ratio = di['ovc'] / di['ovb'] if di['ovb'] else 0.0 fig.text( 0.5, y, - f"{cn_name} : Position = ({pos[2]:.2f}, {pos[1]:.2f}, {pos[0]:.2f}), " + f"{cn_name} : Position = ({pos[2]*sz_mm:.2f}, {pos[1]*sy_mm:.2f}, {pos[0]*sx_mm:.2f}), " f"{segs}, " f"Intersection = {di['inter']}, " f"Score = {di['ovc']:.2f} / {di['ovb']:.2f} / {cb_ratio:.2f}", diff --git a/xfr_cbt_native.py b/xfr_cbt_native.py index 3819003..8ae0e5c 100644 --- a/xfr_cbt_native.py +++ b/xfr_cbt_native.py @@ -280,22 +280,40 @@ def render_xray_projections(cbt_path, ct_path, out_dir, - spine mask：native 分割 label（1-24 = C1..L5，見 config.constant.LABEL_MAP）； label 缺 / grid 不符時退回 HU 300-3000 閾值。 - 骨 μ = clip(HU, 0, 2000)/400（松質骨~0.2-1、皮質/終板~2-5）。 - - 螺絲 voxel μ = 60（金屬等效，最亮白）。 + - 螺絲 voxel μ = 60（金屬等效，最亮白）；螺絲維持白色，每支螺絲以 + 合併圖螺絲體色（darkcyan = L、blue = R）標註，兩面方式不同： + AP 視圖 = bounding box；Lateral 視圖 = 投影外緣 2 px margin 環。 投影视圖 = 沿中心射線 μ 的線積分（AP 沿 y、Lateral 沿 x）；裁到螺絲 bbox 外 margin_mm 的 spine 區域；骨用 percentile(1,99) 獨立視窗，螺絲再疊加裁白。 - 方向：SimpleITK LPS（x→右、y→後、z→上）；AP 頂=上、病人右在畫面左（R 標記）； - Lateral 頂=上、前位在左、後位在右（A/P 標記）。 + 方向：讀 native CT header 的 direction（LPS：世界 +x=左、+y=後、+z=上； + 本 dataset 有 identity 與 x/y 反號兩種 header，flip 依 direction 決定）； + 螢幕方向同合併圖：x-ap = 左下圖（畫面右=病人右、R 標記在右）、 + x-lat = 右下圖（畫面右=前位、A 標記在右）、頂恆=頭端； + 非軸對齊 direction（離對角 ±1 的 off-diagonal）會 warning（投影軸錯位）。 spacing 各軸不等時先重取樣到 min(spacing) 各向同性格，維持投影 aspect ratio。 輸出 x-ap.jpg / x-lat.jpg，回傳 {view: jpg_path}。""" import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt + from scipy.ndimage import binary_erosion vid = os.path.basename(os.path.normpath(out_dir)) ct = sitk.ReadImage(ct_path) cbt = sitk.ReadImage(cbt_path) if ct.GetSize() != cbt.GetSize(): raise ValueError(f'cbt/CT grid 尺寸不符 {cbt.GetSize()} vs {ct.GetSize()}') + D = np.array(ct.GetDirection(), float).reshape(3, 3) # D[i,j] = 影像軸 i 在世界 LPS(x左,y後,z上) 的分量 + # 螢幕方向同合併圖（L*_CBT_*.png 子圖）： + # x-ap = 左下圖（後視：相機在 −y 後側）→ 畫面右 = +x rotated = 病人右（R 標記在右） + # x-lat = 右下圖（矢狀：相機在 +x）→ 畫面右 = +y rotated = 前位（A 標記在右） + # 兩面頂皆 = 頭端（+z） + off = max(abs(D[i, j]) for i in range(3) for j in range(3) if i != j) + if off > 0.5: + logger.warning(f'{vid}: CT direction 非軸對齊 {np.round(D, 2).tolist()}; ' + f'投影视圖方向可能不正確') + flip_rows = D[2, 2] > 0 # z+=上 -> 頂端要放高 z index（翻行） + flip_cols = {'ap': D[0, 0] > 0, # x+=左 -> 畫面右（病人右）要放低 x index（翻欄） + 'lateral': D[1, 1] > 0} # y+=後 -> 畫面右（前位）要放低 y index sp = np.array(ct.GetSpacing(), float) # (x,y,z) native_size = ct.GetSize() # spacing 各軸不等：先重取樣 CT/螺絲（與後面的分割 label）到 @@ -347,16 +365,21 @@ def render_xray_projections(cbt_path, ct_path, out_dir, 'ap': (bone_mu[zs, :, xs].sum(axis=1), metal_mu[zs, :, xs].sum(axis=1)), # 沿 y 積分 'lateral': (bone_mu[zs, ys, :].sum(axis=2), metal_mu[zs, ys, :].sum(axis=2)) # 沿 x 積分 } - flip_x = {'ap': True, 'lateral': False} - markers = {'ap': ('R', 'L'), 'lateral': ('A', 'P')} # 畫面左=前位(A)、右=後位(P) + markers = {'ap': ('L', 'R'), 'lateral': ('P', 'A')} # 依上方 flip_* 對齊後恆成立 + # 螺絲標註色 = 該側合併圖螺絲體色（res_bone_figure 的 _CYL_COLOR）： + # L darkcyan / R blue；兩面標註方式不同： + # AP = 每支螺絲 bounding box；Lateral = 每支螺絲投影外緣 2 px margin 環 + _SCREW_ANNOT_RGB = {'L': (0.0, 139 / 255.0, 139 / 255.0), 'R': (0.0, 0.0, 1.0)} os.makedirs(out_dir, exist_ok=True) paths = {} for view in views: if view not in proj: continue b, m = proj[view] - b = b[::-1, ::-1] if flip_x[view] else b[::-1, :] # 頂=頭端 - m = m[::-1, ::-1] if flip_x[view] else m[::-1, :] + if flip_rows: + b = b[::-1, :]; m = m[::-1, :] # 頂=頭端 + if flip_cols[view]: + b = b[:, ::-1]; m = m[:, ::-1] bnz = b[b > 0] if bnz.size: lo, hi = np.percentile(bnz, 1), np.percentile(bnz, 99) @@ -369,7 +392,57 @@ def render_xray_projections(cbt_path, ct_path, out_dir, ll, rl = markers[view] h, w = img.shape fig, ax = plt.subplots(figsize=(8.0 * w / h, 8.0), dpi=110) - ax.imshow(img, cmap='gray', interpolation='nearest') + # 螺絲（label 1-10 = L1L L1R … L5L L5R，奇=L 偶=R）本身維持白色； + # 標註座標與投影视圖同系，flip 依 CT direction（上方 flip_rows/flip_cols） + n_z = zs.stop - zs.start + n_x = xs.stop - xs.start + n_y = ys.stop - ys.start + if view == 'ap': + # AP：每支螺絲一個 bounding box + ax.imshow(img, cmap='gray', interpolation='nearest') + for val in range(1, 11): + ms = cbt_arr == val + if not ms.any(): + continue + zz, yy, xx = np.where(ms) + keep = ((zz >= zs.start) & (zz < zs.stop) + & (xx >= xs.start) & (xx < xs.stop)) + if not keep.any(): + continue + zpos = zz[keep] - zs.start + r_top, r_bot = ((n_z - 1 - int(zpos.max()), n_z - 1 - int(zpos.min())) + if flip_rows else (int(zpos.min()), int(zpos.max()))) + xpos = n_x - 1 - (xx[keep] - xs.start) if flip_cols['ap'] \ + else xx[keep] - xs.start + c_l, c_r = int(xpos.min()), int(xpos.max()) + pad = 1.0 + ax.add_patch(plt.Rectangle( + (c_l - 0.5 - pad, r_top - 0.5 - pad), + c_r - c_l + 1 + 2 * pad, r_bot - r_top + 1 + 2 * pad, + fill=False, + edgecolor=_SCREW_ANNOT_RGB['L' if val % 2 == 1 else 'R'], + linewidth=1.5, zorder=5)) + else: + # Lateral：每支螺絲投影的外緣 2 px 環（margin）著色 + rgb = np.stack([img] * 3, axis=-1) + for val in range(1, 11): + ms = cbt_arr == val + if not ms.any(): + continue + zz, yy, xx = np.where(ms) + keep = ((zz >= zs.start) & (zz < zs.stop) + & (yy >= ys.start) & (yy < ys.stop)) + if not keep.any(): + continue + zpos = zz[keep] - zs.start + r = (n_z - 1 - zpos) if flip_rows else zpos + c = (n_y - 1 - (yy[keep] - ys.start)) if flip_cols['lateral'] \ + else yy[keep] - ys.start + mask = np.zeros((n_z, n_y), bool) + mask[r.astype(int), c.astype(int)] = True + outline = mask & ~binary_erosion(mask, iterations=2) + rgb[outline] = _SCREW_ANNOT_RGB['L' if val % 2 == 1 else 'R'] + ax.imshow(rgb, interpolation='nearest') ax.set_title(f'{view.upper()} projection (spine + screws) - {vid}', color='white', fontsize=13) ax.text(0.01, 0.98, ll, transform=ax.transAxes, color='cyan', fontsize=13, diff --git a/xfr_de.py b/xfr_de.py new file mode 100644 index 0000000..0baa52a --- /dev/null +++ b/xfr_de.py @@ -0,0 +1,495 @@ +#!/home/xfr/.conda/envs/cbt/bin/python + +import logging +import os +import re +import sys +import time +import queue as queue_module +import subprocess +import multiprocessing as mp +from concurrent.futures import ThreadPoolExecutor + +import SimpleITK as sitk +import torch + +from core.cylinder import create_coordinate_grid +from core.optimizer import run_de_torch_xfr + +import xfr_cbt_native + +standardized_dir = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3/' + +Output_dir = '/mnt/1248/open/cyrou/Output' + +def available_levels(volume_id): + """該 volume 可跑 debug_de 的全部 lumbar level：rotated/ 中輸入三件 + （_cortical / _binary_sdf / _roi）齊全的 L\\d。各 volume 層數不同 + （例：有的只有 L1~L3、有的含 L6），故不用固定 LEVELS 清單。""" + rotated_dir = os.path.join(standardized_dir, volume_id, 'rotated') + if not os.path.isdir(rotated_dir): + return () + levels = [] + for fn in os.listdir(rotated_dir): + m = re.fullmatch(r'(L[1-9]\d*)_cortical\.nii\.gz', fn) + if not m: + continue + level = m.group(1) + if (os.path.exists(os.path.join(rotated_dir, f'{level}_binary_sdf.nii.gz')) + and os.path.exists(os.path.join(rotated_dir, f'{level}_roi.nii.gz'))): + levels.append(level) + return tuple(sorted(levels, key=lambda lv: int(lv[1:]))) + + +LOG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logs') + +# 單側任務的中繼結果（//_.json + plot lock）： +# 同一 level 的 L/R 可跑在不同 GPU，較晚完成的一側讀到兩側結果後跑合併輸出 +# （由 run_de_torch_xfr 寫入，格式與 run_pso_torch_xfr 相同）；per-volume +# cbt 寫出（write_volume_cbt）照舊讀這些檔映回 native +SIDE_RESULT_DIR = os.path.join(LOG_DIR, 'side_results') + +logger = logging.getLogger('xfr_debug') + + +# 目前任務標記（每個 worker 流程各自一份），讓交錯的 log 可以歸屬到 (volume, level) +_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, side=None): + _TASK_TAG['vid'] = volume_id + _TASK_TAG['level'] = level + _TASK_TAG['side'] = side + + +class _Tee: + """同時輸出到 console 與 log 檔，逐行加上如 [0001 L1 LEFT] 的標記。 + volume 取 id 末段（1.3.6.1.4.1.9328.50.4.0001 -> 0001）； + side 由 [LEFT]/[RIGHT]/左側/右側 段落標記判定並沿用給後續行； + 最終結果段屬於整椎，重置沿用值，其中的 Left/Right 摘要行只標該行本身； + np.float64(...) 之類包裹會解包成裸數值。""" + + def __init__(self, console, log_fh): + self.console = console + self.log_fh = log_fh + self.buf = '' + + @staticmethod + def _marker_side(line): + """回傳 (side, persist)：[LEFT]/[RIGHT]/左側/右側 是段落標記（persist=True）， + Left/Right 摘要行只標該行（persist=False）""" + s = line.lstrip() + if s.startswith('[LEFT]') or '左側' in s: + return 'LEFT', True + if s.startswith('[RIGHT]') or '右側' in s: + return 'RIGHT', True + if s.startswith('Left '): + return 'LEFT', False + if s.startswith('Right '): + return 'RIGHT', False + return None, False + + def _prefix(self, side): + if not _TASK_TAG['vid'] or not _TASK_TAG['level']: + return '' + vol = _TASK_TAG['vid'].rsplit('.', 1)[-1] + tag = f'{vol} {_TASK_TAG["level"]}' + if side: + tag += f' {side}' + return f'[{tag}] ' + + def _emit(self, line): + line = _NP_WRAP_RE.sub(r'\1', line) + if '最終結果' in line: + _TASK_TAG['side'] = None + side, persist = self._marker_side(line) + if persist: + _TASK_TAG['side'] = side + eff_side = side if side is not None else _TASK_TAG['side'] + prefix = self._prefix(eff_side) if line.strip() else '' + if prefix: + # 前綴已含 side 時，去掉行首重複的 [LEFT] / [RIGHT] 標記 + marker = f'[{eff_side}] ' + if line.startswith(marker): + line = line[len(marker):] + self.console.write(prefix + line + '\n') + self.log_fh.write(prefix + 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 _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: + max_free = -1 + gpu_id = 0 + for i in range(torch.cuda.device_count()): + free_mem, _ = torch.cuda.mem_get_info(i) + if free_mem > max_free: + max_free = free_mem + gpu_id = i + device = torch.device(f"cuda:{gpu_id}") + logger.info(f"Using GPU {gpu_id}: {torch.cuda.get_device_name(gpu_id)}") + else: + device = torch.device("cpu") + logger.info("CUDA not available, using CPU") + return device + + +def debug_de(volume_id, level, device=None, side='both', run_id=None): + # ====== DE ====== + # swarm_size 會換算成 de_popsize = swarm_size // 6（6 維 bounds，y 由 + # y_indices 表面取），實際族群 ≈ swarm_size，跟 PSO 版同值以便比較 + swarm_size = 100 + max_iter = 100 + + # ====== DEVICE ====== + if device is None: + device = get_device() + + # ====== OTHER ====== + spacing = [0.5, 0.5, 0.5] + 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') + + 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}') + + binary_image = sitk.ReadImage(binary_path) + binary_array = sitk.GetArrayFromImage(binary_image) + image_shape = binary_array.shape + + # 資料層級的快速失敗（例：…9328.50.4.0653 L5 是 1-voxel 寬的退化體積） + if binary_array.sum() == 0: + raise ValueError(f'{volume_id} {level}: empty bone mask') + if min(image_shape) < 8: + raise ValueError(f'{volume_id} {level}: degenerate volume shape {image_shape}') + + grid = create_coordinate_grid(image_shape, device) + + best_l, loss_l, best_r, loss_r, total_time = run_de_torch_xfr( + label_str=f'{volume_id} {level}', + image1_path=cortical_path, + image2_path=binary_path, + image3_path=roi_path, + folder=Output_dir, + swarm_size=swarm_size, + max_iter=max_iter, + spacing=spacing, + CBT=CBT, + device=device, + optimize_size=True, + grid=grid, + side=side, + level=level, + patient_id=volume_id, + side_dir=SIDE_RESULT_DIR, + run_id=run_id, + ) + + +def list_gpu_ids(): + """透過 nvidia-smi 取得 GPU ID（主流程不初始化 CUDA，避免污染 fork/spawn 子流程）""" + try: + out = subprocess.check_output( + ['nvidia-smi', '--query-gpu=index', '--format=csv,noheader'], + text=True, + ) + return [int(line.strip()) for line in out.splitlines() if line.strip()] + except Exception: + return [] + + +def gpu_worker(gpu_id, log_path, run_id, task_queue, result_queue): + """每張 GPU 一個 worker 流程：先鎖死該 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') + logger.info(f'=== [GPU {gpu_id}] DE worker started ===') + + while True: + item = task_queue.get() + if item is None: + break + 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_de(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: + 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)) + logger.info(f'=== [GPU {gpu_id}] DE worker finished ===') + + +def _write_cbt_safe(volume_id, run_id): + """per-volume 收尾：cbt.nii.gz + x-ap/x-lat 投影（錯誤已 log，回傳 ok）""" + try: + xfr_cbt_native.write_volume_cbt(volume_id, run_id) + return True + except Exception as e: + logger.error(f'[CBT-NATIVE] {volume_id}: {e}') + return False + + +def _run_sequential(tasks, run_id): + """沒有（或只有一張）GPU 時的回退：單流程串行。 + tasks 依 (volume, level, side) 分組排列：一個 volume 的 (level, side) + 全部跑完後，立刻寫它的 cbt.nii.gz + 投影（不等其他 volume）。""" + device = get_device() + current_vid = None + for volume_id, level, side in tasks: + if volume_id != current_vid: + if current_vid is not None: + _write_cbt_safe(current_vid, run_id) + current_vid = volume_id + set_task_tag(volume_id, level, 'LEFT' if side == 'L' else 'RIGHT') + try: + debug_de(volume_id, level, device, side=side, run_id=run_id) + except Exception as e: + logger.error(f'Error in {volume_id} {level} {side}: {e}') + if current_vid is not None: + _write_cbt_safe(current_vid, run_id) + + +USAGE = 'Usage: python xfr_de.py [volume_id] [level]' + + +def parse_args(argv): + """volume_id 可用完整 ID 或末段（如 0005）；level 為 L\\d 形式（L1、L2、…， + 不限定 L1~L5；實際執行以各 volume 資料中有的 level 為準，見 available_levels）。 + 兩者可省略（＝全部）；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 not re.fullmatch(r'L[1-9]\d*', level_arg.upper()): + sys.exit(f'{USAGE}\nInvalid level: {level_arg} (expected L1, L2, ...)') + 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(): + # 要處理的 volume 數（並行模式下是「最多嘗試的 volume 數」） + MAX_SUCCESSFUL_VOLUMES = 100 + + # log 檔（console 與檔案同時輸出；各 GPU worker 也會 append 進同一個檔） + os.makedirs(LOG_DIR, exist_ok=True) + run_id = time.strftime("%Y%m%d_%H%M%S") + log_path = os.path.join(LOG_DIR, f'xfr_de_{run_id}.log') + setup_tee(log_path) + _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] + + # 各 volume 的 lumbar level：未指定 level 時，取該 volume 實際有的全部 lumbar level + # （見 available_levels）；指定 level 時只跑該 level（該 volume 缺檔會直接報錯） + vol_levels = {} + for vid in volumes: + vol_levels[vid] = (level_arg,) if level_arg else available_levels(vid) + no_levels = [vid for vid in volumes if not level_arg and not vol_levels[vid]] + if no_levels: + logger.warning(f'{len(no_levels)} volume(s) have no available lumbar level, ' + f'skipped: {", ".join(no_levels)}') + # 任務粒度 = (volume, level, side)：同一 level 的 L/R 是兩個獨立任務， + # 可被不同 GPU 的 worker 領走並行執行（run_de_torch_xfr 單側模式）； + # 順序 L1 L -> L1 R -> L2 L -> L2 R -> ... + tasks = [(vid, level, s) for vid in volumes for level in vol_levels[vid] for s in ('L', 'R')] + expected = {vid: len(vol_levels[vid]) * 2 for vid in volumes} + all_levels = sorted({level for lv in vol_levels.values() for level in lv}, + key=lambda lv: int(lv[1:])) + logger.info(f'Total {len(volumes)} volumes / {len(tasks)} (volume, level, side) tasks ' + f'(levels: {", ".join(all_levels) if all_levels else "(none)"})') + 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: + logger.info(f'Only {len(gpu_ids)} GPU(s) available, running sequentially') + _run_sequential(tasks, run_id) + return + + 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() + result_queue = ctx.Queue() + for t in tasks: + task_queue.put(t) + for _ in gpu_ids: + task_queue.put(None) # 每個 worker 一個結束哨兵 + + # per-volume 收尾：一個 volume 排的 (level, side) 任務全部回報後 + # （成功或失敗），立刻在背景執行緒寫它的 cbt.nii.gz + 投影，不等全部 case + write_pool = ThreadPoolExecutor(max_workers=4, thread_name_prefix='cbt-write') + write_futures = set() + pending_writes = set(volumes) + remaining = {vid: expected[vid] for vid in volumes} + + def _fire_write(vid): + pending_writes.discard(vid) + write_futures.add(write_pool.submit(_write_cbt_safe, vid, run_id)) + + def _on_task_result(msg): + results.append(msg) + vid = msg[2] + if vid in remaining: + remaining[vid] -= 1 + if remaining[vid] <= 0 and vid in pending_writes: + logger.info(f'{vid}: screw tasks complete, writing cbt + projections now') + _fire_write(vid) + + procs = [ctx.Process(target=gpu_worker, + args=(g, log_path, run_id, task_queue, result_queue), + name=f'de-gpu-{g}') + for g in gpu_ids] + for p in procs: + p.start() + + start_time = time.time() + results = [] + finished = 0 + while finished < len(gpu_ids): + if not any(p.is_alive() for p in procs): + logger.warning('A worker exited early; reaping remaining tasks...') + break + try: + msg = result_queue.get(timeout=5) + except queue_module.Empty: + continue + if msg[0] == 'done': + finished += 1 + else: + _on_task_result(msg) + + # 抽乾剩下排進來的結果 + while True: + try: + msg = result_queue.get(timeout=1) + except queue_module.Empty: + break + if msg[0] == 'task': + _on_task_result(msg) + + for p in procs: + p.join(timeout=60) + + # 補漏：worker 提早退出、有任務未回報的 volume 仍照舊嘗試寫 + # （無 side 結果時 write_volume_cbt 會自行 skip） + for vid in volumes: + if vid in pending_writes: + _fire_write(vid) + + # 等待所有 per-volume cbt.nii.gz / 投影寫出完成 + write_ok = sum(1 for f in write_futures if f.result()) + write_fail = len(write_futures) - write_ok + write_pool.shutdown(wait=True) + + total_time = time.time() - start_time + ok = [r for r in results if r[5]] + fail = [r for r in results if not r[5]] + per_volume = {} + 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, side) 任務都執行過且全部成功 + # （expected 為該 volume 實際排入的任務數，各 volume 的 level 數可不同） + n_success_volumes = sum( + 1 for vid in volumes + if per_volume.get(vid, set()) == {True} + and per_volume_n.get(vid, 0) == expected.get(vid, 0)) + + print('=' * 60) + 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: + 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: + logger.info('Failed tasks:') + for _, g, vid, level, side, _, err in fail: + logger.error(f'[GPU {g}] {vid} {level} {side}: {err}') + + # cbt.nii.gz + x-ap.jpg / x-lat.jpg 已於各 volume 的螺絲任務完成後立刻 + # 寫出（_fire_write， Output_dir///；label 1-10 = + # L1L L1R L2L L2R ... L5L L5R），不再等全部 case 跑完才統一收尾 + logger.info(f'CBT writes: {write_ok}/{len(volumes)} volume(s) ok' + + (f', {write_fail} failed (見 [CBT-NATIVE] log)' if write_fail else '')) + + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/xfr_preprocess.py b/xfr_preprocess.py index 20ef746..c678be3 100644 --- a/xfr_preprocess.py +++ b/xfr_preprocess.py @@ -15,7 +15,7 @@ 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) + smooth_mask_sdf, measure_pedicle_width) from visualization.res_bone_figure import (render_bone_figure, compute_normalizing_rotation, rotate_volume_to, rotated_grid, @@ -39,6 +39,9 @@ label_map = { 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 +MAX_PEDICLE_WIDTH_MM = 18.0 # 腰椎椎弓根寬度生理上限；量測超過即判為 + # 橫突 / 肋在 SDF mask 內與弓根融合，記錄值 + # clip 到此上限並標 _clipped（原始值仍記錄 _raw_mm） 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__)), @@ -197,8 +200,11 @@ def _write_rotated_level(vol_dir, level, smd_path, mask_path, roi_path): （1=VBODY、2=棘突、3=other bone、0=background）。 平面參數經 R 剛性旋轉並換元到輸出 grid 的局部座標。 - 完成後把本 level 的旋轉座標（template / R / center / start / size）補寫 - 進 /transform.json（base 部分由 process_single_image 寫入）： + 完成後把本 level 的旋轉座標（template / R / center / start / size）、 + SP / VBODY 分割兩平面與雙側椎弓根寬度（pedicle：L/R mm + （mediolateral 中位數寬度，clip 上限 MAX_PEDICLE_WIDTH_MM）+ raw 值 + + clipped 旗標、量測切片數、AP 窗）補寫進 /transform.json + （base 部分由 process_single_image 寫入）： 供分割後把 rotated mask 反向映射回 data_root 原始 CT 座標 （imaging/transforms.py）。 @@ -340,12 +346,14 @@ def _write_rotated_level(vol_dir, level, smd_path, mask_path, roi_path): # 用旋轉後的平面畫圖（rotated 版 planes），並畫出 VBODY / 棘突著色 #（VBODY=金、棘突=紫，與 label map 對應）；皮質著色由未旋轉 CT + mask - # 現算（未旋轉 _cortical 不再存檔） + # 現算（未旋轉 _cortical 不再存檔）。shift=fstart 讓點雲 / 平面落在 + # 存檔 rotated/*.nii.gz 的局部 index 系——與 CBT 螺絲圖同座標系 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=False, rotation=(R, c_xyz), output_path=p_fig) + planes_only=False, + rotation=(R, c_xyz, fstart), output_path=p_fig) if fig is not None: logger.info(f'[rotated] saved {fig}') else: @@ -375,9 +383,52 @@ def _write_rotated_level(vol_dir, level, smd_path, mask_path, roi_path): f'other={int((label_arr == 3).sum())}, ' f'sp_mode={sp_info["mode"]}, vb_mode={vb_info["mode"]})') - # 記錄 rotated 座標鏈（template 選擇 / R / center / start），補進 - # transform.json；rotated disk a -> 未旋轉 template disk: - # W = R^T (a + start - center) + center（template disk 系，standardize 後） + # 雙側椎弓根寬度（左右 mediolateral，rotated 體上量測）：椎管 AP 窗內、 + # 中線帶外椎弓根區的每側最靠中線 2D 分塊、逐切片 x 向跨度沿 z 取中位數 + #（與上方 SP / VBODY 分割同用 rot_bin_seg / sym_rot，ap / s 定義一致） + ped = measure_pedicle_width(rot_bin_seg, sym_rot, sp_info, vb_info) + sp_x = float(template.GetSpacing()[0]) + + def _ped_mm(res): + """(記錄值 mm（clip 上限）, 原始值 mm, 是否被 clip)。""" + raw = res['width_vox'] * sp_x if res['width_vox'] is not None else None + clipped = raw is not None and raw > MAX_PEDICLE_WIDTH_MM + val = (min(raw, MAX_PEDICLE_WIDTH_MM) if raw is not None else None) + return (round(val, 2) if val is not None else None, + round(raw, 2) if raw is not None else None, bool(clipped)) + L_rec, L_raw, L_clip = _ped_mm(ped['L']) + R_rec, R_raw, R_clip = _ped_mm(ped['R']) + ped_entry = { + 'L_mm': L_rec, 'R_mm': R_rec, + 'L_raw_mm': L_raw, 'R_raw_mm': R_raw, + 'L_clipped': L_clip, 'R_clipped': R_clip, + 'L_slices': ped['L']['n_slices'], 'R_slices': ped['R']['n_slices'], + 'L_z': ped['L']['z'], 'R_z': ped['R']['z'], + 'ap_window': (None if ped['ap_front'] is None + else [round(ped['ap_front'], 1), round(ped['ap_back'], 1)]), + 'band_w': (None if ped['band_w'] is None else round(ped['band_w'], 2)), + 'reason': ped['reason'], + } + + def _fmt_mm(v): + return 'n/a' if v is None else f'{v:.2f}' + clip_note = (f' (raw>{MAX_PEDICLE_WIDTH_MM:g} 判橫突 / 肋融合, clipped)' + if (ped_entry['L_clipped'] or ped_entry['R_clipped']) else '') + logger.info('[rotated] %s %s: pedicle width L=%s, R=%s mm ' + '(raw L=%s, R=%s%s; slices L=%d, R=%d; %s)', + volume_id, level, _fmt_mm(ped_entry['L_mm']), + _fmt_mm(ped_entry['R_mm']), _fmt_mm(ped_entry['L_raw_mm']), + _fmt_mm(ped_entry['R_raw_mm']), clip_note, + ped_entry['L_slices'], ped_entry['R_slices'], ped_entry['reason']) + + # 記錄 rotated 座標鏈（template 選擇 / R / center / start）+ 本 level + # SP / VBODY 分割所用的兩平面（rotated/ 局部系）+ 雙側椎弓根寬度 + #（pedicle，量測於同 rot_bin_seg），補進 transform.json； + # rotated disk a -> 未旋轉 template disk: + # W = R^T (a + start - center) + center（template disk 系，standardize 後）。 + # 平面記錄後 core/optimizer 的 xfr 管線直接讀用、不在影像上重新估算 + # （imaging.transforms.load_rotated_planes）；sym_rot / symp_rot 即本函式 + # 上方 _label map / 平面圖同用的換系平面。 try: merge_rotated_into_transform(vol_dir, level, { 'template': 'smd_resampled' if smd_arr is not None else 'binary_nn', @@ -385,6 +436,9 @@ def _write_rotated_level(vol_dir, level, smd_path, mask_path, roi_path): 'center': [float(v) for v in c_xyz], 'start': [int(v) for v in fstart], 'size': [int(v) for v in fsize], + 'sym': sym_rot, + 'endplate': symp_rot, + 'pedicle': ped_entry, }) logger.info(f'[rotated] transform.json: recorded {level} rotated section ' f'(template={ "smd_resampled" if smd_arr is not None else "binary_nn" }, ' diff --git a/xfr_debug.py b/xfr_pso.py similarity index 89% rename from xfr_debug.py rename to xfr_pso.py index 719491f..85afa70 100644 --- a/xfr_debug.py +++ b/xfr_pso.py @@ -1,5 +1,6 @@ #!/home/xfr/.conda/envs/cbt/bin/python +import argparse import logging import os import re @@ -61,6 +62,9 @@ _TASK_TAG = {'vid': None, 'level': None, 'side': None} _NP_WRAP_RE = re.compile(r'\bnp\.[A-Za-z_][A-Za-z0-9_]*\(([^()]*)\)') +# 平面對齊偵錯行（[PLANE] mirror / endplate）：只出 console，不寫進 log 檔 +_PLANE_DEBUG_RE = re.compile(r'^\[PLANE\] (?:mirror|endplate)') + def set_task_tag(volume_id, level, side=None): _TASK_TAG['vid'] = volume_id @@ -73,7 +77,8 @@ class _Tee: volume 取 id 末段（1.3.6.1.4.1.9328.50.4.0001 -> 0001）； side 由 [LEFT]/[RIGHT]/左側/右側 段落標記判定並沿用給後續行； 最終結果段屬於整椎，重置沿用值，其中的 Left/Right 摘要行只標該行本身； - np.float64(...) 之類包裹會解包成裸數值。""" + np.float64(...) 之類包裹會解包成裸數值。 + [PLANE] mirror / endplate 偵錯行只寫 console、不寫 log 檔。""" def __init__(self, console, log_fh): self.console = console @@ -119,7 +124,8 @@ class _Tee: if line.startswith(marker): line = line[len(marker):] self.console.write(prefix + line + '\n') - self.log_fh.write(prefix + line + '\n') + if not _PLANE_DEBUG_RE.match(line): + self.log_fh.write(prefix + line + '\n') def write(self, data): if not data: @@ -201,7 +207,7 @@ def debug_orientation(volume_id, level): # print(f'Azimuth: {azi}, Alt: {alt}') logger.info(f'Alt: {alt}') -def debug_pso(volume_id, level, device=None, side='both', run_id=None): +def debug_pso(volume_id, level, device=None, side='both', run_id=None, screw_shape='cone'): # ====== PSO ====== swarm_size = 100 max_iter = 100 @@ -274,6 +280,7 @@ def debug_pso(volume_id, level, device=None, side='both', run_id=None): patient_id=volume_id, side_dir=SIDE_RESULT_DIR, run_id=run_id, + screw_shape=screw_shape, ) # exit() @@ -291,7 +298,7 @@ def list_gpu_ids(): return [] -def gpu_worker(gpu_id, log_path, run_id, task_queue, result_queue): +def gpu_worker(gpu_id, log_path, run_id, screw_shape, task_queue, result_queue): """每張 GPU 一個工作流程：先鎖死該 GPU，再從共享隊列領 (volume, level) 任務""" # worker 是獨立流程，要自己把輸出 tee 到 log 檔 setup_tee(log_path) @@ -311,7 +318,8 @@ def gpu_worker(gpu_id, log_path, run_id, task_queue, result_queue): # 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, side=side, run_id=run_id) + debug_pso(volume_id, level, device, side=side, run_id=run_id, + screw_shape=screw_shape) result_queue.put(('task', gpu_id, volume_id, level, side, True, '')) except Exception as e: logger.error(f'[GPU {gpu_id}] Error in {volume_id} {level} {side}: {e}') @@ -331,7 +339,7 @@ def _write_cbt_safe(volume_id, run_id): return False -def _run_sequential(tasks, run_id): +def _run_sequential(tasks, run_id, screw_shape='cone'): """沒有（或只有一張）GPU 時的回退：單流程串行。 tasks 依 (volume, level, side) 分組排列：一個 volume 的 (level, side) 全部跑完後，立刻寫它的 cbt.nii.gz + 投影（不等其他 volume）。""" @@ -344,29 +352,41 @@ def _run_sequential(tasks, run_id): current_vid = volume_id set_task_tag(volume_id, level, 'LEFT' if side == 'L' else 'RIGHT') try: - debug_pso(volume_id, level, device, side=side, run_id=run_id) + debug_pso(volume_id, level, device, side=side, run_id=run_id, + screw_shape=screw_shape) except Exception as e: logger.error(f'Error in {volume_id} {level} {side}: {e}') if current_vid is not None: _write_cbt_safe(current_vid, run_id) -USAGE = 'Usage: python xfr_debug.py [volume_id] [level]' +USAGE = 'Usage: python xfr_pso.py [volume_id] [level] [--screw-shape {cone,cylinder}]' def parse_args(argv): """volume_id 可用完整 ID 或末段（如 0005）；level 為 L\\d 形式（L1、L2、…， 不限定 L1~L5；實際執行以各 volume 資料中有的 level 為準，見 available_levels）。 - 兩者可省略（＝全部）；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') + 兩者可省略（＝全部）；level 必須搭配 volume_id 使用。 + --screw-shape 為最終合併圖的螺絲體外形：cone＝錐體（預設，入口端全直徑 + 線性收細至尖端）、cylinder＝舊式圓柱。只影響繪製，不影響評分 / CSV。""" + parser = argparse.ArgumentParser(prog='xfr_pso.py', + description=USAGE) + parser.add_argument('volume_id', nargs='?', default=None, + help='volume ID（完整 UID 或末段，如 0005）；省略＝全部') + parser.add_argument('level', nargs='?', default=None, + help='level（L1、L2、…）；省略＝該 volume 全部可用 level；' + '須搭配 volume_id') + parser.add_argument('--screw-shape', choices=['cone', 'cylinder'], + default='cone', + help='合併圖螺絲體外形（預設 cone；cylinder 舊式圓柱；' + '只影響繪製，不影響評分 / CSV）') + ns = parser.parse_args(argv) + level_arg = ns.level if level_arg and not re.fullmatch(r'L[1-9]\d*', level_arg.upper()): - sys.exit(f'{USAGE}\nInvalid level: {level_arg} (expected L1, L2, ...)') - 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) + parser.error(f'Invalid level: {level_arg} (expected L1, L2, ...)') + if level_arg and not ns.volume_id: + parser.error('level requires volume_id') + return ns.volume_id, (level_arg.upper() if level_arg else None), ns.screw_shape def main(): @@ -395,7 +415,7 @@ def main(): 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:]) + vid_arg, level_arg, screw_shape = parse_args(sys.argv[1:]) if vid_arg is not None: key = vid_arg.lower() vols = [v for v in volumes @@ -424,12 +444,13 @@ def main(): f'(levels: {", ".join(all_levels) if all_levels else "(none)"})') if vid_arg or level_arg: logger.info(f'Filter: volume_id={vid_arg!r} level={level_arg!r}') + logger.info(f'Screw shape: {screw_shape}') gpu_ids = list_gpu_ids() if len(gpu_ids) <= 1: logger.info(f'Only {len(gpu_ids)} GPU(s) available, running sequentially') - _run_sequential(tasks, run_id) + _run_sequential(tasks, run_id, screw_shape) return logger.info(f'Found {len(gpu_ids)} GPUs: {gpu_ids}, starting {len(gpu_ids)} workers (one per GPU)') @@ -462,7 +483,7 @@ def main(): logger.info(f'{vid}: screw tasks complete, writing cbt + projections now') _fire_write(vid) - procs = [ctx.Process(target=gpu_worker, args=(g, log_path, run_id, task_queue, result_queue), name=f'cbt-gpu-{g}') + procs = [ctx.Process(target=gpu_worker, args=(g, log_path, run_id, screw_shape, task_queue, result_queue), name=f'cbt-gpu-{g}') for g in gpu_ids] for p in procs: p.start() diff --git a/xfr_record_planes.py b/xfr_record_planes.py new file mode 100644 index 0000000..364f1de --- /dev/null +++ b/xfr_record_planes.py @@ -0,0 +1,214 @@ +#!/home/xfr/.conda/envs/cbt/bin/python +"""補記既有 volume 的 transform.json rotated 段平面（sym / endplate）。 + +core/optimizer 的 xfr 管線（run_pso_torch_xfr / run_de_torch_xfr）不再於影像 +上重新估算兩平面（best_symmetry_plane / best_upper_endplate_plane），改直接 +讀 /transform.json levels[L*]['rotated']['sym' / 'endplate'] +（imaging.transforms.load_rotated_planes_from_path）： + - 新世代：xfr_preprocess._write_rotated_level 旋轉時即記錄（rotated/ 局部系）； + - 舊世代（進度 progress.json 已標完成、重跑 preprocess 會整支跳過）：跑本腳本 + 補記——用與舊 optimizer 相同函式、相同預設值，對 + rotated/_binary_sdf.nii.gz（完整 mask、SP 移除前）重新估算， + 併入既有 rotated 段（只加 sym / endplate 兩鍵，不動 template / R / center / + start / size 等既有鍵、不重寫任何 nii）。補記後 optimizer 讀到的平面與 + 改前行為一致（同輸入、同演算法）。 + +預設跳過已有 sym 記錄的 level（部分失敗後重跑只補缺者）；--force 全數重記。 + +Usage: + python xfr_record_planes.py [--force] [--dir ] [volume_id ...] + volume_id 可用完整 ID 或末段（如 0001）；省略 = 全部 volume。 +""" +import argparse +import logging +import os +import re +import sys +import time + +import numpy as np +import SimpleITK as sitk + +from imaging.orientation import best_symmetry_plane, best_upper_endplate_plane +from imaging.transforms import load_transform, save_transform + + +STANDARDIZED_DIR = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr-3/' + +LOG_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'logs') +logger = logging.getLogger('xfr_record_planes') + +# 基準系健全閾值：對齊良好時 theta_v / tau_y 應近似 0（xfr_preprocess 的 R +# 已把鏡稱面法線轉到 +x、終板法線轉到 z>0） +SANITY_DEG = 5.0 + + +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): + log_fh = open(log_path, 'a', buffering=1) + sys.stdout = _Tee(sys.stdout, log_fh) + sys.stderr = _Tee(sys.stderr, log_fh) + + +def rotated_levels(vol_dir): + """該 volume rotated/ 中有 _binary_sdf.nii.gz 的 lumbar level（排序）。""" + rotated_dir = os.path.join(vol_dir, 'rotated') + if not os.path.isdir(rotated_dir): + return () + levels = set() + for fn in os.listdir(rotated_dir): + m = re.fullmatch(r'(L[1-9]\d*)_binary_sdf\.nii\.gz', fn) + if m: + levels.add(m.group(1)) + return tuple(sorted(levels, key=lambda lv: int(lv[1:]))) + + +def record_level(vol_dir, level, sdf_path): + """估算該 level 兩平面併入 transform.json rotated 段；回傳 (sym, endplate)。""" + arr = sitk.GetArrayFromImage(sitk.ReadImage(sdf_path)) + sym = best_symmetry_plane(arr) + endplate = best_upper_endplate_plane(arr) + + meta = load_transform(vol_dir) + lv = meta['levels'].setdefault(level, {}) + rot = lv.setdefault('rotated', {}) + rot['sym'] = sym + rot['endplate'] = endplate + save_transform(vol_dir, meta) + + s_nx, s_ny, _ = sym['normal'] + theta_v = float(np.degrees(np.arctan2(s_ny, s_nx))) + if endplate is not None: + e_nx, e_ny, e_nz = endplate['normal'] + tau_y = float(np.degrees(np.arctan2(e_ny, e_nz))) + else: + tau_y = None + note = f'endplate={endplate is not None}' + if abs(theta_v) > SANITY_DEG or (tau_y is not None and abs(tau_y) > SANITY_DEG): + logger.warning(f'{os.path.basename(vol_dir)} {level}: 基準系殘差偏大 ' + f'(theta_v={theta_v:+.2f} deg, tau_y=' + f'{tau_y if tau_y is None else round(tau_y, 2)}) — ' + f'該 level 的 xfr_preprocess 對齊可能失效，建議人工核對') + else: + logger.info(f'{os.path.basename(vol_dir)} {level}: recorded ' + f'(mirror ratio={sym["ratio"]:.3f}, theta_v={theta_v:+.2f} deg, ' + f'tau_y={tau_y if tau_y is None else round(tau_y, 2)}, {note})') + return sym, endplate + + +def main(): + parser = argparse.ArgumentParser( + description='Record sym/endplate planes into transform.json for ' + 'existing volumes (old generation, before xfr_preprocess ' + 'recording them at rotation time).') + parser.add_argument('volume_ids', nargs='*', + help='Full ID or trailing segment (e.g. 0001); ' + 'default: all volumes') + parser.add_argument('--dir', default=STANDARDIZED_DIR, + help='Standardized output dir (default: standardized-xfr-3)') + parser.add_argument('--force', action='store_true', + help='Re-record levels that already have sym/endplate') + args = parser.parse_args() + + os.makedirs(LOG_DIR, exist_ok=True) + log_path = os.path.join( + LOG_DIR, f'xfr_record_planes_{time.strftime("%Y%m%d_%H%M%S")}.log') + setup_tee(log_path) + 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)}') + + if not os.path.isdir(args.dir): + sys.exit(f'Standardized dir not found: {args.dir}') + volumes = [d for d in sorted(os.listdir(args.dir)) + if os.path.isdir(os.path.join(args.dir, d))] + if args.volume_ids: + keys = [v.lower() for v in args.volume_ids] + vols = [v for v in volumes + if v.lower() in keys or v.rsplit('.', 1)[-1].lower() in keys] + if not vols: + sys.exit(f'No volume matches: {", ".join(args.volume_ids)}') + volumes = vols + logger.info(f'{len(volumes)} volume(s) under {args.dir}' + + (' (force)' if args.force else '')) + + t0 = time.time() + n_rec, n_skip_done, n_err, n_no_levels = 0, 0, 0, 0 + no_ep = [] + for vid in volumes: + vol_dir = os.path.join(args.dir, vid) + levels = rotated_levels(vol_dir) + if not levels: + n_no_levels += 1 + continue + try: + meta = load_transform(vol_dir) + except FileNotFoundError: + logger.warning(f'{vid}: transform.json 缺失，跳過 ' + f'（舊世代可先在 xfr_inverse_transform.py 加 --rebuild）') + n_err += 1 + continue + for level in levels: + rot = meta['levels'].get(level, {}).get('rotated') or {} + if not args.force and rot.get('sym') is not None: + n_skip_done += 1 + continue + sdf_path = os.path.join(vol_dir, 'rotated', f'{level}_binary_sdf.nii.gz') + try: + _, ep = record_level(vol_dir, level, sdf_path) + n_rec += 1 + if ep is None: + no_ep.append(f'{vid} {level}') + except Exception as e: + n_err += 1 + logger.error(f'{vid} {level}: {e}') + dt = time.time() - t0 + logger.info(f'Done in {dt / 60:.1f} min: recorded {n_rec} level(s), ' + f'skipped(already) {n_skip_done}, no rotated levels ' + f'{n_no_levels} volume(s), errors {n_err}') + if no_ep: + logger.warning(f'{len(no_ep)} level(s) 無上終板面記錄（資料不足），' + f'optimizer 對應 level 終板入口剪除停用: {", ".join(no_ep[:10])}' + f'{" ..." if len(no_ep) > 10 else ""}') + + +if __name__ == '__main__': + main() \ No newline at end of file