Metadata-Version: 2.4
Name: tso-sensorium
Version: 0.1.0
Summary: Sensor data recording, post-processing, and episodic dataset creation for robot learning
Author-email: Lorenzo Mazza <lorenzo.mazza@nct-dresden.de>, Ariel Rodriguez <ariel.rodriguezjimenez@nct-dresden.de>
Maintainer-email: Lorenzo Mazza <lorenzo.mazza@nct-dresden.de>, Ariel Rodriguez <ariel.rodriguezjimenez@nct-dresden.de>
License-Expression: MIT
Project-URL: Repository, https://github.com/nct-tso-robotics/tso_sensorium
Keywords: robotics,imitation-learning,dataset,recording,ros
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy>=1.19.0
Requires-Dist: opencv-python>=4.5.0
Requires-Dist: pandas>=1.3.0
Requires-Dist: scipy>=1.7.0
Requires-Dist: joblib>=1.3.0
Requires-Dist: tqdm>=4.62.0
Requires-Dist: pydantic>=2.7
Requires-Dist: pyyaml-include<2.0,>=1.3
Provides-Extra: lerobot
Requires-Dist: lerobot>=0.4.4; extra == "lerobot"
Provides-Extra: gui
Requires-Dist: flask>=3.0; extra == "gui"
Provides-Extra: test
Requires-Dist: pytest>=8.0; extra == "test"
Dynamic: license-file

# TSO Sensorium

A library for sensor data recording, post-processing, and episodic dataset creation for robot learning.

## Package layout

```
tso_sensorium/
├── recording/     # sensor capture: core + ros1/ (Noetic) and ros2/ (Jazzy) adapters
├── processing/    # stereo rectification, deinterlacing, timestamp alignment
├── episodes/      # schema, episode assembly, parallel dataset building
├── export/        # dataset writers: CSV folders, LeRobot (v3.0 format)
└── scripts/       # command-line entry points
```

Processing, episode assembly, and export have no ROS dependency and run anywhere. Recording adapters import the ROS client libraries and are only usable inside the corresponding ROS environment.

## Installation

### With pixi (recommended)

