stenoz-coronary-stenosis-unet

Stenoz — coronary stenosis detection U-Nets (X-ray angiography)

Two small PyTorch U-Nets (from-scratch architecture, no pretrained backbone) trained on the ARCADE X-ray coronary angiography dataset (MICCAI 2023) — one segments the vessel tree (test Dice 0.702), the other detects stenosis lesions directly (test F1 0.645, lesion recall 95.7%, image-level sensitivity 95.7%). Full code, training kernels, a Streamlit app, and the paper are on GitHub: https://github.com/uzbtrust/stenoz

Models

Both share the same encoder–decoder U-Net architecture (4 down/up-sampling stages, double 3×3 conv + BatchNorm + ReLU, skip connections), single-channel 512×512 grayscale input. The only difference is the training target: one is supervised on vessel masks, the other on stenosis-lesion masks. Architecture code: unet.py (in this repo).

Training: Kaggle Notebooks, 1× NVIDIA Tesla T4, 60 epochs, AdamW with a cosine LR schedule, mixed precision (fp16). Full details in the paper on GitHub (report/Stenoz_Ilmiy_Maqola.pdf), sections 4–5.

Dataset

ARCADE (MICCAI 2023) — X-ray coronary angiography (XCA), COCO-format polygon annotations. unet_vessel.pt was trained on the syntax split (vessel segmentation), unet_stenosis.pt on the stenosis split (lesion localization).

Usage

import torch
from huggingface_hub import hf_hub_download

# unet.py must be downloaded from this repo (or from
# https://github.com/uzbtrust/stenoz/blob/main/src/unet.py) and importable
from unet import UNet

weights = hf_hub_download("uzbtrust/stenoz-coronary-stenosis-unet", "unet_stenosis.pt")
ckpt = torch.load(weights, map_location="cpu")
model = UNet(base=ckpt.get("base", 32))
model.load_state_dict(ckpt["state"])
model.eval()

# img: float32 numpy array normalized to [0,1], H and W padded to a multiple of 16
with torch.no_grad():
    prob = torch.sigmoid(model(img_tensor))  # [1,1,H,W] — per-pixel stenosis probability

Full inference code (padding, connected-component detection, thresholding): src/dlstenosis.py and src/dlseg.py.

Limitations

• Trained on X-ray coronary angiography (XCA) only — not validated on MRI/MRA.
• At the default operating point, roughly half of the lesions detected by unet_stenosis.pt don't overlap a ground-truth annotation (precision ≈ 48.6%); recall (95.7%) was deliberately prioritized via pos_weight=10.
• Built for research/demonstration purposes. Not validated, certified, or intended for clinical diagnosis.

License

MIT (code and weights). The ARCADE dataset is distributed under its own original license.