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.
298 lines
11 KiB
Python
298 lines
11 KiB
Python
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,
|
||
cylinder_torch: torch.Tensor,
|
||
cylinder_o_torch: torch.Tensor,
|
||
intersections: int,
|
||
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 獎勵
|
||
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
|
||
# VBODY 獎勵：螺絲落在 (cortical + VBODY) 內的 voxel 每個 100 分
|
||
# （100 分項由純 cortical 擴展到 cortical∪VBODY；cortical voxel 分數不變），
|
||
# 其中落在 VBODY 的 voxel 每個再加 10 分
|
||
if vbody_tensor is not None:
|
||
in_vbody = ((vbody_tensor == 1) & (cylinder_torch == 1)).sum().item()
|
||
in_corti_vb = (((cortical_tensor == 1) | (vbody_tensor == 1)) & (cylinder_torch == 1)).sum().item()
|
||
else:
|
||
in_vbody = 0
|
||
in_corti_vb = overlap
|
||
null_vox = ((cortical_tensor == 0) & (cylinder_torch == 1)).sum().item()
|
||
null_vox2 = ((spine_tensor == 1) & (cylinder_o_torch == 1)).sum().item()
|
||
|
||
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) # 極差的情況
|
||
|
||
overlap_ratio = overlap / cyl_total
|
||
|
||
# if cyl_total < 1000:
|
||
# return float((1000 - cyl_total)*10000)
|
||
|
||
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 += 200 * in_corti_vb # (cortical + VBODY) 每 voxel 再加分
|
||
score += 50 * in_vbody # VBODY 每 voxel 再加分
|
||
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)
|
||
|
||
# === 階段 1：首要目標是找到骨頭（overlap > 0） ===
|
||
if overlap == 0:
|
||
# 完全沒有 overlap 是最糟糕的情況
|
||
score -= 500000 # 超大懲罰
|
||
# 如果連 spine 都沒穿過，更糟
|
||
if intersections == 0:
|
||
score -= 500000
|
||
return float(-score)
|
||
|
||
# === 階段 2：有找到骨頭了，開始改善品質 ===
|
||
|
||
# 1. Overlap 獎勵（非線性，鼓勵快速提升）
|
||
if overlap_ratio < 0.1:
|
||
# 0-10%：每增加 1% 給大量獎勵（鼓勵探索）
|
||
score += overlap * 5000 # 很高的單位獎勵
|
||
elif overlap_ratio < 0.3:
|
||
# 10-30%：中等獎勵
|
||
score += overlap * 3000
|
||
elif overlap_ratio < 0.5:
|
||
# 30-50%：正常獎勵
|
||
score += overlap * 2000
|
||
else:
|
||
# 50%+：獎勵 + 額外比例獎勵
|
||
score += overlap * 2000
|
||
score += (overlap_ratio - 0.5) * 100000 # 超過 50% 額外大獎
|
||
|
||
# 2. Intersection 控制（稍微放寬）
|
||
if intersections == 1:
|
||
score += 20000 # 完美
|
||
elif intersections == 0:
|
||
score -= 200000 # 嚴重錯誤（但比完全沒 overlap 好）
|
||
elif intersections == 2:
|
||
score -= 10000 # 可接受但不理想
|
||
else:
|
||
score -= intersections * 15000
|
||
|
||
# 3. Null voxel 懲罰（漸進式）
|
||
null_ratio = null_vox / cyl_total
|
||
|
||
if overlap_ratio < 0.2:
|
||
# 如果 overlap 還很少，對 null voxel 寬容一點
|
||
score -= null_vox * 300
|
||
elif overlap_ratio < 0.4:
|
||
score -= null_vox * 600
|
||
else:
|
||
# overlap 夠高了，開始嚴格要求
|
||
if null_ratio > 0.5:
|
||
score -= null_vox * 1500
|
||
else:
|
||
score -= null_vox * 800
|
||
|
||
# 4. 反向圓柱懲罰
|
||
score -= null_vox2 * 1000
|
||
|
||
# 5. 尺寸合理性（放寬）
|
||
if diameter is not None and length is not None:
|
||
if diameter < 2.5 or diameter > 6.0: # 放寬從 (3.0, 5.5) 到 (2.5, 6.0)
|
||
score -= 3000
|
||
if length < 25 or length > 60: # 放寬從 (30, 55) 到 (25, 60)
|
||
score -= 3000
|
||
|
||
# 6. 尖端 breach 懲罰
|
||
if cylinder_tip_torch is not None:
|
||
tip_total = cylinder_tip_torch.sum().item()
|
||
if tip_total > 0:
|
||
tip_breach = ((cortical_tensor == 0) & (cylinder_tip_torch == 1)).sum().item()
|
||
tip_breach_ratio = tip_breach / tip_total
|
||
if tip_breach_ratio > 0:
|
||
score -= tip_breach * 5000 # 尖端出界懲罰要比一般 null_vox 重很多
|
||
|
||
return float(-score)
|
||
|
||
def cl_score_torch(
|
||
cortical_tensor: torch.Tensor,
|
||
spine_tensor: torch.Tensor,
|
||
cylinder_torch: torch.Tensor,
|
||
cylinder_o_torch: torch.Tensor,
|
||
intersections: int,
|
||
diameter: float = None,
|
||
length: float = None,
|
||
cylinder_tip_torch: torch.Tensor = None # 新增：尖端 mask
|
||
) -> float:
|
||
"""
|
||
漸進式評分：優先確保找到骨頭，再改善細節
|
||
"""
|
||
cyl_total = cylinder_torch.sum().item()
|
||
overlap = ((cortical_tensor == 1) & (cylinder_torch == 1)).sum().item()
|
||
null_vox = ((cortical_tensor == 0) & (cylinder_torch == 1)).sum().item()
|
||
null_vox2 = ((spine_tensor == 1) & (cylinder_o_torch == 1)).sum().item()
|
||
|
||
if cyl_total == 0:
|
||
return float(1e9) # 極差的情況
|
||
|
||
overlap_ratio = overlap / cyl_total
|
||
|
||
score = 0
|
||
|
||
# === 階段 1：首要目標是找到骨頭（overlap > 0） ===
|
||
if overlap == 0:
|
||
# 完全沒有 overlap 是最糟糕的情況
|
||
score -= 500000 # 超大懲罰
|
||
# 如果連 spine 都沒穿過，更糟
|
||
if intersections == 0:
|
||
score -= 500000
|
||
return float(-score)
|
||
|
||
# === 階段 2：有找到骨頭了，開始改善品質 ===
|
||
|
||
# 1. Overlap 獎勵（非線性，鼓勵快速提升）
|
||
if overlap_ratio < 0.1:
|
||
# 0-10%：每增加 1% 給大量獎勵（鼓勵探索）
|
||
score += overlap * 5000 # 很高的單位獎勵
|
||
elif overlap_ratio < 0.3:
|
||
# 10-30%：中等獎勵
|
||
score += overlap * 3000
|
||
elif overlap_ratio < 0.5:
|
||
# 30-50%：正常獎勵
|
||
score += overlap * 2000
|
||
else:
|
||
# 50%+：獎勵 + 額外比例獎勵
|
||
score += overlap * 2000
|
||
score += (overlap_ratio - 0.5) * 100000 # 超過 50% 額外大獎
|
||
|
||
# 2. Intersection 控制（稍微放寬）
|
||
if intersections == 1:
|
||
score += 20000 # 完美
|
||
elif intersections == 0:
|
||
score -= 200000 # 嚴重錯誤（但比完全沒 overlap 好）
|
||
elif intersections == 2:
|
||
score -= 10000 # 可接受但不理想
|
||
else:
|
||
score -= intersections * 15000
|
||
|
||
# 3. Null voxel 懲罰（漸進式）
|
||
null_ratio = null_vox / cyl_total
|
||
|
||
if overlap_ratio < 0.2:
|
||
# 如果 overlap 還很少，對 null voxel 寬容一點
|
||
score -= null_vox * 300
|
||
elif overlap_ratio < 0.4:
|
||
score -= null_vox * 600
|
||
else:
|
||
# overlap 夠高了，開始嚴格要求
|
||
if null_ratio > 0.5:
|
||
score -= null_vox * 1500
|
||
else:
|
||
score -= null_vox * 800
|
||
|
||
# 4. 反向圓柱懲罰
|
||
score -= null_vox2 * 1000
|
||
|
||
# 5. 尺寸合理性（放寬）
|
||
if diameter is not None and length is not None:
|
||
if diameter < 2.5 or diameter > 6.0: # 放寬從 (3.0, 5.5) 到 (2.5, 6.0)
|
||
score -= 3000
|
||
if length < 25 or length > 60: # 放寬從 (30, 55) 到 (25, 60)
|
||
score -= 3000
|
||
|
||
# 6. 尖端 breach 懲罰
|
||
if cylinder_tip_torch is not None:
|
||
tip_total = cylinder_tip_torch.sum().item()
|
||
if tip_total > 0:
|
||
tip_breach = ((cortical_tensor == 0) & (cylinder_tip_torch == 1)).sum().item()
|
||
tip_breach_ratio = tip_breach / tip_total
|
||
if tip_breach_ratio > 0:
|
||
score -= tip_breach * 5000 # 尖端出界懲罰要比一般 null_vox 重很多
|
||
|
||
return float(-score)
|
||
|
||
def get_overlap_ratio(
|
||
position_params: list,
|
||
diameter: float,
|
||
length: float,
|
||
cortical_tensor: torch.Tensor,
|
||
image_shape: tuple,
|
||
spacing: list,
|
||
device: torch.device,
|
||
grid=None
|
||
) -> float:
|
||
"""
|
||
計算 Cylinder 與 Cortical Bone 的重疊比例 (%)
|
||
"""
|
||
# 生成 Cylinder Mask
|
||
cyl_mask = generate_cylinder_n_torch(
|
||
diameter, length,
|
||
position_params[0], position_params[1], position_params[2],
|
||
position_params[3], position_params[4],
|
||
image_shape, spacing, device, grid
|
||
)
|
||
|
||
# 計算體積 (Voxel count)
|
||
cyl_vol = torch.sum(cyl_mask).item()
|
||
|
||
if cyl_vol == 0:
|
||
return 0.0
|
||
|
||
# 計算重疊部分
|
||
# 注意：這裡使用 cortical_tensor (與 cl_score_torch 邏輯一致)
|
||
overlap_count = ((cortical_tensor == 1) & (cyl_mask == 1)).sum().item()
|
||
|
||
return (overlap_count / cyl_vol) * 100.0
|
||
|
||
def compute_overlap_ratio_from_cylinder_mask(cyl_mask: torch.Tensor,
|
||
spine_mask: torch.Tensor,
|
||
eps: float = 1e-6) -> float:
|
||
"""
|
||
一個常見定義： overlap = intersection / cylinder_volume
|
||
你也可以改成 intersection / spine_volume 或 Dice，依你論文/需求一致即可。
|
||
cyl_mask, spine_mask: uint8/bool tensor, same shape
|
||
"""
|
||
cyl = cyl_mask.bool()
|
||
sp = spine_mask.bool()
|
||
inter = (cyl & sp).sum().item()
|
||
denom = cyl.sum().item()
|
||
return float(inter) / float(denom + eps)
|
||
|
||
def is_solution_ok(loss: float, overlap: float, overlap_thresh: float = OVERLAP_THRESH) -> bool:
|
||
return (loss <= 0) and (overlap >= overlap_thresh)
|
||
|