# AnyData

The new TRI CV/ML unified dataset format and library!

[Design document](https://docs.google.com/document/d/1O8FEbZfrxS5zswaHYGfE9DfGOgMLh-LaNU2zXfA0ags/edit?tab=t.0#heading=h.2p2oc3b1yy9k) | [Dataset status](https://docs.google.com/spreadsheets/d/1xJXuunAccKdf7UW5WEzmN7McC7tAdGlqKTPOp-QbzyI/edit?gid=606832075#gid=606832075)

## S3 dataset storage ##

All our unified and webbed datasets currently live in the `tri-ml-sandbox-16011-us-west-2-datasets` bucket, organized by source format (`videos` vs `frames`):

- **Unified** (= converted from raw): `s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_unified/videos/<DATASET>/` (browse: [AnyFile](https://anyfile.us-west-2.awsinternal.tri.global/browse/s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_unified/videos?sort_by=date_desc))
- **Webbed** (= webdatasetized for streaming): `s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_webbed/videos/<DATASET>/<SPLIT>__<IKEDOCLABSN>_<STORE>_res.../` (browse: [AnyFile](https://anyfile.us-west-2.awsinternal.tri.global/browse/s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_webbed/videos?sort_by=date_desc))

## Installation ##

### Option A: Editable install into an existing environment

Use this when integrating AnyData as a library into another codebase (e.g. [Any4D](https://github.com/TRI-ML/Any4D)) that already has its own environment or container. The default `pyproject.toml` has no version pins, so pip will only install what's missing:

```
git clone --recursive git@github.com:TRI-ML/AnyData.git
cd AnyData
pip install -e .
pip install -e externals/webdataset
pip install -e externals/camviz_internal
pip install -e ".[viz,droid]"   # with optional extras
```

Note: AnyData ships forked externals in `externals/`:
- **webdataset** (`externals/webdataset/`): built into this repo, replaces the stock tar iteration pipeline with multi-camera-aware extraction (threaded parallel reads from multiple camera tar streams per sample).
- **camviz** (`externals/camviz_internal/`): git submodule ([TRI-ML/camviz_internal](https://github.com/TRI-ML/camviz_internal)). 3D dataset visualization (unified + webbed).

With **uv**, both are installed as editable automatically via `[tool.uv.sources]` in `pyproject.toml`, but with `pip` you need the explicit `pip install -e` commands above.

### Option B: Standalone with uv (pinned versions)

Use this when developing AnyData itself or running it standalone. Swap in `pyproject_pinned.toml` for exact dependency versions in an isolated venv:

```
curl -LsSf https://astral.sh/uv/install.sh | sh   # if uv not installed
git clone --recursive git@github.com:TRI-ML/AnyData.git
cd AnyData
mv pyproject.toml pyproject_nopins.toml
cp pyproject_pinned.toml pyproject.toml
uv venv --python 3.11
uv sync
source .venv/bin/activate
```

For visualization tools (includes camviz, matplotlib, pygame, PyOpenGL):
```
uv sync --extra viz
```

Specific converters have their own optional groups (see `pyproject.toml`):

```
uv sync --extra dl3dv10k  # DL3DV10K
uv sync --extra hypersim  # Hypersim

# OmnidataStarter (must be installed manually - requires building from source)
uv pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable"
```

## Downloading and Visualizing ##

If you are in a hurry and want to try something out quickly, just run

```
uv run anydata/inspect/display.py DDAD --subset 2
uv run anydata/inspect/display.py LBM --subset 2
```

This will download 2 episodes of the DDAD or LBM datasets and visualize them using camviz. 
Specifics for each dataset configuration are described in `anydata/inspect/unified.yaml`. 
If you remove `--subset` you will download the entire dataset!

## Dataset Configuration Files ##

```
dataset:                                                        ### Dataset parameters
    name: [Unified]                                             # Name of the dataset
    path: [/data/cv_unified/Spartan_tiny/split_all.json]        # Path to the dataset split 
    cameras: [[scene_right_0,scene_left_0]]                     # Which cameras to consider
    labels: [rgb,extrinsics,intrinsics,depth]                   # Which labels to consider
    length: [30]                                               # Temporal context
    context_stride: [2]                                         # Context stride
    augmentation:                                               ### Augmentations (see some examples below)
        resize: [256,256]                                       # Resize sample to a given resolution
        resize: [-1,256,s16]                                    # Resize sample to a largest dimension of 256, as a multiple of 16
        preserve_depth: True                                    # If true, keep density of valid depth pixels (use when downsampling sparse depth)
```

## Unifying dataset ##

Our official raw datasets are stored in `s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_downloaded`. Taking LBM as an example, first download a few Spartan episodes:

```
uv run anydata/sync/sync_from.py s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_downloaded/spartan_tiny /data/cv_downloaded/spartan_tiny
```

These raw episodes are now stored in `/data/cv_downloaded/spartan_tiny`.
 
You can optionally visualize any episode of the raw spartan data for sanity-check like so:
```
uv run anydata/visualize/vis_spartan.py \
    --episode_path /data/cv_downloaded/spartan_tiny/efs/data/tasks/BimanualPutRedBellPepperInBin/riverway/sim/bc/teleop/2025-01-06T08-58-31-05-00/diffusion_spartan/episode_131/processed/ \
    --output_dir ./validation_videos \
    --max_frames 200
```

Then, unify the dataset! 
This is done by running the converter specific to that dataset (`lbm.py` in this case), and pointing to the raw dataset (usually stored in `/data/cv_downloaded/<DATASET_NAME>`):

```
uv run anydata/converters/lbm.py /data/cv_downloaded/spartan_tiny
```

The unified dataset is now stored in `/data/cv_unified/spartan_tiny`. 
It contains a `split_all.json` file which points to all available sequences.

You can optionally visualize any episode of the LBM unified format (converted from spartan) for sanity-check like so:
```
uv run anydata/visualize/vis_lbm_unified.py \
    --episode_path /data/cv_unified/spartan_tiny/BimanualPutRedBellPepperInBin/riverway/sim/teleop/2025-01-06T08-58-31-05-00/episode_131 \
    --output_dir ./validation_videos \
    --camera scene_right_0 \
    --max_frames 200
```
Next, you can upload this dataset to s3 with:

```
uv run anydata/sync/upload_uni.py spartan_tiny
```
The uploaded dataset will be stored in `s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_unified/spartan_tiny`. 
Alternatively, you can download an already unified dataset using:
```
# Downloads the whole dataset
uv run anydata/sync/download_uni.py spartan_tiny
# Downloads only sequences from an specific split (recommended)
uv run anydata/sync/download_uni.py spartan_tiny/split_all.json     
```

## Webbing dataset ##

Finally, you can now web the dataset! There are two ways to web the dataset, one is on your local instance/machine, the other is using Sagemaker.
For local webbing, this is done using `anydata/webdataset/web_local.py <DATASET>/<RECIPE> <LENGTH>b<STRIDE>` and choosing which dataset + recipe you are going to use, as well as the sequence length and number of contiguous snippets (samples) per tarfile. Each `<RECIPE>` is stored in `anydata/webdataset/configs/<DATASET>.yaml`. For example:
```
uv run anydata/webdataset/web_local.py LBM/spartan 41b4
```

will create a webdataset using the `spartan` recipe in `anydata/webdataset/configs/LBM.yaml`. Each tarfile contains one snippet of 44 frames (41 length + 4 buffer stride - 1), which supports 4 different actual samples of length 41 each, and consecutive snippets start 4 frames apart.

Once the process is done, you can upload the resulting dataset to s3 with:
```
uv run anydata/sync/upload_web.py <PATH/TO/DATASET>
```
For Sagemaker webbing, it first downloads unified data from S3 to the instance, runs `web_local.py`, and uploads the result back to S3 under `cv_webbed`. Create one config file per dataset under `anydata/sagemaker/webdatasetize_config/` with the webdataset recipe, download split, store, and worker settings.


Example sagemaker config for DDAD:
```
datasets: DDAD/split_all.json         ### Link to webdatasetize job configuration in webdataset/configs folder
s3_download_data: DDAD/split_all.json ### This is the json file that contains the list of data samples to be downloaded from s3. It should be the same as the one used in the webdatasetize job configuration.
store: 24                             ### pass store parameter to webbing job
official: true                        ### pass official parameter to webbing job
download_num_procs: 16                ### num_procs used to download data from s3 to sagemaker instance.
num_procs: 4                          ### pass num_procs parameter to webbing job
```


Then submit a single SageMaker webdatasetize job. The second argument is the job name, and the fifth argument is the username used for the SageMaker/ECR/S3 namespace. Usually `g6e-small` is enough for most datasets to get rid of the queue for p5 instances, but with larger datasets or creating webbing data with latent, `p5en` can also be used with an AWS Batch queue name provided, e.g., `cv-p5en`. For frame-only webbing, `r7i` (mapped to `ml.r7i.12xlarge`) is also supported as a CPU-only option. For `mode: latents` or `mode: both`, prefer GPU instances because the latent tokenizer path is CUDA-oriented and CPU execution will be much slower and may not be reliable.

```
bash anydata/sagemaker/run_sagemaker_webdatasetize.sh \
  <dataset> <job-name> [instance_type] [queue] <username> [split_json]
```

Example:
```
bash anydata/sagemaker/run_sagemaker_webdatasetize.sh DDAD fzc-DDAD g6e-small none fzc
```

The launcher also handles Docker automatically. It builds the SageMaker container from `anydata/sagemaker/Dockerfile_280_data`, tags it as `<user>-anydata-280:latest`, and then pushes it to AWS ECR in `us-west-2`. By default `BUILD_TYPE=full`; set `BUILD_TYPE=none` to reuse the existing image.

To distribute one large webbing job across multiple SageMaker instances, submit split jobs with one command:
```
bash anydata/sagemaker/run_sagemaker_webdatasetize_splits.sh \
  DDAD <username> g6e-small none 8
```
The split launcher signature is:
```
bash anydata/sagemaker/run_sagemaker_webdatasetize_splits.sh \
  <dataset> <username> [instance_type] [queue] [num_splits] [first_build_type] [src_split_json]
```

It reads `mode` / `from_mode` from `webdatasetize_config/<dataset>.yaml`, generates `sm_split_0.json` ... `sm_split_<N-1>.json`, uploads them to `s3://.../cv_unified/<from_mode>/<dataset>/`, and submits one SageMaker job per split using the provided username for job names, ECR image names, and S3 output paths. By default, `first_build_type=full` and `REST_BUILD_TYPE=none`, so the first job builds/pushes the image and the rest reuse it. Pass `none` as the sixth argument to skip the first build too:
```
bash anydata/sagemaker/run_sagemaker_webdatasetize_splits.sh \
  DDAD fzc r7i none 4 full
```

The optional seventh argument overrides the source split JSON. Use this when the input split is not the default `s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_unified/<from_mode>/<dataset>/split_all.json`:
```
bash anydata/sagemaker/run_sagemaker_webdatasetize_splits.sh \
  DDAD fzc g6e-small none 4 none \
  s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_unified/frames/DDAD/split_all.json
```

Each split job downloads its corresponding `sm_split_<i>.json` and writes the matching outputs `split_all_<i>.json` and `stats_all_<i>.txt`. After all jobs finish, merge them into consolidated `split_all.json` and `stats_all.txt`:
```
python anydata/sagemaker/merge_webdatasetize_split_outputs.py \
  <output-prefix> --num_splits <N> --delete_shards
```


## Inspecting a Unified dataset ##

You can inspect the dataset by instantiating it and printing batch information:

```
uv run anydata/inspect/batch.py spartan_tiny
```

Alternatively, you can visualize the batch using camviz with:

```
uv run anydata/inspect/display.py spartan_tiny
```

The `spartan` suffix points to a particular configuration inside `anydata/inspect/unified.yaml`.

## Splitting dataset

You can create a `split_all` split for an unified dataset with:

```
python anydata/converters/misc/create_split.py <DATASET> --official (not debug) --upload (upload to s3) 
```

If the dataset is not available locally, you can download only the metadata necessary with:

```
python anydata/converters/misc/download_metadata.py <DATASET>
```

Alternatively, you can create splits with specific filtering criteria with the following options:
```
python anydata/converters/misc/create_split.py <DATASET> --name <NAME_OF_THE_SPLIT> 
--num_frames X (e.g., 40+ for 40 frames or more, 80- for 80 frames or less, or 50 for exactly 50 frames)
--num_cameras X (e.g., 3+ for 3 cameras or more, 5- for 5 cameras or less, or 2 for exactly 2 cameras)
--with_cameras cam1 cam2 ... (only select sequences with those cameras)
--without_cameras cam1 cam2 ... (only select sequences without those cameras)
--with_labels lab1 lab2 ... (only select sequences with those labels)
--without_labels lab1 lab2 ... (only select sequences without those labels)
--with_tags cam1 cam2 ... (only select sequences with those tags)
--without_tags cam1 cam2 ... (only select sequences without those tags)
--resolution height width (only select sequences with that resolution across all cameras)
```
More filtering conditions can be added easily on the `create_split.py` file 

## Webdataset naming

Webdataset splits will often contain some subset of the letters `IKEDOCLABSN` to denote available modalities:

| Code | Description        |
|------|-------------------|
| I    | Image (RGB)       |
| K    | Intrinsics        |
| E    | Extrinsics        |
| D    | Depth             |
| O    | Optical flow      |
| C    | Scene flow        |
| L    | Language          |
| A    | Action            |
| B    | Bounding Box      |
| S    | Semantic          |
| N    | Surface normals   |

For example, for `s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_webbed/frames/DROID/DROID101__IKELA_41b8_res[-16,384]`:
- `IKELA` means it contains images, intrinsics, extrinsics, language, actions.
- `41b8` means the tarfiles contain snippets of length 48, because it supports video clips of 41 frames, and the shifting window stride between tarfiles is 8.
- `res[-16,384]` means images were resized to a maximum dimension of 384, with the other dimension rounded to the closest multiple of 16 (to support easy video tokenization and patchification without having to do weird padding tricks) that maximally respects the original aspect ratio.

## Contribution & Coding Standards

AnyData is a shared library used across multiple projects and divisions. Changes here may have a wide blast radius, so please be careful and follow these guidelines / best practices:

* **Clean & modular**: Write readable, self-contained functions. Keep data processing steps composable and independently testable.
* **Design doc consistency**: Keep the [design document](https://docs.google.com/document/d/1O8FEbZfrxS5zswaHYGfE9DfGOgMLh-LaNU2zXfA0ags/edit?tab=t.0#heading=h.2p2oc3b1yy9k) up-to-date and flag any major outdated or missing aspects.
* **Backward compatibility**: New features must not change default behavior or break existing callers.
Prefer using optional parameters that default to the prior behavior to avoid surprising other users.
* **Defensive coding**: Check for `None` before accessing attributes, validate dict keys exist, etc.
Avoid fragile if-checks that depend on low-level implementation details or string matching.
Code should fail gracefully and print clear warning or error messages rather than swallowing exceptions quietly.
* **No hardcoded paths**: Paths must come from CLI arguments, config files, or environment variables.
* **Documentation**: Brief but precise. Docstrings for APIs; inline comments where the logic is non-obvious. Example YAML config fields should have a comment documenting their meaning and defaults.
* **Branch hygiene**: When merging or rebasing, please reconcile changes carefully.
* **Robustness**: Long-running scripts should be robust to occasional per-sample failures: catch, log with traceback, skip, and continue. Print progress verbosely.
* **Testing**: At minimum, verify that changes do not break core functionality (e.g. unifying and webbing of at least one dataset, or construction and iteration of at least one dataloader).
* **PR reviewing**: Please habitually add Copilot as reviewer to every PR, and do not merge PRs if they have not been reviewed yet by Copilot. It tends to catch a lot of pertinent bugs and issues. I would also recommend emailing IT to ask to get yourself added to the [TRI-Tools](https://github.com/TRI-Tools) org, such that Copilot code reviews will be triggered automatically.
