95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
|
from typing import List, Dict, Optional
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from qwen_vl_utils.vision_process import MIN_PIXELS, MAX_PIXELS, IMAGE_FACTOR
|
||
|
|
|
||
|
|
|
||
|
|
# Tactile sensor file mapping for data processing
|
||
|
|
TACTILE_FILE_MAPPING = {
|
||
|
|
"tactile_data_left": "left_tactile",
|
||
|
|
"tactile_data_right": "right_tactile"
|
||
|
|
}
|
||
|
|
|
||
|
|
# Supported action datasets
|
||
|
|
ACTION_DATASET_NAMES = [
|
||
|
|
"x2_normal", "agibotworld_alpha", "droid", "fractal", "bridge_data_v2",
|
||
|
|
"DobbE", "RH20T", "UMI-biarm", "austin_buds", "austin_sailor", "austin_sirius",
|
||
|
|
"bc_z", "berkeley_autolab_ur5", "berkeley_cable_routing", "berkeley_fanuc_manipulation",
|
||
|
|
"dlr_edan_shared_control", "fmb", "furniture_bench", "jaco_play", "nyu_rot",
|
||
|
|
"stanford_hydra", "stanford_kuka_multimodal", "taco_play", "utaustin_mutex", "viola"
|
||
|
|
]
|
||
|
|
|
||
|
|
# Supported multimodal datasets
|
||
|
|
MULTIMODAL_DATASET_NAMES = [
|
||
|
|
"x2_multimodal_from_action", "x2_multimodal", "x2_subtask_generation",
|
||
|
|
"multimodal_CapsFusion", "multimodal_Robo2VLM", "multimodal_RoboPoint",
|
||
|
|
"multimodal_EQA", "multimodal_Cambrian", "multimodal_pixmo",
|
||
|
|
"multimodal_VQAv2", "multimodal_COCO"
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class X2RDataProcessingConfig:
|
||
|
|
"""Configuration class for X2R data processing pipeline.
|
||
|
|
|
||
|
|
This class contains all the necessary parameters for processing robotic data
|
||
|
|
including camera mappings, tactile sensor configurations, action predictions,
|
||
|
|
and various processing options.
|
||
|
|
"""
|
||
|
|
|
||
|
|
# Action prediction configuration
|
||
|
|
predict_action_keys: List[str] = field(default_factory=list)
|
||
|
|
obs_action_keys: List[str] = field(default_factory=list)
|
||
|
|
|
||
|
|
# Image resolution settings for different views
|
||
|
|
resolution: Dict[str, int] = field(
|
||
|
|
default_factory=lambda: {
|
||
|
|
"face_view": -1,
|
||
|
|
"left_wrist_view": 128,
|
||
|
|
"right_wrist_view": 128
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
# Dataset splitting
|
||
|
|
train_test_split: float = 0.9
|
||
|
|
split_seed: int = 42
|
||
|
|
|
||
|
|
# Instruction handling
|
||
|
|
priority_order: Optional[Dict[str, float]] = None
|
||
|
|
|
||
|
|
# Vision model parameters
|
||
|
|
model_type: str = "qwen2_5"
|
||
|
|
max_pixels: int = MAX_PIXELS
|
||
|
|
min_pixels: int = MIN_PIXELS
|
||
|
|
image_factor: int = IMAGE_FACTOR
|
||
|
|
|
||
|
|
generate_subtask_ratio: float = 0.0
|
||
|
|
|
||
|
|
def __post_init__(self):
|
||
|
|
"""Post-initialization validation and setup."""
|
||
|
|
# Validate train/test split
|
||
|
|
if not 0 < self.train_test_split < 1:
|
||
|
|
raise ValueError(f"train_test_split must be between 0 and 1, got {self.train_test_split}")
|
||
|
|
|
||
|
|
def as_dict(self) -> Dict:
|
||
|
|
"""Convert configuration to dictionary format.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Dict: Configuration as dictionary
|
||
|
|
"""
|
||
|
|
return self.__dict__
|
||
|
|
|
||
|
|
def update(self, **kwargs) -> 'X2RDataProcessingConfig':
|
||
|
|
"""Update configuration parameters.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
**kwargs: Key-value pairs to update
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
X2RDataProcessingConfig: Updated configuration instance
|
||
|
|
"""
|
||
|
|
for key, value in kwargs.items():
|
||
|
|
if hasattr(self, key):
|
||
|
|
setattr(self, key, value)
|
||
|
|
else:
|
||
|
|
raise ValueError(f"Unknown configuration parameter: {key}")
|
||
|
|
return self
|