ImageBatch

Represents an ordered collection of images.

Parameters

Field	Type	Description
data	Sequence[Image | np.ndarray]	Images to store. Bare np.ndarrays are converted via Image(...).
compression	ImageCompression | int | None	If given, overrides every item's codec (each image is re-encoded with it). If None (default), each Image keeps its own codec, and bare ndarrays default to ImageCompression.NONE. See ImageCompression.

Raises

Exception	Condition
TypeError	An item isn't an Image/np.ndarray, or compression isn't a valid ImageCompression/matching int
ValueError	An item has an unsupported dtype or shape (per Image._validate)

Attributes

Attribute	Type	Description
shapes	list[tuple[int, ...]]	Each image's pixel-array shape, in order.
dtypes	list[str]	Each image's NumPy dtype name, in order.
compressions	list[ImageCompression]	Each image's compression codec, in order.

Methods

Method	Description
to_numpy(copy=True)	Returns each image's pixel array as a list[np.ndarray]; copy=False returns the batch's internal arrays directly.
ImageBatch.coerce(value)	Returns value unchanged if already an ImageBatch; otherwise wraps a list/tuple of Image/np.ndarray. Raises TypeError for anything else.

Operators

Operation	Behavior
len(batch)	Number of images.
batch[i]	int returns a materialized Image (defensive copy; IndexError if out of range); slice or boolean np.ndarray mask returns an ImageBatch sub-batch. Anything else raises TypeError.
batch == other	True only if other is an ImageBatch with the same number of images and element-equal pixel arrays in order; compression is not compared. NotImplemented if other isn't an ImageBatch.
hash(batch)	Not supported.

There's no __array__ -- images in a batch can have different shapes/dtypes, so the batch as a whole can't be viewed as one NumPy array. Use to_numpy() (a list) or index into one Image instead.

Serialization

Arrow layout:

StructArray length 1
├── data:        list<binary>  (length 1, inner length N; raw or ZSTD-framed per row)
├── height:      list<int32>   (length 1, inner length N)
├── width:       list<int32>   (length 1, inner length N)
├── channels:    list<int32>   (length 1, inner length N)
├── dtype:       list<string>  (length 1, inner length N)
└── compression: list<int8>    (length 1, inner length N)

One ImageBatch is serialized as one Arrow row, with all N images packed into the single list element of each field.

Visualization

datatypes.visualize(batch, entity_path=...) logs each image under its own indexed child path ({entity_path}/{i}) as rr.Image.

Example

"""Demonstrates the Telekinesis ImageBatch datatype."""

import time
from pathlib import Path

import numpy as np
from loguru import logger
import rerun as rr

from telekinesis import datatypes

def image_batch_example():
    """Demonstrate creation, inspection, visualization, indexing, grayscale conversion, rebuilding, and serialization."""

    # ======================= Create ============================================
    ROOT_PATH = Path(__file__).parent
    image_1 = np.random.randint(0, 255, (512, 512, 3), dtype=np.uint8)
    image_2 = datatypes.Image.from_path(ROOT_PATH / "data/sample.jpg").to_numpy()
    images = [image_1, image_2]
    image_batch = datatypes.ImageBatch(images)
    logger.info(f"Original ImageBatch: {image_batch}")

    # ======================= Inspect ===========================================
    dtypes = image_batch.dtypes
    shapes = image_batch.shapes
    compressions = image_batch.compressions
    numpy_array = image_batch.to_numpy()

    logger.info(f"dtypes={dtypes}, shapes={shapes}, compressions={compressions}")
    logger.info(f"NumPy array: {numpy_array}")

    # ======================= Visualize =========================================
    rr.init("image_batch_example", spawn=True)
    datatypes.visualize(image_batch, entity_path="/ImageBatch")

    # ======================= Index =============================================
    index = 1
    image_at_index = image_batch[index]
    logger.info(f"Image at index {index}: {image_at_index}")
    datatypes.visualize(image_at_index, entity_path="/ImageBatch/Image_1")

    # ======================= Grayscale =========================================
    gray_image = image_at_index.to_grayscale()
    logger.info(f"Grayscale image at index {index}: {gray_image}")
    datatypes.visualize(gray_image, entity_path="/ImageBatch/Image_1/Grayscale")
    gray_image.save_to_path(ROOT_PATH / "data/grayscale_image.jpg")

    # ======================= Rebuild ===========================================
    index = 0
    updated_image = np.random.randint(0, 255, (1907, 512, 3), dtype=np.uint8)
    images[index] = updated_image
    image_batch = datatypes.ImageBatch(images)
    logger.info(f"Rebuilt ImageBatch at index 0: {image_batch}")

    # ======================= Serialize / Deserialize ===========================
    start = time.perf_counter()
    serialized = datatypes.serialize(image_batch)
    serialization_ms = (time.perf_counter() - start) * 1000

    start = time.perf_counter()
    deserialized = datatypes.deserialize(serialized)["param_0"]
    deserialization_ms = (time.perf_counter() - start) * 1000

    logger.info(f"Deserialized ImageBatch: {deserialized}")
    logger.info(f"Round-trip successful: {deserialized == image_batch}")
    logger.info(f"Serialization time: {serialization_ms:.3f} ms")
    logger.info(f"Deserialization time: {deserialization_ms:.3f} ms")


if __name__ == "__main__":
    image_batch_example()