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.

[imstevenpmwork]

Thanks for the contribution.

I wonder if copilot suggestions makes sense. It seems that we could call datasets.Dataset.from_parquet() with the features if we have them already available. WDYT ?

Some other notes for the future -unrelated to this PR-:

• aggregate hot path seems to be changing frequently between panda frames, parquet files and dataset representations, probably not ideal.
• I wonder if we should do this also for non-image datasets.

Thanks again!

[imstevenpmwork]

LGTM, thanks!