From 4204d2cd4c93da2db1b7734e1b6c8ef24eb10a91 Mon Sep 17 00:00:00 2001 From: Xiao Furen Date: Sun, 13 Sep 2026 09:14:55 +0800 Subject: [PATCH] feat(core): improve scoring logic and add X-ray projection rendering Implement a more robust scoring mechanism for screw optimization and add functionality to generate synthetic X-ray projections (AP and lateral views) from CT data. Key changes: - core: add `generate_cylinder_butt_torch` to create a mask for the screw entrance (0.25mm) to exempt it from bone-breaching penalties. - core: update `cl_score_torch_xfr` to include a diameter preference bonus and utilize the entrance mask. - core: adjust optimizer bounds and scoring weights to favor larger diameter screws and improve convergence. - xfr_cbt_native: implement `render_xray_projections` to generate synthetic AP and lateral X-ray images for visualization. - visualization: enhance `render_bone_figure` with semi-transparent spinous process rendering and improved depth sorting for screws. - xfr_debug: improve level detection to support arbitrary lumbar levels (L1-L9) and add safe volume-level cleanup for CBT writing. - config: update allowed diameters and lengths constants. --- config/constant.py | 6 +- core/cylinder.py | 55 +++++++ core/objective.py | 14 +- core/optimizer.py | 12 +- core/scoring.py | 36 ++++- visualization/res_bone_figure.py | 75 ++++++++-- xfr_cbt_native.py | 125 ++++++++++++++++ xfr_debug.py | 117 ++++++++++++--- xfr_rerender_spinous.py | 242 +++++++++++++++++++++++++++++++ 9 files changed, 631 insertions(+), 51 deletions(-) create mode 100644 xfr_rerender_spinous.py diff --git a/config/constant.py b/config/constant.py index 8e70484..db2016e 100644 --- a/config/constant.py +++ b/config/constant.py @@ -19,17 +19,17 @@ ALLOWED_DIAMETERS = [ 5.0, 5.5, 6.0, - 6.5, + # 6.5, # 7.0, # 7.5, ] ALLOWED_LENGTHS = [ # 25, - 30, + # 30, 35, 40, 45, - # 50, + 50, # 60, # 70, # 80, diff --git a/core/cylinder.py b/core/cylinder.py index 19b43c7..ee0db8b 100644 --- a/core/cylinder.py +++ b/core/cylinder.py @@ -275,6 +275,61 @@ def generate_cylinder_numpy(diameter, length, position_z, position_y, position_x return cylinder_mask +def generate_cylinder_butt_torch( + diameter, + position_z, position_y, position_x, + azimuth, altitude, + shape, spacing, device, grid=None, + butt_mm=0.25 # 入口端最後 butt_mm（mm） +) -> torch.Tensor: + """生成圓柱「最後端」mask：入口端（z_rot=0，遠離 VBODY 的一端） + 最靠近的 butt_mm 圓柱短柱（與 generate_cylinder_n_torch 同慣例、同直徑）。 + 尖端在 z_rot=length（靠 VBODY 端），不在此 mask。 + 此 mask 是完整圓柱的子集，可直接用於 not_in_bone 豁免。""" + + 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) + ) + + # 與 generate_cylinder_n_torch 相同的 spacing/單位處理； + # 長度固定為 butt_mm（mm）→ voxel + if spacing == [1, 1, 1]: + radius = diameter / 2.0 + butt_len = butt_mm + elif spacing == [0.5, 0.5, 0.5]: + radius = (diameter / 2.0) * 2 + butt_len = butt_mm * 2 + else: + raise ValueError(f"Unsupported spacing: {spacing}") + + mask = ( + (x_rot**2 + y_rot**2 <= radius**2) + & (z_rot >= 0) + & (z_rot < butt_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 0b3da0f..b22ffed 100644 --- a/core/objective.py +++ b/core/objective.py @@ -8,7 +8,7 @@ from scipy.ndimage import map_coordinates import numpy as np import torch -from core.cylinder import generate_cylinder_n_torch, generate_cylinder_o_torch, snap_to_discrete_values, generate_cylinder_tip_torch, snap_to_discrete_values_xfr +from core.cylinder import generate_cylinder_n_torch, generate_cylinder_o_torch, snap_to_discrete_values, generate_cylinder_tip_torch, generate_cylinder_butt_torch, snap_to_discrete_values_xfr from core.intersection import center_line_intersections_torch from core.scoring import cl_score_torch, cl_score_torch_xfr @@ -120,12 +120,22 @@ def cylinder_circle_line_intersection_loss_deductions_torch( ctx.image2_shape, ctx.spacing, ctx.device, ctx.grid ) + # 最後端（入口端，遠離 VBODY 的一端）0.25mm 豁免 mask + cyl_butt = generate_cylinder_butt_torch( + diameter, + position_z, position_y, position_x, + float(azimuth), float(altitude), + ctx.image2_shape, ctx.spacing, ctx.device, ctx.grid + ) + # loss_value = cl_score_torch( loss_value = cl_score_torch_xfr( ctx.cortical_tensor, ctx.spine_tensor, cyl_fwd, cyl_opp, intersections, + diameter=diameter, length=length, cylinder_tip_torch=cyl_tip, - vbody_tensor=ctx.vbody_tensor + vbody_tensor=ctx.vbody_tensor, + cylinder_butt_torch=cyl_butt ) return loss_value diff --git a/core/optimizer.py b/core/optimizer.py index 5e30f0d..66b4dea 100644 --- a/core/optimizer.py +++ b/core/optimizer.py @@ -10,7 +10,7 @@ from imaging.orientation import (azimuth_rotation, analyze_vertebral_tilt_contou 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 +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 config.constant import OVERLAP_THRESH @@ -63,9 +63,13 @@ def refine_lateral_longer( if cyl.sum().item() == 0: return None inter, _ = center_line_intersections_torch(z_c, y_c, x_c, az_c, alt_c, - L_c, spine_tensor, spacing, device) + 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, - vbody_tensor=vbody_tensor) + 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} @@ -364,7 +368,7 @@ def run_pso_torch_xfr( # z_bounds = (0, image_shape[0]-1) # z_bounds = (z1, (z1+z2)/2) # z_bounds = (.1*image_shape[0], .8*image_shape[0]) - z_bounds = (0, z1+z_height*.8) + z_bounds = (0, z1+z_height*.7) # x_bounds_right = (image_shape[2]/2 + image_shape[2]/10, image_shape[2] - 1) # x_bounds_left = (0, image_shape[2]/2 - image_shape[2]/10 - 1) diff --git a/core/scoring.py b/core/scoring.py index cd47c62..e15100d 100644 --- a/core/scoring.py +++ b/core/scoring.py @@ -2,6 +2,11 @@ import torch 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 直徑，讓結果在安全前提下偏好較粗螺絲。 +DIAMETER_BONUS_PER_MM = 20000.0 + def cl_score_torch_xfr( cortical_tensor: torch.Tensor, spine_tensor: torch.Tensor, @@ -11,10 +16,13 @@ def cl_score_torch_xfr( diameter: float = None, length: float = None, cylinder_tip_torch: torch.Tensor = None, # 新增：尖端 mask - vbody_tensor: torch.Tensor = None # VBODY（椎體）mask (z,y,x) 0/1，None = 不計 VBODY 獎勵 + vbody_tensor: torch.Tensor = None, # VBODY（椎體）mask (z,y,x) 0/1，None = 不計 VBODY 獎勵 + cylinder_butt_torch: torch.Tensor = None # 最後端（入口端，遠離 VBODY 的一端）0.25mm mask ) -> float: """ - 漸進式評分：優先確保找到骨頭，再改善細節 + 漸進式評分：優先確保找到骨頭，再改善細節。 + diameter 提供時附加直徑偏好 bonus（偏好較粗螺絲）。 + cylinder_butt_torch 提供時，該 mask 內 not_in_bone 的 voxel 免除 10000/voxel 扣分。 """ cyl_total = cylinder_torch.sum().item() overlap = ((cortical_tensor == 1) & (cylinder_torch == 1)).sum().item() # in cortical @@ -32,7 +40,12 @@ 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 扣分 + if cylinder_butt_torch is not None: + not_in_bone -= ((cylinder_butt_torch == 1) & (cylinder_torch == 1) + & (spine_tensor == 0)).sum().item() + if cyl_total == 0: return float(1e9) # 極差的情況 @@ -43,17 +56,24 @@ def cl_score_torch_xfr( score = cyl_total + # 直徑偏好 bonus：相同幾何下偏好較大直徑（diameter 單位 mm）。 + # bonus 遠小於 breaching 懲罰（10000 分/voxel），不會把螺絲推出骨頭。 + if diameter is not None and diameter > 0: + score += DIAMETER_BONUS_PER_MM * diameter + + allowed_error = cyl_total/100 + # if in_bone == 0: # return float(not_in_bone*200) score += 10 * in_bone # 10 實在太低 score += 100 * overlap # in cortical - score += 100 * in_corti_vb # (cortical + VBODY) 每 voxel 再加分 + score += 200 * in_corti_vb # (cortical + VBODY) 每 voxel 再加分 score += 50 * in_vbody # VBODY 每 voxel 再加分 - # score -= 1000 * not_in_bone - score -= 1000 * max(0, not_in_bone-10) - # score -= 2000 * null_vox2 - score -= 2000 * max(0, null_vox2-10) + score -= 10000 * not_in_bone + # score -= 10000 * max(0, not_in_bone-allowed_error) + score -= 20000 * null_vox2 + # score -= 2000 * max(0, null_vox2-10) return float(-score) diff --git a/visualization/res_bone_figure.py b/visualization/res_bone_figure.py index 4e8ddfd..5baa1cd 100644 --- a/visualization/res_bone_figure.py +++ b/visualization/res_bone_figure.py @@ -24,6 +24,8 @@ BONE_MU_CORTICAL = 0.02 # 1/mm BONE_MU_TRABECULAR = 0.005 # 1/mm BONE_MARKER_SIZE = 3.0 # 骨散點點面積 (pt^2) BONE_SUBSAMPLE = 1 # 抽稀（1 = 全畫） +SPINOUS_MU = 0.075 # 1/mm，棘突吸收係數（半透明紫色層，密度約皮質骨 3 倍： + # 看得見紫色、後方螺絲路徑仍透見） def set_axes_equal_3d(ax): @@ -265,10 +267,13 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, """統一骨頭 X-ray 四視角圖（合併原 render_bone_figure + res_plt_2_torch）。 四視角：預設 / axial（俯視 XY）/ 冠狀（後視）/ 矢狀。 - 內容：皮質 vs 鬆質吸收骨、椎體(gold)、棘突(purple，獨立層覆蓋椎體)、 + 內容：皮質 vs 鬆質吸收骨、椎體(gold)、棘突(purple 半透明，獨立層覆蓋椎體)、 中矢狀鏡稱面(orange)、上終板面(green)；螺絲模式另畫中心線(紅) + 圓柱（L darkcyan / R blue，o 層粉）。 - 繪製採固定分層（不依深度排序）：基底骨 < VBODY < 棘突 < 終板 < 鏡稱面 < 螺絲。 + 繪製採固定分層（不依深度排序）： + 基底骨 < VBODY < 螺絲(棘突後方) < 棘突(半透明) < 終板 < 鏡稱面 < 螺絲(棘突前方) + 螺絲與棘突另依各視角相機深度拆分：比棘突中位深度深的螺絲畫在棘突之下， + 被半透明紫色正確遮擋（仍可透見螺絲路徑）；較淺的螺絲照舊畫在最上層。 骨骼輸入： binary_path 骨頭遮罩：path (nifti) 或 (z,y,x) ndarray / torch tensor @@ -339,6 +344,7 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, voxel_mm = float(spacing[0]) alpha_cortical = 1.0 - np.exp(-BONE_MU_CORTICAL * voxel_mm) alpha_trabecular = 1.0 - np.exp(-BONE_MU_TRABECULAR * voxel_mm) + alpha_spinous = 1.0 - np.exp(-SPINOUS_MU * voxel_mm) # 骨頭 voxel 拆成皮質 / 鬆質兩組（同 res_plt_2_torch） z_corti, y_corti, x_corti = np.where((spine) & (cortical)) @@ -451,7 +457,8 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, # x_bone 保留完整點雲（含 VBODY / SP）供下方平面 patch 算範圍 #（VBODY 前側是整顆骨最前緣，剔除後綠色終板 patch 會縮小）； # mpl 3D scatter 在同一 collection 內依深度排序 markers， - # 「棘突覆蓋椎體、螺絲覆蓋骨頭」改以固定 zorder 分層達成（見 _fill_ax） + # 「棘突覆蓋椎體」以固定 zorder 分層達成；螺絲與棘突的遮擋則依各視角 + # 相機深度拆分（半透明棘突正確遮擋其後方的螺絲，見 _fill_ax） base_idx = ~(vb_flag | sp_flag) x_base, y_base, z_base = x_bone[base_idx], y_bone[base_idx], z_bone[base_idx] rgba_base, size_base = bone_rgba[base_idx], bone_size[base_idx] @@ -473,7 +480,8 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, spacing = list(spacing) # core.cylinder 以 list 比對 spacing import torch from core.cylinder import (generate_cylinder_n_torch, - generate_cylinder_o_torch) + generate_cylinder_o_torch, + generate_cylinder_butt_torch) from core.intersection import center_line_intersections_torch from core.scoring import cl_score_torch, cl_score_torch_xfr @@ -538,9 +546,14 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, loss = cl_score_torch(cortical_tensor, spine_tensor, cyl_n, cyl_o, inter) else: + cyl_butt = generate_cylinder_butt_torch( + d, pos[0], pos[1], pos[2], pos[3], pos[4], + image_shape, spacing, device, grid) loss = cl_score_torch_xfr(cortical_tensor, spine_tensor, cyl_n, cyl_o, inter, - vbody_tensor=vbody_tensor) + diameter=d, + vbody_tensor=vbody_tensor, + cylinder_butt_torch=cyl_butt) cyl_points = int(torch.sum(cyl_n).item()) ovc = (100.0 * int(((cortical_tensor == 1) & (cyl_n == 1)).sum().item()) / cyl_points) if cyl_points else 0.0 @@ -633,9 +646,21 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, if sp_corti is not None: legend_handles.append(Line2D([], [], marker="o", ls="", ms=6, color="purple", label="SpinousProcess")) - def _fill_ax(ax): + def _view_dir(elev_d, azim_d): + """本視角相機位置方向（從場景中心指向相機），(x, y, z) 資料座標系。 + view_init(elev, azim) 的球座標慣例：azim 繞 +z、自 +x 起算。""" + e, a = np.radians(elev_d), np.radians(azim_d) + return np.array([np.cos(e) * np.cos(a), + np.cos(e) * np.sin(a), + np.sin(e)]) + + def _fill_ax(ax, elev, azim): # 固定分層（關 depth zorder，否則半透明骨頭會被重繪到螺絲上方）： - # 基底骨(5) < VBODY gold(6) < 棘突 purple(6.5) < 終板(7) < 鏡稱面(8) < 螺絲(10) + # 基底骨(5) < VBODY gold(6) < 螺絲「棘突後方」(6.2) + # < 棘突 purple(6.5, 半透明) < 終板(7) < 鏡稱面(8) < 螺絲「棘突前方」(10) + # 螺絲與棘突按本視角相機深度拆分：depth = P·cam（大 = 靠近相機）。 + # 比棘突中位深度深（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) @@ -644,8 +669,28 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, 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) + + # 螺絲拆「棘突前方 / 後方」（front = 本視角較淺或等深）：無螺絲時 + # None；無棘突時全 True（維持舊行為）；有棘突時依相機深度 vs 棘突 + # 中位深度拆分 + if x_screw is None: + front = None + elif 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] + front = d_screw >= np.median(d_sp) + else: + front = np.ones(len(x_screw), dtype=bool) + + if front is not None and (~front).any(): + m = ~front + 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: - sc_sp = ax.scatter(x_sp, y_sp, z_sp, c=np.array([to_rgba("purple", 0.95)]), + 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: @@ -655,19 +700,19 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, plane = ax.plot_surface(_Xp, _Yp, _Zp, color="orange", alpha=0.30, linewidth=1.0, edgecolor="orange", rstride=1, cstride=1) plane.set_zorder(8) - if x_screw is not None: - sc_screw = ax.scatter(x_screw, y_screw, z_screw, - c=screw_rgba, s=screw_size, marker="o") + if 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") sc_screw.set_zorder(10) ax1 = fig.add_subplot(221, projection="3d") - _fill_ax(ax1) + _fill_ax(ax1, 30, -60) # matplotlib 3D 預設視角 ax1.set_xlabel("X-axis"); ax1.set_ylabel("Y-axis"); ax1.set_zlabel("Z-axis") set_axes_equal_3d(ax1) ax2 = fig.add_subplot(222, projection="3d") ax2.view_init(elev=90, azim=-90, roll=0) - _fill_ax(ax2) + _fill_ax(ax2, 90, -90) ax2.set_xlabel("X-axis"); ax2.set_ylabel("Y-axis"); ax2.set_zlabel("Z-axis") set_axes_equal_3d(ax2) if legend_handles: @@ -676,13 +721,13 @@ def render_bone_figure(volume_id, level, binary_path, cortical_path, ax3 = fig.add_subplot(223, projection="3d") # 後視圖：相機在 −y 後側，x 軸畫面左小右大 ax3.view_init(elev=0, azim=-90, roll=0) - _fill_ax(ax3) + _fill_ax(ax3, 0, -90) ax3.set_xlabel("X-axis"); ax3.set_ylabel("Y-axis"); ax3.set_zlabel("Z-axis") set_axes_equal_3d(ax3) ax4 = fig.add_subplot(224, projection="3d") ax4.view_init(elev=0, azim=0, roll=0) - _fill_ax(ax4) + _fill_ax(ax4, 0, 0) ax4.set_xlabel("X-axis"); ax4.set_ylabel("Y-axis"); ax4.set_zlabel("Z-axis") set_axes_equal_3d(ax4) diff --git a/xfr_cbt_native.py b/xfr_cbt_native.py index e9076ba..3819003 100644 --- a/xfr_cbt_native.py +++ b/xfr_cbt_native.py @@ -4,6 +4,10 @@ （L1-L5 x L/R，最多 10 支）存成單一 label 體積： Output_dir///cbt.nii.gz + Output_dir///x-ap.jpg 合成前後(AP) X 光投影视圖 + Output_dir///x-lat.jpg 合成側位 X 光投影视圖 + （只投影脊椎骨 + 螺絲、不含軟組織； + 見 render_xray_projections） label 值：L1L=1 L1R=2 L2L=3 L2R=4 L3L=5 L3R=6 L4L=7 L4R=8 L5L=9 L5R=10 （0 = 背景）。 @@ -74,6 +78,7 @@ META_DB = os.path.join(_PROJ_DIR, 'xfr_image_metadata.json') LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5') LEVEL_LABEL_VAL = {v: int(k) for k, v in LABEL_MAP.items() if v in LEVELS} # {'L1': 20, ...} +PROJ_FILENAME = {'ap': 'x-ap.jpg', 'lateral': 'x-lat.jpg'} logger = logging.getLogger('xfr_cbt_native') @@ -258,9 +263,129 @@ def write_volume_cbt(volume_id, run_id, output_root=Output_dir, date=None): logger.warning(f'{volume_id}: failed sides: {", ".join(skipped)}') logger.info(f'{volume_id}: {n_screws}/10 screws -> {out_path} ' f'(transform.json={n_tf}, re-est={n_fb}, ap_flip={ap_flip})') + + # 收尾：原 CT + 螺絲 -> 合成 AP / Lateral X 光投影视圖 (x-ap.jpg / x-lat.jpg) + try: + render_xray_projections(out_path, ct_path, out_dir) + except Exception as e: + logger.warning(f'{volume_id}: X-ray projection render failed: {e}') return out_path, n_screws +def render_xray_projections(cbt_path, ct_path, out_dir, + views=('ap', 'lateral'), margin_mm=50.0): + """原 CT + cbt.nii.gz（native 同一 grid 的螺絲 label）-> 合成「只有骨頭」X 光投影视圖。 + + 只投影脊椎骨 + 螺絲，不含軟組織： + - 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（金屬等效，最亮白）。 + 投影视圖 = 沿中心射線 μ 的線積分（AP 沿 y、Lateral 沿 x）；裁到螺絲 bbox 外 + margin_mm 的 spine 區域；骨用 percentile(1,99) 獨立視窗，螺絲再疊加裁白。 + 方向：SimpleITK LPS（x→右、y→後、z→上）；AP 頂=上、病人右在畫面左（R 標記）； + Lateral 頂=上、前位在左、後位在右（A/P 標記）。 + spacing 各軸不等時先重取樣到 min(spacing) 各向同性格，維持投影 aspect ratio。 + 輸出 x-ap.jpg / x-lat.jpg，回傳 {view: jpg_path}。""" + import matplotlib + matplotlib.use('Agg') + import matplotlib.pyplot as plt + + 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()}') + sp = np.array(ct.GetSpacing(), float) # (x,y,z) + native_size = ct.GetSize() + # spacing 各軸不等：先重取樣 CT/螺絲（與後面的分割 label）到 + # min(spacing) 的各向同性格，維持投影视圖正確的 aspect ratio + min_sp = float(sp.min()) + isosp = isosize = None + if not np.allclose(sp, min_sp): + isosp = [min_sp, min_sp, min_sp] + isosize = [max(1, int(round(s * v / min_sp))) for s, v in zip(sp, ct.GetSize())] + o, d = ct.GetOrigin(), ct.GetDirection() + ct = sitk.Resample(sitk.Cast(ct, sitk.sitkFloat32), isosize, sitk.Transform(), + sitk.sitkLinear, o, isosp, d, 0.0) + cbt = sitk.Resample(cbt, isosize, sitk.Transform(), + sitk.sitkNearestNeighbor, o, isosp, d, 0) + ct_arr = sitk.GetArrayFromImage(ct).astype(np.float32) # (z,y,x) + cbt_arr = sitk.GetArrayFromImage(cbt) + sp = np.array(ct.GetSpacing(), float) + + metal = cbt_arr > 0 + try: + _, lb_path, _ = find_native_paths(vid) + lb_img = sitk.ReadImage(lb_path) + # 比對「重取樣前」的 native grid（label 與原 CT 同格） + if lb_img.GetSize() != native_size: + raise ValueError(f'label/CT grid 尺寸不符 {lb_img.GetSize()} vs {native_size}') + if isosp is not None: + lb_img = sitk.Resample(lb_img, isosize, sitk.Transform(), + sitk.sitkNearestNeighbor, + ct.GetOrigin(), isosp, ct.GetDirection(), 0) + lb = sitk.GetArrayFromImage(lb_img) + spine = (lb >= 1) & (lb <= 24) # LABEL_MAP: C1..L5 全為脊椎 + except Exception as e: + logger.warning(f'{vid}: native 分割 label 不可用（{e}）；' + f'退回 HU 300-3000 閾值當脊椎骨') + spine = (ct_arr >= 300) & (ct_arr <= 3000) + + bone_mu = np.where(spine & ~metal, np.clip(ct_arr, 0, 2000) / 400.0, 0.0) + metal_mu = 60.0 * metal + + z, y, x = np.where(metal) + if z.size == 0: + raise ValueError(f'{vid}: cbt 無螺絲 voxel') + mz, my, mx = (int(margin_mm / s) for s in (sp[2], sp[1], sp[0])) + zs = slice(max(0, z.min() - mz), min(ct_arr.shape[0], z.max() + mz + 1)) + ys = slice(max(0, y.min() - my), min(ct_arr.shape[1], y.max() + my + 1)) + xs = slice(max(0, x.min() - mx), min(ct_arr.shape[2], x.max() + mx + 1)) + + proj = { + '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) + 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, :] + bnz = b[b > 0] + if bnz.size: + lo, hi = np.percentile(bnz, 1), np.percentile(bnz, 99) + bone = np.clip((b - lo) / (hi - lo + 1e-9), 0, 1) ** 0.7 + else: + bone = np.zeros_like(b) + mmax = float(m.max()) + metal_img = (m / (mmax + 1e-9)) ** 0.5 if mmax > 0 else m * 0 + img = np.clip(bone + metal_img, 0, 1) + 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') + 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, + va='top', ha='left', fontweight='bold') + ax.text(0.99, 0.98, rl, transform=ax.transAxes, color='cyan', fontsize=13, + va='top', ha='right', fontweight='bold') + ax.axis('off') + fig.tight_layout(pad=0.5) + p = os.path.join(out_dir, PROJ_FILENAME[view]) + fig.savefig(p, format='jpg', facecolor='black') + plt.close(fig) + paths[view] = p + logger.info(f'{vid}: {view} projection -> {p}') + return paths + + def main(): parser = argparse.ArgumentParser( description='Write ///cbt.nii.gz (screws mapped to native space) ' diff --git a/xfr_debug.py b/xfr_debug.py index 8999504..719491f 100644 --- a/xfr_debug.py +++ b/xfr_debug.py @@ -8,6 +8,7 @@ import time import queue as queue_module import subprocess import multiprocessing as mp +from concurrent.futures import ThreadPoolExecutor import SimpleITK as sitk import torch @@ -27,7 +28,24 @@ azimuth_rotation_dir = '/mnt/1248/open2/cyrou/azimuth_rotation' tilt_contour_dir = '/mnt/1248/open2/cyrou/tilt_contour' Output_dir = '/mnt/1248/open/cyrou/Output' -LEVELS = ('L1', 'L2', 'L3', 'L4', 'L5') +def available_levels(volume_id): + """該 volume 可跑 debug_pso 的全部 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') @@ -303,29 +321,49 @@ def gpu_worker(gpu_id, log_path, run_id, task_queue, result_queue): logger.info(f'=== [GPU {gpu_id}] 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 時的回退：單流程串行""" + """沒有（或只有一張）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_pso(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_debug.py [volume_id] [level]' def parse_args(argv): - """volume_id 可用完整 ID 或末段（如 0005）；level 為 LEVELS 之一（L1~L5）。 + """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 level_arg.upper() not in LEVELS: - sys.exit(f'{USAGE}\nInvalid level: {level_arg} (choose from {"/".join(LEVELS)})') + 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) @@ -342,7 +380,7 @@ def main(): # 要處理的 volume 數（並行模式下是「最多嘗試的 volume 數」） MAX_SUCCESSFUL_VOLUMES = 100 - MAX_SUCCESSFUL_VOLUMES = 10 + # MAX_SUCCESSFUL_VOLUMES = 10 # log 檔（console 與檔案同時輸出；各 GPU worker 也會 append 進同一個檔） os.makedirs(LOG_DIR, exist_ok=True) @@ -367,12 +405,23 @@ def main(): volumes = vols volumes = volumes[:MAX_SUCCESSFUL_VOLUMES] - levels = (level_arg,) if level_arg else LEVELS + # 各 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 領走並行執行；順序 L1 L -> L1 R -> L2 L -> L2 R -> ... - tasks = [(vid, level, s) for vid in volumes for level in levels for s in ('L', 'R')] + 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(levels)})') + 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}') @@ -393,6 +442,26 @@ def main(): 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'cbt-gpu-{g}') for g in gpu_ids] for p in procs: @@ -412,7 +481,7 @@ def main(): if msg[0] == 'done': finished += 1 else: - results.append(msg) + _on_task_result(msg) # 抽乾剩下排進來的結果 while True: @@ -421,11 +490,22 @@ def main(): except queue_module.Empty: break if msg[0] == 'task': - results.append(msg) + _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]] @@ -437,10 +517,11 @@ def main(): 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) == len(levels) * 2) + and per_volume_n.get(vid, 0) == expected.get(vid, 0)) print('=' * 60) logger.info(f'Finished in {total_time / 60:.1f} min | ' @@ -453,13 +534,11 @@ def main(): for _, g, vid, level, side, _, err in fail: logger.error(f'[GPU {g}] {vid} {level} {side}: {err}') - # 收尾：螺絲位置映回原 CT 空間 -> Output_dir///cbt.nii.gz - # （label 1-10 = L1L L1R L2L L2R ... L5L L5R；無 side 結果的 volume 跳過） - for vid in volumes: - try: - xfr_cbt_native.write_volume_cbt(vid, run_id) - except Exception as e: - logger.error(f'[CBT-NATIVE] {vid}: {e}') + # 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__': diff --git a/xfr_rerender_spinous.py b/xfr_rerender_spinous.py new file mode 100644 index 0000000..775cbad --- /dev/null +++ b/xfr_rerender_spinous.py @@ -0,0 +1,242 @@ +#!/home/xfr/.conda/envs/cbt/bin/python +""" +從 optimizer 存的 output.csv 重新渲染既有螺絲模式四視角圖（X-ray 圖）。 + +用途：render_bone_figure 改了純顯示層（例如棘突改半透明 + 依深度正確遮擋 +螺絲路徑）之後，不需要重跑優化，直接從 CSV 的 best_position 重新出圖。 +原圖先備份到 Output/{date}/backup_opaque_spinous/{vol}/，再重渲染同檔名覆蓋。 + +mask 取法與 xfr_debug / xfr_plot_level 相同（level_file_path： +binary_sdf 優先、缺則 binary；cortical 同），spacing 讀自 mask 檔。 + +Usage: + python xfr_rerender_spinous.py # 預設 20260912，多 GPU 並行 + python xfr_rerender_spinous.py 20260911 # 指定日期 + python xfr_rerender_spinous.py 20260912 0001 # 只該 volume（全名或末段） + python xfr_rerender_spinous.py --dry-run 20260912 # 只列出會重新渲染的圖 + python xfr_rerender_spinous.py --cpus 20260912 # 純 CPU 循序 +""" + +import argparse +import csv +import multiprocessing +import os +import re +import shutil +import sys + +import numpy as np +import SimpleITK as sitk + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from imaging.transforms import level_file_path +from visualization.res_bone_figure import render_bone_figure + +OUTPUT_BASE = '/mnt/1248/open2/cyrou/Output' +MASK_DIR = '/mnt/1248/open2/cyrou/CBT/Seg/Resample/standardized-xfr' +BACKUP = 'backup_opaque_spinous' + +# {level}_{way}_L{d}_{l}_R{d}_{l}_{swarm}_{iter}.png（單側跑法該側可為空） +PNG_RE = re.compile( + r'^(?P[A-Z]\d+?)_(?PCBT|TPS)' + r'_L(?P
\d+(?:\.\d+)?|)_(?P\d+(?:\.\d+)?|)' + r'_R(?P\d+(?:\.\d+)?|)_(?P\d+(?:\.\d+)?|)' + r'_(?P\d+|)_(?P\d+|)\.png$') + + +def _f(s): + s = (s or '').strip() + return float(s) if s else None + + +def _parse_pos(s): + """CSV '(x, y, z)' -> (z, y, x)（best_position 慣例）。""" + v = re.findall(r'[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?', s or '') + if len(v) < 3: + return None + x, y, z = (float(t) for t in v[:3]) + return z, y, x + + +def _pos(row): + if row is None: + return None + p = _parse_pos(row.get('Position_XYZ')) + if p is None: + return None + az, alt = _f(row.get('Raw_Azimuth')), _f(row.get('Raw_Altitude')) + if az is None or alt is None: + return None + return (p[0], p[1], p[2], az, alt) + + +def read_csv_rows(vol_out_dir): + path = os.path.join(vol_out_dir, 'output.csv') + if not os.path.isfile(path): + return None + with open(path, newline='') as f: + return [r for r in csv.DictReader(f) if any(c.strip() for c in r.values())] + + +def find_row(rows, side, d, l): + if d is None or l is None: + return None + for r in rows: + if r.get('Side') != side: + continue + rd, rl = _f(r.get('Diameter')), _f(r.get('Length')) + if rd is not None and rl is not None \ + and abs(rd - d) < 1e-6 and abs(rl - l) < 1e-6: + return r + return None + + +def collect_figs(date_dir, vol_filter=None): + """回傳 (tasks, skipped)：tasks 為可重新渲染的圖（dict），skipped 為 (原因, 路徑)。""" + vols = sorted(d for d in os.listdir(date_dir) + if os.path.isdir(os.path.join(date_dir, d)) and d != BACKUP) + if vol_filter: + vols = [v for v in vols + if v == vol_filter or v.rsplit('.', 1)[-1] == vol_filter] + if not vols: + raise SystemExit(f'volume not found: {vol_filter}') + tasks, skipped = [], [] + for vol in vols: + vdir = os.path.join(date_dir, vol) + rows = read_csv_rows(vdir) or [] + for name in sorted(os.listdir(vdir)): + m = PNG_RE.match(name) + if not m: + continue + png = os.path.join(vdir, name) + level, way = m['level'], m['way'] + d_l, l_l = _f(m['dl']), _f(m['ll']) + d_r, l_r = _f(m['dr']), _f(m['lr']) + row_l = find_row(rows, 'L', d_l, l_l) + row_r = find_row(rows, 'R', d_r, l_r) + if (d_l is not None and row_l is None) or (d_r is not None and row_r is None) \ + or (d_l is None and d_r is None): + skipped.append((vol, name, 'CSV 無對應 L/R 行')) + continue + mask_dir = os.path.join(MASK_DIR, vol) + sdf = level_file_path(mask_dir, level, 'binary_sdf') + binary_path = sdf if os.path.exists(sdf) \ + else level_file_path(mask_dir, level, 'binary') + if not os.path.exists(binary_path): + skipped.append((vol, name, f'無 bone mask: {binary_path}')) + continue + cortical_path = level_file_path(mask_dir, level, 'cortical') + row_any = row_l or row_r + tt = _f(row_any.get('Total_Time')) + tt = None if (tt is None or not np.isfinite(tt)) else tt + tasks.append({ + 'vol': vol, 'name': name, 'png': png, 'level': level, 'way': way, + 'd_l': d_l, 'l_l': l_l, 'd_r': d_r, 'l_r': l_r, + 'pos_l': _pos(row_l), 'pos_r': _pos(row_r), + 'binary': binary_path, + 'cortical': cortical_path if os.path.exists(cortical_path) else None, + 'swarm': int(_f(m['sw']) or 0), 'iter': int(_f(m['it']) or 0), + 'time': tt, + }) + return tasks, skipped + + +def _run_batch(batch): + """一個 worker 固定綁一個 GPU，依序渲染分到的圖（CUDA_VISIBLE_DEVICES + 須在 torch 首次 init CUDA 前設好，故綁定後不再變）。""" + gpu, g_tasks = batch + os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu) + os.environ.setdefault('OMP_NUM_THREADS', '4') + import matplotlib + matplotlib.use('Agg') + return [render_one(t) for t in g_tasks] + + +def render_one(task): + import torch + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + img = sitk.ReadImage(task['binary'], sitk.sitkUInt8) + spacing = list(img.GetSpacing()) + + # 備份原圖（同目錄同檔名會重複跑時，備份檔保留第一份） + backup_dir = os.path.join(os.path.dirname(os.path.dirname(task['png'])), BACKUP, task['vol']) + bpath = os.path.join(backup_dir, task['name']) + os.makedirs(backup_dir, exist_ok=True) + if not os.path.exists(bpath): + shutil.move(task['png'], bpath) + + try: + out = render_bone_figure( + task['vol'], task['level'], task['binary'], task['cortical'], + base_folder=OUTPUT_BASE, spacing=spacing, way=task['way'], + best_position_l=task['pos_l'], best_position_r=task['pos_r'], + diameter_l=task['d_l'], length_l=task['l_l'], + diameter_r=task['d_r'], length_r=task['l_r'], + image2_path=None, device=device, + swarm_size=task['swarm'], max_iter=task['iter'], total_time=task['time'], + write_csv=False, output_path=task['png']) + except Exception as e: + if not os.path.exists(task['png']) and os.path.exists(bpath): + shutil.move(bpath, task['png']) + return (task['vol'], task['name'], 'fail', f'{type(e).__name__}: {e}') + if out is None: + if os.path.exists(bpath): + shutil.move(bpath, task['png']) + return (task['vol'], task['name'], 'fail', 'render 回傳 None（無/空遮罩）') + if out != task['png']: + os.replace(out, task['png']) + return (task['vol'], task['name'], 'ok', f'{out} ({device})') + + +def main(): + ap = argparse.ArgumentParser(description='從 output.csv 重新渲染螺絲模式四視角圖') + ap.add_argument('date', nargs='?', default='20260912') + ap.add_argument('volume', nargs='?', default=None) + ap.add_argument('--dry-run', action='store_true') + ap.add_argument('--cpus', action='store_true', help='CPU 循序（不佔 GPU）') + args = ap.parse_args() + + date_dir = os.path.join(OUTPUT_BASE, args.date) + if not os.path.isdir(date_dir): + raise SystemExit(f'no such date dir: {date_dir}') + tasks, skipped = collect_figs(date_dir, args.volume) + print(f'{args.date}: {len(tasks)} figure(s) to re-render, {len(skipped)} skip(s)') + for vol, name, why in skipped: + print(f' [skip] {vol}/{name}: {why}') + if args.dry_run: + for t in tasks: + print(f" [dry] {t['vol']}/{t['name']} " + f"L=({t['d_l']},{t['l_l']}) R=({t['d_r']},{t['l_r']})") + return + + if not tasks: + return + + if args.cpus or not _cuda_count(): + os.environ['CUDA_VISIBLE_DEVICES'] = '' + for i, t in enumerate(tasks, 1): + r = render_one(t) + print(f'[{i:3d}/{len(tasks)}] {r[2]:4s} {r[0]}/{r[1]} {r[3]}', flush=True) + else: + n = min(_cuda_count(), 4) + ctx = multiprocessing.get_context('spawn') + batches = [(i, tasks[i::n]) for i in range(n)] + with ctx.Pool(n) as pool: + for results in pool.map(_run_batch, batches): + for vol, name, status, detail in results: + print(f'[{status:4s}] {vol}/{name} {detail}', flush=True) + print('=' * 60) + print('Done.') + + +def _cuda_count(): + try: + import torch + return torch.cuda.device_count() + except Exception: + return 0 + + +if __name__ == '__main__': + main() \ No newline at end of file