Xiao Furen
4204d2cd4c
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.
316 lines
No EOL
10 KiB
Python
316 lines
No EOL
10 KiB
Python
from dataclasses import dataclass
|
||
from typing import Optional
|
||
|
||
import random
|
||
|
||
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, 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
|
||
|
||
|
||
# =====================================================================
|
||
# OptimizationContext: the single explicit container for all shared state
|
||
# needed to evaluate a candidate cylinder. New code should build one of
|
||
# these and bind it to an objective via make_objective_function[_xfr];
|
||
# no module globals are read during optimization.
|
||
# =====================================================================
|
||
|
||
@dataclass
|
||
class OptimizationContext:
|
||
"""Shared state for evaluating candidate cylinder placements.
|
||
|
||
Required:
|
||
cortical_tensor: cortical bone mask (uint8, 0/1)
|
||
spine_tensor: bone (binary) mask (uint8, 0/1)
|
||
image2_shape: (Z, Y, X) volume shape
|
||
spacing: voxel spacing, [sx, sy, sz] (mm)
|
||
device: torch device for tensor ops
|
||
|
||
Optional:
|
||
grid: precomputed coordinate grid (z_t, y_t, x_t)
|
||
use_tip_penalty: add the tip-cylinder penalty to the loss
|
||
vbody_tensor: VBODY (vertebral body) mask 0/1 — adds the VBODY
|
||
voxel rewards in cl_score_torch_xfr
|
||
The *_array / spine_roi_tensor / diameter / length fields are kept
|
||
for compatibility and debugging; the loss itself does not use them.
|
||
"""
|
||
cortical_tensor: torch.Tensor
|
||
spine_tensor: torch.Tensor
|
||
image2_shape: tuple
|
||
spacing: list
|
||
device: torch.device
|
||
grid: Optional[tuple] = None
|
||
use_tip_penalty: bool = False
|
||
vbody_tensor: Optional[torch.Tensor] = None
|
||
spine_roi_tensor: Optional[torch.Tensor] = None
|
||
image1_array: Optional[np.ndarray] = None
|
||
image2_array: Optional[np.ndarray] = None
|
||
image3_array: Optional[np.ndarray] = None
|
||
diameter: Optional[float] = None
|
||
length: Optional[float] = None
|
||
|
||
|
||
def cylinder_circle_line_intersection_loss_deductions_torch(
|
||
ctx: OptimizationContext,
|
||
diameter: float,
|
||
length: float,
|
||
params: list[float],
|
||
) -> float:
|
||
"""
|
||
Computes the loss for a given set of cylinder params in PyTorch,
|
||
returning a Python float for PSO consumption.
|
||
|
||
All shared state (tensors, shape, spacing, device, grid, tip penalty)
|
||
is taken from `ctx`; no module globals are involved.
|
||
"""
|
||
position_z, position_y, position_x, azimuth, altitude = params
|
||
|
||
cyl_fwd = generate_cylinder_n_torch(
|
||
diameter,
|
||
length,
|
||
position_z,
|
||
position_y,
|
||
position_x,
|
||
float(azimuth),
|
||
float(altitude),
|
||
ctx.image2_shape,
|
||
ctx.spacing,
|
||
ctx.device,
|
||
ctx.grid
|
||
)
|
||
|
||
cyl_opp = generate_cylinder_o_torch(
|
||
diameter,
|
||
length,
|
||
position_z,
|
||
position_y,
|
||
position_x,
|
||
float(azimuth),
|
||
float(altitude),
|
||
ctx.image2_shape,
|
||
ctx.spacing,
|
||
ctx.device,
|
||
ctx.grid
|
||
)
|
||
|
||
# We call the center_line_intersections in Torch mode
|
||
intersections, _ = center_line_intersections_torch(
|
||
position_z,
|
||
position_y,
|
||
position_x,
|
||
azimuth,
|
||
altitude,
|
||
length,
|
||
ctx.spine_tensor,
|
||
ctx.spacing,
|
||
ctx.device
|
||
)
|
||
|
||
cyl_tip = None
|
||
if ctx.use_tip_penalty:
|
||
cyl_tip = generate_cylinder_tip_torch(
|
||
diameter, length,
|
||
position_z, position_y, position_x,
|
||
float(azimuth), float(altitude),
|
||
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,
|
||
cylinder_butt_torch=cyl_butt
|
||
)
|
||
|
||
return loss_value
|
||
|
||
|
||
# =====================================================================
|
||
# Context-bound objective builders (preferred API)
|
||
# =====================================================================
|
||
|
||
def _evaluate(params: list[float], ctx: OptimizationContext) -> float:
|
||
"""
|
||
Core objective: params = [z, y, x, azimuth, altitude, diameter_raw, length_raw]
|
||
"""
|
||
position_params = params[:5] # [z, y, x, azimuth, altitude]
|
||
diameter_raw = params[5]
|
||
length_raw = params[6]
|
||
|
||
# 將連續值轉換為離散值
|
||
diameter_discrete, length_discrete = snap_to_discrete_values(diameter_raw, length_raw)
|
||
|
||
return cylinder_circle_line_intersection_loss_deductions_torch(
|
||
ctx,
|
||
diameter_discrete,
|
||
length_discrete,
|
||
position_params
|
||
)
|
||
|
||
|
||
def _evaluate_xfr(params: list[float], ctx: OptimizationContext, y_indices) -> float:
|
||
"""
|
||
Core xfr objective: params = [z, x, azimuth, altitude, diameter_raw, length_raw];
|
||
y is derived from the per-column `y_indices` surface.
|
||
"""
|
||
z, x, azimuth, altitude, diameter_raw, length_raw = params
|
||
y = y_indices[round(z), round(x)] #+ random.uniform(-0.5, 0.5)
|
||
|
||
position_params = [z, y, x, azimuth, altitude]
|
||
|
||
# 將連續值轉換為離散值
|
||
diameter_discrete, length_discrete = snap_to_discrete_values_xfr(diameter_raw, length_raw)
|
||
|
||
diameter_loss = .9*diameter_discrete + .1*diameter_raw
|
||
length_loss = .9* length_discrete + .1* length_raw
|
||
|
||
return cylinder_circle_line_intersection_loss_deductions_torch(
|
||
ctx,
|
||
diameter_loss,
|
||
length_loss,
|
||
position_params
|
||
)
|
||
|
||
|
||
def make_objective_function(ctx: OptimizationContext):
|
||
"""
|
||
Return an objective function bound to `ctx`.
|
||
Pass the returned callable directly to pso / differential_evolution / minimize.
|
||
"""
|
||
def objective(params: list[float]) -> float:
|
||
return _evaluate(params, ctx)
|
||
return objective
|
||
|
||
|
||
def make_objective_function_xfr(ctx: OptimizationContext, y_indices):
|
||
"""
|
||
Return an xfr objective function bound to `ctx` and the `y_indices` surface.
|
||
Pass the returned callable directly to pso / differential_evolution / minimize.
|
||
"""
|
||
def objective(params: list[float]) -> float:
|
||
return _evaluate_xfr(params, ctx, y_indices)
|
||
return objective
|
||
|
||
|
||
# =====================================================================
|
||
# Legacy API (backward compatible)
|
||
#
|
||
# The old flow mutated module attributes on this file from the optimizers
|
||
# ("跨檔案注入變數"), which was fragile: any forgotten attribute showed up
|
||
# deep inside an optimizer callback as a NoneType/NameError. It is kept so
|
||
# existing callers (set_global_context + objective_function) keep working,
|
||
# but new code should use OptimizationContext + the make_* factories.
|
||
# =====================================================================
|
||
|
||
# Module-level state for the legacy path only.
|
||
cortical_tensor = None
|
||
spine_tensor = None
|
||
image1_array = None # cortical_nii.gz
|
||
image2_array = None # binarynii.gz
|
||
image2_shape = None
|
||
image3_array = None # roi2.nii.gz
|
||
diameter = None
|
||
length = None
|
||
spacing = [0.5, 0.5, 0.5]
|
||
device = None
|
||
grid = None
|
||
USE_TIP_PENALTY = None
|
||
|
||
_current_context: Optional[OptimizationContext] = None
|
||
|
||
|
||
def set_global_context(
|
||
cortical,
|
||
spine,
|
||
shape,
|
||
spacing_,
|
||
device_,
|
||
grid_,
|
||
use_tip_penalty=False
|
||
):
|
||
"""
|
||
Legacy: set the process-wide context used by the legacy
|
||
objective_function / objective_function_xfr wrappers.
|
||
|
||
Returns the built OptimizationContext for convenience.
|
||
"""
|
||
global _current_context
|
||
global cortical_tensor, spine_tensor, image2_shape, spacing, device, grid, USE_TIP_PENALTY
|
||
|
||
_current_context = OptimizationContext(
|
||
cortical_tensor=cortical,
|
||
spine_tensor=spine,
|
||
image2_shape=shape,
|
||
spacing=spacing_,
|
||
device=device_,
|
||
grid=grid_,
|
||
use_tip_penalty=use_tip_penalty,
|
||
)
|
||
|
||
# mirror into legacy module attributes for any code that reads them
|
||
cortical_tensor = cortical
|
||
spine_tensor = spine
|
||
image2_shape = shape
|
||
spacing = spacing_
|
||
device = device_
|
||
grid = grid_
|
||
USE_TIP_PENALTY = use_tip_penalty
|
||
|
||
return _current_context
|
||
|
||
|
||
def _active_context() -> OptimizationContext:
|
||
"""Resolve the context for the legacy wrappers, with a clear error."""
|
||
if _current_context is not None:
|
||
return _current_context
|
||
|
||
# Fallback: a caller (old-style optimizers) may have set the legacy
|
||
# module attributes directly — rebuild a context from them.
|
||
if cortical_tensor is not None and spine_tensor is not None:
|
||
return OptimizationContext(
|
||
cortical_tensor=cortical_tensor,
|
||
spine_tensor=spine_tensor,
|
||
image2_shape=image2_shape,
|
||
spacing=spacing,
|
||
device=device,
|
||
grid=grid,
|
||
use_tip_penalty=bool(USE_TIP_PENALTY),
|
||
)
|
||
|
||
raise RuntimeError(
|
||
"No optimization context available. Either call set_global_context(...) "
|
||
"or (preferred) build an OptimizationContext and use "
|
||
"make_objective_function / make_objective_function_xfr."
|
||
)
|
||
|
||
|
||
def objective_function(params: list[float]) -> float:
|
||
"""
|
||
Legacy wrapper: evaluates against the context set by set_global_context.
|
||
params = [z, y, x, azimuth, altitude, diameter_raw, length_raw]
|
||
"""
|
||
return _evaluate(params, _active_context())
|
||
|
||
|
||
def objective_function_xfr(params: list[float], y_indices) -> float:
|
||
"""
|
||
Legacy wrapper: evaluates against the context set by set_global_context.
|
||
params = [z, x, azimuth, altitude, diameter_raw, length_raw]
|
||
"""
|
||
return _evaluate_xfr(params, _active_context(), y_indices) |