tso

tso

Dataset schema for TSO surgical CSV datasets.

TSODatasetSchema

TSODatasetSchema(dataset_folders, zarr_path, metadata, dataset_type=value, use_rectified_images=True, rgb_image_crops=None)

Bases: CsvDatasetSchema

Schema for TSO zarr datasets with synchronized CSV and stereo images.

These datasets can contain: - 3D cartesian position in robot and camera frames - Binary gripper state (open/close) - Optional task phase labels - Stereo camera images (left, right) with optional depth

Initialize and validate the TSO dataset schema.

Parameters:

Name	Type	Description	Default
dataset_folders	list[str]	List of folders containing episode CSVs.	required
zarr_path	str	Path to save/load the zarr store.	required
metadata	DatasetMetadata	Metadata to use for creating the zarr store from the raw data.	required
dataset_type	str	Type of dataset. Must be 'tso'.	value
use_rectified_images	bool	Whether to read rectified image path columns.	True
rgb_image_crops	dict[str, dict[str, int]] | None	Optional per-camera Albumentations crop params.	None

Source code in src/versatil/data/raw/schemas/custom/tso.py

def __init__(
    self,
    dataset_folders: list[str],
    zarr_path: str,
    metadata: DatasetMetadata,
    dataset_type: str = DatasetType.TSO.value,
    use_rectified_images: bool = True,
    rgb_image_crops: dict[str, dict[str, int]] | None = None,
):
    """Initialize and validate the TSO dataset schema.

    Args:
        dataset_folders: List of folders containing episode CSVs.
        zarr_path: Path to save/load the zarr store.
        metadata: Metadata to use for creating the zarr store from the raw data.
        dataset_type: Type of dataset. Must be 'tso'.
        use_rectified_images: Whether to read rectified image path columns.
        rgb_image_crops: Optional per-camera Albumentations crop params.
    """
    if dataset_type != DatasetType.TSO.value:
        raise ValueError(
            f"TSODatasetSchema only supports dataset_type='{DatasetType.TSO.value}', "
            f"got '{dataset_type}'"
        )
    self._validate_metadata(metadata)
    self.left_image_csv_key = TSO_LEFT_IMAGE_KEY
    self.right_image_csv_key = TSO_RIGHT_IMAGE_KEY
    self.rectified_left_image_key = TSO_RECTIFIED_LEFT_IMAGE_KEY
    self.rectified_right_image_key = TSO_RECTIFIED_RIGHT_IMAGE_KEY
    self.use_rectified_images = use_rectified_images
    self.rgb_image_crops = self._normalize_rgb_image_crops(
        rgb_image_crops or DEFAULT_TSO_IMAGE_CROPS
    )
    self.depth_dir_pattern = "depth"
    self.depth_file_pattern = r"depth_\1.npy"
    self.left_dir_pattern = "framesLeft"
    self.rectified_left_dir_pattern = "framesLeftRectified"
    super().__init__(
        dataset_folders=dataset_folders,
        episode_filename=TSO_EPISODE_FILENAME,
        zarr_path=zarr_path,
        metadata=metadata,
        dataset_type=dataset_type,
    )

extract_episode

extract_episode(episode)

Extract all data from a TSO episode, optionally cropping/resizing images.

Parameters:

Name	Type	Description	Default
episode	DataFrame	DataFrame with episode data.	required

Returns:

Type	Description
dict[str, ndarray]	Dictionary mapping zarr keys to numpy arrays.

Source code in src/versatil/data/raw/schemas/custom/tso.py

def extract_episode(
    self,
    episode: pd.DataFrame,
) -> dict[str, np.ndarray]:
    """Extract all data from a TSO episode, optionally cropping/resizing images.

    Args:
        episode: DataFrame with episode data.

    Returns:
        Dictionary mapping zarr keys to numpy arrays.
    """
    data = {}
    for zarr_key, observation in self.metadata.observations.items():
        if isinstance(observation, CameraMetadata):
            continue
        elif isinstance(observation, ObservationMetadata):
            data[zarr_key] = episode[
                observation.raw_data_column_keys
            ].values.astype(observation.dtype)

    for zarr_key, action in self.metadata.precomputed_actions.items():
        data[zarr_key] = episode[action.raw_data_column_keys].values.astype(
            action.dtype
        )

    for zarr_key, cam_metadata in self.metadata.cameras.items():
        camera = cam_metadata.raw_camera_key
        camera_resizer = build_camera_resizer(camera_metadata=cam_metadata)
        if cam_metadata.is_depth:
            left_column = self._get_rgb_column(Cameras.LEFT.value)
            paths = [self._compute_depth_path(p) for p in episode[left_column]]
            images = [
                camera_resizer(image=np.load(p))["image"][..., np.newaxis]
                for p in paths
            ]  # (H, W, 1)
        else:
            column = self._get_rgb_column(camera)
            images = [
                camera_resizer(image=self._read_rgb_image(p, camera))["image"]
                for p in episode[column]
            ]
        data[zarr_key] = np.stack(images).astype(cam_metadata.dtype)

    return data

get_image_path_column

get_image_path_column(camera)

Get CSV column name for a camera image path.

Source code in src/versatil/data/raw/schemas/custom/tso.py

def get_image_path_column(self, camera: str) -> str:
    """Get CSV column name for a camera image path."""
    return self._get_rgb_column(camera)

compute_depth_path

compute_depth_path(base_image_path)

Compute depth file path from a left RGB image path.

Source code in src/versatil/data/raw/schemas/custom/tso.py

def compute_depth_path(self, base_image_path: str) -> str:
    """Compute depth file path from a left RGB image path."""
    return self._compute_depth_path(base_image_path)