AImageLab-Zip/US_Cond-UNet

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:

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