zea.models.ultrapips

UltraPIPS: perceptual similarity for B-mode ultrasound.

UltraPIPS is an LPIPS-style perceptual metric whose backbone is an ultrasound foundation model instead of a network pretrained on natural images. The features of both images are unit-normalized along the channel axis, compared with a squared difference, averaged over space, summed over channels and summed over all feature levels. Unlike LPIPS there is no learned linear head: every channel counts equally.

To try this model, simply load one of the available presets:

from zea.models.ultrapips import UltraPIPS

model = UltraPIPS.from_preset("ultrapips-tusa")
distance = model([image1, image2])  # images in [-1, 1], shape (B, H, W, C)

Backbone

The default (and currently only) backbone is the Swin transformer encoder of TUSA (Grutman et al., 2026), a self-supervised texture model trained on open B-mode datasets from many organs. It is the swinViT of a 2D MONAI SwinUNETR (v2) with feature_size=36, depths=(2, 3, 4, 5) and num_heads=(3, 6, 12, 24), run on 128 x 128 single-channel images. The five hidden states it returns (patch embedding plus the four stages) are the feature levels UltraPIPS compares.

Important

This is a zea implementation of the model. For the original code, see here and here for the TUSA backbone.

Note

The original UltraPIPS and TUSA code and the pretrained TUSA weights are released by their authors under the GNU General Public License v3.0. The ultrapips-tusa preset holds those weights, converted to Keras, and remains subject to that license. The encoder architecture follows MONAI’s SwinUNETR (Apache 2.0).

Reference

T. Grutman and T. Ilovitsh, “UltraPIPS: Improving Model Perception in B-mode Ultrasound with Foundation Models,” in MICCAI Workshop on Advances in Simplifying Medical Ultrasound (ASMUS), 2026. arXiv:2608.26033

Functions

load_torch_state_dict(net, state_dict[, prefix])	Copy the swinViT part of a TUSA (MONAI SwinUNETR) state dict into a SwinViT.
relative_position_index(window_size)	Index into the relative position bias table for every pair of window tokens.
shifted_window_mask(height, width, ...)	Attention mask that stops shifted windows from mixing wrapped-around regions.
window_partition(x, window_size)	Split (B, H, W, C) into windows of shape (B * num_windows, wh * ww, C).
window_reverse(windows, window_size, height, ...)	Inverse of window_partition(), back to (B, H, W, C).

Classes

PatchMerging(*args, **kwargs)	Halve the resolution and double the channels by merging 2x2 neighborhoods.
SwinStage(*args, **kwargs)	One stage of the Swin encoder: alternating regular and shifted blocks, then merging.
SwinTransformerBlock(*args, **kwargs)	Swin transformer block: (shifted) window attention followed by an MLP.
SwinViT(*args, **kwargs)	2D Swin transformer encoder (MONAI SwinUNETR.swinViT, use_v2=True).
UltraPIPS(*args, **kwargs)	Perceptual image patch similarity for B-mode ultrasound (UltraPIPS).
UnetResBlock(*args, **kwargs)	Residual block of two 3x3 convolutions with instance norm and leaky ReLU.
WindowAttention(*args, **kwargs)	Window-based multi-head self-attention with a relative position bias.

class zea.models.ultrapips.PatchMerging(*args, **kwargs)

Bases: Layer

Halve the resolution and double the channels by merging 2x2 neighborhoods.

Initialize the layer.

Parameters:

dim (int) – Number of input channels.

build(input_shape)

call(x)

Merge a (B, H, W, C) map into (B, ceil(H / 2), ceil(W / 2), 2C).

class zea.models.ultrapips.SwinStage(*args, **kwargs)

Bases: Layer

One stage of the Swin encoder: alternating regular and shifted blocks, then merging.

Initialize the stage.

Parameters:

• dim (int) – Number of input channels.

• depth (int) – Number of transformer blocks.

• num_heads (int) – Number of attention heads.

• window_size (tuple[int, int]) – Window size.

