atlas_gui.datasets package

Submodules

atlas_gui.datasets.dataset module

class DatasetBase[source]

Bases: ABC

Abstract base class for dataset handlers used in the segmentation GUI.

This interface enforces a standard structure for different dataset types (e.g., REASSEMBLE, RLDS), including methods for loading data, accessing segments, and reading/writing annotations.

abstract get_segment(segment_idx)[source]

Retrieve a specific segment from the dataset.

Parameters:

segment_idx (int) – Index of the segment to retrieve.

Returns:

A dataset-specific segment.

In REASSEMBLE, this is a high-level action within one H5 file. In RLDS, this is an episode.

Return type:

Any

abstract load_data(file_path)[source]

Load dataset from the given file path.

Parameters:

file_path (str) – Path to the dataset.

This method should also prepare the dataset for iteration and load the first sample, if applicable.

abstract write_annot_data(segment_idx, annots)[source]

Write annotation data back to the dataset or external storage.

Parameters:

  • segment_idx (int) – Index of the segment to annotate.

  • annots (dict) – Dictionary containing annotation data.

Note

  • In REASSEMBLE, this modifies the dataset file directly.

  • In RLDS, annotations are stored in external JSON files.

abstract load_annot_data(segment_idx)[source]

Load annotation data for a given segment.

Parameters:

segment_idx (int) – Index of the segment.

Returns:

Loaded annotation data for the segment.

Return type:

dict

abstract load_segments_info(file_path)[source]

Load segment metadata for the given file.

Parameters:

file_path (str) – Path to the dataset file.

This populates self.segments_info with segment metadata (e.g., start/end times, text descriptions).

abstract get_max_timestamp()[source]

Return the maximum timestamp for the current segment.

Returns:

Maximum timestamp value in seconds.

Return type:

float

atlas_gui.datasets.frames module

class Frames(config)[source]

Bases: DatasetBase

Dataset handler for image frame sequences stored in folders.

Supports three folder structures (auto-detected):

  • Flat: images directly in the selected folder (single segment, single camera)

  • Subfolder: subfolders each containing images (each subfolder = one segment, single camera)

  • Multi-camera: subfolders containing camera subfolders with images (each top-level subfolder = one segment, camera subfolders = cameras)

VALID_EXTENSIONS = {'.bmp', '.jpeg', '.jpg', '.png', '.tif', '.tiff'}
load_data(file_path)[source]

Load dataset from the given file path.

Parameters:

file_path (str) – Path to the dataset.

This method should also prepare the dataset for iteration and load the first sample, if applicable.

load_segments_info(file_path)[source]

Load segment metadata for the given file.

Parameters:

file_path (str) – Path to the dataset file.

This populates self.segments_info with segment metadata (e.g., start/end times, text descriptions).

get_segment(segment_idx)[source]

Retrieve a specific segment from the dataset.

Parameters:

segment_idx (int) – Index of the segment to retrieve.

Returns:

A dataset-specific segment.

In REASSEMBLE, this is a high-level action within one H5 file. In RLDS, this is an episode.

Return type:

Any

get_max_timestamp()[source]

Return the maximum timestamp for the current segment.

Returns:

Maximum timestamp value in seconds.

Return type:

float

write_annot_data(segment_idx, annots)[source]

Write annotation data back to the dataset or external storage.

Parameters:

  • segment_idx (int) – Index of the segment to annotate.

  • annots (dict) – Dictionary containing annotation data.

Note

  • In REASSEMBLE, this modifies the dataset file directly.

  • In RLDS, annotations are stored in external JSON files.

load_annot_data(segment_idx)[source]

Load annotation data for a given segment.

Parameters:

segment_idx (int) – Index of the segment.

Returns:

Loaded annotation data for the segment.

Return type:

dict

atlas_gui.datasets.reassemble module

class Reassemble(config)[source]

Bases: DatasetBase

Dataset class for handling REASSEMBLE HDF5-formatted data.

This class supports loading, caching, and annotating segments of data stored in REASSEMBLE format. It uses a moving window cache to efficiently load previous, current, and next segments, and provides thread-safe file access.

load_data(file_path)[source]

Load the HDF5 file and initialize segment metadata and cache.

Parameters:

file_path (str) – Path to the H5 file.

load_segments_info(file_path)[source]

Load segment metadata.

Parameters:

file_path (str) – Path to the H5 file.

set_segments_info(segments_info)[source]

Manually set the segment metadata.

Parameters:

segments_info (dict) – Dictionary of segment info.

get_segment(segment_idx)[source]

Retrieve a segment by index using the cache and support for sequential navigation.

Parameters:

segment_idx (int) – Index of the segment to retrieve.

Returns:

Segment data.

Return type:

dict

get_max_timestamp()[source]

Get the maximum timestamp from current segment after offsetting to start at zero.

Returns:

Maximum relative timestamp.

