Skip to content

[python] Batch TorchCodec frame decoding for LeRobot datasets - #9836

Open
QuakeWang wants to merge 3 commits into
apache:masterfrom
QuakeWang:lerobot-batch-decode
Open

QuakeWang wants to merge 3 commits into
apache:masterfrom
QuakeWang:lerobot-batch-decode

Conversation

@QuakeWang

Copy link
Copy Markdown
Member

Purpose

LeRobot already groups video frame requests by payload and reuses decoders, but still decodes each frame individually.

Use TorchCodec's get_frames_at(...).data once per full payload group within a read batch. Preserve sample order, bounded decoder caching, and existing PyAV/custom callback behavior.

Tests

  • Related suites: 61 passed, 31 skipped due to missing optional dependencies.
  • Verified real MP4 batch/single-frame parity, delta windows, payload isolation, cache reuse, and backend fallback.
  • Flake8 4.0.1, license headers, and git diff --check passed.

Decode each payload group with TorchCodec get_frames_at instead of repeated single-frame calls. Preserve sample order, bounded decoder reuse, and PyAV/custom callback behavior.

Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com>
output_column="frame",
max_open_videos=8,
collate_fn=None):
collate_fn=None,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Maybe we can make decode_fn optional when decode_batch_fn is provided?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Thanks, updated decode_fn to default to None and require at least one of the two callbacks. Removed the redundant single-frame callback from LeRobot and added tests for batch-only usage and invalid arguments.

Make decode_fn optional when decode_batch_fn is provided, while validating that at least one callable is supplied. Remove the redundant LeRobot single-frame callback and cover missing and invalid callbacks.

Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com>
@XiaoHongbo-Hope

Copy link
Copy Markdown
Contributor

+1

Resolve the LeRobot test import conflict while preserving camera parallelism and TorchCodec batch decoding.

Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants