From 87b3198e26f18f74620bafbf666cb02ef58b881b Mon Sep 17 00:00:00 2001 From: hofee Date: Tue, 2 Sep 2025 17:53:04 +0800 Subject: [PATCH] finish task_templates_divider --- .gitignore | 4 ++ app.py | 37 ++++++++++++++++++ configs/divide_task_config.yaml | 10 +++++ configs/generate_task_config.yaml | 6 +++ runners/data_generator.py | 8 ++++ runners/data_recorder.py | 8 ++++ runners/task_generator.py | 8 ++++ runners/task_templates_divider.py | 65 +++++++++++++++++++++++++++++++ 8 files changed, 146 insertions(+) create mode 100644 .gitignore create mode 100644 app.py create mode 100644 configs/divide_task_config.yaml create mode 100644 configs/generate_task_config.yaml create mode 100644 runners/data_generator.py create mode 100644 runners/data_recorder.py create mode 100644 runners/task_generator.py create mode 100644 runners/task_templates_divider.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..edcc8da --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +__pycache__/ +workspace/ +logs/ +cache/ diff --git a/app.py b/app.py new file mode 100644 index 0000000..1850a86 --- /dev/null +++ b/app.py @@ -0,0 +1,37 @@ +from pyapp.application import PyApplication +from runners.task_generator import TaskGenerator +from runners.data_generator import DataGenerator +from runners.data_recorder import DataRecorder +from runners.task_templates_divider import TaskTemplatesDivider + +@PyApplication("agent") +class AgentApp(): + @staticmethod + def start(): + TaskGenerator(config_path="configs/task_generator.json").run() + DataGenerator(config_path="configs/data_generator.json").run() + DataRecorder(config_path="configs/data_recorder.json").run() + +@PyApplication("divide_task") +class DivideTaskApp(): + @staticmethod + def start(): + TaskTemplatesDivider(config_path="configs/divide_task_config.yaml").run() + +@PyApplication("generate_task") +class GenerateTaskApp(): + @staticmethod + def start(): + TaskGenerator(config_path="configs/task_generator.json").run() + +@PyApplication("generate_data") +class GenerateDataApp(): + @staticmethod + def start(): + DataGenerator(config_path="configs/data_generator.json").run() + +@PyApplication("record_data") +class RecordDataApp(): + @staticmethod + def start(): + DataRecorder(config_path="configs/data_recorder.json").run() \ No newline at end of file diff --git a/configs/divide_task_config.yaml b/configs/divide_task_config.yaml new file mode 100644 index 0000000..3490bda --- /dev/null +++ b/configs/divide_task_config.yaml @@ -0,0 +1,10 @@ +runner: + workspace: + 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" + divide_num: 100 + total_nums: 10000 diff --git a/configs/generate_task_config.yaml b/configs/generate_task_config.yaml new file mode 100644 index 0000000..83ef6d3 --- /dev/null +++ b/configs/generate_task_config.yaml @@ -0,0 +1,6 @@ +runner: + workspace: + name: "generate_task" + root_dir: "workspace" + generate: + target_task_templates: "task_template_target_0.json" diff --git a/runners/data_generator.py b/runners/data_generator.py new file mode 100644 index 0000000..0d43681 --- /dev/null +++ b/runners/data_generator.py @@ -0,0 +1,8 @@ +from pyapp.runner import Runner + +class DataGenerator(Runner): + def __init__(self, config_path: str): + super().__init__(config_path) + + def run(self): + pass \ No newline at end of file diff --git a/runners/data_recorder.py b/runners/data_recorder.py new file mode 100644 index 0000000..8cbaca8 --- /dev/null +++ b/runners/data_recorder.py @@ -0,0 +1,8 @@ +from pyapp.runner import Runner + +class DataRecorder(Runner): + def __init__(self, config_path: str): + super().__init__(config_path) + + def run(self): + pass \ No newline at end of file diff --git a/runners/task_generator.py b/runners/task_generator.py new file mode 100644 index 0000000..798e762 --- /dev/null +++ b/runners/task_generator.py @@ -0,0 +1,8 @@ +from pyapp.runner import Runner + +class TaskGenerator(Runner): + def __init__(self, config_path: str): + super().__init__(config_path) + + def run(self): + pass \ No newline at end of file diff --git a/runners/task_templates_divider.py b/runners/task_templates_divider.py new file mode 100644 index 0000000..a1d06ed --- /dev/null +++ b/runners/task_templates_divider.py @@ -0,0 +1,65 @@ +import os +import shutil +import json + +from pyapp.runner import Runner +from pyapp.utils.log import Log + +class TaskTemplatesDivider(Runner): + def __init__(self, config_path: str): + 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.divide_num = self.divide_config["divide_num"] + self.total_nums = self.divide_config["total_nums"] + + self.task_list = self.load_all_task_templates() + self.new_task_list = [] + + + def run(self): + self.generate_task_templates_dir() + self.generate_divided_task_templates_json() + + 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)): + if file.endswith('.json'): + task_list.append(os.path.join(self.task_templates_root_dir, task_template_dir, file)) + Log.success(f"Loaded {len(task_list)} tasks") + return task_list + + def generate_task_templates_dir(self): + for task_id in range(len(self.task_list)): + os.makedirs(self.output_task_templates_dir, exist_ok=True) + if not os.path.exists(os.path.join(self.output_task_templates_dir, f"task_template_{task_id}.json")): + shutil.copy(self.task_list[task_id], os.path.join(self.output_task_templates_dir, f"task_template_{task_id}.json")) + self.new_task_list.append(os.path.join(self.output_task_templates_dir, f"task_template_{task_id}.json")) + + def generate_divided_task_templates_json(self): + tasks_num = len(self.new_task_list) + if tasks_num < self.total_nums: + Log.error(f"tasks_num < total_nums, tasks_num: {tasks_num}, total_nums: {self.total_nums}", terminate=True) + return + task_templates_num_per_target = self.total_nums // self.divide_num + task_templates_targets_jsons = {} + for i in range(self.divide_num): + task_templates_targets_jsons[f"task_template_target_{i}.json"] = {} + + for task_id in range(self.total_nums): + task_json_id = task_id // task_templates_num_per_target + task_templates_targets_jsons[f"task_template_target_{task_json_id}.json"][f"task_template_{task_id}"] = self.new_task_list[task_id] + + if not os.path.exists(self.output_template_targets_dir): + os.makedirs(self.output_template_targets_dir, exist_ok=True) + for key, value in task_templates_targets_jsons.items(): + with open(os.path.join(self.output_template_targets_dir, key), "w") as f: + json.dump(value, f) + Log.success(f"Divide {self.total_nums} task templates to {self.divide_num} task templates targets") + + \ No newline at end of file