opentau.datasets.image_writer

Asynchronous image writing utilities for high-frequency data recording.

This module provides functionality for writing images to disk asynchronously using multithreading or multiprocessing, which is critical for controlling robots and recording data at high frame rates without blocking the main process.

The module supports two execution models:

Threading mode (num_processes=0): Creates a pool of worker threads for concurrent image writing within a single process.

Multiprocessing mode (num_processes>0): Creates multiple processes, each with their own thread pool, for maximum parallelism.

Key Features:

Asynchronous writing: Images are queued and written in background workers, preventing I/O blocking of the main process.
Multiple input formats: Supports torch Tensors, numpy arrays, and PIL Images with automatic conversion.
Format flexibility: Handles both channel-first (C, H, W) and channel-last (H, W, C) image formats.
Type conversion: Automatically converts float arrays in [0, 1] to uint8 in [0, 255] for PIL Image compatibility.
Safe cleanup: Decorator ensures image writers are properly stopped even when exceptions occur.

Classes:

AsyncImageWriter
Main class for asynchronous image writing with configurable threading or multiprocessing backends.

Functions:

image_array_to_pil_image: Convert numpy array to PIL Image with format and type conversion.
write_image: Write an image (numpy array or PIL Image) to disk.
worker_thread_loop: Worker thread loop for processing image write queue.
worker_process: Worker process that manages multiple threads for image writing.
safe_stop_image_writer: Decorator to safely stop image writer on exceptions.

Example

Create an async image writer with threading:

>>> writer = AsyncImageWriter(num_processes=0, num_threads=4)
>>> writer.save_image(image_array, Path("output/image.jpg"))
>>> writer.wait_until_done()  # Wait for all images to be written
>>> writer.stop()  # Clean up resources

Functions

`image_array_to_pil_image`(image_array[, ...])	Convert a numpy array to a PIL Image.
`safe_stop_image_writer`(func)	Decorator to safely stop image writer on exceptions.
`worker_process`(queue, num_threads)	Worker process that manages multiple threads for image writing.
`worker_thread_loop`(queue)	Worker thread loop for asynchronous image writing.
`write_image`(image, fpath)	Write an image to disk.

Classes

AsyncImageWriter([num_processes, num_threads])

This class abstract away the initialisation of processes or/and threads to save images on disk asynchrounously, which is critical to control a robot and record data at a high frame rate.

class opentau.datasets.image_writer.AsyncImageWriter(num_processes: int = 0, num_threads: int = 1)[source]

Bases: object

This class abstract away the initialisation of processes or/and threads to save images on disk asynchrounously, which is critical to control a robot and record data at a high frame rate.

When num_processes=0, it creates a threads pool of size num_threads. When num_processes>0, it creates processes pool of size num_processes, where each subprocess starts their own threads pool of size num_threads.

The optimal number of processes and threads depends on your computer capabilities. We advise to use 4 threads per camera with 0 processes. If the fps is not stable, try to increase or lower the number of threads. If it is still not stable, try to use 1 subprocess, or more.

__init__(num_processes: int = 0, num_threads: int = 1)[source]

save_image(image: Tensor | ndarray | Image, fpath: Path) → None[source]

Queue an image for asynchronous writing.

Converts torch tensors to numpy arrays and adds the image to the write queue.

Parameters:

image – Image to save (torch Tensor, numpy array, or PIL Image).
fpath – Path where the image will be saved.

stop() → None[source]

Stop all worker threads/processes and clean up resources.

Sends sentinel values to all workers and waits for them to finish. Terminates processes if they don’t respond.

wait_until_done() → None[source]: Wait until all queued images have been written to disk.

opentau.datasets.image_writer.image_array_to_pil_image(image_array: ndarray, range_check: bool = True) → Image[source]

Convert a numpy array to a PIL Image.

Supports channel-first (C, H, W) and channel-last (H, W, C) formats. Converts float arrays in [0, 1] to uint8 in [0, 255].

Parameters:

image_array – Input image array of shape (C, H, W) or (H, W, C).
range_check – If True, validates that float arrays are in [0, 1] range. Defaults to True.

Returns:

PIL Image object.

Raises:

ValueError – If array has wrong number of dimensions, wrong number of channels, or float values are outside [0, 1] range.
NotImplementedError – If image doesn’t have 3 channels.

Note

TODO(aliberts): handle 1 channel and 4 for depth images

opentau.datasets.image_writer.safe_stop_image_writer(func)[source]

Decorator to safely stop image writer on exceptions.

Ensures that the image writer is properly stopped if an exception occurs during function execution.

Parameters:: func – Function to wrap.
Returns:: Wrapped function that stops image writer on exceptions.

opentau.datasets.image_writer.worker_process(queue: Queue, num_threads: int) → None[source]

Worker process that manages multiple threads for image writing.

Creates and manages a pool of worker threads that process items from the queue.

Parameters:

queue – Queue containing (image_array, file_path) tuples or None sentinels.
num_threads – Number of worker threads to create in this process.

opentau.datasets.image_writer.worker_thread_loop(queue: Queue) → None[source]

Worker thread loop for asynchronous image writing.

Continuously processes items from the queue until receiving None (sentinel). Each item should be a tuple of (image_array, file_path).

Parameters:: queue – Queue containing (image_array, file_path) tuples or None sentinel.

opentau.datasets.image_writer.write_image(image: ndarray | Image, fpath: Path) → None[source]

Write an image to disk.

Converts numpy arrays to PIL Images if needed, then saves to the specified path.

Parameters:

image – Image to save (numpy array or PIL Image).
fpath – Path where the image will be saved.

Raises:

TypeError – If image type is not supported.