diff --git a/src/lerobot/policies/groot/processor_groot.py b/src/lerobot/policies/groot/processor_groot.py index a02049a2b..5856f5ff1 100644 --- a/src/lerobot/policies/groot/processor_groot.py +++ b/src/lerobot/policies/groot/processor_groot.py @@ -143,6 +143,7 @@ class _GrootN17CheckpointProcessorAssets: shortest_image_edge: int | None crop_fraction: float | None use_albumentations: bool + letter_box_transform: bool @dataclass(frozen=True) @@ -199,6 +200,9 @@ def _load_n1_7_checkpoint_processor_assets(config: GrootConfig) -> _GrootN17Chec use_albumentations = processor_kwargs.get("use_albumentations", False) if not isinstance(use_albumentations, bool): use_albumentations = False + letter_box_transform = processor_kwargs.get("letter_box_transform", False) + if not isinstance(letter_box_transform, bool): + letter_box_transform = False valid_action_horizon = _load_n1_7_checkpoint_action_horizon(processor_kwargs, config.embodiment_tag) video_horizon = _load_n1_7_checkpoint_video_horizon(processor_kwargs, config.embodiment_tag) @@ -225,6 +229,7 @@ def _load_n1_7_checkpoint_processor_assets(config: GrootConfig) -> _GrootN17Chec shortest_image_edge=as_optional_int(processor_kwargs.get("shortest_image_edge")), crop_fraction=as_optional_float(processor_kwargs.get("crop_fraction")), use_albumentations=use_albumentations, + letter_box_transform=letter_box_transform, ) @@ -1058,6 +1063,7 @@ def _build_n1_7_relative_action_processor_assets( shortest_image_edge=base_assets.shortest_image_edge if base_assets is not None else None, crop_fraction=base_assets.crop_fraction if base_assets is not None else None, use_albumentations=base_assets.use_albumentations if base_assets is not None else False, + letter_box_transform=base_assets.letter_box_transform if base_assets is not None else False, ) @@ -1179,6 +1185,7 @@ def make_groot_pre_post_processors( shortest_image_edge = None crop_fraction = None use_albumentations = checkpoint_assets.use_albumentations if checkpoint_assets is not None else False + letter_box_transform = checkpoint_assets.letter_box_transform if checkpoint_assets is not None else False input_steps: list[ProcessorStep] = [ RenameObservationsProcessorStep(rename_map={}), @@ -1191,6 +1198,7 @@ def make_groot_pre_post_processors( shortest_image_edge=shortest_image_edge, crop_fraction=crop_fraction, use_albumentations=use_albumentations, + letter_box_transform=letter_box_transform, device=config.device, ), DeviceProcessorStep(device=config.device), @@ -1315,6 +1323,7 @@ def _transform_n1_7_image_for_vlm_albumentations( image_target_size: list[int] | None, shortest_image_edge: int | None, crop_fraction: float | None, + letter_box_transform: bool = False, ) -> np.ndarray: """cv2/INTER_AREA eval transform mirroring Isaac-GR00T's albumentations preprocessing. @@ -1339,6 +1348,18 @@ def _transform_n1_7_image_for_vlm_albumentations( if not image_np.flags.c_contiguous: image_np = np.ascontiguousarray(image_np) + if letter_box_transform: + height, width = image_np.shape[:2] + if height != width: + square_edge = max(height, width) + pad_h = square_edge - height + pad_w = square_edge - width + top = pad_h // 2 + bottom = pad_h - top + left = pad_w // 2 + right = pad_w - left + image_np = cv2.copyMakeBorder(image_np, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0) + resize_edge = shortest_image_edge or target_h def resize_shortest_edge(frame: np.ndarray) -> np.ndarray: @@ -1377,9 +1398,12 @@ def _transform_n1_7_image_for_vlm_torch( image_target_size: list[int] | None, shortest_image_edge: int | None, crop_fraction: float | None, + letter_box_transform: bool = False, ) -> torch.Tensor: - """Default (non-albumentations) N1.7 image transform: pad-to-square, resize to - ``shortest_image_edge``, center-crop by ``crop_fraction``, resize to ``image_target_size``. + """Default (non-albumentations) N1.7 image transform. + + Optionally pads to square, then resizes to ``shortest_image_edge``, center-crops + by ``crop_fraction``, and resizes to ``image_target_size``. Operates on a ``(C, H, W)`` uint8 tensor and keeps the result on the input tensor's device so the resize/crop run on GPU when the tensor is. Bicubic @@ -1394,13 +1418,14 @@ def _transform_n1_7_image_for_vlm_torch( target_h, target_w = image_target_size _, height, width = image.shape - square_edge = max(height, width) - if height != width: - left = (square_edge - width) // 2 - top = (square_edge - height) // 2 - image = tv_functional.pad( - image, [left, top, square_edge - width - left, square_edge - height - top], fill=0 - ) + if letter_box_transform: + square_edge = max(height, width) + if height != width: + left = (square_edge - width) // 2 + top = (square_edge - height) // 2 + image = tv_functional.pad( + image, [left, top, square_edge - width - left, square_edge - height - top], fill=0 + ) resize_edge = shortest_image_edge or target_h image = tv_functional.resize( @@ -1945,6 +1970,7 @@ class GrootN17VLMEncodeStep(ProcessorStep): shortest_image_edge: int | None = None crop_fraction: float | None = None use_albumentations: bool = False + letter_box_transform: bool = False device: str | None = None _proc: ProcessorMixin | None = field(default=None, init=False, repr=False) @@ -1986,6 +2012,7 @@ class GrootN17VLMEncodeStep(ProcessorStep): image_target_size=self.image_target_size, shortest_image_edge=self.shortest_image_edge, crop_fraction=self.crop_fraction, + letter_box_transform=self.letter_box_transform, ) for timestep in range(video_np.shape[1]) for view_idx in range(video_np.shape[2]) @@ -2010,6 +2037,7 @@ class GrootN17VLMEncodeStep(ProcessorStep): image_target_size=self.image_target_size, shortest_image_edge=self.shortest_image_edge, crop_fraction=self.crop_fraction, + letter_box_transform=self.letter_box_transform, ) for timestep in range(sample.shape[0]) for view_idx in range(sample.shape[1]) @@ -2083,6 +2111,7 @@ class GrootN17VLMEncodeStep(ProcessorStep): "shortest_image_edge": self.shortest_image_edge, "crop_fraction": self.crop_fraction, "use_albumentations": self.use_albumentations, + "letter_box_transform": self.letter_box_transform, "device": self.device, } diff --git a/tests/policies/groot/test_groot_n1_7.py b/tests/policies/groot/test_groot_n1_7.py index d1cd6610b..61233f778 100644 --- a/tests/policies/groot/test_groot_n1_7.py +++ b/tests/policies/groot/test_groot_n1_7.py @@ -46,6 +46,7 @@ from lerobot.policies.groot.processor_groot import ( N1_7_NATIVE_ACTION_HORIZON, _make_relative_action_training_stats, _transform_n1_7_image_for_vlm_albumentations, + _transform_n1_7_image_for_vlm_torch, make_groot_pre_post_processors, ) from lerobot.processor import ( @@ -245,6 +246,7 @@ def _write_raw_n1_7_libero_checkpoint(path): "shortest_image_edge": 256, "crop_fraction": 0.95, "use_albumentations": True, + "letter_box_transform": False, "max_action_horizon": 40, "max_state_dim": 132, "max_action_dim": 132, @@ -609,6 +611,7 @@ def test_raw_n1_7_libero_checkpoint_processors_use_checkpoint_assets(tmp_path): assert vlm_encode.shortest_image_edge == 256 assert vlm_encode.crop_fraction == 0.95 assert vlm_encode.use_albumentations is True + assert vlm_encode.letter_box_transform is False assert decode_actions.raw_stats["action"]["gripper"]["q99"] == [115.0] assert decode_actions.env_action_dim == 7 assert decode_actions.use_percentiles is True @@ -682,6 +685,7 @@ def test_groot_n1_7_saved_processors_round_trip_checkpoint_specific_fields(tmp_p config_filename="policy_postprocessor.json", ) pack_inputs = next(step for step in loaded_preprocessor.steps if isinstance(step, GrootN17PackInputsStep)) + vlm_encode = next(step for step in loaded_preprocessor.steps if isinstance(step, GrootN17VLMEncodeStep)) decode_actions = next( step for step in loaded_postprocessor.steps if isinstance(step, GrootN17ActionDecodeStep) ) @@ -690,6 +694,7 @@ def test_groot_n1_7_saved_processors_round_trip_checkpoint_specific_fields(tmp_p assert pack_inputs.action_horizon == 40 assert pack_inputs.video_modality_keys == ["image", "wrist_image"] assert pack_inputs.clip_outliers is True + assert vlm_encode.letter_box_transform is False torch.testing.assert_close( pack_inputs.stats[OBS_STATE]["min"], torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]), @@ -1858,6 +1863,58 @@ def test_groot_n1_7_vlm_image_transform_matches_albumentations_eval_path(): np.testing.assert_array_equal(np.asarray(transformed), expected) +def test_groot_n1_7_albumentations_letterbox_is_opt_in(): + pytest.importorskip("cv2", exc_type=ImportError) + + image = np.full((3, 5, 3), 255, dtype=np.uint8) + + default = _transform_n1_7_image_for_vlm_albumentations( + image, + image_crop_size=None, + image_target_size=[10, 10], + shortest_image_edge=10, + crop_fraction=None, + ) + letterboxed = _transform_n1_7_image_for_vlm_albumentations( + image, + image_crop_size=None, + image_target_size=[10, 10], + shortest_image_edge=10, + crop_fraction=None, + letter_box_transform=True, + ) + + assert default.shape == (10, 17, 3) + assert default.min() == 255 + assert letterboxed.shape == (10, 10, 3) + assert letterboxed.min() < 255 + + +def test_groot_n1_7_torch_letterbox_is_opt_in(): + image = torch.full((3, 3, 5), 255, dtype=torch.uint8) + + default = _transform_n1_7_image_for_vlm_torch( + image, + image_crop_size=None, + image_target_size=[10, 10], + shortest_image_edge=10, + crop_fraction=None, + ) + letterboxed = _transform_n1_7_image_for_vlm_torch( + image, + image_crop_size=None, + image_target_size=[10, 10], + shortest_image_edge=10, + crop_fraction=None, + letter_box_transform=True, + ) + + assert tuple(default.shape) == (3, 10, 10) + assert int(default.min()) == 255 + assert tuple(letterboxed.shape) == (3, 10, 10) + assert int(letterboxed.min()) < 255 + + def test_groot_n1_7_vlm_encode_transforms_non_square_two_camera_sample_like_core_albumentations(): cv2 = pytest.importorskip("cv2", exc_type=ImportError) @@ -1928,6 +1985,7 @@ def test_groot_n1_7_vlm_encode_config_round_trips_model_name(): shortest_image_edge=256, crop_fraction=0.95, use_albumentations=True, + letter_box_transform=True, ) restored = GrootN17VLMEncodeStep(**step.get_config()) @@ -1938,6 +1996,7 @@ def test_groot_n1_7_vlm_encode_config_round_trips_model_name(): assert restored.shortest_image_edge == 256 assert restored.crop_fraction == 0.95 assert restored.use_albumentations is True + assert restored.letter_box_transform is True def test_groot_n1_7_processor_uses_qwen_component_assets(monkeypatch):