The preferred workflow is to enter the environment once and run commands
normally inside it. Pixi creates the environment on first use and installs this
package in editable mode automatically. Install [Pixi](https://pixi.sh/) first
if it is not already available.

For ROS 1 development, start a clean terminal and enter the ROS 1 environment:

```bash
pixi shell -e ros1
```

Confirm that Python is using the RoboStack ROS installation from this
repository's Pixi environment:

```bash
python -c "import rospy; print(rospy.__file__)"
```

The path must be under `.pixi/envs/ros1/`, not `/opt/ros/noetic/`. Once inside
the shell, run ROS commands, application commands, and tests directly.

The `tso_testbed` configurations also require the custom messages generated by
the testbed catkin workspace. Add only that workspace's generated Python
packages to the Pixi environment:

```bash
export TESTBED_WORKSPACE=/path/to/robot_testbed
export PYTHONPATH="$TESTBED_WORKSPACE/devel/lib/python3/dist-packages"
python -c "import rospy, testbed_msgs; print(rospy.__file__); print(testbed_msgs.__file__)"
```

`rospy` must still resolve under `.pixi/envs/ros1/`, while `testbed_msgs` must
resolve under the testbed workspace. Do not source the workspace's `setup.bash`
inside the Pixi shell: it also adds `/opt/ros/noetic` and replaces RoboStack's
ROS Python packages with the system installation.

The following commands can then use the real testbed configuration:

```bash
python -m tso_sensorium.scripts.record \
    --config_path configs/recording/tso_testbed.yaml \
    --output_folder "$HOME/tso_sensorium_recordings"
pytest tests/recording/test_ros1_adapter.py
```

Use a separate terminal for each foreground process and run `pixi shell -e
ros1` in each one. Leave an environment with `exit` or Ctrl-D. The other
environments work the same way:

```bash
pixi shell             # core development and non-ROS tests
pixi shell -e ros2     # ROS 2 recording and tests
```

Do not automatically source `/opt/ros/noetic/setup.bash` or a catkin workspace
from `.bashrc`. Those scripts prepend the complete system ROS installation to
`PYTHONPATH`, which can make Pixi's Python import the wrong ROS packages. If the
current terminal has already sourced system ROS, open a clean terminal before
entering the Pixi shell, then add only the generated message path as shown
above.

For scripts and CI, where an interactive shell is inconvenient, use the
equivalent one-shot form:

```bash
pixi run test
pixi run -e ros1 pytest tests/recording/test_ros1_adapter.py
pixi run -e ros2 pytest tests/recording/test_ros2_adapter.py
```

The environments are declared in `pyproject.toml` and locked in `pixi.lock`.
On machines with a small home quota, set `PIXI_CACHE_DIR` to storage with more
space before creating an environment. LeRobot export remains a pip extra to
avoid duplicating a multi-gigabyte torch installation per checkout.

### Alternative installations

#### Pip

```bash
pip install -e .              # core: processing, episodes, export
pip install -e ".[gui]"       # + browser dashboard (Flask)
pip install -e ".[lerobot]"   # + LeRobot dataset export (Python >= 3.10)
pip install -e ".[test]"      # + pytest
```

Python requirements: 3.9+ (the pydantic-backed config layer); LeRobot
export additionally needs 3.10+. Extras combine as usual, e.g.
`pip install -e ".[gui,test]"`. After install the scripts are available
as console commands (`tso-record`, `tso-record-ui`, `tso-annotate`,
`tso-generate-dataset`, `tso-label-phases`, and the `*-ros2` variants),
equivalent to the `python -m tso_sensorium.scripts.*` forms used below.

#### Mamba for ROS recording

For recording without Pixi, ROS client libraries must come from a system ROS
installation or a dedicated [RoboStack](https://robostack.github.io/)
environment because `rospy` and `rclpy` are not available from PyPI. The
environment files under `environments/` provide the latter:

```bash
mamba env create -f environments/ros1-noetic.yml   # or ros2-jazzy.yml
mamba activate tso-sensorium-ros1
pip install -e . --no-deps                          # deps already come from conda
```

Note: the Noetic robot PC runs Python 3.8 and the LeRobot exporter needs
Python 3.10+, so recording and LeRobot export are not meant to share one
environment. Record on the robot, export on the processing machine.

## User guide

### Configuration

Every script is driven by a [pydantic](https://docs.pydantic.dev)-validated
YAML file under `configs/`, selected with `--config_path`. The same rules
apply everywhere:

- **Unknown keys fail loudly.** A typo'd key aborts with a validation error
  naming the key, instead of being silently ignored.
- **Polymorphic entries pick their variant with a `type:` key** — recorder
  kinds (`topic`, `video`), frame transforms (`deinterlace`, `resize`,
  `rectify`), table transforms (`parse_vector3`, `sum_columns`, ...),
  writers (`csv`, `lerobot`), phase labelers and triggers. The fields next
  to `type` are the variant's own options.
- **Configs compose with `!include`**, resolved relative to the including
  file. This keeps shared pieces in one place, e.g.
  `configs/recording/mock_ui.yaml` includes
  `configs/dataset/mock.yaml` as its `generation:` section, and the same
  file is included by `configs/annotation/mock.yaml`.
- **Calibration files packaged with the library** can be referenced as
  `package://<asset name>`.

**CLI overrides**: any field can be overridden with `--dot.path value`
pairs on top of the YAML. Values are parsed as YAML scalars, so `true`,
`3.5`, and `[a, b]` become their typed equivalents. List elements are
addressed by index, and setting a `type` key resets that section to the
new variant (follow-up flags then fill its fields):

```bash
python -m tso_sensorium.scripts.generate_dataset \
    --config_path configs/dataset/bowel_retraction.yaml \
    --recordings_root /data/recordings \
    --save_frames true \
    --videos.0.frame_column framePath \
    --writer.type lerobot --writer.output_root /data/lerobot
```

All modules require Python 3.9+ (pydantic-backed config layer).

### Recording

When the endoscope is mounted on a moving robot, publish its calibrated live
camera orientation from the robot-state topic in a separate ROS 1 Pixi shell:

```bash
python -m tso_sensorium.scripts.publish_robot_camera_transform \
    --subscribe /ur5e_rcm_twist_controller/RobotState \
    --topic /robot_camera_transform \
    --calibration_path configs/calibration/tso_endoscope_mount.yaml
```

The calibration file contains the fixed EE-to-camera axis mapping and optical
tilt. Replace it for a different holder or endoscope geometry. The equivalent
installed command is `tso-publish-camera-transform`.

Record ROS topics and video streams in one shot (inside a ROS 1
environment):

```bash
export TESTBED_WORKSPACE=/path/to/robot_testbed
export PYTHONPATH="$TESTBED_WORKSPACE/devel/lib/python3/dist-packages"
python -m tso_sensorium.scripts.record \
    --config_path configs/recording/tso_testbed.yaml \
    --output_folder "$HOME/tso_sensorium_recordings"
```

For interactive sessions, `record_ui` runs the browser UI (requires the `gui`
extra):

```bash
export TESTBED_WORKSPACE=/path/to/robot_testbed
export PYTHONPATH="$TESTBED_WORKSPACE/devel/lib/python3/dist-packages"
python -m tso_sensorium.scripts.record_ui \
    --config_path configs/recording/tso_testbed_ui.yaml \
    --session.output_folder "$HOME/tso_sensorium_recordings"
```

For the force-sensing setup, the bundled UI config records UR5e state,
the left, right, and combined endoscope streams, the robot-camera transform,
and the Bota wrench and IMU topics. It saves under
`~/tso_sensorium_recordings` by default:

```bash
python -m tso_sensorium.scripts.record_ui \
    --config_path configs/recording/force_session_ui.yaml
```

Video readiness uses an explicitly configured lightweight per-frame status
topic, not raw image subscriptions. Each video entry in the recording YAML
must specify `liveness_topic` when used with the UI; topic names are never
inferred. The bundled setups use `sensor_msgs/CameraInfo` messages.
Recording still subscribes to the configured image topic when an episode
starts and releases that subscription when it stops. The configured live
preview is a separate, intentional image subscription.

Set the status topic and its message type alongside the recorded image topic:

```yaml
- type: video
  file_name: camera
  topic_name: /camera/image_raw
  liveness_topic: /camera/camera_info
  liveness_message_type: sensor_msgs.msg.CameraInfo
```

Other per-frame status messages, such as `std_msgs.msg.Header`, can be selected
with `liveness_message_type`. Command-line recording does not use these fields.
The status topic must publish with each acquired frame; latched calibration
alone does not establish ongoing camera freshness. A missing or stale status
topic keeps the source unready rather than silently falling back to expensive
raw-image subscriptions. If a camera can only provide raw images for status,
explicitly select that image topic and `sensor_msgs.msg.Image`, accepting its
bandwidth cost. This affects readiness only, not the recording format.

Open `http://<host>:8080` from any machine on the network. The dashboard
has two sections: **Record** (live camera feed, per-sensor liveness,
start/stop of demonstrations, per-episode topic selection) and **Library**
(episode browsing with in-browser replay, phase annotation, and dataset
generation — see below). The server is unauthenticated; expose it on
trusted networks only.

To try the full record → browse → annotate → generate loop without any
hardware, first enter `pixi shell -e ros1`, then run the bundled synthetic
sensors:

First check whether the current `ROS_MASTER_URI` already has a running master:

```bash
rostopic list
```

If that succeeds, reuse the existing master. If it fails, start `roscore` in
that terminal. Do not start a second master:

```bash
roscore
```

Open a second terminal, enter `pixi shell -e ros1`, and start the sensors:

```bash
python -m tso_sensorium.scripts.mock_sensors
```

Open a third terminal, enter `pixi shell -e ros1`, and start the dashboard:

```bash
python -m tso_sensorium.scripts.record_ui \
    --config_path configs/recording/mock_ui.yaml
```

ROS may warn when `~/.ros/log` exceeds 1 GB. Inspect it with `rosclean check`;
run `rosclean purge` only if deleting old ROS logs is acceptable.

The same dashboard runs on ROS 2 (`record_ui_ros2`,
`mock_sensors_ros2`) with identical configs — recorder entries reference
message types by dotted path (`std_msgs.msg.Bool`), which resolve to the
ROS 2 classes inside a ROS 2 environment:

```bash
python -m tso_sensorium.scripts.mock_sensors_ros2 &
python -m tso_sensorium.scripts.record_ui_ros2 \
    --config_path configs/recording/mock_ui.yaml
```

### Annotation and dataset studio (no ROS required)

The Library section also runs as a standalone app on any machine — a
processing workstation without ROS, pointed at a folder of recordings:

```bash
python -m tso_sensorium.scripts.annotate \
    --config_path configs/annotation/mock.yaml
# or directly:
python -m tso_sensorium.scripts.annotate \
    --recordings_root /data/recordings --port 8090
```

From the dashboard you can:

- **Switch the recordings folder** being browsed (the path field at the
  top; the episode list, annotations, and legend reload for that folder).
- **Edit the dataset metadata and phase legend**: dataset name, task, and
  the mapping of integer phase labels to a phase name plus its language
  instruction variants. One instruction per phase is deterministic;
  several lines make a stochastic mapping — one variant is sampled per
  episode, seeded by the episode name so regeneration is reproducible.
  The legend is saved to `dataset_metadata.json` at the recordings root.
- **Annotate episodes on a timeline** (pencil icon next to an episode):
  colored phase segments under the video, click to seek and select, drag
  segment edges, split at the playhead, and assign phases from the legend.
  Edits are saved to `annotations.json` inside the episode folder with
  `source: manual`, so re-running an automatic labeler never overwrites
  them.
- **Generate datasets** with per-run options (format, output root, frame
  extraction, sync tolerance); discarded episodes are listed with the
  reason.

### Phase legends

Legends map phase labels to names and instruction variants. Select one in
your dataset-generation config:

```yaml
annotations:
  legend: package://instructions/endoscope_guidance.yaml
  legend_source: config
  language_source: phase_legend
```

`legend` also accepts a YAML filename or an inline metadata mapping.

- `legend_source`: `config` uses the supplied legend; `auto` (default) prefers
  the phase legend in the recordings root's `dataset_metadata.json`, falling
  back to the configured legend when none is saved.
- `language_source`: `phase_legend` samples instructions from the legend;
  `annotation` (default) uses segment-specific text when present, otherwise
  falling back to the legend.

See the [annotation configuration](tso_sensorium/episodes/generation_config.py)
for all fields and defaults.

### Automatic phase labeling

Labelers segment episodes from recorded signals and write automatic
segments into each episode's `annotations.json`:

```bash
python -m tso_sensorium.scripts.label_phases \
    --config_path configs/labeling/bowel_retraction.yaml \
    --recordings_root /data/recordings
```

`configs/labeling/` shows both labeler kinds: `column_threshold` (two
phases split by a signal threshold) and `sequential_trigger` (an ordered
phase sequence where each trigger hands over to the next phase — gripper
state changes, motion starting or settling). The typical workflow is
auto-label → correct visually in the dashboard → generate.

### Dataset generation

```bash
python -m tso_sensorium.scripts.generate_dataset \
    --config_path configs/dataset/bowel_retraction.yaml \
    --recordings_root /data/recordings

# Same recordings, exported as a LeRobot dataset instead:
python -m tso_sensorium.scripts.generate_dataset \
    --config_path configs/dataset/bowel_retraction.yaml \
    --recordings_root /data/recordings \
    --writer.type lerobot --writer.output_root /data/lerobot

# Recompute transformed actions without re-encoding observations or videos:
python -m tso_sensorium.scripts.generate_dataset \
    --config_path configs/dataset/bowel_retraction.yaml \
    --recordings_root /data/recordings \
    --writer.type lerobot_action_update \
    --writer.dataset_root /data/lerobot
```

When the config has an `annotations:` section, each episode's phase
segments are joined onto the aligned table as an integer phase column and
a language instruction column. Every writer stores generation statistics
and any configured `schema.coordinate_frame_features` in
`dataset_metadata.json`. Coordinate-frame entries list their exact component
columns, named frame, and whether that frame is fixed, moving, or unknown
across timesteps. LeRobot exports can also carry named
`schema.auxiliary_features`, such as an integer phase label, without folding
them into the policy action vector. Action-only updates require an existing
LeRobot v3 dataset whose episode order, lengths, schema, observations,
auxiliary values, and task strings exactly match the regenerated episodes.
They atomically replace only actions, derived action statistics, and
`dataset_metadata.json`; a validation failure or cancellation leaves the
existing dataset unchanged.

### Library usage

Generate an episode table programmatically:

```python
from tso_sensorium.episodes import EpisodeGenerator

episode = (
    EpisodeGenerator()
    .add_video(
        video_path="episode/left.mp4",
        timestamps_path="episode/left_timestamps.csv",
        sync_col_name="timestamp",
        frames_output_path="episode/frames/left",
        frame_col_name="left_frame",
        save_frames=True,
    )
    .add_state(
        state_data_path="episode/robot_state.csv",
        sync_col_name="timestamp",
        dataset_cols=["x", "y", "z", "roll"],
    )
    .generate_dataset()
    .save_dataset("episode/episode.csv")
)
```

Rectify and deinterlace stereo images:

```python
from tso_sensorium.processing import Rectifier, deinterlace_cv_image

rectifier = Rectifier(calibration_file_path="calibration.yml")
left, right = deinterlace_cv_image(image=interlaced_image)
left, right = rectifier.rectify(left=left, right=right)
```

## Development

Enter the default Pixi shell for framework-independent development:

```bash
pixi shell
pytest
pytest -m ""
ruff format tso_sensorium/ tests/
ruff check tso_sensorium/ tests/
```

ROS adapter tests require the matching environment and skip in the default
environment:

```bash
exit
pixi shell -e ros1
pytest tests/recording/test_ros1_adapter.py
```

For non-interactive use, the equivalents are `pixi run test`, `pixi run
test-all`, `pixi run lint`, and `pixi run -e ros1 pytest
tests/recording/test_ros1_adapter.py`.

CI checks pull requests and `main`, including ROS 1/2 tests and package builds.
Published GitHub releases upload to PyPI after those checks pass.
