README.md

---
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}
}
```