Files
2026-07-13 13:05:14 +08:00

68 lines
1.5 KiB
Python

from __future__ import annotations
from typing import Any
import numpy as np
import pytest
import rerun as rr
import torch
rng = np.random.default_rng(12345)
RANDOM_IMAGE_SOURCE = rng.integers(0, 255, size=(10, 20))
IMAGE_INPUTS: list[Any] = [
RANDOM_IMAGE_SOURCE,
RANDOM_IMAGE_SOURCE,
]
def segmentation_image_image_expected() -> Any:
return rr.SegmentationImage(RANDOM_IMAGE_SOURCE)
def test_image() -> None:
expected = segmentation_image_image_expected()
for img in IMAGE_INPUTS:
arch = rr.SegmentationImage(img)
assert arch == expected
GOOD_IMAGE_INPUTS: list[Any] = [
# Mono
rng.integers(0, 255, (10, 20)),
# Assorted Extra Dimensions
rng.integers(0, 255, (1, 10, 20)),
rng.integers(0, 255, (10, 20, 1)),
# Torch tensors
torch.randint(0, 255, (10, 20)),
]
BAD_IMAGE_INPUTS: list[Any] = [
rng.integers(0, 255, (10, 20, 3)),
rng.integers(0, 255, (10, 20, 4)),
rng.integers(0, 255, (10,)),
rng.integers(0, 255, (1, 10, 20, 3)),
rng.integers(0, 255, (1, 10, 20, 4)),
rng.integers(0, 255, (10, 20, 3, 1)),
rng.integers(0, 255, (10, 20, 4, 1)),
rng.integers(0, 255, (10, 20, 2)),
rng.integers(0, 255, (10, 20, 5)),
rng.integers(0, 255, (10, 20, 3, 2)),
]
def test_segmentation_image_shapes() -> None:
import rerun as rr
rr.set_strict_mode(True)
for img in GOOD_IMAGE_INPUTS:
rr.SegmentationImage(img)
for img in BAD_IMAGE_INPUTS:
with pytest.raises(ValueError):
rr.SegmentationImage(img)