Fixes aggregation of image datasets

[fracapuano]

Title

Fixes aggregation of image datasets

Type / Scope

• Type: Bug
• Scope: Bug fix for aggregate_datasets. Also affects lerobot-edit-dataset

Summary / Motivation

• Fixes aggregate_datasets loses Image feature schema for image datasets #2715. Also, adds tests to ensure this edge case is properly covered.
• As a side note: I'd advocate for removing image datasets altogether :)

Related issues

• Fixes / Closes: aggregate_datasets loses Image feature schema for image datasets #2715

What changed

• Short, concrete bullets of the modifications (files/behaviour).
• Short note if this introduces breaking changes and migration steps.

How was this tested

• Tests added: list new tests or test files.

How to run locally (reviewer)

• Run the relevant tests:

  pytest tests/datasets/test_aggregate.py::test_aggregate_image_datasets

Run these tests to confirm no breaking changes on closely related parts of the library.

pytest tests/datasets/test_aggregate.py
========================================== test session starts ===========================================
platform darwin -- Python 3.10.13, pytest-8.4.1, pluggy-1.6.0
rootdir: /Users/fracapuano/Desktop/personal/lerobot
configfile: pyproject.toml
plugins: timeout-2.4.0, anyio-4.10.0, cov-6.2.1, mock-serial-0.0.1, hydra-core-1.3.2
collected 4 items                                                                                        

Testing with DEVICE='cpu'

tests/datasets/test_aggregate.py ....                                                              [100%]

=========================================== 4 passed in 15.24s ===========================================
(lerobot) ➜  lerobot git:(fix/fracapuano-aggregate-image-datasets) ✗ pytest tests/datasets/test_dataset_tools.py
========================================== test session starts ===========================================
platform darwin -- Python 3.10.13, pytest-8.4.1, pluggy-1.6.0
rootdir: /Users/fracapuano/Desktop/personal/lerobot
configfile: pyproject.toml
plugins: timeout-2.4.0, anyio-4.10.0, cov-6.2.1, mock-serial-0.0.1, hydra-core-1.3.2
collected 40 items                                                                                       

Testing with DEVICE='cpu'

tests/datasets/test_dataset_tools.py ........................................                      [100%]

========================================== 40 passed in 44.80s ===========================================

Checklist (required before merge)

• Linting/formatting run (pre-commit run -a)
• All tests pass locally (pytest) (run test_aggregate.py, test_dataset_tools.py)
• NA [] Documentation updated
• CI is green

Reviewer notes

• Anything the reviewer should focus on (performance, edge-cases, specific files) or general notes.
• Anyone in the community is free to review the PR.

[Copilot]

Pull request overview

This PR fixes a bug in image dataset aggregation where HuggingFace Image() feature types were being lost during the aggregation process, causing images to be stored with generic struct schemas instead. The fix ensures proper preservation of image schemas by passing feature metadata through the aggregation pipeline.

Key changes:

• Modified aggregate_data() and append_or_create_parquet_file() to retrieve and pass HuggingFace features schema for image datasets
• Added special handling for reading and writing parquet files containing images using datasets.Dataset.from_parquet() to preserve image format
• Added comprehensive test coverage with test_aggregate_image_datasets() to verify schema preservation and data integrity

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
tests/datasets/test_aggregate.py	Added comprehensive test for image dataset aggregation including schema validation and data integrity checks
src/lerobot/datasets/aggregate.py	Updated aggregation logic to retrieve and pass HuggingFace features schema when processing image datasets, ensuring proper Image() type preservation in parquet files

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The function to_parquet_with_hf_images is being called with a features parameter, but the current function signature in utils.py only accepts (df: pandas.DataFrame, path: Path) and does not have a features parameter. This will cause a TypeError at runtime. The function signature needs to be updated to accept and use the features parameter to properly preserve the HuggingFace Image schema.

[Copilot]

The function to_parquet_with_hf_images is being called with a features parameter, but the current function signature in utils.py only accepts (df: pandas.DataFrame, path: Path) and does not have a features parameter. This will cause a TypeError at runtime. The function signature needs to be updated to accept and use the features parameter to properly preserve the HuggingFace Image schema.

[Copilot]

When reading image datasets using datasets.Dataset.from_parquet, the features parameter should be passed to ensure image columns are properly loaded with the correct schema. Without this, the image data might not be correctly preserved during the read-update-write cycle. Consider using datasets.Dataset.from_parquet(str(src_path), features=hf_features) to maintain schema consistency.

Suggested change

	src_ds = datasets.Dataset.from_parquet(str(src_path))
	src_ds = datasets.Dataset.from_parquet(str(src_path), features=hf_features)

[Copilot]

When reading existing image datasets using datasets.Dataset.from_parquet, the features parameter should be passed to ensure image columns are properly loaded with the correct schema. Without this, the image data might not be correctly preserved during the read-merge-write cycle. Consider using datasets.Dataset.from_parquet(str(dst_path), features=hf_features) to maintain schema consistency.

Suggested change

	existing_ds = datasets.Dataset.from_parquet(str(dst_path))
	existing_ds = datasets.Dataset.from_parquet(str(dst_path), features=hf_features)