• mlp_ratio (float) – Hidden size of the MLP relative to dim.

• qkv_bias (bool) – Whether the query, key and value projection has a bias.

build(input_shape)

call(x)

Apply the stage to a (B, H, W, C) feature map.

class zea.models.ultrapips.SwinTransformerBlock(*args, **kwargs)

Bases: Layer

Swin transformer block: (shifted) window attention followed by an MLP.

Initialize the block.

Parameters:

• dim (int) – Number of feature channels.

• num_heads (int) – Number of attention heads.

• window_size (tuple[int, int]) – Window size.

• shift_size (tuple[int, int]) – Cyclic shift, (0, 0) for regular windows.

• mlp_ratio (float) – Hidden size of the MLP relative to dim.

• qkv_bias (bool) – Whether the query, key and value projection has a bias.

build(input_shape)

call(x, attn_mask=None)

Apply the block to a (B, H, W, C) feature map.

Parameters:

• x (Tensor) – Feature map.

• attn_mask (np.ndarray, optional) – Shifted window mask, ignored for blocks without a shift.

class zea.models.ultrapips.SwinViT(*args, **kwargs)

Bases: Layer

2D Swin transformer encoder (MONAI SwinUNETR.swinViT, use_v2=True).

Returns the five hidden states of the encoder: the patch embedding and the output of each of the four stages, each optionally layer-normalized.

Initialize the encoder.

Parameters:

• in_channels (int) – Number of input image channels.

• feature_size (int) – Embedding dimension; doubles at every stage.

• depths (tuple[int]) – Number of transformer blocks per stage.

• num_heads (tuple[int]) – Number of attention heads per stage.

• window_size (tuple[int, int]) – Attention window size.

• patch_size (int) – Patch size of the embedding.

• mlp_ratio (float) – Hidden size of the MLPs relative to their input.

• qkv_bias (bool) – Whether the query, key and value projections have a bias.

build(input_shape)

call(x, normalize=True)

Encode a (B, H, W, C) image.

Parameters:

• x (Tensor) – Input image, normalized the way the weights expect.

• normalize (bool) – Whether to layer-normalize (without affine) each returned hidden state over its channels.

Returns:

Five (B, H_i, W_i, C_i) feature maps.

Return type:

list[Tensor]

get_config()

Returns the config of the object.

An object config is a Python dictionary (serializable) containing the information needed to re-instantiate it.

class zea.models.ultrapips.UltraPIPS(*args, **kwargs)

Bases: BaseModel

Perceptual image patch similarity for B-mode ultrasound (UltraPIPS).

Initialize the UltraPIPS model.

Parameters:

• backbone (str, optional) – Feature extractor. Only "tusa" (the default of the reference implementation) is ported. Defaults to "tusa".

• image_size (int, optional) – Side length the images are resized to before they are encoded. The TUSA backbone was trained at 128. Defaults to 128.

• disable_checks (bool, optional) – Disable the input value checks, e.g. to trace the metric in a TensorFlow graph. Defaults to False.

BACKBONES = ('tusa',)

build(input_shape=None)

Build the backbone (its input shape is fixed by image_size).

call(inputs)

Compute the UltraPIPS distance.

Parameters:

inputs (list) – Two images of shape (B, H, W, C) or (H, W, C) with 1 or 3 channels and values in [-1, 1].

Returns:

Distance of shape (B,), or a scalar without a batch dimension.

Return type:

Tensor

custom_load_weights(preset, backend='keras', **kwargs)

Load weights from a preset (Hugging Face handle or local directory).

Parameters:

• preset (str) – Preset identifier passed from from_preset().

• backend (str) – "keras" loads model.weights.h5; "torch" converts the original TUSA unet.pt checkpoint and needs PyTorch.

features(image)

Encode an image into its (layer-normalized) backbone feature maps.

Parameters:

image (Tensor) – (B, H, W, C) image in [-1, 1] with 1 or 3 channels.

Returns:

