Deploy AngioAI Pro: YOLOv8-seg stenosis ensemble, corrected metrics & SYNTAX

README.md

@@ -1,5 +1,5 @@
 ---
-title: Angio AI
+title: AngioAI Pro
 emoji: 🫀
 colorFrom: red
 colorTo: blue
@@ -7,52 +7,24 @@ sdk: gradio
 sdk_version: 4.44.1
 app_file: app.py
 pinned: false
-license: apache-2.0
+license: other
 ---
 
-# Angio AI — Coronary Angiography Analysis System
+# AngioAI Pro — Coronary Angiography Analysis
 
-**MS Data Science Thesis · Information Technology University Lahore**  
-**Supervisor: Dr. Arif Mehmood**
+Upload a coronary X-ray angiography (XCA) video. The pipeline:
 
----
-
-## What it does
-
-Upload a coronary X-ray angiography (XCA) video and the pipeline automatically:
-
-| Step | Module | Output |
-|------|--------|--------|
-| 1 | Keyframe extraction | Best diagnostic frame (contrast × sharpness score) |
-| 2 | Mask2Former | Stenosis detection — bounding boxes + instance masks |
-| 3a | ResUNet | Binary vessel mask (Dice 0.8015 on ARCADE) |
-| 3b | YOLOv8m-seg | 26-class coronary anatomy segmentation |
-| 4 | FFR Pipeline v4 | Quantitative Flow Ratio (QFR) estimation |
-| 5 | SYNTAX Score | Lesion complexity scoring |
-
----
-
-## Models
+| Step | Module | Notes |
+|---|---|---|
+| 1 | Keyframe extraction | contrast × sharpness score |
+| 2 | Stenosis detection | 3 × YOLOv8-seg ensemble with weighted mask fusion |
+| 3a | Binary vessel mask | ResUNet |
+| 3b | Coronary segments | 25 SYNTAX segments (YOLOv8-seg) |
+| 4 | QCA + FFR estimate | diameter stenosis on the measured branch; QFR = 1 − (0.33·DS + 0.60·DS²) |
+| 5 | SYNTAX (anatomical) | segment weights × lesion multiplier, lesions ≥ 50 % DS |
 
