core.py

"""
Core inference สำหรับ Bone Age Space
- โมเดลเป็น soft classification 240 bins (แต่ละ bin = 1 เดือน)
  => expected value ของ softmax = bone age (เดือน)
  => std ของ softmax distribution = ความไม่แน่นอน (SD, เดือน)  [variance-based]
- Grad-CAM: hook เอา activation/gradient จาก output ของ backbone (N,768,16,16)
  ใช้ expected value เป็น scalar target ในการ backprop
"""
import os
import cv2
import torch
import numpy as np
try:
    import segmentation_models_pytorch as smp
    _SMP_OK = True
except Exception as _e:
    smp = None
    _SMP_OK = False
    print(f"[seg] segmentation_models_pytorch ไม่พร้อมใช้ ({_e!r}) -> ใช้ mask แบบ threshold แทน")

from model import BoneAgeModel
from configuration import BoneAgeConfig
from preprocess_infer import preprocess_image

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
WEIGHTS_PATH = os.path.join(os.path.dirname(__file__), "best_model.pth")

_config = BoneAgeConfig(backbone="convnextv2_tiny", num_classes=240, in_chans=2)
_model = BoneAgeModel(
    backbone=_config.backbone,
    feature_dim=_config.feature_dim,
    dropout=_config.dropout,
    num_classes=_config.num_classes,
    in_chans=_config.in_chans,
)
_model.load_state_dict(torch.load(WEIGHTS_PATH, map_location=DEVICE))
_model = _model.eval().to(DEVICE)

_NUM_BINS = _config.num_classes  # 240
_BINS = torch.arange(_NUM_BINS, device=DEVICE).float()

# =============================================================
# DeepLabV3+ สำหรับ segment มือ (ใช้มาส์กพื้นหลังออกจาก Grad-CAM)
# วางไฟล์ weight ชื่อ best_deeplabv3plus_png_weights.pth ไว้ในโฟลเดอร์นี้
# =============================================================
_SEG_MODEL = None
SEG_WEIGHTS_PATH = os.path.join(os.path.dirname(__file__), "best_deeplabv3plus_png_weights.pth")

# ตัดพื้นหลังเป็นสีดำในภาพ Grad-CAM
#   True  = ตัดพื้นหลังทิ้ง (สวยถ้า mask ดี แต่ถ้า mask เพี้ยนจะเห็นมือโดนตัด/รูดำ)
#   False = ไม่ตัดพื้นหลัง เก็บภาพ X-ray ไว้ แค่ให้ heatmap อยู่บนมือ (ปลอดภัยสุด ไม่มีภาพพัง)
BLACKOUT_BG = True


def _get_seg_model():
    global _SEG_MODEL
    if not _SMP_OK:
        return None
    if _SEG_MODEL is None:
        if not os.path.exists(SEG_WEIGHTS_PATH):
            print(f"[seg] ⚠️ ไม่พบ {SEG_WEIGHTS_PATH} -> ใช้ mask แบบ threshold แทน")
            return None
        m = smp.DeepLabV3Plus(encoder_name="resnet50", encoder_weights=None,
                              in_channels=3, classes=1)
        m.load_state_dict(torch.load(SEG_WEIGHTS_PATH, map_location=DEVICE, weights_only=False))
        _SEG_MODEL = m.eval().to(DEVICE)
    return _SEG_MODEL


@torch.no_grad()
def _hand_mask(proc: np.ndarray) -> np.ndarray:
    """คืน mask มือ (H,W) float 0/1 จาก DeepLabV3+ ; ถ้าไม่มีโมเดล fallback เป็น Otsu"""
    seg = _get_seg_model()
    if seg is None:
        _, mm = cv2.threshold(proc, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        m = (mm > 0).astype(np.uint8)
    else:
        rgb = cv2.cvtColor(proc, cv2.COLOR_GRAY2RGB).astype(np.float32) / 255.0
        t = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0).to(DEVICE)
        prob = torch.sigmoid(seg(t)).cpu().numpy()[0, 0]
        m = (prob > 0.5).astype(np.uint8)
    # เก็บเฉพาะก้อนใหญ่สุด (= มือ) ลบจุดขาวหลง — ปลอดภัย ไม่ทำให้มือหาย
    n, lbl, st, _ = cv2.connectedComponentsWithStats(m, connectivity=8)
    if n > 1:
        big = 1 + int(np.argmax(st[1:, cv2.CC_STAT_AREA]))
        m = (lbl == big).astype(np.uint8)
    k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
    m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, k)
    m = cv2.morphologyEx(m, cv2.MORPH_OPEN, k)
    return m.astype(np.float32)


def _build_inputs(image: np.ndarray, sex: str):
    """return (img_tensor (1,1,512,512), female_tensor (1,), proc_uint8 (512,512))"""
    proc = preprocess_image(image)  # (512,512) uint8
    img_t = torch.from_numpy(proc).unsqueeze(0).unsqueeze(0).float().to(DEVICE)
    is_female = 1.0 if str(sex).lower().startswith("f") else 0.0
    female_t = torch.tensor([is_female], dtype=torch.float32).to(DEVICE)
    return img_t, female_t, proc


# ช่วงความคลาดเคลื่อนที่ยอมรับได้ (±เดือน) สำหรับคำนวณ "ความมั่นใจ"
CONF_TAU_MONTHS = 12  # ±12 เดือน (1 ปี) ; อยากเข้มขึ้นใช้ 6


