261 lines
11 KiB
Python
261 lines
11 KiB
Python
# 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)
|
|
|
|
|
|
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": task_template["scene"],
|
|
"task_name": task_template["task"],
|
|
"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": [],
|
|
"recording_setting": task_template.get('recording_setting', {})
|
|
}
|
|
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"]
|
|
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]()
|
|
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
|
|
# ------------------------------------------------------------- |