wehub-resource-sync
146 lines
4.1 KiB
Python
146 lines
4.1 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import rerun as rr
|
|
import torch
|
|
from rerun.archetypes.image import Image
|
|
from rerun.datatypes.tensor_data import TensorData
|
|
from rerun.error_utils import RerunWarning
|
|
|
|
rng = np.random.default_rng(12345)
|
|
RANDOM_IMAGE_SOURCE = rng.uniform(0.0, 1.0, (10, 20, 3))
|
|
|
|
|
|
IMAGE_INPUTS: list[Any] = [
|
|
{"image": RANDOM_IMAGE_SOURCE},
|
|
{"image": RANDOM_IMAGE_SOURCE, "width": 20, "height": 10},
|
|
{"image": RANDOM_IMAGE_SOURCE, "color_model": "RGB", "width": 20, "height": 10},
|
|
{"image": RANDOM_IMAGE_SOURCE, "color_model": rr.datatypes.ColorModel.RGB, "width": 20, "height": 10},
|
|
{
|
|
"bytes": RANDOM_IMAGE_SOURCE.tobytes(),
|
|
"color_model": "RGB",
|
|
"datatype": "f64",
|
|
"width": 20,
|
|
"height": 10,
|
|
},
|
|
{
|
|
"bytes": RANDOM_IMAGE_SOURCE.tobytes(),
|
|
"color_model": "RGB",
|
|
"datatype": rr.datatypes.ChannelDatatype.F64,
|
|
"width": 20,
|
|
"height": 10,
|
|
},
|
|
{
|
|
"bytes": RANDOM_IMAGE_SOURCE.tobytes(),
|
|
"color_model": "RGB",
|
|
"datatype": np.float64,
|
|
"width": 20,
|
|
"height": 10,
|
|
},
|
|
# This was allowed in 0.17
|
|
{"image": TensorData(array=RANDOM_IMAGE_SOURCE)},
|
|
]
|
|
|
|
|
|
def image_data_expected() -> Any:
|
|
return Image(RANDOM_IMAGE_SOURCE, color_model="RGB", width=20, height=10)
|
|
|
|
|
|
def test_image() -> None:
|
|
expected = image_data_expected()
|
|
|
|
for input in IMAGE_INPUTS:
|
|
arch = rr.Image(**input)
|
|
|
|
assert arch.buffer == expected.buffer
|
|
assert arch.format == expected.format
|
|
|
|
|
|
GOOD_IMAGE_INPUTS: list[Any] = [
|
|
# Mono
|
|
rng.uniform(0.0, 1.0, (10, 20)),
|
|
# RGB
|
|
rng.uniform(0.0, 1.0, (10, 20, 3)),
|
|
# RGBA
|
|
rng.uniform(0.0, 1.0, (10, 20, 4)),
|
|
# Assorted Extra Dimensions
|
|
rng.uniform(0.0, 1.0, (1, 10, 20)),
|
|
rng.uniform(0.0, 1.0, (1, 10, 20, 3)),
|
|
rng.uniform(0.0, 1.0, (1, 10, 20, 4)),
|
|
rng.uniform(0.0, 1.0, (10, 20, 1)),
|
|
rng.uniform(0.0, 1.0, (10, 20, 3, 1)),
|
|
rng.uniform(0.0, 1.0, (10, 20, 4, 1)),
|
|
# Torch tensors
|
|
torch.rand(10, 20, 1),
|
|
torch.rand(10, 20, 3),
|
|
]
|
|
|
|
BAD_IMAGE_INPUTS: list[Any] = [
|
|
rng.uniform(0.0, 1.0, (10,)),
|
|
rng.uniform(0.0, 1.0, (10, 20, 2)),
|
|
rng.uniform(0.0, 1.0, (10, 20, 5)),
|
|
rng.uniform(0.0, 1.0, (10, 20, 3, 2)),
|
|
]
|
|
|
|
|
|
def test_image_shapes() -> None:
|
|
import rerun as rr
|
|
|
|
rr.set_strict_mode(True)
|
|
|
|
for img in GOOD_IMAGE_INPUTS:
|
|
rr.Image(img)
|
|
|
|
for img in BAD_IMAGE_INPUTS:
|
|
with pytest.raises(ValueError):
|
|
rr.Image(img)
|
|
|
|
|
|
def test_image_compress() -> None:
|
|
rr.set_strict_mode(False)
|
|
|
|
# RGB Supported
|
|
image_data = np.asarray(rng.uniform(0, 255, (10, 20, 3)), dtype=np.uint8)
|
|
|
|
compressed = rr.Image(image_data).compress(jpeg_quality=80)
|
|
assert type(compressed) is rr.EncodedImage
|
|
|
|
# Mono Supported
|
|
image_data = np.asarray(rng.uniform(0, 255, (10, 20)), dtype=np.uint8)
|
|
|
|
compressed = rr.Image(image_data).compress(jpeg_quality=80)
|
|
assert type(compressed) is rr.EncodedImage
|
|
|
|
# RGBA Not supported
|
|
with pytest.warns(RerunWarning) as warnings:
|
|
image_data = np.asarray(rng.uniform(0, 255, (10, 20, 4)), dtype=np.uint8)
|
|
compressed = rr.Image(image_data, "RGBA").compress(jpeg_quality=80)
|
|
|
|
assert len(warnings) == 1
|
|
assert "Cannot JPEG compress an image of type" in str(warnings[0])
|
|
|
|
assert type(compressed) is rr.Image
|
|
|
|
# 16-bit Not supported
|
|
with pytest.warns(RerunWarning) as warnings:
|
|
image_data = np.asarray(rng.uniform(0, 255, (10, 20, 3)), dtype=np.uint16)
|
|
compressed = rr.Image(image_data).compress(jpeg_quality=80)
|
|
|
|
assert len(warnings) == 1
|
|
assert "Cannot JPEG compress an image of datatype" in str(warnings[0])
|
|
|
|
assert type(compressed) is rr.Image
|
|
|
|
# Floating point not supported
|
|
with pytest.warns(RerunWarning) as warnings:
|
|
image_data = np.asarray(rng.uniform(0, 255, (10, 20)), dtype=np.float32)
|
|
compressed = rr.Image(image_data).compress(jpeg_quality=80)
|
|
|
|
assert len(warnings) == 1
|
|
assert "Cannot JPEG compress an image of datatype" in str(warnings[0])
|
|
|
|
assert type(compressed) is rr.Image
|