CBT_project/core/scoring.py

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)