Xiao Furen
264e17b13a
- Implemented xfr_pso.py to handle PSO optimization tasks for lumbar levels, including GPU management and logging. - Introduced xfr_record_planes.py to supplement existing volumes with symmetry and endplate plane data in transform.json. - Added functionality to skip already recorded levels unless forced, ensuring efficient processing. - Enhanced logging for better traceability of operations and errors during execution.
316 lines
12 KiB
Python
316 lines
12 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 個 voxel 的舊軟性 10000/voxel 懲罰，
|
||
# 現已改由 HARD_BREACH_PER_VOX 硬約束管制）；bonus 只在同 breach 數
|
||
# （＝同安全等級）的可行解之間讓優化偏好較粗螺絲。
|
||
DIAMETER_BONUS_PER_MM = 20000.0
|
||
# DIAMETER_BONUS_PER_MM = 5000000.0
|
||
# DIAMETER_BONUS_PER_MM = 2000000.0
|
||
# 硬約束的每 breach voxel 懲罰（不計 butt 豁免內）。
|
||
# 取 1e9：大於其餘所有評分項可能造成的總分差
|
||
# （d≤6.0 / L≤45 / 0.5mm voxel 時 |score| 展佈約 ≤ 2.1e8），
|
||
# 使 optimizer 先「把 breach voxel 數降到最低」（lexicographic），
|
||
# 再在同 breach 數下比品質。結果：完全在骨內（除 butt 豁免帶）的
|
||
# 螺絲永遠優於任何有 breach 的螺絲；PSO 找不到可行解時才回
|
||
# 最少 breach 的候選。
|
||
HARD_BREACH_PER_VOX = 1e9
|
||
|
||
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 的一端）5mm 豁免 mask
|
||
) -> float:
|
||
"""
|
||
漸進式評分：優先確保找到骨頭，再改善細節。
|
||
diameter 提供時附加直徑偏好 bonus（偏好較粗螺絲）。
|
||
cylinder_butt_torch 提供時，該 mask 內 not_in_bone 的 voxel 免除扣分
|
||
（butt 豁免帶）；豁免帶外的 breach voxel 以 HARD_BREACH_PER_VOX
|
||
硬約束處理——breach 數少的解永遠優於 breach 數多的解，與品質分無關。
|
||
"""
|
||
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 的一端）5mm 豁免：
|
||
# 這些 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）。
|
||
# breach 是硬約束（HARD_BREACH_PER_VOX ≫ 任何品質項），
|
||
# bonus 永遠不會把螺絲推出骨頭（butt 豁免帶除外）。
|
||
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 += 20 * in_vbody # VBODY 每 voxel 再加分
|
||
score -= 10000 * not_in_bone
|
||
# 硬約束：butt 豁免外的 breach 每 voxel 扣 HARD_BREACH_PER_VOX，
|
||
# 任何一個 breach voxel 的懲罰都超過其餘所有項的總分差 →
|
||
# 「breach=0」的解永遠優於「breach>0」的解（lexicographic）。
|
||
# score -= HARD_BREACH_PER_VOX * not_in_bone
|
||
# score -= 10000 * max(0, not_in_bone-allowed_error)
|
||
score -= 10000 * 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)
|
||
|