Five (B, H_i, W_i, C_i) feature maps.

Return type:

list[Tensor]

classmethod from_pth(pth_path, **kwargs)

Create an UltraPIPS model from the original TUSA PyTorch checkpoint.

Parameters:

• pth_path (str) – Path to the TUSA unet.pt state dict.

• **kwargs – Passed to the constructor.

Returns:

Model with the converted weights.

Return type:

UltraPIPS

get_config()

Serialize the model arguments.

preprocess_input(image)

Convert images to the input the backbone expects.

Converts RGB to grayscale, resizes to image_size (bilinear with antialiasing, like torchvision) and keeps the [-1, 1] range, which is the (x - 0.5) / 0.5 normalization TUSA was trained with.

Warning

Images are resized to a square image_size x image_size, so non-square inputs are stretched and lose their aspect ratio, as in the reference implementation. A warning is logged once per input shape; crop or pad to a square beforehand to avoid the distortion.

Parameters:

image (Tensor) – (B, H, W, C) image in [-1, 1] with 1 or 3 channels.

Returns:

(B, image_size, image_size, 1) tensor.

Return type:

Tensor

class zea.models.ultrapips.UnetResBlock(*args, **kwargs)

Bases: Layer

Residual block of two 3x3 convolutions with instance norm and leaky ReLU.

Initialize the block.

Parameters:

• filters (int) – Number of input and output channels.

• negative_slope (float) – Slope of the leaky ReLU.

build(input_shape)

call(x)

Apply the block to a (B, H, W, C) feature map.

class zea.models.ultrapips.WindowAttention(*args, **kwargs)

Bases: Layer

Window-based multi-head self-attention with a relative position bias.

Initialize the attention layer.

Parameters:

• dim (int) – Number of feature channels.

• num_heads (int) – Number of attention heads.

• window_size (tuple[int, int]) – Window size the bias table is sized for.

• qkv_bias (bool) – Whether the query, key and value projection has a bias.

build(input_shape)

call(x, attn_mask=None)

Attend within each window.

Parameters:

• x (Tensor) – Windows of shape (B * num_windows, N, C).

• attn_mask (np.ndarray, optional) – (num_windows, N, N) additive mask for shifted windows.

Returns:

Same shape as x.

Return type:

Tensor

zea.models.ultrapips.load_torch_state_dict(net, state_dict, prefix='swinViT.')

Copy the swinViT part of a TUSA (MONAI SwinUNETR) state dict into a SwinViT.

Kernels are permuted from PyTorch’s (out, in, h, w) to Keras’ (h, w, in, out) and dense weights are transposed. The relative position index buffers are not copied: they are recomputed from the window size.

Parameters:

• net (SwinViT) – Built encoder to load the weights into.

• state_dict (dict) – {key: tensor} mapping from PyTorch, e.g. the full SwinUNETR state dict.

• prefix (str) – Prefix of the encoder keys in state_dict.

Raises:

KeyError – If the state dict does not match the encoder’s architecture.

zea.models.ultrapips.relative_position_index(window_size)

Index into the relative position bias table for every pair of window tokens.

Returns:

Integer array of shape (wh * ww, wh * ww).

Return type:

np.ndarray

zea.models.ultrapips.shifted_window_mask(height, width, window_size, shift_size)

Attention mask that stops shifted windows from mixing wrapped-around regions.

Parameters:

• height (int) – Padded feature map height.

• width (int) – Padded feature map width.

• window_size (tuple[int, int]) – Window size.

• shift_size (tuple[int, int]) – Cyclic shift.

Returns:

(num_windows, N, N) float32 mask with 0 where tokens may attend to each other and -100 where they may not.

Return type:

np.ndarray

zea.models.ultrapips.window_partition(x, window_size)

Split (B, H, W, C) into windows of shape (B * num_windows, wh * ww, C).

zea.models.ultrapips.window_reverse(windows, window_size, height, width)

Inverse of window_partition(), back to (B, H, W, C).