@torch.no_grad()
def predict_stats(image: np.ndarray, sex: str, tau_months: float = CONF_TAU_MONTHS):
    """คืน mean (เดือน), sd (เดือน), confidence (%) จาก softmax distribution
    confidence = ผลรวมความน่าจะเป็นที่อายุจริงอยู่ในช่วง ±tau_months รอบค่าที่ทำนาย"""
    img_t, female_t, _ = _build_inputs(image, sex)
    with torch.cuda.amp.autocast(enabled=(DEVICE == "cuda")):
        logits = _model(img_t, female_t, return_logits=True)
    probs = logits.float().softmax(1)[0]          # (240,)
    mean = (probs * _BINS).sum()
    var = (probs * (_BINS - mean) ** 2).sum()
    sd = var.clamp(min=0).sqrt()

    # ความมั่นใจ (%) = ผลรวม prob ในหน้าต่าง ±tau รอบค่าที่ทำนาย
    m = float(mean.item())
    lo = int(max(0, round(m - tau_months)))
    hi = int(min(_NUM_BINS - 1, round(m + tau_months)))
    confidence = float(probs[lo:hi + 1].sum().item()) * 100.0
    return m, float(sd.item()), confidence


def _get_cam_layer():
    """เลือก layer สำหรับ Grad-CAM: stages[2] (384ch @32x32) ละเอียดกว่าชั้นสุดท้าย
    ถ้าโครงสร้างไม่ตรง (ไม่ใช่ convnext) fallback ไปที่ทั้ง backbone (16x16)"""
    bb = _model.backbone
    try:
        return bb.stages[2]
    except Exception:
        return bb


def gradcam_overlay(image: np.ndarray, sex: str, alpha: float = 0.6):
    """
    คืน (overlay_rgb uint8 (512,512,3), mean_months)
    Grad-CAM++ ที่ stages[2] ของ ConvNeXt โดยใช้ expected value เป็น target
    """
    img_t, female_t, proc = _build_inputs(image, sex)

    store = {}

    def fwd_hook(_m, _inp, out):
        out.retain_grad()      # ให้เก็บ .grad ของ activation ได้หลัง backward
        store["A"] = out

    target_layer = _get_cam_layer()
    h = target_layer.register_forward_hook(fwd_hook)

    try:
        _model.zero_grad(set_to_none=True)
        # ไม่ใช้ autocast เพื่อให้ gradient เป็น fp32 เสถียร
        logits = _model(img_t, female_t, return_logits=True)
        probs = logits.softmax(1)
        expected = (probs * _BINS).sum(1).squeeze()  # scalar (เดือน)
        expected.backward()

        A = store["A"][0]          # (C, H, W)
        G = store["A"].grad[0]     # (C, H, W)

        # ---- Grad-CAM++ weighting ----
        g2 = G.pow(2)
        g3 = G.pow(3)
        sum_a = A.sum(dim=(1, 2), keepdim=True)              # (C,1,1)
        denom = 2.0 * g2 + sum_a * g3
        denom = torch.where(denom != 0, denom, torch.ones_like(denom))
        alpha_kij = g2 / denom
        weights = (alpha_kij * torch.relu(G)).sum(dim=(1, 2))  # (C,)
        cam = torch.relu((weights[:, None, None] * A).sum(0))  # (H, W)
        cam = cam.detach().cpu().numpy().astype(np.float32)
    finally:
        h.remove()

    # upscale -> 512 (cubic ให้เนียนกว่า)
    cam = cv2.resize(cam, (proc.shape[1], proc.shape[0]), interpolation=cv2.INTER_CUBIC)
    cam = np.clip(cam, 0, None)

    # percentile clip กัน outlier เดี่ยว ๆ ครองภาพ แล้ว min-max normalize
    hi = np.percentile(cam, 99)
    if hi > 1e-6:
        cam = np.clip(cam, 0, hi)
    cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)

    # Gaussian smooth ลบขอบบล็อกจากการ upscale ให้ heatmap ดูนวลขึ้น
    cam = cv2.GaussianBlur(cam, (0, 0), sigmaX=9)
    cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
    # gamma เน้นยอดพีคให้เด่น
    cam = np.power(cam, 1.4)

    # mask เฉพาะบริเวณมือด้วย DeepLabV3+ (heatmap ขึ้นเฉพาะบนมือ)
    hand = _hand_mask(proc)
    soft = cv2.GaussianBlur(
        cv2.dilate(hand, np.ones((7, 7), np.uint8), iterations=1), (0, 0), sigmaX=5)
    cam = cam * soft

    # เบลนด์แบบถ่วงน้ำหนักด้วยความแรงของ CAM (TURBO สวยกว่า JET)
    heat = cv2.applyColorMap((cam * 255).astype(np.uint8), cv2.COLORMAP_TURBO)  # BGR
    base = cv2.cvtColor(proc, cv2.COLOR_GRAY2BGR).astype(np.float32)
    w = (alpha * cam)[..., None]  # (H, W, 1)
    overlay = base * (1 - w) + heat.astype(np.float32) * w
    if BLACKOUT_BG:
        # ตัดพื้นหลังออก -> นอกบริเวณมือเป็นสีดำ
        out_mask = cv2.GaussianBlur((hand > 0.5).astype(np.float32), (0, 0), sigmaX=2)[..., None]
        overlay = overlay * out_mask
    overlay = overlay.astype(np.uint8)
    overlay_rgb = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)

    mean = float((probs.detach() * _BINS).sum(1).item())
    return overlay_rgb, mean