Xiao Furen
2ae08ac2cd
Refactor the preprocessing and segmentation pipeline to handle AP orientation variations and improve anatomical boundary detection. Key changes include: - Implement automated AP orientation detection in `process_single_image` to handle prone scans by flipping CT and labels when necessary. - Enhance `segment_spinous_process` using a gap-based approach to identify the spinal canal, providing more stable thresholds for spinous process and vertebral body segmentation. - Improve optimization search space by using the vertebral body (VBODY) projection for x/z bounding box calculation instead of the whole bone. - Refactor `render_bone_figure` to unify 2D/3D visualization and support detailed anatomical coloring (VBODY, spinous process). - Update `cl_score_torch_xfr` with more robust penalty handling for out-of-bone and null-voxel regions. - Add `retry_robust` utility to handle transient NFS file system errors. - Update `xfr_preprocess.py` to include anatomical segmentation coloring in rotated level visualizations.
91 lines
No EOL
2.9 KiB
Python
91 lines
No EOL
2.9 KiB
Python
import torch
|
||
import numpy as np
|
||
|
||
from core.cylinder import generate_cylinder_n_torch, snap_to_discrete_values
|
||
from core.scoring import compute_overlap_ratio_from_cylinder_mask
|
||
|
||
|
||
def set_axes_equal_3d(ax):
|
||
"""
|
||
Make axes of 3D plot have equal scale so that spheres appear as spheres,
|
||
cubes are 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, cortical_tensor, image_shape, image2_path,
|
||
base_folder, label_str, diameter_l, length_l,
|
||
diameter_r, length_r, best_position_l, best_position_r,
|
||
swarm_size, max_iter, total_time, spacing, CBT, device,
|
||
grid):
|
||
"""相容入口（舊簽名）：實作已合併進 res_bone_figure.render_bone_figure。
|
||
label_str / image_shape 不再使用（volume / level 由 image2_path 反推）。"""
|
||
from visualization.res_bone_figure import render_bone_figure
|
||
return render_bone_figure(
|
||
None, None, spine_tensor, cortical_tensor, base_folder,
|
||
spacing=spacing, way='CBT' if CBT else 'TPS',
|
||
best_position_l=best_position_l, best_position_r=best_position_r,
|
||
diameter_l=diameter_l, length_l=length_l,
|
||
diameter_r=diameter_r, length_r=length_r,
|
||
image2_path=image2_path, device=device, grid=grid,
|
||
swarm_size=swarm_size, max_iter=max_iter, total_time=total_time,
|
||
)
|
||
|
||
|
||
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 |