简介

LeRobot v3.0 数据集图像观测的 ROI 裁剪算子，按给定或自动推导的 ROI 裁剪每帧图像并 resize，产出一个新的 v3.0 数据集（源数据集只读）。移植 lerobot rl/crop_dataset_roi.py 的 convert_lerobot_dataset_to_cropped_lerobot_dataset，把上游的 cv2 窗口交互选框换成非交互配置。

功能描述

• 两种 ROI 模式：
• ◦manual：ROI 来自 crop_params，形如 {image_key: (top, left, height, width)}
• ◦center_square：对首帧里键名含 image 且形状为 3 维（C, H, W）的每个观测键，自动推导居中正方形 ROI（边长 min(H, W)，偏移 (H-side)//2 / (W-side)//2），是上游 select_square_roi_for_images 的非交互替代
• 裁剪后统一 resize 到 resize_size（(height, width)），并把 meta/info.json 里对应特征的 shape 同步改成 [3, height, width]
• 逐帧走 lerobot writer API（LeRobotDataset.create + add_frame + save_episode），episode 按 episode_index 变化切分；输出 repo_id 为 <源 repo_id>_cropped，fps / robot_type / features 沿用源数据集，源有 video 特征时输出也写视频
• task 保留：task=None（默认）时按 task_index 逐帧回填原 task 文本（上游强制所有帧写同一个字符串）；显式给 task 则整体覆盖
• dtype 归一：dataset[idx] 把特征统一成 float32 张量，写回时按目标 features 声明的 dtype（float32 / float64）转回；裁剪 resize 后的图像 clamp 到 [0, 1]
• 非图像特征（state / action 等）原样透传；task_index / timestamp / episode_index / frame_index / index / task 这些由 writer 自己生成的列不复制
• 输入校验：源目录必须是完成态 v3.0（meta/info.json 的 codebase_version == v3.0）；manual 模式缺 crop_params、ROI 的 top/left 为负或 height/width 非正都报错；center_square 一个图像键都没命中时报错
• 幂等与崩溃安全：输出已是完成态 v3.0 且 force=False 时跳过并返回 SKIPPED；转换中途失败会删掉半成品输出目录；写完再核一次输出是不是 v3.0，不是就报错
• 行级容错：dataset_path / output_path 缺失或转换抛异常都返回 Failed: ...，不打断整批

算子参数

输入

输出

参数

注意事项

• crop_params 的键必须与 meta/info.json 的特征名完全一致（如 observation.images.cam）。写错的键不会报错，也不会裁剪，只是静默不生效。
• ROI 只校验 top/left ≥ 0、height/width > 0，不校验是否落在源图像内；超界部分由 torchvision F.crop 处理，不会报错。
• 输出目录会被清空重建：只要没走 SKIPPED 分支，已存在的 output_path 会被整个 rmtree 掉，不要把输出指向还有别的数据的目录。
• 图像 shape 元数据只对 crop_params / 自动推导命中的键更新。手工改动 ROI 或 resize 尺寸后必须让算子重跑（force=True），否则元数据与实际帧尺寸会不一致。
• 全量逐帧解码 + 逐帧重编码，耗时随总帧数线性增长；一行任务处理一个数据集，并发靠行数（batch_size=1），不要指望单行内并行。
• state / action 不随裁剪变化，下游训练配置需要自己对齐新的图像尺寸。

调用示例

1from __future__ import annotations
2
3import os
4
5import daft
6from daft import col
7
8from daft.aihc.common.udf import aihc_udf
9from daft.aihc.functions.embodied.lerobot.crop_dataset_roi_udf import CropDatasetROI
10
11if __name__ == "__main__":
12    if os.getenv("DAFT_RUNNER", "native") == "ray":
13        import ray
14        ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
15        daft.set_runner_ray()
16    daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
17
18samples = {
19        # 源目录必须是完成态 v3.0 数据集（meta/info.json 的 codebase_version == v3.0）
20        "dataset_path": ["/path/to/lerobot_dataset"],
21        "output_path": ["/tmp/lerobot_crop_out"],
22    }
23    ds = daft.from_pydict(samples)
24    ds = ds.with_column(
25        "result",
26        aihc_udf(
27            CropDatasetROI,
28            construct_args={
29                # 手动 ROI：{"roi_mode": "manual",
30                #            "crop_params": {"observation.images.cam": (0, 0, 24, 32)},
31                #            "resize_size": (16, 16)}
32                "roi_mode": "center_square",
33                "resize_size": (32, 32),
34                "task": None,
35                "force": False,
36            },
37            num_cpus=2,
38            concurrency=1,
39            batch_size=1,
40        )(col("dataset_path"), col("output_path")),
41    )
42    ds.show()