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

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
) -> 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()

    in_bone= ((spine_tensor == 1) & (cylinder_torch == 1)).sum().item()
    not_in_bone= ((spine_tensor == 0) & (cylinder_torch == 1)).sum().item()
    
    # if cyl_total == 0:
    #     return float(1000*1000)
        # return float(1e9)  # 極差的情況
    
    overlap_ratio = overlap / cyl_total

    # if cyl_total < 1000:
    #     return float((1000 - cyl_total)*10000)

    score = cyl_total

    # if in_bone == 0:
    #     return float(not_in_bone*200)

    score += overlap*30
    score += in_bone*10
    score -= not_in_bone*1000
    score -= null_vox2*1000

    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)