finish task_templates_divider

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+workspace/
+logs/
+cache/

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()

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

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"

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

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

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

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")
+
+