Skip to content

feat(datasets): improve image transform support#2885

Merged
CarolinePascal merged 4 commits into
huggingface:mainfrom
reeceomahoney:feat/improve-transform-support
Feb 5, 2026
Merged

feat(datasets): improve image transform support#2885
CarolinePascal merged 4 commits into
huggingface:mainfrom
reeceomahoney:feat/improve-transform-support

Conversation

@reeceomahoney
Copy link
Copy Markdown
Contributor

Type / Scope

  • Type: Feature | Test
  • Scope: Datasets

Summary / Motivation

  • The motivation for this PR was wanting to add a resize transform to my training pipeline. The LeRobotDataset class supports all torchvision transforms but this was restricted to a hardcoded list in the ImageTransformsConfig

Related issues

  • Fixes / Closes:
  • Related:

What changed

  • Updated the make_transform_from_config function to accept any valid torchvision v2 transform rather than just the hardcoded list of options that existed
  • Added 2 new tests to support this new feature

How was this tested (or how to run locally)

pytest -sv tests/ 
lerobot-train \
  --dataset.repo_id=lerobot/pusht \
  --policy.type=act \
  --job_name=test \
  --policy.device=cuda \
  --wandb.enable=false \
  --policy.repo_id=${HF_USER}/my_policy \
  --dataset.image_transforms.enable=true \
  --dataset.image_transforms.tfs={"resize": {"weight": 1.0, "type": "Resize", "kwargs": {"size": [224, 224]}}}

Checklist (required before merge)

  • Linting/formatting run (pre-commit run -a)
  • All tests pass locally (pytest)
  • Documentation updated
  • CI is green

Reviewer notes

Copilot AI review requested due to automatic review settings January 30, 2026 12:10
@github-actions github-actions Bot added dataset Issues regarding data inputs, processing, or datasets tests Problems with test coverage, failures, or improvements to testing labels Jan 30, 2026
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR enhances the image transform support in the LeRobot datasets module by allowing any torchvision v2 transform to be used, replacing the previously hardcoded list of supported transforms.

Changes:

  • Refactored make_transform_from_config to dynamically load any valid torchvision v2 transform using hasattr and getattr
  • Added tests for the new dynamic transform loading feature (Resize transform) and error handling for invalid transform names

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File Description
src/lerobot/datasets/transforms.py Replaced hardcoded transform type checks with dynamic attribute lookup in torchvision.transforms.v2, while maintaining special handling for SharpnessJitter
tests/datasets/test_image_transforms.py Added two new tests: one for validating a new v2 transform (Resize) and one for error handling with invalid transform types

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

Comment thread src/lerobot/datasets/transforms.py Outdated
@CarolinePascal
Copy link
Copy Markdown
Collaborator

Hi @reeceomahoney,

Thank you for your contribution ! Before merging your PR, would It be possible to tackle these two issues :

  • As mentioned by Copilot, the hasattr is a bit too permissive. We would like to check if the return value of getattr is indeed a Transform. Calling : issubclass(my_transform,Transform) should do.
  • I'm wondering if could add another test along the Resize, simply testing Identity could be enough.

Once done, feel free to ping me for the merge !

Best,

Caroline.

@reeceomahoney reeceomahoney force-pushed the feat/improve-transform-support branch from 936acea to 30d108a Compare February 4, 2026 12:17
@reeceomahoney
Copy link
Copy Markdown
Contributor Author

@CarolinePascal these changes have been made

@CarolinePascal CarolinePascal self-requested a review February 4, 2026 20:08
CarolinePascal
CarolinePascal previously approved these changes Feb 4, 2026
Copy link
Copy Markdown
Collaborator

@CarolinePascal CarolinePascal left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me, thanks !

Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

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


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

Comment thread src/lerobot/datasets/transforms.py Outdated
Comment thread tests/datasets/test_image_transforms.py
@CarolinePascal
Copy link
Copy Markdown
Collaborator

Hi @reeceomahoney,

I left a comment under Copilot's review ;)

@reeceomahoney
Copy link
Copy Markdown
Contributor Author

reeceomahoney commented Feb 5, 2026

@CarolinePascal fixed, have removed the need for the hasattr check too

@CarolinePascal CarolinePascal self-requested a review February 5, 2026 12:52
Copy link
Copy Markdown
Collaborator

@CarolinePascal CarolinePascal left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Awesome, let's merge !

@CarolinePascal CarolinePascal merged commit 97e7e0f into huggingface:main Feb 5, 2026
8 checks passed
@reeceomahoney reeceomahoney deleted the feat/improve-transform-support branch February 5, 2026 16:10
Zaraxxx pushed a commit to Zaraxxx/lerobot that referenced this pull request Feb 20, 2026
* improve image transform support

* add tests

* Add stricter transform check and extra test

* improve subclass check
xianglunkai pushed a commit to xianglunkai/lerobot that referenced this pull request Feb 28, 2026
* improve image transform support

* add tests

* Add stricter transform check and extra test

* improve subclass check
lu391see pushed a commit to lu391see/lerobot_tactile that referenced this pull request Mar 24, 2026
* improve image transform support

* add tests

* Add stricter transform check and extra test

* improve subclass check
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

dataset Issues regarding data inputs, processing, or datasets tests Problems with test coverage, failures, or improvements to testing

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants