File size: 2,491 Bytes
1cc903b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """
Dataset management tool for operations NOT requiring Isaac Sim (SimulationApp).
This script handles dataset operations that do not require the Isaac Sim application:
- inspect: Inspect dataset metadata
- read: Read dataset states
- augment: Add end-effector pose to dataset
- merge: Merge multiple datasets
For operations requiring Isaac Sim (record, replay), use dataset_sim.py instead."""
import argparse
from pathlib import Path
from .utils import (
dataset_inspection,
dataset_processing,
setup_inspect_parser,
setup_read_parser,
setup_augment_parser,
setup_merge_parser,
)
def main():
"""Main entry point for dataset management tool (no Isaac Sim required)."""
parser = argparse.ArgumentParser(
description="LeHome dataset management tool (no Isaac Sim required)",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="Note: For operations requiring Isaac Sim (record, replay), use dataset_sim.py instead.",
)
subparsers = parser.add_subparsers(
dest="command", help="Available commands", required=True
)
setup_inspect_parser(subparsers)
setup_read_parser(subparsers)
setup_augment_parser(subparsers)
setup_merge_parser(subparsers)
args = parser.parse_args()
if args.command == "inspect":
dataset_inspection.inspect(
Path(args.dataset_root),
show_frames=args.show_frames,
show_stats=args.show_stats,
)
elif args.command == "read":
dataset_inspection.read_states(
Path(args.dataset_root),
num_frames=args.num_frames,
episode=args.episode,
output_csv=args.output_csv,
show_stats=args.show_stats,
)
elif args.command == "augment":
dataset_processing.augment_ee_pose(
Path(args.dataset_root).resolve(),
Path(args.urdf_path).resolve(),
state_unit=args.state_unit,
output_root=Path(args.output_root).resolve() if args.output_root else None,
overwrite=args.overwrite,
)
elif args.command == "merge":
import ast
source_roots = [Path(p) for p in ast.literal_eval(args.source_roots)]
dataset_processing.merge_datasets(
source_roots,
Path(args.output_root),
output_repo_id=args.output_repo_id,
merge_custom_meta=args.merge_custom_meta,
)
if __name__ == "__main__":
main()
|