bdfc77147d
Signed-off-by: Meng <meng_chen99@163.com> Signed-off-by: Yangshen Deng <yangshen.d@outlook.com> Signed-off-by: Suli Wang <50828270+wsl2000@users.noreply.github.com> Co-authored-by: Meng <meng_chen99@163.com> Co-authored-by: Suli Wang <50828270+wsl2000@users.noreply.github.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""RoboArena dataset transform.
|
|
|
|
RoboArena uses 0-indexed exterior cameras, 3 views total (OXE_DROID embodiment).
|
|
Same stitching layout as DROID — both map to OXE_DROID in DreamZero.
|
|
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from vllm_omni.diffusion.models.dreamzero.transform.base import (
|
|
register_transform,
|
|
)
|
|
from vllm_omni.diffusion.models.dreamzero.transform.droid import (
|
|
DroidTransform,
|
|
)
|
|
|
|
|
|
class RoboArenaTransform(DroidTransform):
|
|
"""Transform for RoboArena dataset.
|
|
|
|
Same embodiment as DROID (OXE_DROID), same stitching and template.
|
|
Only difference: 0-indexed exterior camera keys.
|
|
|
|
RoboArena observation keys (0-indexed):
|
|
observation/exterior_image_0_left → left exterior
|
|
observation/exterior_image_1_left → right exterior
|
|
observation/wrist_image_left → wrist
|
|
|
|
"""
|
|
|
|
IMAGE_KEY_MAP = {
|
|
"observation/exterior_image_0_left": "images/exterior_0",
|
|
"observation/exterior_image_1_left": "images/exterior_1",
|
|
"observation/wrist_image_left": "images/wrist",
|
|
}
|
|
|
|
|
|
register_transform("roboarena", RoboArenaTransform())
|