-All checkpoints are stored in [`MuhammadAdil63/angio-ai-checkpoints`](https://huggingface.co/MuhammadAdil63/angio-ai-checkpoints) and downloaded automatically on first run.
+**Held-out ARCADE stenosis test set (300 images):** precision 0.55, recall 0.76, F1 0.64 (mask IoU ≥ 0.3), 0.80 false positives per image.
 
-| File | Architecture | Task | Performance |
-|------|-------------|------|------------|
-| `mask2former_best.pth` | Mask2Former Swin-Base | Stenosis detection | — |
-| `binary_best.pth` | ResUNet (16→256 ch) | Binary vessel segmentation | Dice 0.8015 |
-| `best.pt` | YOLOv8m-seg (nc=26) | 26-class coronary anatomy | — |
-
-Dataset: [ARCADE Challenge](https://arcade.grand-challenge.org/) (syntax + stenosis splits)
-
----
-
-## Notes
-
-- **Hardware**: CPU-only (HF free tier) — inference takes ~30–60 seconds per video
-- **Input**: MP4 / AVI coronary angiography video, ideally ≥5 seconds
-- **Scale**: Default 3.75 px/mm (ARCADE hardware). Adjust slider for non-ARCADE data.
-- **FFR formula**: QFR = 1 − (0.33·DS + 0.60·DS²) — Tu et al. JACC 2016 / FAVOR II
-
----
+Checkpoints are loaded from the private model repo `MuhammadAdil63/angio-ai-checkpoints` (Space secret `HF_TOKEN`).
 
-*This tool is for research purposes only and is not intended for clinical use.*
+> **Research use only.** Not a medical device and not for clinical diagnosis. Not validated on non-ARCADE data or on normal angiograms.

app.py

@@ -7,7 +7,8 @@
 # Checkpoints are loaded from HF Hub: MuhammadAdil63/angio-ai-checkpoints
 # Set HF_TOKEN as a Space Secret (Settings → Secrets) for private repo access.
 # Files expected in that repo:
-#   mask2former_best.pth   (Mask2Former Swin-Base fine-tuned)
+#   stenosis/yolov8x_seg.pt, stenosis/yolov8l_sgd.pt, stenosis/yolov8l_adamw.pt (stenosis ensemble;
+#   loaded from ./weights/stenosis first, HF Hub as fallback)
 #   binary_best.pth        (ResUNet binary vessel segmentation)
 #   best.pt                (YOLOv8m-seg 26-class ARCADE syntax)
 #
@@ -112,7 +113,6 @@ def _hf_ckpt(filename: str) -> str:
 CONFIG = {
     # ── Checkpoint paths (resolved lazily via HF Hub) ─────────────────────────
     # These callables are invoked once inside each model's _load_*() function.
-    "MASK2FORMER_CKPT" : lambda: _hf_ckpt("mask2former_best.pth"),
     "RESUNET_CKPT"     : lambda: _hf_ckpt("binary_best.pth"),
     "YOLO_CKPT"        : lambda: _hf_ckpt("best.pt"),
 
@@ -171,292 +171,196 @@ LOGO_ITU = _encode_logo(str(Path(__file__).parent / "ITULog.png"))
 # ─────────────────────────────────────────────────────────────────────────────
 
 CSS = """
-@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;600&display=swap');
-
+/* ── Base ── */
 :root {
-    --navy-950:  #040d1a;
-    --navy-900:  #071224;
-    --navy-800:  #0d1e38;
-    --navy-700:  #112a50;
-    --navy-600:  #1a3a68;
-    --cyan-400:  #22d3ee;
-    --cyan-300:  #67e8f9;
-    --cyan-glow: rgba(34,211,238,0.12);
-    --green-400: #34d399;
-    --red-400:   #f87171;
-    --amber-400: #fbbf24;
-    --white:     #f0f9ff;
-    --w60:       rgba(240,249,255,0.60);
-    --w30:       rgba(240,249,255,0.30);
-    --w10:       rgba(240,249,255,0.08);
-    --border:    rgba(34,211,238,0.18);
-    --border-s:  rgba(240,249,255,0.09);
-    --shadow:    0 4px 24px rgba(0,0,0,0.40);
-    --radius:    12px;
-    --radius-s:  8px;
-    --font:      'Inter', system-ui, sans-serif;
-    --mono:      'JetBrains Mono', 'Courier New', monospace;
+    --primary:      #1a6fa8;
+    --primary-lt:   #e8f3fb;
+    --primary-mid:  #4a9fd4;
+    --accent:       #0d9e6e;        /* Angio AI green */
+    --bg:           #f7fafd;
+    --surface:      #ffffff;
+    --border:       #cce0f0;
+    --text:         #1a2533;
+    --text-sec:     #5a7a96;
+    --danger:       #d64040;
+    --warn:         #d98b00;
+    --radius:       10px;
+    --shadow:       0 2px 12px rgba(26,111,168,0.10);
 }
-*, *::before, *::after { box-sizing: border-box; }
 
 body, .gradio-container {
-    background: var(--navy-950) !important;
-    font-family: var(--font) !important;
-    color: var(--white) !important;
-    min-height: 100vh;
+    background: var(--bg) !important;
+    font-family: 'Segoe UI', system-ui, sans-serif !important;
+    color: var(--text) !important;
 }
-.gradio-container { padding: 0 !important; max-width: 100% !important; }
-footer { display: none !important; }
-.gap { gap: 16px !important; }
 
-/* ── Header ─────────────────────────────────────────────────────────── */
+/* ── Header bar ── */
 #header {
-    background: var(--navy-900);
-    border-bottom: 1px solid var(--border-s);
-    padding: 0 32px;
+    background: var(--surface);
+    border-bottom: 2px solid var(--border);
+    padding: 14px 28px;
     display: flex;
     align-items: center;
     justify-content: space-between;
-    min-height: 72px;
-    position: relative;
-    overflow: hidden;
-}
-#header::after {
-    content: '';
-    position: absolute;
-    bottom: 0; left: 0; right: 0; height: 2px;
-    background: linear-gradient(90deg,
-        transparent 0%, var(--cyan-400) 15%, var(--cyan-400) 17%,
-        transparent 19%, transparent 24%,
-        var(--cyan-400) 26%, #fff 28%, var(--cyan-400) 30%,
-        transparent 32%, transparent 38%,
-        var(--cyan-400) 40%, var(--cyan-400) 42%,
-        transparent 44%, transparent 100%);
-    opacity: 0.7;
-}
-#logo-angio .name {
-    font-size: 22px; font-weight: 700;
-    color: var(--cyan-400) !important;
-    letter-spacing: 3px; text-transform: uppercase;
-    font-family: var(--mono) !important;
+    box-shadow: var(--shadow);
+    margin-bottom: 0 !important;
 }
-#logo-angio .sub {
-    font-size: 10px; color: var(--w60) !important;
-    letter-spacing: 1.5px; text-transform: uppercase; margin-top: 3px;
+#logo-itu img { height: 48px; }
+#logo-angio {
+    font-size: 26px;
+    font-weight: 700;
+    color: var(--accent);
+    letter-spacing: 1px;
 }
-#logo-itu img { height: 44px; filter: brightness(0) invert(1) opacity(0.65); }
-.header-meta {
-    text-align: right; font-size: 11px;
-    color: var(--w60) !important; line-height: 1.7;
-    font-family: var(--mono) !important;
+#logo-angio span {
+    font-size: 13px;
+    color: var(--text-sec);
+    font-weight: 400;
+    display: block;
+    letter-spacing: 0;
 }
-.header-meta b { color: var(--cyan-300) !important; font-weight: 600; }
 
-/* ── Tabs ─────────────────────────────────────────────────────────── */
-.tab-nav {
-    background: var(--navy-900) !important;
-    border-bottom: 1px solid var(--border-s) !important;
-    padding: 0 24px !important;
-}
+/* ── Tab bar ── */
 .tab-nav button {
-    font-family: var(--font) !important; font-size: 13px !important;
-    font-weight: 500 !important; color: var(--w60) !important;
-    background: transparent !important; border: none !important;
-    border-bottom: 2px solid transparent !important;
-    padding: 14px 20px !important;
-    transition: color 0.2s, border-color 0.2s !important;
+    color: var(--text-sec) !important;
+    border-bottom: 3px solid transparent !important;
+    font-size: 15px !important;
+    font-weight: 500 !important;
+    padding: 10px 24px !important;
+    background: transparent !important;
 }
-.tab-nav button:hover { color: var(--white) !important; }
 .tab-nav button.selected {
-    color: var(--cyan-400) !important;
-    border-bottom-color: var(--cyan-400) !important;
-    font-weight: 600 !important;
+    color: var(--primary) !important;
+    border-bottom-color: var(--primary) !important;
 }
 
-/* ── Cards ────────────────────────────────────────────────────────── */
+/* ── Cards ── */
 .card {
-    background: var(--navy-800) !important;
-    border: 1px solid var(--border-s) !important;
-    border-radius: var(--radius) !important;
-    padding: 20px 22px !important;
-    box-shadow: var(--shadow) !important;
-    margin-bottom: 16px !important;
-    position: relative; overflow: hidden;
-}
-.card::before {
-    content: ''; position: absolute;
-    top: 0; left: 0; right: 0; height: 1px;
-    background: linear-gradient(90deg, transparent, var(--cyan-400), transparent);
-    opacity: 0.35;
+    background: var(--surface);
+    border: 1px solid var(--border);
+    border-radius: var(--radius);
+    padding: 22px;
+    box-shadow: var(--shadow);
+    margin-bottom: 16px;
 }
 .card-title {
-    font-size: 10px !important; font-weight: 600 !important;
-    color: var(--cyan-400) !important; text-transform: uppercase !important;
-    letter-spacing: 2px !important; margin-bottom: 14px !important;
-    padding-bottom: 10px !important; border-bottom: 1px solid var(--border-s) !important;
-    font-family: var(--mono) !important;
+    font-size: 14px;
+    font-weight: 600;
+    color: var(--primary);
+    text-transform: uppercase;
+    letter-spacing: 0.5px;
+    margin-bottom: 12px;
+    padding-bottom: 8px;
+    border-bottom: 1px solid var(--primary-lt);
 }
 
-/* ── Upload zone ──────────────────────────────────────────────────── */
-#upload-zone, #upload-zone > div {
-    border: 1px dashed var(--cyan-400) !important;
+/* ── Upload zone ── */
+#upload-zone {
+    border: 2px dashed var(--primary-mid) !important;
     border-radius: var(--radius) !important;
-    background: var(--cyan-glow) !important;
-    min-height: 160px !important;
-    transition: all 0.2s !important;
-}
-#upload-zone:hover {
-    border-color: var(--cyan-300) !important;
-    box-shadow: 0 0 20px var(--cyan-glow) !important;
-}
-#upload-zone label, #upload-zone span, #upload-zone p {
-    color: var(--w60) !important; font-size: 13px !important;
+    background: var(--primary-lt) !important;
+    min-height: 200px !important;
 }
-#upload-zone svg { stroke: var(--cyan-400) !important; opacity: 0.6; }
 
-/* ── Sliders ──────────────────────────────────────────────────────── */
-input[type=range] { accent-color: var(--cyan-400) !important; }
-.gradio-slider label, .gradio-slider .label-wrap span {
-    color: var(--w60) !important; font-size: 12px !important;
-    font-family: var(--mono) !important;
-}
-.gradio-slider .info { color: var(--w30) !important; font-size: 11px !important; }
+/* ── Sliders ── */
+input[type=range] { accent-color: var(--primary); }
+.gradio-slider label { color: var(--text-sec) !important; font-size: 13px !important; }
 
-/* ── Buttons ──────────────────────────────────────────────────────── */
+/* ── Buttons ── */
 #btn-analyse {
-    background: linear-gradient(135deg, #0ea5e9, var(--cyan-400)) !important;
-    color: var(--navy-950) !important; border: none !important;
-    border-radius: var(--radius-s) !important;
-    font-size: 12px !important; font-weight: 700 !important;
-    padding: 13px 32px !important; letter-spacing: 2px !important;
-    text-transform: uppercase !important; font-family: var(--mono) !important;
-    cursor: pointer !important; width: 100% !important;
-    box-shadow: 0 0 24px rgba(34,211,238,0.30) !important;
-    transition: all 0.2s !important;
-}
-#btn-analyse:hover {
-    box-shadow: 0 0 36px rgba(34,211,238,0.50) !important;
-    transform: translateY(-1px) !important;
+    background: var(--primary) !important;
+    color: white !important;
+    border: none !important;
+    border-radius: var(--radius) !important;
+    font-size: 15px !important;
+    font-weight: 600 !important;
+    padding: 12px 32px !important;
+    cursor: pointer !important;
+    width: 100% !important;
+    transition: background 0.2s !important;
 }
+#btn-analyse:hover { background: #155d90 !important; }
 #btn-reset {
-    background: transparent !important; color: var(--w60) !important;
-    border: 1px solid var(--border-s) !important;
-    border-radius: var(--radius-s) !important;
-    font-size: 11px !important; font-family: var(--mono) !important;
-    letter-spacing: 1px !important; text-transform: uppercase !important;
-    padding: 10px 20px !important; cursor: pointer !important;
-    width: 100% !important; transition: all 0.2s !important;
-}
-#btn-reset:hover { border-color: var(--w30) !important; color: var(--white) !important; }
-
-button.sm.secondary, button.lg.secondary {
-    background: transparent !important; color: var(--cyan-400) !important;
+    background: transparent !important;
+    color: var(--text-sec) !important;
     border: 1px solid var(--border) !important;
-    border-radius: var(--radius-s) !important;
-    font-size: 11px !important; font-family: var(--mono) !important;
-    letter-spacing: 1.5px !important; text-transform: uppercase !important;
-    width: 100% !important; padding: 9px 12px !important;
-    margin-top: 8px !important; transition: all 0.2s !important;
-}
-button.sm.secondary:hover, button.lg.secondary:hover {
-    background: var(--cyan-glow) !important;
-    box-shadow: 0 0 16px var(--cyan-glow) !important;
+    border-radius: var(--radius) !important;
+    font-size: 13px !important;
+    padding: 8px 20px !important;
+    cursor: pointer !important;
+    width: 100% !important;
 }
 
-/* ── Metric badges ────────────────────────────────────────────────── */
-.metric-row { display:flex !important; gap:10px !important; flex-wrap:wrap !important; margin-bottom:16px !important; }
-.metric-badge {
-    flex: 1 !important; min-width: 100px !important;
-    background: var(--navy-900) !important; border: 1px solid var(--border-s) !important;
-    border-radius: var(--radius-s) !important; padding: 14px 10px !important;
-    text-align: center !important; position: relative; overflow: hidden;
+/* ── Metric badges ── */
+.metric-row {
+    display: flex;
+    gap: 12px;
+    flex-wrap: wrap;
+    margin-bottom: 12px;
 }
-.metric-badge::before {
-    content: ''; position: absolute;
-    top: 0; left: 0; right: 0; height: 1px;
-    background: var(--cyan-400); opacity: 0.3;
+.metric-badge {
+    flex: 1;
+    min-width: 110px;
+    background: var(--primary-lt);
+    border: 1px solid var(--border);
+    border-radius: 8px;
+    padding: 10px 14px;
+    text-align: center;
 }
 .metric-badge .val {
-    font-size: 22px !important; font-weight: 700 !important;
-    color: var(--cyan-400) !important; font-family: var(--mono) !important; line-height: 1.2;
+    font-size: 22px;
+    font-weight: 700;
+    color: var(--primary);
 }
-.metric-badge .val.ischemic  { color: var(--red-400) !important; text-shadow: 0 0 12px rgba(248,113,113,0.4); }
-.metric-badge .val.borderline{ color: var(--amber-400) !important; }
-.metric-badge .val.ok        { color: var(--green-400) !important; text-shadow: 0 0 12px rgba(52,211,153,0.4); }
+.metric-badge .val.ischemic { color: var(--danger); }
+.metric-badge .val.borderline { color: var(--warn); }
+.metric-badge .val.ok { color: var(--accent); }
 .metric-badge .lbl {
-    font-size: 9px !important; color: var(--w30) !important;
-    text-transform: uppercase !important; letter-spacing: 1.5px !important;
-    margin-top: 4px !important; font-family: var(--mono) !important;
+    font-size: 11px;
+    color: var(--text-sec);
+    text-transform: uppercase;
+    letter-spacing: 0.4px;
 }
 
-/* ── Status box ───────────────────────────────────────────────────── */
-#status-box {
-    padding: 12px 18px !important; border-radius: var(--radius-s) !important;
-    font-size: 12px !important; font-weight: 500 !important;
-    border: 1px solid var(--border) !important;
-    background: var(--cyan-glow) !important;
-    color: var(--cyan-300) !important;
-    font-family: var(--mono) !important; letter-spacing: 0.5px;
+/* ── Result images ── */
+#result-images img {
+    border-radius: var(--radius);
+    border: 1px solid var(--border);
 }
 
-/* ── Images ───────────────────────────────────────────────────────── */
-.gradio-image img {
-    border-radius: var(--radius-s) !important;
-    border: 1px solid var(--border-s) !important;
-}
-.gradio-image .label-wrap { display: none !important; }
-
-/* ── Code (JSON) ──────────────────────────────────────────────────── */
-.gradio-code, .cm-editor, .cm-scroller {
-    background: var(--navy-950) !important;
-    border: 1px solid var(--border-s) !important;
-    font-family: var(--mono) !important; font-size: 12px !important;
+/* ── Status ── */
+#status-box {
+    padding: 10px 16px;
+    border-radius: 8px;
+    font-size: 13px;
+    font-weight: 500;
+    border: 1px solid var(--border);
+    background: var(--primary-lt);
+    color: var(--primary);
 }
 
-/* ── Footer ───────────────────────────────────────────────────────── */
+/* ── Footer ── */
 #footer {
-    background: var(--navy-950) !important;
-    border-top: 1px solid var(--border-s) !important;
-    padding: 14px 32px !important; text-align: center !important;
-    font-size: 11px !important; color: var(--w30) !important;
-    font-family: var(--mono) !important; letter-spacing: 0.5px;
-}
-#footer b { color: var(--cyan-400) !important; }
-
-/* ── Gradio overrides ─────────────────────────────────────────────── */
-.gradio-container .prose, .gradio-container h1, .gradio-container h2,
-.gradio-container h3, .gradio-container p, .gradio-container label,
-.gradio-container span { color: var(--white) !important; }
-.gradio-container input, .gradio-container textarea, .gradio-container select {
-    background: var(--navy-800) !important;
-    border-color: var(--border-s) !important; color: var(--white) !important;
+    text-align: center;
+    font-size: 12px;
+    color: var(--text-sec);
+    padding: 20px;
+    border-top: 1px solid var(--border);
+    margin-top: 24px;
 }
-::-webkit-scrollbar { width: 5px; height: 5px; }
-::-webkit-scrollbar-track { background: var(--navy-950); }
-::-webkit-scrollbar-thumb { background: var(--navy-600); border-radius: 3px; }
-::-webkit-scrollbar-thumb:hover { background: var(--cyan-400); }
 """
 
-
 # ─────────────────────────────────────────────────────────────────────────────
 #  HEADER HTML
 # ─────────────────────────────────────────────────────────────────────────────
 
 def _header_html() -> str:
-    itu_tag = (f'<img src="{LOGO_ITU}" alt="ITU">'
-               if LOGO_ITU else
-               '<span style="font-size:11px;color:rgba(240,249,255,0.4);font-family:JetBrains Mono,monospace">ITU LAHORE</span>')
-    return f"""
+    return """
 <div id="header">
   <div id="logo-angio">
-    <div class="name">Angio&thinsp;AI</div>
-    <div class="sub">Coronary Angiography Analysis System</div>
+    AngioAI Pro
+    <span>Coronary Angiography Analysis · Research use only</span>
   </div>
-  <div class="header-meta">
-    <b>Muhammad Adil</b> &nbsp;&middot;&nbsp; MS Data Science<br>
-    Supervisor: Dr. Arif Mehmood &nbsp;&middot;&nbsp; ITU Lahore
-  </div>
-  <div id="logo-itu">{itu_tag}</div>
 </div>"""
 
 # ─────────────────────────────────────────────────────────────────────────────
@@ -504,63 +408,24 @@ def extract_best_frame(video_path: str, n: int = 1) -> tuple:
     return rgb, best_idx, best_score, total
 
 # ─────────────────────────────────────────────────────────────────────────────
-#  MODULE 2 — MASK2FORMER INFERENCE
+#  MODULE 2 — STENOSIS DETECTION (YOLOv8-seg ensemble)
 # ─────────────────────────────────────────────────────────────────────────────
 
-_m2f_model  = None
-_m2f_proc   = None
-
-def _load_mask2former():
-    global _m2f_model, _m2f_proc
-    if _m2f_model is not None:
-        return
-    from transformers import (AutoImageProcessor,
-                               Mask2FormerForUniversalSegmentation)
-    ckpt = CONFIG["MASK2FORMER_CKPT"]()   # resolves via HF Hub download
-    if not os.path.exists(ckpt):
-        raise FileNotFoundError(f"Mask2Former checkpoint not found:\n{ckpt}")
-
-    _m2f_proc  = AutoImageProcessor.from_pretrained(
-        "facebook/mask2former-swin-base-coco-instance")
-    _m2f_model = Mask2FormerForUniversalSegmentation.from_pretrained(
-        "facebook/mask2former-swin-base-coco-instance",
-        num_labels=2, ignore_mismatched_sizes=True)
-
-    state = torch.load(ckpt, map_location="cpu", weights_only=False)
-    weights = state.get("model_state", state)
-    _m2f_model.load_state_dict(weights, strict=False)
-    _m2f_model.to(CONFIG["DEVICE"]).eval()
+# Stenosis detection: 3 x YOLOv8-seg ensemble (see stenosis_ensemble.py).
+# Replaces Mask2Former, which scored worst on the held-out ARCADE test set (F1@0.3 0.485 vs 0.637).
+_sten_ens = None
 
 
-def run_mask2former(frame_rgb: np.ndarray, conf_thr: float) -> list:
+def run_stenosis(frame_rgb: np.ndarray, conf_thr: float) -> list:
     """
-    Returns list of dicts: [{bbox:[x1,y1,x2,y2], score:float, mask:np.ndarray}, ...]
+    Returns list of dicts: [{bbox:[x1,y1,x2,y2], score:float, mask:np.ndarray, n_models:int}, ...]
     """
-    _load_mask2former()
-    from PIL import Image
-    pil = Image.fromarray(frame_rgb)
-    inputs = _m2f_proc(images=pil, return_tensors="pt")
-    inputs = {k: v.to(CONFIG["DEVICE"]) for k, v in inputs.items()}
-
-    with torch.no_grad():
-        outputs = _m2f_model(**inputs)
+    global _sten_ens
+    if _sten_ens is None:
+        from stenosis_ensemble import StenosisEnsemble
+        _sten_ens = StenosisEnsemble()
+    return _sten_ens.predict(frame_rgb, score_thr=conf_thr)
 
-    results = _m2f_proc.post_process_instance_segmentation(
-        outputs, target_sizes=[(512, 512)])[0]
-
-    detections = []
-    for seg_info in results["segments_info"]:
-        score = float(seg_info["score"])
-        if score < conf_thr:
-            continue
-        mask  = (results["segmentation"] == seg_info["id"]).cpu().numpy().astype(np.uint8)
-        ys, xs = np.where(mask)
-        if len(ys) == 0:
-            continue
-        x1, y1, x2, y2 = int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())
-        detections.append({"bbox": [x1, y1, x2, y2], "score": score, "mask": mask})
-
-    return detections
 
 # ─────────────────────────────────────────────────────────────────────────────
 #  MODULE 3 — RESUNET BINARY SEGMENTATION
@@ -815,34 +680,20 @@ def run_resunet(frame_rgb: np.ndarray):
 # ─────────────────────────────────────────────────────────────────────────────
 
 # ARCADE 26-class name list (index = YOLO cls_id + 1)
-_YOLO_CLASS_NAMES = [
-    "background",  # 0 (background, not predicted by YOLO)
-    "LMCA",        # 1
-    "LAD prox",    # 2
-    "LAD mid",     # 3
-    "LAD dist",    # 4
-    "D1",          # 5
-    "D2",          # 6
-    "D3",          # 7
-    "Septal",      # 8
-    "LAD var",     # 9
-    "LCX prox",    # 10
-    "LCX mid",     # 11
-    "LCX dist",    # 12
-    "OM1",         # 13
-    "OM2",         # 14
-    "OM3",         # 15
-    "LCX var",     # 16
-    "RCA prox",    # 17
-    "RCA mid",     # 18
-    "RCA dist",    # 19
-    "PDA",         # 20
-    "PLV",         # 21
-    "AM",          # 22
-    "Conus",       # 23
-    "SAN",         # 24
-    "Other",       # 25
-]
+# YOLO cls_id + 1 == ARCADE category id == SYNTAX segment (verified from best.pt names:
+# '1'..'9','9a','10','10a','11','12','12a','13','14','14a','15','16','16a','16b','16c','12b','14b','stenosis').
+# cls_id 25 ('stenosis') is not an anatomical segment and is ignored in the seg map.
+SYNTAX_SEGMENTS = {
+    1:"1 RCA prox",   2:"2 RCA mid",    3:"3 RCA dist",   4:"4 R-PDA",
+    5:"5 Left main",  6:"6 LAD prox",   7:"7 LAD mid",    8:"8 LAD apical",
+    9:"9 D1",         10:"9a D1a",      11:"10 D2",       12:"10a D2a",
+    13:"11 LCx prox", 14:"12 Intermediate", 15:"12a OM1",  16:"13 LCx dist",
+    17:"14 L-PL",     18:"14a L-PL a",  19:"15 L-PDA",    20:"16 R-PL",
+    21:"16a R-PL a",  22:"16b R-PL b",  23:"16c R-PL c",  24:"12b OM2",
+    25:"14b L-PL b",
+}
+_YOLO_CLASS_NAMES = ["background"] + [SYNTAX_SEGMENTS[i] for i in range(1, 26)]
+_YOLO_STENOSIS_CLS = 25
 
 _yolo_model = None   # singleton -- loaded once on first call, reused every frame
 
@@ -930,15 +781,17 @@ def run_yolo_seg(frame_rgb: np.ndarray, seg_conf: float) -> np.ndarray:
     n_written = 0
     for i in order:
         cls_id = int(cls_ids[i])
+        if cls_id >= _YOLO_STENOSIS_CLS:      # 'stenosis' class -> not a segment
+            continue
         mask_resized = cv2.resize(
             masks[i], (w, h), interpolation=cv2.INTER_LINEAR
         )
         binary    = mask_resized > 0.5
-        label_val = min(cls_id + 1, 25)   # clamp to valid 1-25 range
+        label_val = cls_id + 1
         seg_map[binary] = label_val
         n_written += int(binary.sum())
 
-    detected_cls = [_YOLO_CLASS_NAMES[min(int(c) + 1, 25)] for c in cls_ids]
+    detected_cls = [_YOLO_CLASS_NAMES[int(c) + 1] for c in cls_ids if int(c) < _YOLO_STENOSIS_CLS]
     print(f"[SEG] Detections: {len(cls_ids)}  |  classes: {detected_cls}  |  "
           f"vessel px: {n_written} ({100 * n_written / (h * w):.1f}%)")
     print(f"[SEG] seg_map unique labels: {np.unique(seg_map)}")
@@ -975,6 +828,7 @@ FFR_CFG = {
     # If auto-computed ref exceeds this, clamp it
     "ref_diam_max_mm"    : 5.0,
     "ref_diam_min_mm"    : 1.0,
+    "max_offline_px"     : 12,    # lesion farther than this from the measured centreline is not quantified
     "px_per_mm"          : 3.75,
 }
 
@@ -1157,6 +1011,7 @@ def run_ffr(frame_rgb, binary_mask, detections):
         centers_arr = np.array(centers)
         print(f"[FFR] After noise filter: {len(diams_sm)} measurements kept")
 
+    centers_arr = np.array(centers)
     stenoses=[]; unloc=[]
 
     for det in detections:
@@ -1169,17 +1024,25 @@ def run_ffr(frame_rgb, binary_mask, detections):
         x1,y1,x2,y2=bbox
         inside=np.where((centers_arr[:,0]>=y1)&(centers_arr[:,0]<=y2)&
                         (centers_arr[:,1]>=x1)&(centers_arr[:,1]<=x2))[0]
-        idx = int(inside[np.argmin(diams_sm[inside])]) if len(inside)>0 else \
-              int(np.argmin(np.linalg.norm(centers_arr-np.array([(y1+y2)/2,(x1+x2)/2]),axis=1)))
+        if len(inside)>0:
+            idx = int(inside[np.argmin(diams_sm[inside])])
+        else:
+            dist = np.linalg.norm(centers_arr-np.array([(y1+y2)/2,(x1+x2)/2]),axis=1)
+            idx = int(np.argmin(dist))
+            # lesion is not on the measured (dominant) branch -> do not measure it there
+            if dist[idx] > FFR_CFG["max_offline_px"]:
+                unloc.append({**det,"reason":"not on measured branch"}); continue
         if idx>=len(diams_sm): continue
         d_sten=float(diams_sm[idx]); ds=max(0.,1.-d_sten/ref_px)
         if ds < FFR_CFG["min_ds_fraction"]:
-            unloc.append({**det,"reason":"no-narrowing"}); continue
+            unloc.append({**det,"reason":"no measurable narrowing"}); continue
         ffr_est=_ds_to_ffr(ds)
         stenoses.append({"index":idx,"position_mm":float(arc_pos[idx]/px_mm),
                          "d_mm":float(d_sten/px_mm),"ref_mm":float(ref_mm),
                          "pct_ds":float(ds*100),"ffr":ffr_est,
-                         "score":score,"significant":ffr_est<=0.80})
+                         "score":score,"significant":ffr_est<=0.80,
+                         "segment":det.get("segment","—"),"segment_id":det.get("segment_id",0),
+                         "bbox":[int(v) for v in bbox],"n_models":det.get("n_models",0)})
 
     global_ffr=1.
     for s in stenoses: global_ffr*=s["ffr"]
@@ -1195,7 +1058,7 @@ def run_ffr(frame_rgb, binary_mask, detections):
         b=det.get("bbox",[]); sc=det.get("score",0)
         if len(b)<4: continue
         x1,y1,x2,y2=b
-        col="red" if sc>=0.5 else "orange"
+        col="red" if sc>=0.5 else "orange"   # fused ensemble score
         axes[0].add_patch(mpatches.Rectangle(
             (x1,y1),x2-x1,y2-y1,fill=False,edgecolor=col,linewidth=2))
         axes[0].text(x1,y1-5,f"{sc:.2f}",color=col,fontsize=8,fontweight="bold")
@@ -1240,9 +1103,14 @@ def run_ffr(frame_rgb, binary_mask, detections):
         axes[2].legend(fontsize=8); axes[2].grid(alpha=0.3); axes[2].set_ylim(0)
     axes[2].set_title("Lumen diameter profile", fontsize=11, color="#1a6fa8")
 
-    ffr_label = ("⚠ ISCHEMIC" if global_ffr<=0.80 else "✓ Non-ischemic")
-    fig.suptitle(f"FFR Analysis — Global FFR: {global_ffr:.3f}  {ffr_label}  |  "
-                 f"Stenoses: {len(stenoses)}  |  Ref: {ref_mm:.2f} mm",
+    if not detections:
+        ffr_title = "No stenosis detected"
+    elif not stenoses:
+        ffr_title = f"{len(detections)} detected, none quantifiable on the measured branch"
+    else:
+        ffr_label = ("⚠ ISCHEMIC" if global_ffr<=0.80 else "✓ Non-ischemic")
+        ffr_title = f"Branch FFR: {global_ffr:.3f}  {ffr_label}  |  Quantified: {len(stenoses)}/{len(detections)}"
+    fig.suptitle(f"FFR Analysis — {ffr_title}  |  Ref: {ref_mm:.2f} mm",
                  fontsize=12,fontweight="bold",color="#1a2533",y=1.01)
     plt.tight_layout()
 
@@ -1267,6 +1135,10 @@ def run_ffr(frame_rgb, binary_mask, detections):
         "ref_diam_mm"  : round(ref_mm,2),
         "vessel_len_mm": round(float(arc_pos[-1]/px_mm),1) if len(arc_pos)>0 else 0,
         "n_unlocalized": len(unloc),
+        "n_detected"   : len(detections),
+        "unlocalized"  : [{"segment":u.get("segment","—"),"score":round(float(u.get("score",0)),3),
+                          "bbox":[int(v) for v in u.get("bbox",[])],"reason":u.get("reason","")} for u in unloc],
+        "n_significant": sum(1 for s_ in stenoses if s_["pct_ds"]>=50),
     }
     return result, buf
 
@@ -1307,16 +1179,7 @@ ARCADE_COLOURS = {
     25: (180, 180, 180),   # Other           — silver grey
 }
 
-ARCADE_LABELS = {
-    1:"LMCA",     2:"LAD prox",  3:"LAD mid",   4:"LAD dist",
-    5:"D1",       6:"D2",        7:"D3",         8:"Septal",
-    9:"LAD var",  10:"LCX prox", 11:"LCX mid",  12:"LCX dist",
-    13:"OM1",     14:"OM2",      15:"OM3",       16:"LCX var",
-    17:"RCA prox",18:"RCA mid",  19:"RCA dist",  20:"PDA",
-    21:"PLV",     22:"AM",       23:"Conus",     24:"SAN",
-    25:"Other",
-}
-
+ARCADE_LABELS = dict(SYNTAX_SEGMENTS)
 
 def render_nnunet_overlay(frame_rgb, seg_map, alpha=0.55):
     """
@@ -1383,36 +1246,51 @@ def render_nnunet_overlay(frame_rgb, seg_map, alpha=0.55):
 #  MODULE 7 — SYNTAX SCORE  (simplified anatomical estimate)
 # ─────────────────────────────────────────────────────────────────────────────
 
-def compute_syntax_score(stenoses, seg_map):
+# SYNTAX segment weighting factors (Sianos et al., EuroIntervention 2005, Table 1).
+# Keys = ARCADE label ids (see SYNTAX_SEGMENTS). Verify against syntaxscore.org before clinical use.
+_SYNTAX_W_RIGHT = {1:1, 2:1, 3:1, 4:1, 5:5, 6:3.5, 7:2.5, 8:1, 9:1, 10:1, 11:0.5, 12:0.5,
+                   13:1.5, 14:1, 15:1, 16:0.5, 17:0.5, 18:0.5, 19:0, 20:0.5, 21:0.5, 22:0.5, 23:0.5,
+                   24:1, 25:0.5}
+_SYNTAX_W_LEFT  = {1:0, 2:0, 3:0, 4:0, 5:6, 6:3.5, 7:2.5, 8:1, 9:1, 10:1, 11:0.5, 12:0.5,
+                   13:2.5, 14:1, 15:1, 16:1.5, 17:1, 18:1, 19:1, 20:0, 21:0, 22:0, 23:0,
+                   24:1, 25:1}
+
+
+def segment_for_detection(det, seg_map, dilate_px=7):
+    """Majority anatomical label under the (slightly dilated) stenosis mask."""
+    m = det.get("mask")
+    if m is None:
+        x1,y1,x2,y2 = [int(v) for v in det["bbox"]]
+        m = np.zeros_like(seg_map, dtype=np.uint8); m[y1:y2+1, x1:x2+1] = 1
+    m = cv2.dilate(m.astype(np.uint8), np.ones((dilate_px, dilate_px), np.uint8)) > 0
+    labs = seg_map[m]; labs = labs[labs > 0]
+    if labs.size == 0:
+        return 0, "unassigned"
+    sid = int(np.bincount(labs).argmax())
+    return sid, SYNTAX_SEGMENTS.get(sid, "unassigned")
+
+
+def compute_syntax_score(stenoses, dominance="right"):
     """
-    Simplified SYNTAX estimate:
-      Each stenosis contributes a weight based on DS% and
-      anatomical territory (estimated from nnUNet class ID).
-    Reference: Sianos et al. EuroIntervention 2005.
-    Note: true SYNTAX requires angiographic expertise; this is
-    an automated surrogate for display purposes.
+    Anatomical SYNTAX estimate: sum over lesions with DS >= 50% of
+    segment weight x 2 (non-occlusive) or x 5 (occlusion, DS >= 99%).
+    NOT scored yet: adverse lesion characteristics (bifurcation, ostial, length > 20 mm,
+    tortuosity, calcification, thrombus, diffuse disease) -> true SYNTAX is >= this value.
+    Lesions without an assigned segment are listed but not scored.
     """
-    if not stenoses:
-        return 0.0, "Low (0)"
-
-    base_score = 0.
+    W = _SYNTAX_W_LEFT if dominance == "left" else _SYNTAX_W_RIGHT
+    score, scored, unassigned = 0.0, 0, 0
     for st in stenoses:
-        ds = st["pct_ds"] / 100
-        # Weight by position: proximal vessels score higher
-        pos_w = max(0.5, 1.5 - st["position_mm"] / 100)
-        # Occlusion multiplier
-        occ_w = 5.0 if ds >= 0.99 else (2.5 if ds >= 0.70 else 1.5 if ds >= 0.50 else 1.0)
-        base_score += pos_w * occ_w * (1 + ds)
-
-    # Add segment count from nnUNet
-    n_segs = len(np.unique(seg_map)) - 1  # exclude background
-    base_score += n_segs * 0.3
-
-    score = round(base_score, 1)
-    if score < 23:  tier = "Low (<23)"
-    elif score < 33: tier = "Intermediate (23–32)"
-    else:            tier = "High (≥33)"
-    return score, tier
+        if st["pct_ds"] < 50:
+            st["syntax_points"] = 0.0; continue
+        sid = st.get("segment_id", 0)
+        if not sid:
+            st["syntax_points"] = None; unassigned += 1; continue
+        pts = W.get(sid, 0) * (5 if st["pct_ds"] >= 99 else 2)
+        st["syntax_points"] = pts; score += pts; scored += 1
+    score = round(score, 1)
+    tier = "Low (≤22)" if score <= 22 else ("Intermediate (23–32)" if score <= 32 else "High (≥33)")
+    return score, tier, scored, unassigned
 
 # ─────────────────────────────────────────────────────────────────────────────
 #  JSON ENCODER — handles numpy float32/int64 etc.
@@ -1435,7 +1313,7 @@ def _to_json(obj) -> str:
 #  MAIN PIPELINE FUNCTION
 # ─────────────────────────────────────────────────────────────────────────────
 
-def analyse(video_path, sten_conf, seg_conf, px_per_mm_override, progress=gr.Progress()):
+def analyse(video_path, sten_conf, seg_conf, px_per_mm_override, dominance="right", progress=gr.Progress()):
     """
     Full pipeline. Called by Gradio on button click.
     Returns: (status_html, frame_img, seg_img, ffr_img, metrics_html, json_str)
@@ -1451,17 +1329,8 @@ def analyse(video_path, sten_conf, seg_conf, px_per_mm_override, progress=gr.Pro
         frame_rgb, frame_idx, frame_score, total_frames = extract_best_frame(video_path)
 
         # Step 2 — stenosis detection
-        progress(0.30, desc="Running Mask2Former…")
-        detections = run_mask2former(frame_rgb, sten_conf)
-
-        # Compute vessel overlap for each detection
-        binary_mask_for_overlap = np.ones((512,512),dtype=np.uint8)  # placeholder until ResUNet
-        for det in detections:
-            b=det.get("bbox",[]); 
-            if len(b)<4: det["overlap"]=0.; continue
-            x1,y1,x2,y2=[int(v) for v in b]
-            area=(x2-x1)*(y2-y1)
-            det["overlap"]=min(float(binary_mask_for_overlap[y1:y2,x1:x2].sum())/max(area,1),1.)
+        progress(0.30, desc="Detecting stenoses (YOLOv8-seg ensemble)…")
+        detections = run_stenosis(frame_rgb, sten_conf)
 
         # Step 3a — binary segmentation (ResUNet)
         progress(0.50, desc="Running ResUNet…")
@@ -1469,7 +1338,7 @@ def analyse(video_path, sten_conf, seg_conf, px_per_mm_override, progress=gr.Pro
         print(f"[ResUNet] binary_mask sum={binary_mask.sum()}  "
               f"unique={np.unique(binary_mask)}  "
               f"vessel%={100*binary_mask.mean():.1f}")
-        # Recompute overlap with real mask
+        # Vessel overlap of each detection (used by FFR to drop off-vessel detections)
         for det in detections:
             b=det.get("bbox",[])
             if len(b)<4: continue
@@ -1481,26 +1350,37 @@ def analyse(video_path, sten_conf, seg_conf, px_per_mm_override, progress=gr.Pro
         progress(0.65, desc="Running seg…")
         seg_map = run_yolo_seg(frame_rgb, seg_conf)
         seg_overlay = render_nnunet_overlay(frame_rgb, seg_map)
+        for det in detections:
+            det["segment_id"], det["segment"] = segment_for_detection(det, seg_map)
 
         # Step 4 — FFR
         progress(0.80, desc="Computing FFR…")
         ffr_result, ffr_fig = run_ffr(frame_rgb, binary_mask, detections)
 
         if "error" in ffr_result:
-            return (_status("warn", ffr_result["error"]),
+            ffr_result["n_detected"] = len(detections)
+            ffr_result["detections"] = [{"segment":d.get("segment","—"),"score":round(d["score"],3),
+                                          "bbox":d["bbox"]} for d in detections]
+            return (_status("warn", "Vessel centreline could not be traced — FFR/QCA not available. "
+                                    + ffr_result["error"]),
                     frame_rgb, binary_overlay, seg_overlay, None,
                     _metrics_html(ffr_result, {}, frame_idx, total_frames),
                     _to_json(ffr_result))
 
         # Step 5 — SYNTAX score
         progress(0.92, desc="Computing SYNTAX score…")
-        syntax_score, syntax_tier = compute_syntax_score(ffr_result["stenoses"], seg_map)
+        syntax_score, syntax_tier, n_scored, n_unassigned = compute_syntax_score(ffr_result["stenoses"], dominance)
         ffr_result["syntax_score"] = syntax_score
         ffr_result["syntax_tier"]  = syntax_tier
+        ffr_result["syntax_lesions_scored"] = n_scored
+        ffr_result["syntax_lesions_unassigned"] = n_unassigned
+        ffr_result["dominance"] = dominance
 
         progress(1.00, desc="Done.")
+        n_det = len(detections)
         status = _status("ok",
-            f"Analysis complete — Frame {frame_idx}/{total_frames} "
+            ("No stenosis detected — " if n_det == 0 else f"{n_det} stenosis candidate(s) detected — ")
+            + f"Frame {frame_idx}/{total_frames} "
             f"(quality score {frame_score:.1f})")
         metrics = _metrics_html(ffr_result, {"score":frame_score,"idx":frame_idx},
                                 frame_idx, total_frames)
@@ -1533,85 +1413,128 @@ def _status(kind: str, msg: str) -> str:
             f'{ico} {msg}</div>')
 
 
+def _badge(val, lbl, cls="", small=False):
+    st = ' style="font-size:14px"' if small else ""
+    return (f'<div class="metric-badge"><div class="val {cls}"{st}>{val}</div>'
+            f'<div class="lbl">{lbl}</div></div>')
+
+
 def _metrics_html(ffr_result, frame_info, frame_idx, total) -> str:
-    if not ffr_result or "error" in ffr_result:
+    if not ffr_result:
         return '<div class="card"><div class="card-title">Metrics</div><p style="color:#5a7a96">Run analysis to see results.</p></div>'
 
-    gffr  = ffr_result.get("global_ffr", "—")
-    isch  = ffr_result.get("ischemic", None)
-    ns    = ffr_result.get("n_stenoses", 0)
-    ref   = ffr_result.get("ref_diam_mm","—")
-    vlen  = ffr_result.get("vessel_len_mm","—")
-    syn   = ffr_result.get("syntax_score","—")
-    stier = ffr_result.get("syntax_tier","—")
-
-    if isch is True:  ffr_cls,ffr_lbl = "ischemic","⚠ ISCHEMIC"
-    elif isch is False: ffr_cls,ffr_lbl = "ok","✓ NON-ISCHEMIC"
-    else: ffr_cls,ffr_lbl = "","—"
-
-    gffr=0.82
-    ns=1
-    syn = 0.4
-    syn_val = float(syn) if isinstance(syn,float) else 0
-    syn_cls = "ischemic" if syn_val>=33 else ("borderline" if syn_val>=23 else "ok")
-
-    stenosis_rows = ""
-    for i,st in enumerate(ffr_result.get("stenoses",[])):
+    n_det = int(ffr_result.get("n_detected", 0))
+    err   = ffr_result.get("error")
+
+    # ── Case A: no stenosis detected ──────────────────────────────────────────
+    if n_det == 0:
+        extra = ("" if err else
+                 f'<div style="font-size:12px;color:#5a7a96;margin-top:6px">Measured branch: '
+                 f'ref diameter {ffr_result.get("ref_diam_mm","—")} mm, length {ffr_result.get("vessel_len_mm","—")} mm '
+                 f'&nbsp;|&nbsp; Frame {frame_idx}/{total} &nbsp;|&nbsp; Scale {FFR_CFG["px_per_mm"]} px/mm</div>')
+        return f"""
+<div class="card">
+  <div class="card-title">Key metrics</div>
+  <div class="metric-row">
+    {_badge("0", "Stenoses detected", "ok")}
+    {_badge("✓ None", "Stenosis", "ok", small=True)}
+    {_badge("N/A", "FFR (no lesion)")}
+    {_badge("0", "SYNTAX (anatomical)", "ok")}
+  </div>
+  <div style="font-size:13px;color:#0d9e6e;margin-top:6px"><b>No stenosis detected</b> at the current
+  confidence threshold. This is not a clinical exclusion of coronary artery disease.</div>
+  {extra}
+</div>"""
+
+    # ── Case B: detections, but QCA/FFR failed (no centreline) ───────────────
+    if err:
+        rows = "".join(f'<tr><td style="padding:4px 8px">#{i+1}</td><td style="padding:4px 8px">{d["segment"]}</td>'
+                       f'<td style="padding:4px 8px">{d["score"]:.2f}</td></tr>'
+                       for i, d in enumerate(ffr_result.get("detections", [])))
+        return f"""
+<div class="card">
+  <div class="card-title">Key metrics</div>
+  <div class="metric-row">
+    {_badge(n_det, "Stenoses detected", "borderline")}
+    {_badge("N/A", "FFR / %DS", "", small=True)}
+    {_badge("N/A", "SYNTAX", "", small=True)}
+  </div>
+  <div style="font-size:12px;color:#d98b00;margin-top:4px">Vessel centreline could not be traced, so
+  diameter stenosis, FFR and SYNTAX could not be computed.</div>
+</div>
+<div class="card"><div class="card-title">Detections</div><table style="width:100%;border-collapse:collapse;font-size:13px">
+<thead><tr style="background:#e8f3fb;color:#1a6fa8"><th style="padding:6px 8px;text-align:left">#</th>
+<th style="padding:6px 8px;text-align:left">Segment</th><th style="padding:6px 8px;text-align:left">Confidence</th></tr></thead>
+<tbody>{rows}</tbody></table></div>"""
+
+    # ── Case C: normal result ────────────────────────────────────────────────
+    sts   = ffr_result.get("stenoses", [])
+    n_q   = len(sts)
+    n_sig = int(ffr_result.get("n_significant", 0))
+    gffr  = ffr_result.get("global_ffr")
+    syn   = ffr_result.get("syntax_score", 0.0)
+    stier = ffr_result.get("syntax_tier", "—")
+
+    if n_q == 0:
+        ffr_val, ffr_cls, isch_lbl, isch_cls = "N/A", "", "Not quantified", ""
+    else:
+        ffr_val = f"{gffr:.2f}"
+        isch = ffr_result.get("ischemic")
+        ffr_cls = isch_cls = "ischemic" if isch else "ok"
+        isch_lbl = "⚠ ISCHEMIC" if isch else "✓ NON-ISCHEMIC"
+    syn_cls = "ischemic" if syn >= 33 else ("borderline" if syn >= 23 else "ok")
+    sig_cls = "ischemic" if n_sig else "ok"
+
+    rows = ""
+    for i, st in enumerate(sts):
         sig_col = "#d64040" if st.get("significant") else "#0d9e6e"
-        stenosis_rows += (
-            f'<tr>'
-            f'<td style="padding:4px 8px">#{i+1}</td>'
-            f'<td style="padding:4px 8px">{st["position_mm"]:.1f} mm</td>'
-            f'<td style="padding:4px 8px">{st["pct_ds"]:.1f}%</td>'
-            f'<td style="padding:4px 8px">{st["d_mm"]:.2f} mm</td>'
-            f'<td style="padding:4px 8px;color:{sig_col};font-weight:600">{st["ffr"]:.3f}</td>'
-            f'</tr>'
-        )
+        pts = st.get("syntax_points")
+        pts_s = "—" if pts is None else (f"{pts:g}" if pts else "0")
+        rows += (f'<tr><td style="padding:4px 8px">#{i+1}</td>'
+                 f'<td style="padding:4px 8px">{st.get("segment","—")}</td>'
+                 f'<td style="padding:4px 8px">{st["score"]:.2f}</td>'
+                 f'<td style="padding:4px 8px">{st["pct_ds"]:.0f}%</td>'
+                 f'<td style="padding:4px 8px">{st["d_mm"]:.2f} mm</td>'
+                 f'<td style="padding:4px 8px;color:{sig_col};font-weight:600">{st["ffr"]:.2f}</td>'
+                 f'<td style="padding:4px 8px">{pts_s}</td>'
+                 f'<td style="padding:4px 8px">quantified</td></tr>')
+    for j, u in enumerate(ffr_result.get("unlocalized", [])):
+        rows += (f'<tr style="color:#7a9ab6"><td style="padding:4px 8px">#{n_q+j+1}</td>'
+                 f'<td style="padding:4px 8px">{u["segment"]}</td><td style="padding:4px 8px">{u["score"]:.2f}</td>'
+                 f'<td style="padding:4px 8px">—</td><td style="padding:4px 8px">—</td><td style="padding:4px 8px">—</td>'
+                 f'<td style="padding:4px 8px">—</td><td style="padding:4px 8px">{u["reason"]}</td></tr>')
+    th = "".join(f'<th style="padding:6px 8px;text-align:left">{h}</th>' for h in
+                 ["#", "Segment", "Conf.", "DS%", "Min diam", "FFR", "SYNTAX pts", "Status"])
 
     return f"""
 <div class="card">
   <div class="card-title">Key metrics</div>
   <div class="metric-row">
-    <div class="metric-badge">
-      <div class="val {ffr_cls}">{gffr if gffr!="—" else "—"}</div>
-      <div class="lbl">Global FFR</div>
-    </div>
-    <div class="metric-badge">
-      <div class="val {ffr_cls}" style="font-size:14px">{ffr_lbl}</div>
-      <div class="lbl">Ischaemia</div>
-    </div>
-    <div class="metric-badge">
-      <div class="val">{ns}</div>
-      <div class="lbl">Stenoses</div>
-    </div>
-    <div class="metric-badge">
-      <div class="val {syn_cls}">{syn}</div>
-      <div class="lbl">SYNTAX score</div>
-    </div>
-    <div class="metric-badge">
-      <div class="val">{ref} mm</div>
-      <div class="lbl">Ref diameter</div>
-    </div>
-    <div class="metric-badge">
-      <div class="val">{vlen} mm</div>
-      <div class="lbl">Vessel length</div>
-    </div>
+    {_badge(n_det, "Stenoses detected", "borderline")}
+    {_badge(f"{n_sig}", "Significant (DS ≥ 50%)", sig_cls)}
+    {_badge(ffr_val, "Branch FFR (QFR est.)", ffr_cls)}
+    {_badge(isch_lbl, "Ischaemia", isch_cls, small=True)}
+    {_badge(f"{syn:g}", "SYNTAX (anatomical)", syn_cls)}
+    {_badge(f'{ffr_result.get("ref_diam_mm","—")} mm', "Ref diameter (measured branch)")}
+    {_badge(f'{ffr_result.get("vessel_len_mm","—")} mm', "Measured branch length")}
   </div>
   <div style="font-size:12px;color:#5a7a96;margin-top:4px">
-    SYNTAX tier: <b>{stier}</b> &nbsp;|&nbsp;
-    Frame {frame_idx}/{total} &nbsp;|&nbsp;
-    Scale: {FFR_CFG["px_per_mm"]} px/mm
+    Quantified on measured branch: <b>{n_q}/{n_det}</b> &nbsp;|&nbsp;
+    SYNTAX tier: <b>{stier}</b> ({ffr_result.get("dominance","right")} dominance; lesions ≥50% DS only;
+    adverse lesion features not yet scored{"; " + str(ffr_result.get("syntax_lesions_unassigned")) + " lesion(s) without segment" if ffr_result.get("syntax_lesions_unassigned") else ""}) &nbsp;|&nbsp;
+    Frame {frame_idx}/{total} &nbsp;|&nbsp; Scale: {FFR_CFG["px_per_mm"]} px/mm
   </div>
 </div>
-{'<div class="card"><div class="card-title">Stenosis table</div><table style="width:100%;border-collapse:collapse;font-size:13px"><thead><tr style="background:#e8f3fb;color:#1a6fa8"><th style="padding:6px 8px;text-align:left">#</th><th style="padding:6px 8px;text-align:left">Position</th><th style="padding:6px 8px;text-align:left">DS%</th><th style="padding:6px 8px;text-align:left">Min diam</th><th style="padding:6px 8px;text-align:left">FFR</th></tr></thead><tbody>' + stenosis_rows + '</tbody></table></div>'
- if stenosis_rows else ""}
+<div class="card"><div class="card-title">Stenosis table</div>
+<table style="width:100%;border-collapse:collapse;font-size:13px">
+<thead><tr style="background:#e8f3fb;color:#1a6fa8">{th}</tr></thead><tbody>{rows}</tbody></table></div>
 """
 
 # ─────────────────────────────────────────────────────────────────────────────
 #  GRADIO APP
 # ─────────────────────────────────────────────────────────────────────────────
 
-with gr.Blocks(css=CSS, title="Angio AI") as demo:
+with gr.Blocks(css=CSS, title="AngioAI Pro") as demo:
 
     # Header
     gr.HTML(_header_html())
@@ -1637,8 +1560,8 @@ with gr.Blocks(css=CSS, title="Angio AI") as demo:
 
                     gr.HTML('<div class="card"><div class="card-title">Model controls</div>')
                     sten_conf = gr.Slider(
-                        minimum=0.05, maximum=0.95, value=0.30, step=0.05,
-                        label="Stenosis confidence threshold  (Mask2Former)",
+                        minimum=0.05, maximum=0.95, value=0.20, step=0.05,
+                        label="Stenosis confidence threshold  (YOLOv8-seg ensemble)",
                         info="Detections below this confidence are discarded",
                     )
                     seg_conf = gr.Slider(
@@ -1651,50 +1574,41 @@ with gr.Blocks(css=CSS, title="Angio AI") as demo:
                         label="Scale  (px / mm)",
                         info="ARCADE default = 3.75 px/mm (0.267 mm/px). Adjust if using non-ARCADE data.",
                     )
+                    dominance = gr.Radio(
+                        choices=["right", "left"], value="right",
+                        label="Coronary dominance (for SYNTAX weighting)",
+                    )
                     gr.HTML('</div>')
 
                 # Right — info card
                 with gr.Column(scale=1):
                     gr.HTML("""
 <div class="card">
-  <div class="card-title">Analysis Pipeline</div>
-  <div style="display:flex;flex-direction:column;gap:0">
-    <div style="display:flex;align-items:center;gap:10px;padding:8px 0;border-bottom:1px solid rgba(240,249,255,0.07)">
-      <span style="font-family:JetBrains Mono,monospace;font-size:10px;color:#22d3ee;min-width:18px">01</span>
-      <span style="font-size:12px;color:rgba(240,249,255,0.7)">Keyframe extraction</span>
-    </div>
-    <div style="display:flex;align-items:center;gap:10px;padding:8px 0;border-bottom:1px solid rgba(240,249,255,0.07)">
-      <span style="font-family:JetBrains Mono,monospace;font-size:10px;color:#22d3ee;min-width:18px">02</span>
-      <span style="font-size:12px;color:rgba(240,249,255,0.7)">Stenosis detection</span>
-    </div>
-    <div style="display:flex;align-items:center;gap:10px;padding:8px 0;border-bottom:1px solid rgba(240,249,255,0.07)">
-      <span style="font-family:JetBrains Mono,monospace;font-size:10px;color:#22d3ee;min-width:18px">03</span>
-      <span style="font-size:12px;color:rgba(240,249,255,0.7)">26-class coronary segmentation</span>
-    </div>
-    <div style="display:flex;align-items:center;gap:10px;padding:8px 0;border-bottom:1px solid rgba(240,249,255,0.07)">
-      <span style="font-family:JetBrains Mono,monospace;font-size:10px;color:#22d3ee;min-width:18px">04</span>
-      <span style="font-size:12px;color:rgba(240,249,255,0.7)">Binary vessel segmentation</span>
-    </div>
-    <div style="display:flex;align-items:center;gap:10px;padding:8px 0;border-bottom:1px solid rgba(240,249,255,0.07)">
-      <span style="font-family:JetBrains Mono,monospace;font-size:10px;color:#22d3ee;min-width:18px">05</span>
-      <span style="font-size:12px;color:rgba(240,249,255,0.7)">FFR estimation — QFR v4</span>
-    </div>
-    <div style="display:flex;align-items:center;gap:10px;padding:8px 0">
-      <span style="font-family:JetBrains Mono,monospace;font-size:10px;color:#22d3ee;min-width:18px">06</span>
-      <span style="font-size:12px;color:rgba(240,249,255,0.7)">SYNTAX score computation</span>
-    </div>
-  </div>
-  <div style="margin-top:14px;padding-top:12px;border-top:1px solid rgba(240,249,255,0.07);font-size:11px;color:rgba(240,249,255,0.35);font-family:JetBrains Mono,monospace;line-height:1.8">
-    FFR &le; 0.80 &rarr; ischemic &nbsp;|&nbsp; 1&minus;(0.33&middot;DS+0.60&middot;DS&sup2;)<br>
-    Tu et al. JACC 2016 &middot; FAVOR II
+  <div class="card-title">Pipeline overview</div>
+  <ol style="font-size:13px;color:#1a2533;line-height:1.9;padding-left:16px">
+    <li>Best frame extracted from video</li>
+    <li>Stenosis detection — YOLOv8-seg ensemble (3 models)</li>
+    <li>Coronary segmentation — 26-class anatomy labelling</li>
+    <li>Binary vessel mask — ResUNet</li>
+    <li>FFR estimation — QFR v4</li>
+    <li>SYNTAX score computation</li>
+  </ol>
+  <hr style="border-color:#cce0f0;margin:12px 0">
+  <div style="font-size:12px;color:#5a7a96">
+    <b>FFR threshold:</b> ≤ 0.80 → ischemic<br>
+    <b>Formula:</b> 1 − (0.33·DS + 0.60·DS²)<br>
+    <b>Reference:</b> Tu et al. JACC 2016
   </div>
 </div>
 <div class="card">
-  <div class="card-title">Checkpoints</div>
-  <div style="font-family:JetBrains Mono,monospace;font-size:11px;color:rgba(240,249,255,0.45);line-height:2.2">
-    <span style="color:#22d3ee">&#9658;</span> mask2former_best.pth<br>
-    <span style="color:#22d3ee">&#9658;</span> binary_best.pth<br>
-    <span style="color:#22d3ee">&#9658;</span> best.pt
+  <div class="card-title">Model checkpoints</div>
+  <div style="font-size:12px;color:#5a7a96;line-height:1.8">
+    Stenosis: <code>weights/stenosis/</code><br>
+    <code>yolov8x_seg.pt · yolov8l_sgd.pt · yolov8l_adamw.pt</code><br><br>
+    Others auto-downloaded from<br>
+    <code>MuhammadAdil63/angio-ai-checkpoints</code><br>
+    <code>binary_best.pth</code><br>
+    <code>best.pt</code>
   </div>
 </div>
 """)
@@ -1740,15 +1654,14 @@ with gr.Blocks(css=CSS, title="Angio AI") as demo:
     # Footer
     gr.HTML("""
 <div id="footer">
-  Angio AI · MS Data Science · Information Technology University Lahore ·
-  Supervisor: Dr. Arif Mehmood &nbsp;|&nbsp;
-  Image-based QFR — Tu et al. JACC 2016 · FAVOR II trial
+  AngioAI Pro · Research use only — not a medical device, not for clinical diagnosis &nbsp;|&nbsp;
+  Image-based QFR — Tu et al. JACC 2016 · FAVOR II trial &nbsp;|&nbsp; SYNTAX weights — Sianos et al. 2005
 </div>""")
 
     # ── WIRING ────────────────────────────────────────────────────────────
 
-    def run_and_switch(video, sc, sg, px):
-        return analyse(video, sc, sg, px)
+    def run_and_switch(video, sc, sg, px, dom):
+        return analyse(video, sc, sg, px, dom)
 
     def load_demo1():
         return _get_demo_video(DEMO_VIDEO_1_NAME)
@@ -1757,7 +1670,7 @@ with gr.Blocks(css=CSS, title="Angio AI") as demo:
 
     btn_analyse.click(
         fn=run_and_switch,
-        inputs=[video_input, sten_conf, seg_conf, px_per_mm],
+        inputs=[video_input, sten_conf, seg_conf, px_per_mm, dominance],
         outputs=[status_out, frame_out, binary_out, seg_out, ffr_out, metrics_out, json_out],
     ).then(
         # Copy status HTML to results page status box
@@ -1793,5 +1706,5 @@ if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=True,
+        share=os.environ.get("SPACE_ID") is None,   # public share link only when run locally
     )

requirements.txt

@@ -1,10 +1,7 @@
 --extra-index-url https://download.pytorch.org/whl/cpu
 torch==2.7.0+cpu
 torchvision==0.22.0+cpu
-tokenizers==0.21.1
-transformers==4.49.0
-accelerate>=0.34.2
-ultralytics==8.2.18
+ultralytics==8.4.31
 huggingface_hub>=0.26.0
 opencv-python-headless==4.9.0.80
 scikit-image==0.25.2

stenosis_ensemble.py

@@ -0,0 +1,115 @@
+"""
+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]