Files
gen_data_agent/data_gen_dependencies/data_generate.py
2025-09-05 15:49:00 +08:00

30 lines
1.3 KiB
Python

import json
from data_gen_dependencies.omni_robot import IsaacSimRpcRobot
from data_gen_dependencies.omniagent import Agent
def generate_data(task_file, client_host, use_recording=True):
task = json.load(open(task_file))
robot_cfg = task["robot"]["robot_cfg"]
stand = task["robot"]["stand"]
robot_position = task["robot"]["init_position"]
robot_rotation = task["robot"]["init_rotation"]
scene_usd = task["scene_usd"]
robot = IsaacSimRpcRobot(robot_cfg=robot_cfg, scene_usd=scene_usd,
client_host=client_host,
position=robot_position, rotation=robot_rotation, stand_type=stand["stand_type"],
stand_size_x=stand["stand_size_x"], stand_size_y=stand["stand_size_y"])
agent = Agent(robot)
render_semantic = False
if "render_semantic" in task["recording_setting"]:
render_semantic = task["recording_setting"]["render_semantic"]
agent.run(task_file=task_file,
camera_list=task["recording_setting"]["camera_list"],
use_recording=use_recording,
workspaces=task['scene']['function_space_objects'],
fps=task["recording_setting"]["fps"],
render_semantic=render_semantic,
)
print("job done")
robot.client.Exit()