Return type:

float

write_annot_data(segment_idx, annots)[source]

Write annotations either to H5 file (in-place) or external JSON file.

Parameters:

  • segment_idx (int) – Index of the segment to annotate.

  • annots (dict) – Dictionary of annotations to save.

load_annot_data(segment_idx)[source]

Load annotations either from H5 file or external JSON file.

Parameters:

segment_idx (int) – Index of the segment.

Returns:

Dictionary of annotation data.

Return type:

dict

atlas_gui.datasets.rlds module

class RLDS(config, split='train')[source]

Bases: DatasetBase

Dataset handler for RLDS (Reinforcement Learning Datasets) format datasets.

This class loads RLDS datasets via TensorFlow Datasets, and supports segment access, annotation I/O, and automatic structuring of episode information. Each segment corresponds to an entire episode.

load_data(file_path=None)[source]

Load the RLDS dataset.

Parameters:

file_path (str, optional) – Path to the TFDS-formatted dataset directory. Required if download=False, ignored if download=True.

get_segment(segment_idx)[source]

Return a specific segment (episode) from the dataset, stacked into NumPy arrays.

Parameters:

segment_idx (int) – Index of the segment to load.

Returns:

A dictionary containing stacked step-wise data (e.g., observations, actions).

Return type:

dict

load_segments_info(file_path=None)[source]

Load metadata for each episode in the dataset.

Creates a dictionary where each segment entry includes: - index - start time - end time (based on FPS and number of steps) - language instruction (text) - unique ID (derived from episode metadata or index)

get_max_timestamp()[source]

Return the end timestamp of the current episode.

Returns:

The maximum timestamp for the current segment.

Return type:

float

write_annot_data(segment_idx, annots)[source]

Write annotation data to a per-dataset JSON file.

Annotations are stored using the UID of the segment as the key.

Parameters:

  • segment_idx (int) – Index of the segment being annotated.

  • annots (dict) – Dictionary of annotations (must be JSON serializable).

load_annot_data(segment_idx)[source]

Load annotation data for the given segment index, using UID-based lookup.

Parameters:

segment_idx (int) – Index of the segment to load annotations for.

Returns:

Annotation data for the given segment, or an empty dict if none found.

Return type:

dict

atlas_gui.datasets.rosbag_ds module

RosbagDataset: A dataset handler for ROS bag files without requiring ROS installation.

This module uses the rosbags library to read ROS1 (.bag) and ROS2 (.db3) bag files.

Install dependencies:

pip install rosbags numpy

For image decompression support:

pip install opencv-python

class Rosbag(config, split='train')[source]

Bases: DatasetBase

Dataset handler for ROS bag files (both ROS1 .bag and ROS2 .db3 formats).

This class loads rosbag files using the pure-Python rosbags library, requiring no ROS installation. Each bag file in the folder is treated as a separate segment.

Config structure (YAML):

dataset_type: rosbag dataset_name: my_rosbag_dataset fps: 30 annotation_dir: annotations/rosbag/ annotation_group: low_level

low_level_keys:

  • /robot_state/joint_positions

  • /robot_state/joint_velocities

  • /robot_state/ee_pose

camera_keys:

  • /camera/color/image_raw

  • /camera/depth/image_raw

color_format: “BGR”

default_graphs:

  • /robot_state/joint_positions

action_map:

1: Approach 2: Grasp …

# Streaming mode (default: True) - loads camera frames on-demand # Set to False to load everything into memory stream_mode: True

# Frame cache size for streaming mode (default: 30) frame_cache_size: 30

load_data(file_path)[source]

Load rosbag data from a folder containing bag files. Each bag file in the folder is treated as a separate segment.

Parameters:

file_path (str) – Path to a directory containing bag files (.bag or .db3).

load_segments_info(file_path=None)[source]

Load metadata for each bag file (segment) in the dataset. Each bag file is treated as one segment.

Populates segments_info with:

  • index

  • start time

  • end time

  • duration

  • text (empty string, for compatibility)

  • uid (unique identifier from filename)

  • topics (list of available topics)

get_segment(segment_idx)[source]

Load and return a specific segment (bag file).

In stream_mode (default): Only loads low_level_keys data fully. Camera data is indexed but loaded on-demand via get_frame_by_index().

With stream_mode=False: Loads everything into memory (original behavior).

The returned structure:

  • Camera data: FLAT structure with camera key directly containing image array Example: {‘cam1’: np.array (N, H, W, C)}

  • Low-level data: NESTED structure based on topic paths Example: {‘robot_state’: {‘joint_positions’: np.array (N, joints)}}

  • Timestamps: Top-level dict with all timestamps Example: {‘timestamps’: {‘cam1’: np.array (N,), ‘joint_positions’: np.array (N,)}}

Parameters:

segment_idx (int) – Index of the segment to load.

Returns:

