From 4e511582156015d1bc70c7348c14065e4fc0acf0 Mon Sep 17 00:00:00 2001 From: hofee Date: Fri, 5 Sep 2025 11:10:42 +0800 Subject: [PATCH] finish task_gen --- app.py | 4 +-- configs/divide_task_config.yaml | 6 ++-- configs/generate_task_config.yaml | 4 ++- runners/task_generator.py | 29 +++++++++++++++--- runners/task_templates_divider.py | 18 ++++++----- .../cost_function.py | 4 +-- .../layout_2d.py | 0 .../layout_object.py | 17 ++++++----- .../multi_add_util.py | 0 {others => task_gen_dependencies}/object.py | 2 +- {others => task_gen_dependencies}/sdf.py | 2 +- .../solver_2d.py | 12 ++++---- .../solver_3d.py | 4 +-- .../task_generate.py | 30 +++++++++++-------- .../transform_utils.py | 0 .../transforms.py | 0 {others => task_gen_dependencies}/utils.py | 0 17 files changed, 82 insertions(+), 50 deletions(-) rename {others => task_gen_dependencies}/cost_function.py (95%) rename {others => task_gen_dependencies}/layout_2d.py (100%) rename {others => task_gen_dependencies}/layout_object.py (76%) rename {others => task_gen_dependencies}/multi_add_util.py (100%) rename {others => task_gen_dependencies}/object.py (99%) rename {others => task_gen_dependencies}/sdf.py (98%) rename {others => task_gen_dependencies}/solver_2d.py (97%) rename {others => task_gen_dependencies}/solver_3d.py (97%) rename {others => task_gen_dependencies}/task_generate.py (92%) rename {others => task_gen_dependencies}/transform_utils.py (100%) rename {others => task_gen_dependencies}/transforms.py (100%) rename {others => task_gen_dependencies}/utils.py (100%) diff --git a/app.py b/app.py index 6801501..3d2b269 100644 --- a/app.py +++ b/app.py @@ -4,8 +4,8 @@ from runners.data_generator import DataGenerator from runners.data_recorder import DataRecorder from runners.task_templates_divider import TaskTemplatesDivider -@PybootApplication("agent") -class AgentApp: +@PybootApplication("main") +class MainApp: @staticmethod def start(): TaskGenerator(config_path="configs/generate_task_config.yaml").run() diff --git a/configs/divide_task_config.yaml b/configs/divide_task_config.yaml index 3490bda..638089a 100644 --- a/configs/divide_task_config.yaml +++ b/configs/divide_task_config.yaml @@ -3,8 +3,8 @@ runner: name: "divide_task" root_dir: "workspace" divide: - task_templates_root_dir: "/home/ubuntu/Projects/tasks_divide/pick&place" - output_task_templates_dir: "/home/ubuntu/Projects/Refactory_isaac_sim/input/templates/task_templates" - output_template_targets_dir: "/home/ubuntu/Projects/Refactory_isaac_sim/input/templates/template_targets" + input_task_templates_root_dir: "/home/ubuntu/Projects/tasks_divide/pick&place" + output_task_templates_dir: + output_template_targets_dir: divide_num: 100 total_nums: 10000 diff --git a/configs/generate_task_config.yaml b/configs/generate_task_config.yaml index 83ef6d3..b0910ba 100644 --- a/configs/generate_task_config.yaml +++ b/configs/generate_task_config.yaml @@ -3,4 +3,6 @@ runner: name: "generate_task" root_dir: "workspace" generate: - target_task_templates: "task_template_target_0.json" + input_data_root: "/home/ubuntu/Projects/docker_isaac_sim/input/data" + input_target_task_templates_path: "workspace/divide_task/template_targets/task_template_target_0.json" + output_task_root_dir: diff --git a/runners/task_generator.py b/runners/task_generator.py index bb17021..396e688 100644 --- a/runners/task_generator.py +++ b/runners/task_generator.py @@ -1,12 +1,33 @@ +import json +import os + from pyboot.runner import Runner -from others.task_generate import OldTaskGenerator +from pyboot.utils.log import Log +from task_gen_dependencies.task_generate import OldTaskGenerator class TaskGenerator(Runner): def __init__(self, config_path: str): super().__init__(config_path) + self.generate_config = self.config["generate"] + self.input_target_task_templates_path = self.generate_config["input_target_task_templates_path"] + self.input_data_root = self.generate_config["input_data_root"] + self.output_task_root_dir = self.generate_config.get("output_task_root_dir", None) + if self.output_task_root_dir is None: + self.output_task_root_dir = os.path.join(self.workspace_path, "task_root_dir") + self.target_task_templates = json.load(open(self.input_target_task_templates_path, "r")) def run(self): - pass + for task_template_name, task_template_path in self.target_task_templates.items(): + task_template = json.load(open(task_template_path, "r")) + self.generate_from_template(task_template_name, task_template) + Log.success(f"Generated {task_template['recording_setting']['num_of_episode']} tasks from <{task_template_name}>") - def generate_from_template(self, template: dict): - pass \ No newline at end of file + def generate_from_template(self, template_name: str, template: dict): + task_dir = os.path.join(self.output_task_root_dir, template_name) + task_num = template["recording_setting"]["num_of_episode"] + task_name = template["task"] + + old_task_generator = OldTaskGenerator(template, self.input_data_root) + if not os.path.exists(task_dir): + os.makedirs(task_dir) + old_task_generator.generate_tasks(save_path=task_dir, task_num=task_num, task_name=task_name) diff --git a/runners/task_templates_divider.py b/runners/task_templates_divider.py index 352c6aa..3d5de45 100644 --- a/runners/task_templates_divider.py +++ b/runners/task_templates_divider.py @@ -10,9 +10,13 @@ class TaskTemplatesDivider(Runner): super().__init__(config_path) self.divide_config = self.config["divide"] - self.task_templates_root_dir = self.divide_config["task_templates_root_dir"] - self.output_task_templates_dir = self.divide_config["output_task_templates_dir"] - self.output_template_targets_dir = self.divide_config["output_template_targets_dir"] + self.input_task_templates_root_dir = self.divide_config["input_task_templates_root_dir"] + self.output_task_templates_dir = self.divide_config.get("output_task_templates_dir", None) + if self.output_task_templates_dir is None: + self.output_task_templates_dir = os.path.join(self.workspace_path, "task_templates") + self.output_template_targets_dir = self.divide_config.get("output_template_targets_dir", None) + if self.output_template_targets_dir is None: + self.output_template_targets_dir = os.path.join(self.workspace_path, "template_targets") self.divide_num = self.divide_config["divide_num"] self.total_nums = self.divide_config["total_nums"] @@ -26,11 +30,11 @@ class TaskTemplatesDivider(Runner): def load_all_task_templates(self): task_list = [] - for task_template_dir in os.listdir(self.task_templates_root_dir): - if os.path.isdir(os.path.join(self.task_templates_root_dir, task_template_dir)): - for file in os.listdir(os.path.join(self.task_templates_root_dir, task_template_dir)): + for task_template_dir in os.listdir(self.input_task_templates_root_dir): + if os.path.isdir(os.path.join(self.input_task_templates_root_dir, task_template_dir)): + for file in os.listdir(os.path.join(self.input_task_templates_root_dir, task_template_dir)): if file.endswith('.json'): - task_list.append(os.path.join(self.task_templates_root_dir, task_template_dir, file)) + task_list.append(os.path.join(self.input_task_templates_root_dir, task_template_dir, file)) Log.success(f"Loaded {len(task_list)} tasks") return task_list diff --git a/others/cost_function.py b/task_gen_dependencies/cost_function.py similarity index 95% rename from others/cost_function.py rename to task_gen_dependencies/cost_function.py index 64ee8a2..97ab7a9 100644 --- a/others/cost_function.py +++ b/task_gen_dependencies/cost_function.py @@ -1,7 +1,7 @@ import numpy as np -from others.sdf import get_distance_with_sdf -from others.transform_utils import transform_points, random_point +from task_gen_dependencies.sdf import get_distance_with_sdf +from task_gen_dependencies.transform_utils import transform_points, random_point def cost_without_collision(obj_active, obj_passive, threshold=1, **kwargs): diff --git a/others/layout_2d.py b/task_gen_dependencies/layout_2d.py similarity index 100% rename from others/layout_2d.py rename to task_gen_dependencies/layout_2d.py diff --git a/others/layout_object.py b/task_gen_dependencies/layout_object.py similarity index 76% rename from others/layout_object.py rename to task_gen_dependencies/layout_object.py index da92279..74d2107 100644 --- a/others/layout_object.py +++ b/task_gen_dependencies/layout_object.py @@ -5,9 +5,9 @@ import trimesh from scipy.interpolate import RegularGridInterpolator from scipy.spatial.transform import Rotation as R -from others.object import OmniObject -from others.sdf import compute_sdf_from_obj_surface -from others.transform_utils import farthest_point_sampling, random_point +from task_gen_dependencies.object import OmniObject +from task_gen_dependencies.sdf import compute_sdf_from_obj_surface +from task_gen_dependencies.transform_utils import farthest_point_sampling, random_point, get_bott_up_point def load_and_prepare_mesh(obj_path, up_axis, scale=1.0): @@ -63,9 +63,12 @@ class LayoutObject(OmniObject): if mesh_points.shape[0] > N_collision_points: self.collision_points = farthest_point_sampling(mesh_points, N_collision_points) # 碰撞检测点 - - self.anchor_points = {'top': random_point(self.anchor_points['top'], 3)[np.newaxis, :], - 'buttom': random_point(self.anchor_points['buttom'], 3)[np.newaxis, :]} + self.anchor_points = {} + self.anchor_points['top'] = get_bott_up_point(mesh_points, 1.5, descending=False) + self.anchor_points['buttom'] = get_bott_up_point(mesh_points, 1.5, descending=True) - self.size = self.mesh.extents.copy() + self.anchor_points['top'] = random_point(self.anchor_points['top'], 3)[np.newaxis, :] + self.anchor_points['buttom'] = random_point(self.anchor_points['buttom'], 3)[np.newaxis, :] + + self.size = self.mesh.extents.copy() self.up_axis = up_axis[0] \ No newline at end of file diff --git a/others/multi_add_util.py b/task_gen_dependencies/multi_add_util.py similarity index 100% rename from others/multi_add_util.py rename to task_gen_dependencies/multi_add_util.py diff --git a/others/object.py b/task_gen_dependencies/object.py similarity index 99% rename from others/object.py rename to task_gen_dependencies/object.py index 5ff5e42..fceb294 100644 --- a/others/object.py +++ b/task_gen_dependencies/object.py @@ -4,7 +4,7 @@ import pickle import json import os -from others.transforms import transform_coordinates_3d +from task_gen_dependencies.transforms import transform_coordinates_3d from grasp_nms import nms_grasp class OmniObject: diff --git a/others/sdf.py b/task_gen_dependencies/sdf.py similarity index 98% rename from others/sdf.py rename to task_gen_dependencies/sdf.py index ea54e23..e1e45bd 100644 --- a/others/sdf.py +++ b/task_gen_dependencies/sdf.py @@ -1,6 +1,6 @@ import numpy as np import open3d as o3d -from others.transform_utils import transform_points +from task_gen_dependencies.transform_utils import transform_points def compute_sdf_from_obj_surface(mesh, resolution=2): # 2mm if mesh is None or len(mesh.vertices) == 0: diff --git a/others/solver_2d.py b/task_gen_dependencies/solver_2d.py similarity index 97% rename from others/solver_2d.py rename to task_gen_dependencies/solver_2d.py index 5d95571..f04cf68 100644 --- a/others/solver_2d.py +++ b/task_gen_dependencies/solver_2d.py @@ -1,10 +1,10 @@ from scipy.spatial.transform import Rotation as R -from others.layout_2d import DFS_Solver_Floor -from others.multi_add_util import * -from others.utils import axis_to_quaternion, quaternion_rotate, get_rotation_matrix_from_quaternion, \ +from task_gen_dependencies.layout_2d import DFS_Solver_Floor +from task_gen_dependencies.multi_add_util import * +from task_gen_dependencies.utils import axis_to_quaternion, quaternion_rotate, get_rotation_matrix_from_quaternion, \ get_quaternion_from_rotation_matrix, get_xyz_euler_from_quaternion - +from pyboot.utils.log import Log def rotate_along_axis(target_affine, angle_degrees, rot_axis='z', use_local=True): """ @@ -354,9 +354,7 @@ class LayoutSolver2D: # objects = [obj[0] for obj in placed_objects] # visualize_objects(objects, self.plane_width, self.plane_height,str(obj_id)+label_flag+".jpg") if len(fail_objs)>0: - print("*******no solution objects************") - print(fail_objs) - print("******************") + Log.warning("failed objects in layout 2d: " + str(fail_objs)) for obj_id in objs_succ: self.objects[obj_id].obj_pose = rotate_along_axis(self.objects[obj_id].obj_pose, random.uniform(-self.angle_step/2.0, self.angle_step/2.0), rot_axis='z', use_local=False) diff --git a/others/solver_3d.py b/task_gen_dependencies/solver_3d.py similarity index 97% rename from others/solver_3d.py rename to task_gen_dependencies/solver_3d.py index ed7fe41..24a9f7c 100644 --- a/others/solver_3d.py +++ b/task_gen_dependencies/solver_3d.py @@ -3,8 +3,8 @@ from scipy.interpolate import RegularGridInterpolator from scipy.optimize import dual_annealing from scipy.spatial.transform import Rotation as R -from others.cost_function import get_cost_func -from others.transform_utils import unnormalize_vars, \ +from task_gen_dependencies.cost_function import get_cost_func +from task_gen_dependencies.transform_utils import unnormalize_vars, \ normalize_vars, pose2mat, euler2quat np.random.seed(0) diff --git a/others/task_generate.py b/task_gen_dependencies/task_generate.py similarity index 92% rename from others/task_generate.py rename to task_gen_dependencies/task_generate.py index ec4144b..6ff97a0 100644 --- a/others/task_generate.py +++ b/task_gen_dependencies/task_generate.py @@ -3,13 +3,13 @@ import os import json import numpy as np -from others.object import OmniObject -from others.layout_object import LayoutObject -from others.solver_2d import LayoutSolver2D -from others.solver_3d import LayoutSolver3D +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 = {} @@ -54,18 +54,18 @@ class LayoutGenerator(): 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 --') + #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 --') + #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 --') + #print('-- extra layout done --') ''' Check completion ''' res_infos = [] @@ -94,8 +94,8 @@ class LayoutGenerator(): class OldTaskGenerator: - def __init__(self, task_template): - self.data_root = os.path.dirname(os.path.dirname(__file__)) + "/assets" + def __init__(self, task_template, data_root): + self.data_root = data_root self.init_info(task_template) def _load_json(self, relative_path): @@ -135,7 +135,7 @@ class OldTaskGenerator: else: obj_dir = os.path.join(self.data_root, obj["data_info_dir"]) if "metadata" in obj: - print(obj["metadata"]) + #print(obj["metadata"]) info = obj["metadata"]["info"] info["interaction"] = obj["metadata"]["interaction"] else: @@ -215,6 +215,10 @@ class OldTaskGenerator: output_file = os.path.join(save_path, f'{task_name}_%d.json' % (i)) self.task_template['objects'] = [] self.task_template['objects'] += self.fix_obj_infos + self.task_template['robot'] = { + "init_position" : self.robot_init_pose["position"], + "init_rotation" : self.robot_init_pose["quaternion"], + } flag_failed = False for key in self.layouts: @@ -231,10 +235,10 @@ class OldTaskGenerator: self.task_template['objects'].append(fix_pose_dict) break if flag_failed: - print(f"Failed to place key object, skipping") + Log.warning(f"Failed to place key object, skipping") continue - print('Saved task json to %s' % output_file) + Log.info('Saved task json to %s' % output_file) with open(output_file, 'w') as f: json.dump(self.task_template, f, indent=4) diff --git a/others/transform_utils.py b/task_gen_dependencies/transform_utils.py similarity index 100% rename from others/transform_utils.py rename to task_gen_dependencies/transform_utils.py diff --git a/others/transforms.py b/task_gen_dependencies/transforms.py similarity index 100% rename from others/transforms.py rename to task_gen_dependencies/transforms.py diff --git a/others/utils.py b/task_gen_dependencies/utils.py similarity index 100% rename from others/utils.py rename to task_gen_dependencies/utils.py