rerun-io--rerun
wehub-resource-sync
58 行
1.7 KiB
Python
58 行
1.7 KiB
Python
"""Query various image representations."""
|
|
|
|
# region: setup
|
|
from __future__ import annotations
|
|
|
|
from io import BytesIO
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pyarrow as pa
|
|
from datafusion import col
|
|
from PIL import Image
|
|
|
|
import rerun as rr
|
|
|
|
sample_video_path = (
|
|
Path(__file__).parents[4] / "tests" / "assets" / "rrd" / "video_sample"
|
|
)
|
|
|
|
server = rr.server.Server(datasets={"video_dataset": sample_video_path})
|
|
CATALOG_URL = server.url()
|
|
client = rr.catalog.CatalogClient(CATALOG_URL)
|
|
dataset = client.get_dataset(name="video_dataset")
|
|
df = dataset.filter_contents([
|
|
"/compressed_images/**",
|
|
"/raw_images/**",
|
|
]).reader(index="log_time")
|
|
times = pa.table(df.select("log_time"))["log_time"].to_numpy()
|
|
# endregion: setup
|
|
|
|
# region: compressed_image
|
|
column_name = "/compressed_images:EncodedImage:blob"
|
|
row = df.filter(col("log_time") == times[0]).select(column_name)
|
|
image_byte_array = pa.table(row)[column_name].to_numpy()[0][0]
|
|
image = np.asarray(Image.open(BytesIO(image_byte_array.tobytes())))
|
|
print(f"{image.shape=}")
|
|
# endregion: compressed_image
|
|
|
|
# region: raw_image
|
|
content_column = "/raw_images:Image:buffer"
|
|
format_column = "/raw_images:Image:format"
|
|
row = df.filter(col("log_time") == times[0]).select(
|
|
content_column, format_column
|
|
)
|
|
table = pa.table(row)
|
|
format_details = table[format_column][0][0]
|
|
flattened_image = table[content_column].to_numpy()[0][0]
|
|
num_channels = rr.datatypes.color_model.ColorModel.auto(
|
|
int(format_details["color_model"].as_py())
|
|
).num_channels()
|
|
image = flattened_image.reshape(
|
|
format_details["height"].as_py(),
|
|
format_details["width"].as_py(),
|
|
num_channels,
|
|
)
|
|
print(f"{image.shape=}")
|
|
# endregion: raw_image
|