Dictionary containing camera data (flat) and low-level data (nested).

Return type:

dict

get_max_timestamp()[source]

Return the maximum timestamp for the current segment.

Returns:

Maximum timestamp in seconds.

Return type:

float

get_frame_by_index(frame_idx, camera_key=None)[source]

Get a specific frame by index (streaming mode).

This method loads only the requested frame from disk, not all frames.

Parameters:

  • frame_idx (int) – Index of the frame to retrieve.

  • camera_key (Optional[str]) – Camera topic key. If None, uses first camera_key from config.

Return type:

Optional[ndarray]

Returns:

Image as numpy array, or None if not found.

get_num_frames(camera_key=None)[source]

Get the number of frames for a camera topic.

Parameters:

camera_key (Optional[str]) – Camera topic key. If None, uses first camera_key from config.

Return type:

int

Returns:

Number of frames, or 0 if not found.

write_annot_data(segment_idx, annots)[source]

Write annotation data to a per-dataset JSON file.

Annotations are stored using the UID of the segment as the key.

Parameters:

  • segment_idx (int) – Index of the segment being annotated.

  • annots (Dict[str, Any]) – Dictionary of annotations (must be JSON serializable).

load_annot_data(segment_idx)[source]

Load annotation data for the given segment.

Parameters:

segment_idx (int) – Index of the segment.

Returns:

Annotation data, or empty dict if none found.

Return type:

dict

get_topics(segment_idx=None)[source]

Get list of available topics for a segment.

Parameters:

segment_idx (Optional[int]) – Segment index. If None, uses current segment.

Return type:

List[str]

Returns:

List of topic names.

get_camera_frame(camera_key, frame_idx, segment_idx=None)[source]

Get a specific frame from a camera topic.

Parameters:

  • camera_key (str) – Camera key from config (e.g., ‘cam1’ or ‘/cam1/image_raw’).

  • frame_idx (int) – Frame index.

  • segment_idx (Optional[int]) – Segment index. If None, uses current segment.

Return type:

Optional[ndarray]

Returns:

Image as numpy array, or None if not found.

get_frame_at_timestamp(timestamp, camera_key=None, segment_idx=None)[source]

Get the image closest to the given timestamp.

Parameters:

  • timestamp (float) – Target timestamp in seconds.

  • camera_key (Optional[str]) – Specific camera key. If None, uses first camera_key from config.

  • segment_idx (Optional[int]) – Segment index. If None, uses current segment.

Return type:

Optional[ndarray]

Returns:

Image as numpy array, or None if not found.

load_rosbag_dataset(file_path, config)[source]

Convenience function to quickly load a rosbag dataset from a folder. Each bag file in the folder will be treated as a separate segment.

Parameters:

  • file_path (str) – Path to directory containing bag files.

  • config (Dict[str, Any]) – Configuration dictionary.

Return type:

Rosbag

Returns:

Loaded Rosbag instance.

atlas_gui.datasets.video module

class Video(config)[source]

Bases: DatasetBase

Dataset handler for video files (.mp4, .avi, .mkv).

Supports three structures (auto-detected):

  • Single file: a direct file path or folder with one video (1 segment, 1 camera)

  • Folder of videos: each video file = one segment (N segments, 1 camera)

  • Multi-camera: folder of subfolders, each containing multiple video files (each subfolder = one segment, each video = one camera, basename = camera key)

VALID_EXTENSIONS = {'.avi', '.mkv', '.mp4'}
load_data(file_path)[source]

Load dataset from the given file path.

Parameters:

file_path (str) – Path to the dataset.

This method should also prepare the dataset for iteration and load the first sample, if applicable.

load_segments_info(file_path)[source]

Load segment metadata for the given file.

Parameters:

file_path (str) – Path to the dataset file.

This populates self.segments_info with segment metadata (e.g., start/end times, text descriptions).

get_segment(segment_idx)[source]

Retrieve a specific segment from the dataset.

Parameters:

segment_idx (int) – Index of the segment to retrieve.

Returns:

A dataset-specific segment.

In REASSEMBLE, this is a high-level action within one H5 file. In RLDS, this is an episode.

Return type:

Any

get_max_timestamp()[source]

Return the maximum timestamp for the current segment.

Returns:

Maximum timestamp value in seconds.

Return type:

float

write_annot_data(segment_idx, annots)[source]

Write annotation data back to the dataset or external storage.

Parameters:

  • segment_idx (int) – Index of the segment to annotate.

  • annots (dict) – Dictionary containing annotation data.

Note

  • In REASSEMBLE, this modifies the dataset file directly.

  • In RLDS, annotations are stored in external JSON files.

load_annot_data(segment_idx)[source]

Load annotation data for a given segment.

Parameters:

segment_idx (int) – Index of the segment.

Returns:

Loaded annotation data for the segment.

Return type:

dict

Module contents