Spaces:
Sleeping
Sleeping
File size: 5,130 Bytes
f0d9c8b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | """
AngioAI - stenosis detection: 3 x YOLOv8-seg ensemble with weighted mask fusion.
Validated on ARCADE stenosis (threshold/config chosen on VAL, reported on held-out TEST, 300 imgs):
F1@IoU0.3 = 0.637 precision 0.549 recall 0.759 AP30 0.707 F1@IoU0.5 = 0.467 0.80 FP/img
(see D:\\AngioAI - Commercial\\stenosis\\results\\ensemble\\)
Fusion (same as eval_ensemble_stenosis.py):
- predictions from all members clustered by box IoU >= CLUSTER_IOU (one prediction per model per cluster)
- cluster score = sum_m(w_m * s_m) / sum(all w) -> a lesion seen by only one model is down-weighted
- cluster mask = weighted pixel vote >= 0.5
Usage:
ens = StenosisEnsemble() # loads weights from ./weights/stenosis or ANGIOAI_WEIGHTS
dets = ens.predict(frame_rgb, score_thr=0.20)
# -> [{"bbox":[x1,y1,x2,y2], "score":float, "mask":uint8[H,W], "n_models":int, "models":[...]}, ...]
"""
import os
from pathlib import Path
import numpy as np
WEIGHTS_DIR = Path(os.environ.get("ANGIOAI_WEIGHTS", Path(__file__).resolve().parent / "weights")) / "stenosis"
# name: (file, imgsz used in training, fusion weight = val AP@IoU0.3 of the single model)
MEMBERS = {
"yolov8x_seg": ("yolov8x_seg.pt", 640, 0.6407),
"yolov8l_sgd": ("yolov8l_sgd.pt", 512, 0.6173),
"yolov8l_adamw": ("yolov8l_adamw.pt", 512, 0.6107),
}
CLUSTER_IOU = 0.20
DEFAULT_SCORE_THR = 0.20 # chosen on ARCADE val (max F1@IoU0.3)
MEMBER_CONF = 0.01 # keep low-confidence member preds; the fused score does the filtering
MAX_DET = 30
HF_REPO = os.environ.get("ANGIOAI_HF_REPO", "MuhammadAdil63/angio-ai-checkpoints")
def _resolve(fname: str) -> str:
local = WEIGHTS_DIR / fname
if local.exists():
return str(local)
from huggingface_hub import hf_hub_download # fallback for cloud deployment
return hf_hub_download(repo_id=HF_REPO, filename=f"stenosis/{fname}", repo_type="model",
token=os.environ.get("HF_TOKEN"))
def _box_iou(a, b):
ix = max(0, min(a[2], b[2]) - max(a[0], b[0])); iy = max(0, min(a[3], b[3]) - max(a[1], b[1]))
i = ix * iy
u = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - i
return i / u if u else 0.0
def _mask_box(m):
ys, xs = np.nonzero(m)
return (int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1) if len(ys) else None
class StenosisEnsemble:
def __init__(self, device=None, members=None):
from ultralytics import YOLO
import torch
self.device = device if device is not None else (0 if torch.cuda.is_available() else "cpu")
self.members = members or MEMBERS
self.models = {n: YOLO(_resolve(f)) for n, (f, _, _) in self.members.items()}
self.W = sum(w for _, _, w in self.members.values())
def _member_preds(self, frame_rgb):
H, W = frame_rgb.shape[:2]
bgr = np.ascontiguousarray(frame_rgb[..., ::-1]) # ultralytics expects BGR numpy
out = {}
for n, model in self.models.items():
_, imgsz, _ = self.members[n]
r = model.predict(bgr, imgsz=imgsz, conf=MEMBER_CONF, max_det=MAX_DET,
retina_masks=True, verbose=False, device=self.device)[0]
preds = []
if r.masks is not None:
for m, s in zip(r.masks.data.cpu().numpy() > 0.5, r.boxes.conf.cpu().numpy()):
if m.shape != (H, W):
import cv2
m = cv2.resize(m.astype(np.uint8), (W, H), interpolation=cv2.INTER_NEAREST) > 0
b = _mask_box(m)
if b:
preds.append((float(s), b, m))
out[n] = preds
return out
def fuse(self, member_preds):
cand = sorted(((s, b, m, n) for n, ps in member_preds.items() for s, b, m in ps),
key=lambda x: -(x[0] * self.members[x[3]][2]))
clusters = []
for s, b, m, n in cand:
for c in clusters:
if n not in c["m"] and _box_iou(b, c["ref"]) >= CLUSTER_IOU:
c["m"][n] = (s, m); break
else:
clusters.append({"ref": b, "m": {n: (s, m)}})
dets = []
for c in clusters:
score = sum(self.members[n][2] * s for n, (s, _) in c["m"].items()) / self.W
acc = None; tot = 0.0
for n, (s, m) in c["m"].items():
w = self.members[n][2] * s
acc = m * w if acc is None else acc + m * w
tot += w
mask = acc >= 0.5 * tot
b = _mask_box(mask)
if b:
dets.append({"bbox": [b[0], b[1], b[2] - 1, b[3] - 1], "score": float(score),
"mask": mask.astype(np.uint8), "n_models": len(c["m"]),
"models": sorted(c["m"])})
return sorted(dets, key=lambda d: -d["score"])
def predict(self, frame_rgb, score_thr=DEFAULT_SCORE_THR):
return [d for d in self.fuse(self._member_preds(frame_rgb)) if d["score"] >= score_thr]
|