Cond-UNet Attention for Ultrasound Segmentation
Cond-UNet Attention is a binary ultrasound segmentation model based on an attention-conditioned U-Net. It was trained to predict a foreground mask from an RGB ultrasound image.
Model Details
Attention-conditioned U-Net for binary ultrasound segmentation: depth 5, base
width 16, 512 x 512 input, 8px patches, and 768-dimensional attention
embeddings. It has one foreground logit per pixel. Organ conditioning is
optional; omitted IDs use the unknown token (-1). DWT and shape conditioning
are disabled.
Usage
This repository contains custom Transformers code. Pass trust_remote_code=True
when loading it.
from transformers import pipeline
segmenter = pipeline(
"image-segmentation",
model="AImageLab-Zip/US_Cond-UNet",
trust_remote_code=True,
)
result = segmenter("ultrasound.png")
mask = result["mask"]
When organ metadata is known, pass its integer class ID:
result = segmenter("ultrasound.png", organ_id=3)
If organ_id is not provided, the model automatically uses -1, matching the
unknown-organ conditioning used in training.
Use the following IDs when organ metadata is available:
| Organ | organ_id |
|---|---|
| Appendix | 0 |
| Breast | 1 |
| Cardiac | 2 |
| Thyroid | 3 |
| Fetal / Fetal HC | 4 |
| Kidney | 5 |
| Liver | 6 |
| Testicle | 7 |
| Unknown | -1 |
Results and Citation
The model results are reported in the BMVC 2026 paper.
If you use this model, please cite:
@inproceedings{morelli2026new,
title={A New Multicenter Testicular US Dataset and a Lightweight Cond-UNet for Generalization in US Segmentation},
author={Morelli, Nicola and Marchesini, Kevin and Santi, Daniele and Grana, Costantino and Bolelli, Federico and others},
booktitle={Proceedings of the British Machine Vision Conference},
year={2026}
}