Skip to content

pixano_inference.impls._helpers

Torch helpers re-exported for the built-in torch backends.

The implementations live in :mod:pixano_inference.frameworks.torch; they are re-exported here for the built-in (transformers-based) backends that import them.

convert_image_pil_to_tensor(image, device, size=None)

Convert a PIL image to a (C, H, W) float tensor, optionally resizing it.

Parameters:

Name Type Description Default
image Any

PIL image.

required
device 'torch.device'

Torch device.

required
size int | None

Optional target size (both height and width).

None

Returns:

Type Description
'Tensor'

Image as a (C, H, W) float tensor.

Source code in pixano_inference/frameworks/torch.py
def convert_image_pil_to_tensor(image: Any, device: "torch.device", size: int | None = None) -> "Tensor":
    """Convert a PIL image to a ``(C, H, W)`` float tensor, optionally resizing it.

    Args:
        image: PIL image.
        device: Torch device.
        size: Optional target size (both height and width).

    Returns:
        Image as a ``(C, H, W)`` float tensor.
    """
    import torch

    assert_torch_installed()
    image = image.convert("RGB")
    if size is not None:
        image = image.resize((size, size))
    image_np = np.array(image) / 255.0
    return torch.from_numpy(image_np).to(device=device).permute(2, 0, 1)

encode_mask_to_rle(mask)

Encode a binary mask tensor using RLE.

Parameters:

Name Type Description Default
mask 'Tensor'

A binary mask of shape (height, width).

required

Returns:

Type Description
dict[str, list[int]]

RLE encoded mask as a dictionary.

Source code in pixano_inference/frameworks/torch.py
def encode_mask_to_rle(mask: "Tensor") -> dict[str, list[int]]:
    """Encode a binary mask tensor using RLE.

    Args:
        mask: A binary mask of shape (height, width).

    Returns:
        RLE encoded mask as a dictionary.
    """
    import torch

    assert_torch_installed()
    rle: dict[str, Any] = {"counts": [], "size": list(mask.shape)}
    mask = mask.permute(1, 0).flatten()
    diff_arr = torch.diff(mask)
    nonzero_indices = torch.where(diff_arr != 0)[0] + 1
    lengths = torch.diff(torch.concatenate((torch.tensor([0]), nonzero_indices, torch.tensor([len(mask)]))))

    # note that the odd counts are always the numbers of zeros
    if mask[0] == 1:
        lengths = torch.concatenate(([0], lengths))

    rle["counts"] = lengths.tolist()

    return rle

resolve_device(config)

Return the torch device for a deployment config (see :func:resolve_device_from_num_gpus).

Source code in pixano_inference/frameworks/torch.py
def resolve_device(config: ModelDeploymentConfig) -> Any:
    """Return the torch device for a deployment config (see :func:`resolve_device_from_num_gpus`)."""
    return resolve_device_from_num_gpus(config.resources.num_gpus)

resolve_torch_dtype(dtype_str)

Map a dtype string to a torch.dtype.

Parameters:

Name Type Description Default
dtype_str str

One of "float32", "float16", "bfloat16".

required

Returns:

Type Description
Any

Corresponding torch.dtype.

Raises:

Type Description
ValueError

If dtype_str is not recognised.

Source code in pixano_inference/frameworks/torch.py
def resolve_torch_dtype(dtype_str: str) -> Any:
    """Map a dtype string to a ``torch.dtype``.

    Args:
        dtype_str: One of ``"float32"``, ``"float16"``, ``"bfloat16"``.

    Returns:
        Corresponding ``torch.dtype``.

    Raises:
        ValueError: If *dtype_str* is not recognised.
    """
    import torch

    mapping = {
        "float32": torch.float32,
        "float16": torch.float16,
        "bfloat16": torch.bfloat16,
    }
    if dtype_str not in mapping:
        raise ValueError(f"Unsupported torch_dtype '{dtype_str}'. Choose from {list(mapping)}")
    return mapping[dtype_str]

should_compile(device, requested)

Decide whether to torch.compile a model.

torch.compile needs a working compiler and only pays off on GPU; on CPU it is often a slow no-win and can fail outright. So honour an explicit compile param when given, else auto-detect: compile only on CUDA devices.

Parameters:

Name Type Description Default
device Any

The resolved torch device (or a device string).

required
requested bool | None

The user's compile param, or None for auto.

required

Returns:

Type Description
bool

Whether to compile the model.

Source code in pixano_inference/impls/_helpers.py
def should_compile(device: Any, requested: bool | None) -> bool:
    """Decide whether to ``torch.compile`` a model.

    ``torch.compile`` needs a working compiler and only pays off on GPU; on CPU it is often a
    slow no-win and can fail outright. So honour an explicit ``compile`` param when given, else
    auto-detect: compile only on CUDA devices.

    Args:
        device: The resolved torch device (or a device string).
        requested: The user's ``compile`` param, or ``None`` for auto.

    Returns:
        Whether to compile the model.
    """
    if requested is not None:
        return bool(requested)
    device_type = getattr(device, "type", None) or str(device)
    return str(device_type).startswith("cuda")