---
license: mit
tags:
  - medical-imaging
  - segmentation
  - coronary-artery
  - 3d
  - ct
  - cta
  - unet
  - mamba
  - pytorch
  - monai
library_name: pytorch
pipeline_tag: image-segmentation
---

# Coronary Artery Segmentation (3D CTA) — U-Net vs. Attention–Mamba2

Two 3D segmentation models for **voxel-wise binary segmentation of the coronary artery tree** from cardiac CT angiography (CTA) volumes. Released as a portfolio study comparing a convolutional baseline against a state-space/attention hybrid encoder under an identical training and evaluation harness.

- **Repository:** https://github.com/noahschuetz/coronary-artery-segmentation
- **Task:** binary 3D semantic segmentation (foreground = coronary artery)
- **Framework:** PyTorch + MONAI
- **License:** MIT

## Checkpoints

| File | Model | Input channels | Params |
|---|---|---|---|
| `baseline_unet.pth` | 3D U-Net (MONAI) | 1 (CT) | ~31M |
| `att_mamba2_unet.pth` | Attention–Mamba2 U-Net | 2 (CT + Frangi vesselness) | ~31M |

`.pth` files are PyTorch state dicts. Instantiate the architecture from the repo, then load weights.

## Results

Identical training: AdamW (lr 1e-4, weight decay 1e-4, 5-epoch warmup), 100 epochs, 96³ patches, 4 patches/volume. Metrics on a held-out test split. clDice is topology-aware (penalizes centerline breaks).

| Model | Dice | clDice | Precision | Recall | IoU | Time/Epoch |
|---|---|---|---|---|---|---|
| 3D U-Net | 0.788 | 0.864 | 0.820 | 0.762 | 0.653 | ~1.7 min |
| Att-Mamba2 U-Net | 0.791 | 0.865 | 0.812 | 0.775 | 0.657 | ~1 min |

**Finding:** the two encoders are at parity on segmentation quality (~0.003 Dice apart). The state-space hybrid reaches that quality at roughly half the per-epoch training time. This is an efficiency/negative result, not a SOTA claim.

## Usage

```python
import torch
from huggingface_hub import hf_hub_download
from src.models.model_factory import get_model  # from the project repo

model_cfg = {"name": "att_mamba2_unet", "features": [32, 64, 128, 256], "in_channels": 2}
data_cfg = {"patch_size": [96, 96, 96]}
model = get_model(model_cfg, data_cfg)

ckpt = hf_hub_download("noahschuetz/coronary-segmentation", "att_mamba2_unet.pth")
state = torch.load(ckpt, map_location="cpu", weights_only=True)  # .pth is a plain state dict
model.load_state_dict(state)
model.eval()
```

Preprocessing must match training: reorient → resample → HU windowing → intensity normalization → (for the Att-Mamba2 model) append a Frangi vesselness channel. See `src/data/transforms.py`. Inference uses sliding-window patching over the full volume; see `scripts/inference.py`.

## Training data

Developed on the public **ImageCAS** coronary CTA dataset (cardiac CT angiography with expert coronary artery annotations). The dataset is not redistributed with these weights; obtain it from its original source.

## Intended use & limitations

- **Intended use:** research and educational demonstration of 3D vessel segmentation and architecture comparison.
- **Not for clinical use.** These models are not a medical device and have not been validated for diagnosis or treatment.
- **Domain shift:** trained on a single dataset's acquisition characteristics; performance will degrade on scanners, contrast protocols, or populations that differ from ImageCAS.
- **Thin-vessel failure modes:** distal/small branches and severe stenoses are the hardest cases; clDice is reported precisely because voxel Dice under-weights these topology errors.

## Citation

If you use these weights, please cite this repository and the ImageCAS dataset.