robits.dataset.io package

Submodules

robits.dataset.io.reader module

class robits.dataset.io.reader.DatasetReader(dataset_path: Path)

Bases: object

Reads a dataset from a file

load(load_camera_images: bool = True, wrapper: Callable | None = None) → Dataset

Loads the dataset into memory.

Parameters:

• load_camera_images – Whether to load camera images (default: True).

• wrapper – Optional wrapper function for iterating over sequences.

Returns:

A Dataset object containing all entries.

load_camera_data(seq: int, camera_name: str) → CameraData

Loads camera data for a specific sequence index.

Parameters:

• seq – The sequence index of the image.

• camera_name – The name of the camera.

Returns:

A CameraData object containing RGB and depth images.

property num_items: int

Returns:

the number of entries within the loaded dataset

read_entry(seq: int, camera_names: List[str]) → Entry

Reads a dataset entry including proprioception and camera data.

Parameters:

• seq – The sequence index of the entry.

• camera_names – A list of camera names whose data should be loaded.

Returns:

An Entry object containing the loaded data.

validate() → bool

Performs some sanity check on the dataset

Returns:

True if the sanity checks are passed

robits.dataset.io.recorder module

class robits.dataset.io.recorder.DatasetRecorder(robot: UnimanualRobot)

Bases: object

Handles recording of dataset entries from a robot.

do_recording() → None

The main recording loop that collects data at a fixed frequency.

start_recording() → None

Starts the recording process in a separate thread.

stop_recording() → None

Stops the recording process and joins the recording thread.

robits.dataset.io.stats_writer module

class robits.dataset.io.stats_writer.StatsWriter(dataset_path: Path, dataset: Dataset | None = None)

Bases: object

compute_stats() → Dict[str, Dict[str, Any]]

Computes statistics for each proprioceptive key across all entries.

Returns:

Dictionary of statistics per proprioception field.

has_stats() → bool

load_stats() → Dict[str, Dict[str, Any]]

write_stats() → None

Computes and writes statistics to a JSON file.

robits.dataset.io.writer module

class robits.dataset.io.writer.DatasetWriter(output_path, recorder, rgb_format='jpg')

Bases: object

Class for writing dataset recordings asynchronously.

do_write() → None

Continuously writes recorded data until stopped.

finalize_write() → None

Waits for the recorder to finish and finalizes dataset writing.

get_demo_path(output_path) → Path

Generates a new demo path within the given output directory.

Parameters:

output_path – The base output path for storing datasets.

Returns:

The formatted dataset path.

prepare_write() → None

Prepares the directory structure and starts the writing thread.

write_obs(observation) → None

Writes a single observation to the dataset.

Parameters:

observation – The observation data containing sequence number, perception, and proprioception.