thewtex's picture
Add ultrasound foundation-model TorchScript exports (USF-MAE, SAMUS, URFM, UltraFedFM)
5ea26b1 verified
|
Raw
History Blame Contribute Delete
11.3 kB
---
license: other
license_name: per-model
license_link: LICENSE.md
tags:
- ultrasound
- itkimpact
- torchscript
- feature-extractor
- image-registration
library_name: itk-impact
---
# IMPACT TorchScript feature extractors for ultrasound
TorchScript **feature extractors** for 2D B-mode ultrasound, packaged for
[ITKIMPACT](https://github.com/InsightSoftwareConsortium/ITKIMPACT) and
[IMPACT](https://github.com/vboussot/ImpactLoss) — the registration metric that
compares learned features instead of intensities. The layout mirrors
[`VBoussot/impact-torchscript-models`](https://huggingface.co/VBoussot/impact-torchscript-models),
the upstream repository of 30 general-purpose models: one `<Family>/<Name>.pt`
per model plus a [`models.json`](models.json) describing it, so the
same tooling reads both repositories.
Each `.pt` wraps a published ultrasound foundation model's **encoder** in the
IMPACT feature-extractor contract:
```python
forward(
x: torch.Tensor, # (1, C, H, W)
nb_layers: torch.Tensor, # how many feature maps to return
stats: torch.Tensor, # (min, max, mean, std) from IMPACT, or empty
direction: torch.Tensor, # unused by 2D models
) -> list[torch.Tensor] # exactly nb_layers maps, deepest last
```
IMPACT applies no intensity preprocessing, so each wrapper carries the model's
own normalization: a raw B-mode frame goes in unchanged, single- or
three-channel.
> [!IMPORTANT]
> These are **derivative works** — TorchScript exports of weights published by
> the authors cited below, whose licenses are reproduced in
> [`LICENSES/`](LICENSES/) and whose attribution is in
> [`NOTICE`](NOTICE). The weights are unmodified; the
> wrapper adds normalization, layer selection and the forward signature above.
## Models
| Model | Architecture | Feature layers | Dim | Input channels | `ImpactPatchSize` | License | In this repository |
| --- | --- | --- | --- | --- | --- | --- | --- |
| `US/USF-MAE` | MAE ViT-B/16 encoder, blocks 2/5/8/11 tapped as stride-16 feature maps | 4 maps, 768 channels each, stride 16 — taps after transformer blocks 2, 5, 8, 11 | 2D | 3 | `0 0` (whole frame), or any multiple of 16 | [MIT](LICENSES/USF-MAE-MIT.txt) | [`US/USF-MAE.pt`](US/USF-MAE.pt) |
| `US/SAMUS` | SAM ViT-B image encoder at 256 px, patch size 8, global-attention blocks 2/5/8/11 plus the 256-channel neck | 5 maps on a 32x32 grid from the 256 px tile — 4 x 768-channel token maps after the global-attention blocks 2, 5, 8, 11 (the last after the `x + 0.5*cnnx` fusion) plus the 256-channel SAM neck | 2D | 1 | `256 256` **only** — the traced graph is valid at that size | [MIT](LICENSES/SAMUS-MIT.txt) | [`US/SAMUS.pt`](US/SAMUS.pt) |
| `US/URFM` | MAE ViT-B/16 encoder with a BiomedCLIP representation target, blocks 2/5/8/11 tapped as stride-16 feature maps | 4 maps, 768 channels each, stride 16 — taps after transformer blocks 2, 5, 8, 11 | 2D | 3 | `0 0` (whole frame), or any multiple of 16 | [MIT](LICENSES/URFM-MIT.txt) | build it yourself |
| `US/UltraFedFM` | Federated MAE ViT-B/16 encoder, blocks 2/5/8/11 tapped as stride-16 feature maps | 4 maps, 768 channels each, stride 16 — taps after transformer blocks 2, 5, 8, 11 | 2D | 3 | `0 0` (whole frame), or any multiple of 16 | [Apache-2.0](LICENSES/UltraFedFM-Apache-2.0.txt) | build it yourself |
### Provenance
| Model | Source code | Original weights | Paper |
| --- | --- | --- | --- |
| `US/USF-MAE` | <https://github.com/Yusufii9/USF-MAE> | <https://drive.google.com/file/d/1ZPu_7KhMEuaq-XdLhVp2EEgMgLJ4dKhr/view> | <https://doi.org/10.1016/j.bspc.2026.110313> |
| `US/SAMUS` | <https://github.com/xianlin7/SAMUS> | <https://drive.google.com/file/d/1nQjMAvbPeolNpCxQyU_HTiOiB5704pkH/view> | <https://doi.org/10.48550/arXiv.2309.06824> |
| `US/URFM` | <https://huggingface.co/QingboKang/URFM> | <https://huggingface.co/QingboKang/URFM/blob/main/mae_vit_base_patch16_dec768d8b_all_biomedclip_1199.pth> | <https://doi.org/10.1016/j.isci.2025.112917> |
| `US/UltraFedFM` | <https://github.com/yuncheng97/UltraFedFM> | <https://drive.google.com/file/d/13cczqVFk84c_9QDP2OLURU1jWryCky5f/view> | <https://doi.org/10.1038/s41746-025-02085-0> |
### Not redistributed here
Catalogued, licensed and buildable, but not shipped here: `US/URFM`, `US/UltraFedFM`.
Their upstream checkpoints sit behind a sign-in — a gated Hugging Face
repository, and cloud shares scoped to specific accounts — so nothing
here can fetch them for you. Build them yourself with the scripts in
`Data/Models/builds/US/` (see **Reproducing these exports** below); the
exports drop straight into the same `US/` layout.
## Usage — Python
```python
import itk
import torch
from huggingface_hub import hf_hub_download
model_path = hf_hub_download(
repo_id="fideus-labs/impact-torchscript-models",
filename="US/USF-MAE.pt",
local_dir="models",
)
frame = itk.imread("bmode.mha", itk.F) # a 2D itk.Image[itk.F, 2] B-mode frame
# Voxel size = the frame's own spacing, so IMPACT does not resample and the
# feature image lands pixel-for-pixel on the input grid.
config = itk.ModelConfiguration(
model_path, # TorchScript file
2, # dimension (models.json)
3, # numberofchannels (models.json)
[0, 0], # patch size; 0 0 = whole frame
[float(spacing) for spacing in frame.GetSpacing()],
0, # patch overlap
[False, False, False, True], # layers mask: ask for 4, keep the deepest
False, # useMixedPrecision
)
interpolator = itk.BSplineInterpolateImageFunction[type(frame), itk.D, itk.F]
extractor = itk.ImageToFeaturesMap[type(frame), interpolator].New()
extractor.SetModelConfiguration(config)
extractor.SetDevice("cuda:0" if torch.cuda.is_available() else "cpu")
extractor.AddInput(frame)
extractor.Update()
features = extractor.GetOutput(0) # itk.VectorImage, 768 components per pixel
```
The **layers mask does double duty**: its length is how many feature maps
IMPACT asks the model for, and each `1` keeps that map. `[0, 0, 0, 1]` on a
4-layer model therefore means "compute all four, keep the deepest".
For `US/SAMUS`, pass `1` input channel, `[256, 256]` as the patch size, and a
**5**-element mask (`[0, 0, 0, 0, 1]` keeps the 256-channel SAM neck).
As a registration loss in PyTorch, with
[ImpactLoss](https://github.com/vboussot/ImpactLoss):
```python
from IMPACT import IMPACTReg
loss = IMPACTReg(
"US/USF-MAE.pt",
shape=[0, 0],
in_channels=1,
weights=[0, 0, 0, 1],
repo_id="fideus-labs/impact-torchscript-models",
)
```
## Usage — Elastix
One resolution of `ParameterMaps/ParameterMap_US_2D_Static.txt`
(`ImpactVoxelSize` is the spacing the frame is resampled to; one token spans
16 voxels, so this level extracts features at 4 mm per token — roughly the
native spacing of a curvilinear B-mode frame):
```txt
(Metric "Impact")
(ImpactMode "Static")
(ImpactModelsPath0 "/Data/Models/US/USF-MAE.pt")
(ImpactDimension0 2)
(ImpactNumberOfChannels0 3)
(ImpactPatchSize0 0 0)
(ImpactVoxelSize0 0.25 0.25)
(ImpactLayersMask0 "0001")
(ImpactPCA0 0)
(ImpactSubsetFeatures0 64)
(ImpactDistance0 "L2")
(ImpactLayersWeight0 1)
```
`ImpactModelsPath` is resolved inside the IMPACT model directory, so
`US/USF-MAE.pt` sits where `sh Data/Models/download_models.sh` puts it. The
mask is a string here and one bit long per feature layer: 4 bits for the MAE
models, 5 for `US/SAMUS`, which also needs `(ImpactPatchSize 256 256)` and
`(ImpactNumberOfChannels 1)`.
## Reproducing these exports
Every model here is built from its upstream checkpoint by a standalone script
in `Data/Models/builds/US/` of
[ImpactLoss](https://github.com/vboussot/ImpactLoss) — download, wrap, check
the layer contract, `torch.jit` export. The scripts import only `torch`,
`gdown` and `huggingface_hub`, and write `./<Name>.pt` into the working
directory, so a rebuild is:
```bash
cd Data/Models/builds/US
python USF-MAE.py # -> ./USF-MAE.pt
```
The exports are reproducible in *behaviour*, not byte-for-byte: TorchScript
archives differ between runs, but a rebuilt model returns feature maps
identical to these (max absolute difference 0.0, checked across every layer
count, several input shapes and both normalization paths).
## Licensing
There is no single license: each model keeps the one its authors chose.
The texts are in [`LICENSES/`](LICENSES/), one per model,
indexed by [`LICENSE.md`](LICENSE.md), and
[`NOTICE`](NOTICE) names the original authors. The
`license: other` / `license_name: per-model` front matter above says the same
thing in Hugging Face's vocabulary.
Models whose licenses did **not** permit redistribution — several published
ultrasound foundation models are CC BY-NC or carry no license at all — were
excluded rather than repackaged.
## Integrity
[`SHA256SUMS`](SHA256SUMS) covers every other file:
```bash
sha256sum -c SHA256SUMS
```
## Citation
Cite the original work, not this packaging:
```bibtex
@article{megahed2026usfmae,
title = {USF-MAE: Ultrasound self-supervised foundation model with masked autoencoding},
author = {Megahed, Youssef and Ducharme, Robin and Erman, Aylin and
Walker, Mark C. and Hawken, Steven and Chan, Adrian D. C.},
journal = {Biomedical Signal Processing and Control},
volume = {122},
pages = {110313},
year = {2026},
doi = {10.1016/j.bspc.2026.110313}
}
```
```bibtex
@misc{lin2023samus,
title = {SAMUS: Adapting Segment Anything Model for Clinically-Friendly and
Generalizable Ultrasound Image Segmentation},
author = {Lin, Xian and Xiang, Yangyang and Yu, Li and Yan, Zengqiang},
year = {2023},
eprint = {2309.06824},
archivePrefix = {arXiv},
primaryClass = {cs.CV},
doi = {10.48550/arXiv.2309.06824},
note = {Later arXiv versions are titled "Beyond Adapting SAM: Towards
End-to-End Ultrasound Image Segmentation via Auto Prompting"}
}
```
```bibtex
@article{kang2025urfm,
title = {URFM: A general Ultrasound Representation Foundation Model for advancing
ultrasound image diagnosis},
author = {Kang, Qingbo and Lao, Qicheng and Gao, Jun and Bao, Wuyongga and He, Zhu and
Du, Chenlin and Lu, Qiang and Li, Kang},
journal = {iScience},
volume = {28},
number = {8},
pages = {112917},
year = {2025},
doi = {10.1016/j.isci.2025.112917}
}
```
```bibtex
@article{jiang2025ultrafedfm,
title = {From pretraining to privacy: federated ultrasound foundation model with
self-supervised learning},
author = {Jiang, Yuncheng and Feng, Chun-Mei and Ren, Jinke and Wei, Jun and
Zhang, Zixun and Hu, Yiwen and Liu, Yunbi and Sun, Rui and Tang, Xuemei and
Du, Juan and Wan, Xiang and Xu, Yong and Du, Bo and Gao, Xin and
Wang, Guangyu and Zhou, Shaohua and Cui, Shuguang and Li, Zhen},
journal = {npj Digital Medicine},
volume = {8},
number = {1},
year = {2025},
doi = {10.1038/s41746-025-02085-0}
}
```