groot_deployment / Isaac-GR00T /examples /mask-guided-background-suppression /test_extra_augmentation.py
| #!/usr/bin/env python3 | |
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| """Smoke test: apply extra_augmentation_config to raw frames and save comparison images.""" | |
| from __future__ import annotations | |
| import argparse | |
| import importlib | |
| import json | |
| import os | |
| from pathlib import Path | |
| import sys | |
| from gr00t.configs.data.embodiment_configs import MODALITY_CONFIGS | |
| from gr00t.data.dataset.lerobot_episode_loader import LeRobotEpisodeLoader | |
| from gr00t.data.embodiment_tags import EmbodimentTag | |
| from gr00t.model.gr00t_n1d7.image_augmentations import ( | |
| apply_with_replay, | |
| build_image_transformations_albumentations, | |
| ) | |
| import numpy as np | |
| from PIL import Image | |
| def save_comparison(original, augmented, mask, output_path): | |
| orig_arr = np.array(original) | |
| aug_arr = augmented.transpose(1, 2, 0) if augmented.shape[0] == 3 else augmented | |
| panels = [orig_arr, aug_arr] | |
| if mask is not None: | |
| mask_vis = np.where(mask[..., None] > 0, 255, 0).astype(np.uint8) | |
| mask_vis = np.broadcast_to(mask_vis, (*mask.shape[:2], 3)).copy() | |
| if mask_vis.shape[:2] != orig_arr.shape[:2]: | |
| mask_vis = np.array( | |
| Image.fromarray(mask_vis).resize( | |
| (orig_arr.shape[1], orig_arr.shape[0]), Image.NEAREST | |
| ) | |
| ) | |
| panels.append(mask_vis) | |
| h = panels[0].shape[0] | |
| resized = [] | |
| for p in panels: | |
| if p.shape[0] != h: | |
| new_w = int(p.shape[1] * h / p.shape[0]) | |
| p = np.array(Image.fromarray(p).resize((new_w, h), Image.BILINEAR)) | |
| resized.append(p) | |
| Image.fromarray(np.concatenate(resized, axis=1)).save(output_path) | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--dataset_path", required=True) | |
| parser.add_argument("--embodiment_tag", required=True) | |
| parser.add_argument("--modality_config_path", default=None) | |
| parser.add_argument("--extra_augmentation_config", type=str, required=True) | |
| parser.add_argument("--output_dir", type=str, default="/tmp/augmentation_vis") | |
| parser.add_argument("--num_frames", type=int, default=5) | |
| args = parser.parse_args() | |
| if args.modality_config_path: | |
| path = Path(args.modality_config_path) | |
| sys.path.append(str(path.parent)) | |
| importlib.import_module(path.stem) | |
| embodiment_tag = EmbodimentTag[args.embodiment_tag].value | |
| modality_configs = MODALITY_CONFIGS[embodiment_tag] | |
| extra_aug_config = json.loads(args.extra_augmentation_config) | |
| train_transform, _ = build_image_transformations_albumentations( | |
| image_target_size=[224, 224], | |
| image_crop_size=[224, 224], | |
| random_rotation_angle=0, | |
| color_jitter_params=None, | |
| shortest_image_edge=512, | |
| crop_fraction=0.95, | |
| extra_augmentation_config=extra_aug_config, | |
| ) | |
| loader = LeRobotEpisodeLoader( | |
| dataset_path=args.dataset_path, | |
| modality_configs=modality_configs, | |
| ) | |
| episode_df = loader[0] | |
| video_cols = [c for c in episode_df.columns if c.startswith("video.")] | |
| mask_cols = [c for c in episode_df.columns if c.startswith("mask.")] | |
| print(f"Video columns: {video_cols}") | |
| print(f"Mask columns: {mask_cols}") | |
| num_frames = min(args.num_frames, len(episode_df)) | |
| frame_indices = np.linspace(0, len(episode_df) - 1, num_frames, dtype=int) | |
| os.makedirs(args.output_dir, exist_ok=True) | |
| for vcol in video_cols: | |
| view_name = vcol.replace("video.", "") | |
| mcol = f"mask.{view_name}" | |
| has_mask = mcol in mask_cols | |
| view_dir = os.path.join(args.output_dir, view_name.replace(".", "_")) | |
| os.makedirs(view_dir, exist_ok=True) | |
| for fidx in frame_indices: | |
| orig_img = episode_df[vcol].iloc[fidx] | |
| mask_arr = np.array(episode_df[mcol].iloc[fidx]) if has_mask else None | |
| masks_list = [mask_arr] if mask_arr is not None else None | |
| transformed, _ = apply_with_replay(train_transform, [orig_img], masks_list) | |
| aug_arr = transformed[0].numpy() | |
| out_path = os.path.join(view_dir, f"frame_{fidx:04d}.png") | |
| save_comparison(orig_img, aug_arr, mask_arr, out_path) | |
| print(f" Saved: {out_path}") | |
| print(f"\nDone! {num_frames} frames x {len(video_cols)} views saved to {args.output_dir}") | |
| print("\n" + "=" * 60) | |
| print("Testing full training pipeline (processor + dataloader) ...") | |
| print("=" * 60) | |
| from gr00t.configs.base_config import get_default_config | |
| from gr00t.data.dataset.factory import DatasetFactory | |
| from gr00t.model.gr00t_n1d7.processing_gr00t_n1d7 import Gr00tN1d7Processor | |
| config = get_default_config() | |
| config = config.load_dict( | |
| { | |
| "data": { | |
| "download_cache": False, | |
| "datasets": [ | |
| { | |
| "dataset_paths": [args.dataset_path], | |
| "mix_ratio": 1.0, | |
| "embodiment_tag": embodiment_tag, | |
| } | |
| ], | |
| } | |
| } | |
| ) | |
| config.model.extra_augmentation_config = extra_aug_config | |
| config.model.use_albumentations_transforms = True | |
| processor = Gr00tN1d7Processor( | |
| modality_configs=config.data.modality_configs, | |
| statistics=None, | |
| image_crop_size=config.model.image_crop_size, | |
| image_target_size=config.model.image_target_size, | |
| random_rotation_angle=config.model.random_rotation_angle, | |
| color_jitter_params=config.model.color_jitter_params, | |
| model_name=config.model.model_name, | |
| model_type=config.model.backbone_model_type, | |
| formalize_language=config.model.formalize_language, | |
| max_state_dim=config.model.max_state_dim, | |
| max_action_dim=config.model.max_action_dim, | |
| apply_sincos_state_encoding=config.model.apply_sincos_state_encoding, | |
| max_action_horizon=config.model.action_horizon, | |
| use_albumentations=config.model.use_albumentations_transforms, | |
| extra_augmentation_config=config.model.extra_augmentation_config, | |
| shortest_image_edge=config.model.shortest_image_edge, | |
| crop_fraction=config.model.crop_fraction, | |
| use_relative_action=config.model.use_relative_action, | |
| ) | |
| processor.train() | |
| dataset_factory = DatasetFactory(config=config) | |
| train_dataset, _ = dataset_factory.build(processor=processor) | |
| sample = next(iter(train_dataset)) | |
| print(f"Sample keys: {list(sample.keys())}") | |
| print(f"VLM keys: {list(sample['vlm_content'].keys())}") | |
| for k, v in sample.items(): | |
| if hasattr(v, "shape"): | |
| print(f" {k}: shape={v.shape}, dtype={v.dtype}") | |
| print("\nPipeline test PASSED!") | |
| if __name__ == "__main__": | |
| main() | |