import torch import numpy as np import matplotlib.pyplot as plt from matplotlib.colors import to_rgba from matplotlib.lines import Line2D from mpl_toolkits.mplot3d.art3d import Poly3DCollection import os from datetime import datetime import csv from core.cylinder import generate_cylinder_n_torch, generate_cylinder_o_torch, snap_to_discrete_values from core.intersection import center_line_intersections_torch from core.scoring import cl_score_torch, compute_overlap_ratio_from_cylinder_mask, cl_score_torch_xfr from imaging.orientation import (azimuth_rotation, analyze_vertebral_tilt_contour, best_symmetry_plane, best_upper_endplate_plane, segment_spinous_process) from utils.helpers import save_with_unique_name # Volume absorption 渲染(Beer-Lambert):每 voxel 不透明度 = 1 - exp(-mu * voxel_width) # 骨頭核心厚度達 70-90 voxel,沿視線堆疊會使任何 per-voxel alpha 累積成不透明。 # 因此以「抽稀 (SUBSAMPLE) 降低堆疊數量」+「低 mu 控制每點吸收」兩項共同調出淡薄 X-ray 陰影, # 同時保留皮質 / 鬆質的吸收入射差異(mu 比值)。 BONE_MU_CORTICAL = 0.02 # 1/mm → 每 voxel = 1-exp(-0.02*0.5) ~ 0.010 BONE_MU_TRABECULAR = 0.005 # 1/mm → 每 voxel = 1-exp(-0.005*0.5) ~ 0.0025 BONE_MARKER_SIZE = 3.0 # 骨骼散點點面積 (pt^2);略大以補償抽稀後的顆粒感 BONE_SUBSAMPLE = 1 # 每 10 個骨 voxel 畫 1 個,降低堆疊不透明度(1=全畫) def set_axes_equal_3d(ax): """ Make axes of 3D plot have equal scale so that spheres appear as spheres, cubes as cubes, etc. """ x_limits = ax.get_xlim3d() y_limits = ax.get_ylim3d() z_limits = ax.get_zlim3d() x_range = abs(x_limits[1] - x_limits[0]) x_middle = np.mean(x_limits) y_range = abs(y_limits[1] - y_limits[0]) y_middle = np.mean(y_limits) z_range = abs(z_limits[1] - z_limits[0]) z_middle = np.mean(z_limits) plot_radius = 0.5*max([x_range, y_range, z_range]) ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius]) ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius]) ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius]) try: ax.set_box_aspect([1, 1, 1]) except AttributeError: pass def res_plt_2_torch( spine_tensor: torch.Tensor, cortical_tensor: torch.Tensor, image_shape: tuple[int, int, int], image2_path: str, base_folder: str, label_str: str, diameter_l: float, length_l: float, diameter_r: float, length_r: float, best_position_l: list[float], best_position_r: list[float], swarm_size: int, max_iter: int, total_time: float, spacing: list[float], CBT: bool, device: torch.device, grid=None ) -> None: """ Same plotting function as before, but it uses torch-based generation and then moves data to CPU for matplotlib 3D scatter. """ cyl_l = generate_cylinder_n_torch( diameter_l, length_l, best_position_l[0], best_position_l[1], best_position_l[2], best_position_l[3], best_position_l[4], image_shape, spacing, device, grid ) cyl_lo = generate_cylinder_o_torch( diameter_l, length_l, best_position_l[0], best_position_l[1], best_position_l[2], best_position_l[3], best_position_l[4], image_shape, spacing, device, grid ) cyl_r = generate_cylinder_n_torch( diameter_r, length_r, best_position_r[0], best_position_r[1], best_position_r[2], best_position_r[3], best_position_r[4], image_shape, spacing, device, grid ) cyl_ro = generate_cylinder_o_torch( diameter_r, length_r, best_position_r[0], best_position_r[1], best_position_r[2], best_position_r[3], best_position_r[4], image_shape, spacing, device, grid ) intersections_l, line_mask_l = center_line_intersections_torch( best_position_l[0], best_position_l[1], best_position_l[2], best_position_l[3], best_position_l[4], int(length_l), spine_tensor, spacing, device ) loss_l = cl_score_torch(cortical_tensor, spine_tensor, cyl_l, cyl_lo, intersections_l) intersections_r, line_mask_r = center_line_intersections_torch( best_position_r[0], best_position_r[1], best_position_r[2], best_position_r[3], best_position_r[4], int(length_r), spine_tensor, spacing, device ) # loss_r = cl_score_torch(cortical_tensor, spine_tensor, cyl_r, cyl_ro, intersections_r) loss_r = cl_score_torch_xfr(cortical_tensor, spine_tensor, cyl_r, cyl_ro, intersections_r) azi = azimuth_rotation(image2_path) res = analyze_vertebral_tilt_contour(image2_path, edge_type='superior', show_plot=False, debug=False) alt = res['superior']['tilt_angle_deg'] # Move data to CPU for plotting line_mask_l_cpu = line_mask_l.cpu().numpy() line_mask_r_cpu = line_mask_r.cpu().numpy() cyl_l_cpu = cyl_l.cpu().numpy() cyl_lo_cpu = cyl_lo.cpu().numpy() cyl_r_cpu = cyl_r.cpu().numpy() cyl_ro_cpu = cyl_ro.cpu().numpy() spine_cpu = spine_tensor.cpu().numpy() z_lin1, y_lin1, x_lin1 = np.where(line_mask_l_cpu == 1) z_lin2, y_lin2, x_lin2 = np.where(line_mask_r_cpu == 1) z_cyl_l1, y_cyl_l1, x_cyl_l1 = np.where(cyl_l_cpu == 1) z_cyl_l2, y_cyl_l2, x_cyl_l2 = np.where(cyl_lo_cpu == 1) z_cyl_r1, y_cyl_r1, x_cyl_r1 = np.where(cyl_r_cpu == 1) z_cyl_r2, y_cyl_r2, x_cyl_r2 = np.where(cyl_ro_cpu == 1) # 骨頭 voxel 依「體積吸收」分成兩組:皮質(高不透明度)與鬆質(低不透明度) cortical_cpu = cortical_tensor.cpu().numpy() 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) z_corti, y_corti, x_corti = np.where((spine_cpu == 1) & (cortical_cpu == 1)) z_trab, y_trab, x_trab = np.where((spine_cpu == 1) & (cortical_cpu == 0)) # 中矢狀面:骨頭的最佳鏡稱面,一般平面 a·x + b·y + c·z = d(法線方向任意) sym = best_symmetry_plane(spine_cpu) # 棘突:鏡稱面中線帶(|s|<=w)且在 AP 谷底之後側的骨 voxel,換不同顏色標示 sp_mask, sp_th, sp_info = segment_spinous_process(spine_cpu, sym) if sp_mask is not None and sp_mask.any(): sp_corti = sp_mask[z_corti, y_corti, x_corti] sp_trab = sp_mask[z_trab, y_trab, x_trab] sp_n_bone = max(int(spine_cpu.sum()), 1) print(f"[SPINOUS] n={sp_info['n_sp']} " f"({100.0 * sp_info['n_sp'] / sp_n_bone:.1f}% of bone) " f"band=+/-{sp_info['band_w']:.1f} voxel AP>={sp_info['ap_thresh']:.1f} " f"mode={sp_info['mode']}") else: sp_corti = sp_trab = None # X-ray 外觀:骨頭合成一個半透明體積吸收點雲(下方); # 螺絲(中心線 + 圓柱 + 入口軌跡延长)合成一個全不透明點雲,永遠畫在骨頭之上 def _rgba_block(n, color, a): arr = np.empty((n, 4)) arr[:] = to_rgba(color) arr[:, 3] = a return arr x_bone = np.concatenate([x_corti, x_trab]) y_bone = np.concatenate([y_corti, y_trab]) z_bone = np.concatenate([z_corti, z_trab]) bone_rgba = np.concatenate([ _rgba_block(len(x_corti), 'lightblue', float(alpha_cortical)), _rgba_block(len(x_trab), 'lightblue', float(alpha_trabecular)), ]) bone_size = np.full(len(x_bone), BONE_MARKER_SIZE) # 抽稀:降低堆疊不透明度以呈現淡薄 X-ray 陰影 if BONE_SUBSAMPLE > 1: x_bone = x_bone[::BONE_SUBSAMPLE] y_bone = y_bone[::BONE_SUBSAMPLE] z_bone = z_bone[::BONE_SUBSAMPLE] bone_rgba = bone_rgba[::BONE_SUBSAMPLE] bone_size = bone_size[::BONE_SUBSAMPLE] if sp_corti is not None: sp_flag = np.concatenate([sp_corti, sp_trab]) if BONE_SUBSAMPLE > 1: sp_flag = sp_flag[::BONE_SUBSAMPLE] bone_rgba[sp_flag] = to_rgba('purple', 0.95) x_screw = np.concatenate([x_lin1, x_lin2, x_cyl_l1, x_cyl_l2, x_cyl_r1, x_cyl_r2]) y_screw = np.concatenate([y_lin1, y_lin2, y_cyl_l1, y_cyl_l2, y_cyl_r1, y_cyl_r2]) z_screw = np.concatenate([z_lin1, z_lin2, z_cyl_l1, z_cyl_l2, z_cyl_r1, z_cyl_r2]) _a, _b, _c, _d = sym['plane'] _n = np.array([_a, _b, _c]) _u = np.array(sym['u']) _v = np.array(sym['v']) _p0 = _d * _n # 平面上最接近原點的點 xyz_bone = np.stack([x_bone - _p0[0], y_bone - _p0[1], z_bone - _p0[2]], axis=1) _pu = xyz_bone @ _u _pv = xyz_bone @ _v _U_, _V_ = np.meshgrid(np.linspace(_pu.min(), _pu.max(), 8), np.linspace(_pv.min(), _pv.max(), 8)) _Xp = _p0[0] + _U_ * _u[0] + _V_ * _v[0] _Yp = _p0[1] + _U_ * _u[1] + _V_ * _v[1] _Zp = _p0[2] + _U_ * _u[2] + _V_ * _v[2] # 上終板平面:RANSAC 擬合骨頭頂面(前側)的最佳 a·x + b·y + c·z = d symp = best_upper_endplate_plane(spine_cpu) _EX = _EY = _EZ = None if symp is not None: _ea, _eb, _ec, _ed = symp['plane'] _en = np.array([_ea, _eb, _ec]) _eu = np.array(symp['u']) _ev = np.array(symp['v']) _ep0 = _ed * _en xz_ep = np.stack([x_bone - _ep0[0], y_bone - _ep0[1], z_bone - _ep0[2]], axis=1) _pu_ep = xz_ep @ _eu _pv_ep = xz_ep @ _ev _EU, _EV = np.meshgrid(np.linspace(_pu_ep.min(), _pu_ep.max(), 8), np.linspace(_pv_ep.min(), _pv_ep.max(), 8)) _EX = _ep0[0] + _EU * _eu[0] + _EV * _ev[0] _EY = _ep0[1] + _EU * _eu[1] + _EV * _ev[1] _EZ = _ep0[2] + _EU * _eu[2] + _EV * _ev[2] screw_rgba = np.concatenate([ _rgba_block(len(x_lin1), 'r', 1.0), _rgba_block(len(x_lin2), 'r', 1.0), _rgba_block(len(x_cyl_l1), 'darkcyan', 1.0), _rgba_block(len(x_cyl_l2), 'pink', 1.0), _rgba_block(len(x_cyl_r1), 'blue', 1.0), _rgba_block(len(x_cyl_r2), 'pink', 1.0), ]) screw_size = np.concatenate([ np.full(len(x_lin1), 3), np.full(len(x_lin2), 3), np.full(len(x_cyl_l1), 36), np.full(len(x_cyl_l2), 36), np.full(len(x_cyl_r1), 36), np.full(len(x_cyl_r2), 36), ]) fig = plt.figure(figsize=(12, 12)) # 圖例色塊提高到可讀不透明度(實際渲染仍用真實吸收 alpha) _leg_alpha_c = max(float(alpha_cortical), 0.35) _leg_alpha_t = max(float(alpha_trabecular), 0.2) legend_handles = [ Line2D([], [], marker='o', ls='', ms=5, color=to_rgba('lightblue', _leg_alpha_c), label='Spine (cortical)'), Line2D([], [], marker='o', ls='', ms=5, color=to_rgba('lightblue', _leg_alpha_t), label='Spine (trabecular)'), Line2D([], [], marker='o', ls='', ms=2, color='r', label='Centerline'), Line2D([], [], marker='o', ls='', ms=6, color='darkcyan', label='Cylinder(L)'), Line2D([], [], marker='o', ls='', ms=6, color='blue', label='Cylinder(R)'), Line2D([], [], marker='o', ls='', ms=6, color='pink', label='Entry track (outer)'), Line2D([], [], color='orange', lw=2, alpha=0.6, label=f"Mirror plane {sym['plane'][0]:+.2f}x {sym['plane'][1]:+.2f}y {sym['plane'][2]:+.2f}z = {sym['plane'][3]:.1f}"), ] if sp_corti is not None: legend_handles.append(Line2D([], [], marker='o', ls='', ms=6, color='purple', label=f"Spinous process (mirror-plane midline, {sp_info['n_sp']} vox)")) if symp is not None: legend_handles.append(Line2D([], [], color='green', lw=2, alpha=0.7, label=f"Upper endplate plane tilt {symp['tilt_deg']:.1f} deg")) def _fill_ax(ax): # X-ray 外觀:關閉 mplot3d 依深度自動排序 zorder(否則半透明骨頭會被重繪到 # 螺絲上方);改為固定分層:吸收骨頭 zorder=5,全不透明螺絲 zorder=10 ax.computed_zorder = False sc_bone = ax.scatter(x_bone, y_bone, z_bone, c=bone_rgba, s=bone_size, marker='o') sc_bone.set_zorder(5) sc_screw = ax.scatter(x_screw, y_screw, z_screw, c=screw_rgba, s=screw_size, marker='o') sc_screw.set_zorder(10) # 中矢狀面(理論左右對稱切分面):半透明橘色平面 x = x_mid # 平面邊緣畫橘色線,讓 axial / 正視(側看時)也能清楚看到切分線 plane = ax.plot_surface(_Xp, _Yp, _Zp, color='orange', alpha=0.30, linewidth=1.0, edgecolor='orange', rstride=1, cstride=1) plane.set_zorder(8) # 上終板平面:半透明綠色平面(邊緣綠線) if _EX is not None: ep = ax.plot_surface(_EX, _EY, _EZ, color='green', alpha=0.35, linewidth=1.0, edgecolor='green', rstride=1, cstride=1) ep.set_zorder(7) ax1 = fig.add_subplot(221, projection='3d') _fill_ax(ax1) ax1.set_xlabel('X-axis'); ax1.set_ylabel('Y-axis'); ax1.set_zlabel('Z-axis') set_axes_equal_3d(ax1) ax2 = fig.add_subplot(222, projection='3d') ax2.view_init(elev=90, azim=-90, roll=0) _fill_ax(ax2) ax2.set_xlabel('X-axis'); ax2.set_ylabel('Y-axis'); ax2.set_zlabel('Z-axis') set_axes_equal_3d(ax2) ax2.legend(handles=legend_handles) ax3 = fig.add_subplot(223, projection='3d') ax3.view_init(elev=0, azim=90, roll=0) _fill_ax(ax3) ax3.set_xlabel('X-axis'); ax3.set_ylabel('Y-axis'); ax3.set_zlabel('Z-axis') set_axes_equal_3d(ax3) ax4 = fig.add_subplot(224, projection='3d') ax4.view_init(elev=0, azim=0, roll=0) _fill_ax(ax4) ax4.set_xlabel('X-axis'); ax4.set_ylabel('Y-axis'); ax4.set_zlabel('Z-axis') set_axes_equal_3d(ax4) cyl_points_l = torch.sum(cyl_l).item() cyl_points_r = torch.sum(cyl_r).item() overlap_l = ((cortical_tensor == 1) & (cyl_l == 1)).sum().item() overlap_r = ((cortical_tensor == 1) & (cyl_r == 1)).sum().item() overlap_b_l = ((spine_tensor == 1) & (cyl_l == 1)).sum().item() overlap_b_r = ((spine_tensor == 1) & (cyl_r == 1)).sum().item() overlap_cortical_l = (overlap_l / cyl_points_l) * 100 if cyl_points_l else 0.0 overlap_cortical_r = (overlap_r / cyl_points_r) * 100 if cyl_points_r else 0.0 overlap_vertebral_l = (overlap_b_l / cyl_points_l) * 100 if cyl_points_l else 0.0 overlap_vertebral_r = (overlap_b_r / cyl_points_r) * 100 if cyl_points_r else 0.0 cb_ratio_l = overlap_cortical_l/overlap_vertebral_l if overlap_vertebral_l else 0.0 cb_ratio_r = overlap_cortical_r/overlap_vertebral_r if overlap_vertebral_r else 0.0 user_altitude_l = 90 - best_position_l[4] - alt user_altitude_r = 90 - best_position_r[4] - alt user_azimuth_l = 90 - best_position_l[3] - azi user_azimuth_r = 90 - best_position_r[3] - azi date_str = datetime.now().strftime("%Y%m%d") patient_id = os.path.basename(os.path.dirname(image2_path)) output_folder = os.path.join(base_folder, date_str, patient_id) os.makedirs(output_folder, exist_ok=True) csv_path = os.path.join(output_folder, 'output.csv') # 檢查檔案是否存在 (決定是否寫入標題) file_exists = os.path.isfile(csv_path) # 欄位標題 (Header) headers = [ 'Label', 'Side', 'Diameter', 'Length', 'Swarm_Size', 'Max_Iter', 'Position_XYZ', 'Raw_Azimuth', 'Azimuth_Diff', 'Raw_Altitude', 'Altitude_Diff', 'Intersections', 'Best_Loss', 'cyl_points', 'Overlap_Cortical', 'Overlap_Bone', 'Cortical_Bone_Ratio', 'User_Azimuth', 'User_Altitude', 'Total_Time' ] try: with open(csv_path, 'a', newline='') as csvfile: writer = csv.writer(csvfile) # 如果是新檔案,寫入 Header if not file_exists: writer.writerow(headers) # 寫入 Left 數據 writer.writerow([ label_str, 'L', diameter_l, length_l, swarm_size, max_iter, # f"({best_position_l[0]:.2f}, {best_position_l[1]:.2f}, {best_position_l[2]:.2f})", f"({best_position_l[2]:.2f}, {best_position_l[1]:.2f}, {best_position_l[0]:.2f})", f"{best_position_l[3]:.2f}", f"{best_position_l[3]-azi:.2f}", f"{best_position_l[4]:.2f}", f"{best_position_l[4]-alt:.2f}", intersections_l, f"{loss_l:.2f}", cyl_points_l, f"{overlap_cortical_l:.2f}", f"{overlap_vertebral_l:.2f}", f"{(overlap_cortical_l/overlap_vertebral_l if overlap_vertebral_l!=0 else 0):.2f}", f"{user_azimuth_l:.2f}", f"{user_altitude_l:.2f}", f"{total_time:.2f}" ]) # 寫入 Right 數據 writer.writerow([ label_str, 'R', diameter_r, length_r, swarm_size, max_iter, # f"({best_position_r[0]:.2f}, {best_position_r[1]:.2f}, {best_position_r[2]:.2f})", f"({best_position_r[2]:.2f}, {best_position_r[1]:.2f}, {best_position_r[0]:.2f})", f"{best_position_r[3]:.2f}", f"{best_position_r[3]-azi:.2f}", f"{best_position_r[4]:.2f}", f"{best_position_r[4]-alt:.2f}", intersections_r, f"{loss_r:.2f}", cyl_points_r, f"{overlap_cortical_r:.2f}", f"{overlap_vertebral_r:.2f}", f"{(overlap_cortical_r/overlap_vertebral_r if overlap_vertebral_r!=0 else 0):.2f}", f"{user_azimuth_r:.2f}", f"{user_altitude_r:.2f}", f"{total_time:.2f}" ]) print(f"[CSV Saved] {csv_path}") except Exception as e: print(f"[Error] Failed to write CSV: {e}") fig.text(0.5, 0.98, f'{label_str} Best Position', ha='center', fontsize=15) fig.text( 0.5, 0.44, f'L: Diameter = {diameter_l} mm, {length_l} mm, ' f'R: Diameter = {diameter_r} mm, {length_r} mm, ' f'Swarm size = {swarm_size}, Iteration = {max_iter}, Total time = {total_time:.2f} s', ha='center', fontsize=12 ) fig.text( 0.5, 0.03, # f'Left : Position = ({best_position_l[0]:.2f}, {best_position_l[1]:.2f}, {best_position_l[2]:.2f}), ' f'Left : Position = ({best_position_l[2]:.2f}, {best_position_l[1]:.2f}, {best_position_l[0]:.2f}), ' f'Azimuth = {user_azimuth_l:.2f}, Altitude = {user_altitude_l:.2f}, ' f'Intersection = {intersections_l}, Score = {overlap_cortical_l:.2f} / {overlap_vertebral_l:.2f} / {cb_ratio_l:.2f}', ha='center', fontsize=9 ) fig.text( 0.5, 0.01, # f'Right : Position = ({best_position_r[0]:.2f}, {best_position_r[1]:.2f}, {best_position_r[2]:.2f}), ' f'Right : Position = ({best_position_r[2]:.2f}, {best_position_r[1]:.2f}, {best_position_r[0]:.2f}), ' f'Azimuth = {user_azimuth_r:.2f}, Altitude = {user_altitude_r:.2f}, ' f'Intersection = {intersections_r}, Score = {overlap_cortical_r:.2f} / {overlap_vertebral_r:.2f} / {cb_ratio_r:.2f}', ha='center', fontsize=9 ) fig.tight_layout() date_str = datetime.now().strftime("%Y%m%d") file_name = os.path.basename(image2_path) level = file_name.split('_')[0] output_folder = os.path.join(base_folder, date_str, patient_id) os.makedirs(output_folder, exist_ok=True) if CBT == True: way = 'CBT' else: way = 'TPS' path = save_with_unique_name(output_folder, label_str, way, diameter_l, length_l, diameter_r, length_r, swarm_size, max_iter) fig.savefig(path, dpi=200, bbox_inches="tight") print("[Saved figure]", path) plt.close(fig) def eval_overlap_from_position( pos, optimize_size: bool, spine_tensor: torch.Tensor, image_shape, spacing, device: torch.device, grid=None, fixed_diameter: float | None = None, fixed_length: float | None = None, ): """ 根據 position 生成 cylinder mask,再算 overlap ratio """ if optimize_size: d, L = snap_to_discrete_values(pos[5], pos[6]) params_5 = pos[:5] else: if fixed_diameter is None or fixed_length is None: raise ValueError("fixed_diameter and fixed_length must be provided when optimize_size=False") d, L = fixed_diameter, fixed_length params_5 = pos z, y, x, az, alt = params_5 cyl_mask = generate_cylinder_n_torch( d, L, z, y, x, az, alt, image_shape, spacing, device=device, grid=grid ) overlap = compute_overlap_ratio_from_cylinder_mask(cyl_mask, spine_tensor) return overlap, d, L