hofee/gen_data_agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

finish data generate

Browse Source
This commit is contained in:
hofee
2025-09-05 19:18:37 +08:00
parent 21fbd5a323
commit 02440ceeb5
21 changed files with 17866 additions and 679 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
task_gen_dependencies/object.py
View File

@@ -181,16 +181,14 @@ class OmniObject:
grasp_data['grasp_pose'][:, :3, 3] = translation
grasp_data['width'] = width
print(f"Grasp num of {obj.name} after NMS: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
#print(f"Grasp num of {obj.name} after NMS: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
# downsample
if grasp_data['grasp_pose'].shape[0]>max_grasps:
grasp_num = grasp_data['grasp_pose'].shape[0]
index = np.random.choice(grasp_num, max_grasps, replace=False)
grasp_data['grasp_pose'] = grasp_data['grasp_pose'][index]
grasp_data['width'] = grasp_data['width'][index]
print(f"Grasp num of {obj.name} after downsample: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
else:
print(f"Grasp num of {obj.name}: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
#print(f"Grasp num of {obj.name} after downsample: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
action_info[grasp_part] = grasp_data

22
task_gen_dependencies/task_generate.py
View File

@@ -99,7 +99,7 @@ class OldTaskGenerator:
def __init__(self, task_template, data_root):
self.data_root = data_root
self.init_info(task_template)
self.task_template = task_template
def _load_json(self, relative_path):
with open(os.path.join(self.data_root, relative_path), 'r') as file:
@@ -171,15 +171,15 @@ class OldTaskGenerator:
scene_info = task_template["scene"]
self.scene_usd = task_template["scene"]["scene_usd"]
self.task_template = {
"scene_usd": self.scene_usd,
"arm": arm,
"scene": task_template["scene"],
"task_name": task_template["task"],
"robot_id": robot_id,
"robot":{"robot_id": robot_id,"arm": arm},
"stages": task_template['stages'],
"object_with_material": task_template.get('object_with_material', {}),
"lights": task_template.get('lights', {}),
"cameras": task_template.get('cameras', {}),
"objects": []
"objects": [],
"recording_setting": task_template.get('recording_setting', {})
}
constraint = task_template.get("constraints")
robot_init_workspace_id = scene_info["scene_id"].split('/')[-1]
@@ -232,14 +232,10 @@ class OldTaskGenerator:
generated_tasks_path.append(output_file)
task_instance['objects'] = []
task_instance['objects'] += self.fix_obj_infos
task_instance['robot'] = {
"init_position" : self.robot_init_pose["position"],
"init_rotation" : self.robot_init_pose["quaternion"],
"robot_cfg": self.robot_cfg,
"stand": self.stand,
}
task_instance['recording_setting'] = self.task_template["recording_setting"]
task_instance['robot']["init_position"] = self.robot_init_pose["position"]
task_instance['robot']["init_rotation"] = self.robot_init_pose["quaternion"]
task_instance['robot']["robot_cfg"] = self.robot_cfg
task_instance['robot']["stand"] = self.stand
flag_failed = False
for key in self.layouts:
obj_infos = self.layouts[key]()