""" 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