File size: 2,115 Bytes
8966f37 f277c59 8966f37 98c698a 8966f37 59ce18d 517df9a | 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 | ---
pipeline_tag: image-segmentation
library_name: transformers
tags:
- ultrasound
- medical-image-segmentation
- attention-unet
- custom-pipeline
---
# 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.
```python
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:
```python
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](https://federicobolelli.it/media/publications/pdfs/0475.pdf).
If you use this model, please cite:
```bibtex
@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}
}
```
|