# old dependencies ---------------------------------------------------- import os import json from copy import copy, deepcopy import numpy as np from task_gen_dependencies.object import OmniObject from task_gen_dependencies.layout_object import LayoutObject from task_gen_dependencies.solver_2d import LayoutSolver2D from task_gen_dependencies.solver_3d import LayoutSolver3D from utils.pose import PoseUtil from pyboot.utils.log import Log def list_to_dict(data: list): tmp = {} for i in range(len(data)): tmp[str(i)] = data[i] return tmp class LayoutGenerator(): def __init__(self, workspace, obj_infos, objects, key_obj_ids, extra_obj_ids, constraint=None, fix_obj_ids=[]): self.workspace = workspace self.objects = objects self.obj_infos = obj_infos self.key_obj_ids = key_obj_ids self.extra_obj_ids = extra_obj_ids self.fix_obj_ids = fix_obj_ids self.constraint = constraint if constraint is None: self.key_obj_ids_2d = self.key_obj_ids self.key_obj_ids_3d = [] else: self.key_obj_ids_2d = [constraint["passive"]] self.key_obj_ids_3d = [constraint["active"]] self.constraint = constraint workspace_xyz, workspace_size = np.array(workspace['position']), np.array(workspace['size']) workspace_size = workspace_size * 1000 # extra info about workspace self.solver_2d = LayoutSolver2D(workspace_xyz, workspace_size, objects, fix_obj_ids=fix_obj_ids, obj_infos=obj_infos) self.solver_3d = LayoutSolver3D(workspace_xyz, workspace_size, objects, obj_infos=obj_infos) self.succ_obj_ids = [] def __call__(self): ''' Generate Layout ''' # import pdb;pdb.set_trace() if len(self.key_obj_ids_2d) > 0: objs_succ = self.solver_2d(self.key_obj_ids_2d, self.succ_obj_ids, object_extent=30, start_with_edge=True, key_obj=True, initial_angle=0) self.update_obj_info(objs_succ) #print('-- 2d layout done --') if len(self.key_obj_ids_3d) > 0: objs_succ = self.solver_3d(self.key_obj_ids_3d, self.succ_obj_ids, constraint=self.constraint) self.update_obj_info(objs_succ) #print('-- 3d layout done --') if len(self.extra_obj_ids) > 0: objs_succ = self.solver_2d(self.extra_obj_ids, self.succ_obj_ids, object_extent=30, start_with_edge=False, key_obj=False) self.update_obj_info(objs_succ) #print('-- extra layout done --') ''' Check completion ''' res_infos = [] if len(self.key_obj_ids) > 0: for obj_id in self.key_obj_ids: if obj_id not in self.succ_obj_ids: return None res_infos.append(self.obj_infos[obj_id]) if len(self.extra_obj_ids) > 0: if len(self.succ_obj_ids) > 0: for obj_id in self.succ_obj_ids: res_infos.append(self.obj_infos[obj_id]) return res_infos def update_obj_info(self, obj_ids): if not isinstance(obj_ids, list): obj_ids = [obj_ids] for obj_id in obj_ids: pose = self.objects[obj_id].obj_pose xyz, quat = pose[:3, 3], PoseUtil.get_quaternion_wxyz_from_rotation_matrix(pose[:3, :3]) self.obj_infos[obj_id]['position'] = (xyz / 1000).tolist() self.obj_infos[obj_id]['quaternion'] = quat.tolist() self.obj_infos[obj_id]["is_key"] = obj_id in self.key_obj_ids self.succ_obj_ids.append(obj_id) 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: return json.load(file) def init_info(self, task_template): # Load all objects & constraints self.fix_objs = task_template["objects"].get('fix_objects', []) all_objs = task_template["objects"]["task_related_objects"] + task_template["objects"][ "extra_objects"] + self.fix_objs self.fix_obj_ids = [obj["object_id"] for obj in self.fix_objs] self.key_obj_ids, self.extra_obj_ids = {'0': []}, {'0': []} for obj in task_template["objects"]["task_related_objects"]: ws_id = obj.get("workspace_id", "0") if ws_id not in self.key_obj_ids: self.key_obj_ids[ws_id] = [] self.key_obj_ids[ws_id].append(obj["object_id"]) for obj in task_template["objects"]["extra_objects"]: ws_id = obj.get("workspace_id", "0") if ws_id not in self.extra_obj_ids: self.extra_obj_ids[ws_id] = [] self.extra_obj_ids[ws_id].append(obj["object_id"]) obj_infos = {} objects = {} all_key_objs = [obj_id for ws_id in self.key_obj_ids for obj_id in self.key_obj_ids[ws_id]] for obj in all_objs: obj_id = obj["object_id"] if obj_id == "fix_pose": info = dict() info['object_id'] = obj_id info['position'] = obj['position'] info['direction'] = obj['direction'] obj_infos[obj_id] = info objects[obj_id] = OmniObject("fix_pose") else: obj_dir = os.path.join(self.data_root, obj["data_info_dir"]) if "metadata" in obj: #print(obj["metadata"]) info = obj["metadata"]["info"] info["interaction"] = obj["metadata"]["interaction"] else: info = json.load(open(obj_dir + "/object_parameters.json")) info["data_info_dir"] = obj_dir info["obj_path"] = obj_dir + "/Aligned.obj" info['object_id'] = obj_id if "add_particle" in obj: info["add_particle"] = obj["add_particle"] if 'extent' in obj: info['extent'] = obj['extent'] obj_infos[obj_id] = info objects[obj_id] = LayoutObject(info, use_sdf=obj_id in all_key_objs) self.obj_infos, self.objects = obj_infos, objects self.fix_obj_infos = [] for fix_obj in self.fix_objs: fix_obj['is_key'] = True fix_obj.update(obj_infos[fix_obj['object_id']]) self.fix_obj_infos.append(fix_obj) if "robot" not in task_template: arm = "right" robot_id = "A2D" else: arm = task_template["robot"]["arm"] robot_id = task_template["robot"]["robot_id"] scene_info = task_template["scene"] self.scene_usd = task_template["scene"]["scene_usd"] self.task_template = { "scene_usd": self.scene_usd, "arm": arm, "task_name": task_template["task"], "robot_id": robot_id, "stages": task_template['stages'], "object_with_material": task_template.get('object_with_material', {}), "lights": task_template.get('lights', {}), "cameras": task_template.get('cameras', {}), "objects": [] } constraint = task_template.get("constraints") robot_init_workspace_id = scene_info["scene_id"].split('/')[-1] # Retrieve scene information self.scene_usd = scene_info["scene_usd"] if "function_space_objects" in scene_info: workspaces = scene_info["function_space_objects"] if robot_init_workspace_id not in task_template["robot"]["robot_init_pose"]: self.robot_init_pose = task_template["robot"]["robot_init_pose"] else: self.robot_init_pose = task_template["robot"]["robot_init_pose"][robot_init_workspace_id] else: scene_info = self._load_json(scene_info["scene_info_dir"] + "/scene_parameters.json") workspaces = scene_info["function_space_objects"] # Normalize format if isinstance(scene_info["robot_init_pose"], list): scene_info['robot_init_pose'] = list_to_dict(scene_info["robot_init_pose"]) self.robot_init_pose = scene_info["robot_init_pose"][robot_init_workspace_id] if isinstance(workspaces, list): workspaces = list_to_dict(workspaces) workspaces = {'0': workspaces[robot_init_workspace_id]} elif isinstance(workspaces, dict) and 'position' in workspaces: workspaces = {'0': workspaces} self.layouts = {} for key in workspaces: ws, key_ids, extra_ids = workspaces[key], self.key_obj_ids.get(key, []), self.extra_obj_ids.get(key, []) self.layouts[key] = LayoutGenerator(ws, obj_infos, objects, key_ids, extra_ids, constraint=constraint, fix_obj_ids=self.fix_obj_ids) self.stand = { "stand_type": "cylinder", "stand_size_x": 0.1, "stand_size_y": 0.1 } if "robot" not in task_template: self.robot_cfg = "A2D_120s.json" else: self.robot_cfg = task_template["robot"]["robot_cfg"] if "stand" in task_template["robot"]: self.stand = task_template["robot"]["stand"] def generate_tasks(self, save_path, task_num, task_name): os.makedirs(save_path, exist_ok=True) generated_tasks_path = [] for i in range(task_num): task_instance = deepcopy(self.task_template) output_file = os.path.join(save_path, f'{task_name}_%d.json' % (i)) 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"] flag_failed = False for key in self.layouts: obj_infos = self.layouts[key]() if not obj_infos: if obj_infos is None: flag_failed = True break continue task_instance['objects'] += obj_infos for object_info in self.obj_infos: if object_info == 'fix_pose': fix_pose_dict = self.obj_infos['fix_pose'] task_instance['objects'].append(fix_pose_dict) break if flag_failed: Log.warning(f"Failed to place key object, skipping") continue Log.info('Saved task json to %s' % output_file) with open(output_file, 'w') as f: json.dump(task_instance, f, indent=4) return generated_tasks_path # -------------------------------------------------------------