68 lines
1.5 KiB
Python
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)
|