solve merge

data_generator.py

@@ -18,6 +18,9 @@ class DataGenerator:
         self.random_config = config["runner"]["generate"]["random_config"]
         self.light_and_camera_config = config["runner"]["generate"]["light_and_camera_config"]
         self.obj_dir = config["runner"]["generate"]["object_dir"]
+        self.use_list = config["runner"]["generate"]["use_list"]
+        self.object_list_path = config["runner"]["generate"]["object_list_path"]
+        
         self.max_views = config["runner"]["generate"]["max_views"]
         self.min_views = config["runner"]["generate"]["min_views"]
         self.min_diag = config["runner"]["generate"]["min_diag"]
@@ -30,7 +33,18 @@ class DataGenerator:
         self.to_idx = config["runner"]["generate"]["to"]
         self.set_status_path = f"http://localhost:{self.port}/project/set_status"
         self.log_path = f"http://localhost:{self.port}/project/add_log"
-        self.obj_name_list = os.listdir(self.obj_dir)[self.from_idx: self.to_idx]
+        self.origin_obj_name_list = os.listdir(self.obj_dir)[self.from_idx: self.to_idx]
+        if not os.path.exists(self.output_dir):
+            os.makedirs(self.output_dir)
+        self.obj_name_list = []
+        if self.use_list:
+            self.target_name_list = [line.strip() for line in open(self.object_list_path).readlines()]
+            for obj_name in self.target_name_list:
+                if obj_name in self.origin_obj_name_list:
+                    self.obj_name_list.append(obj_name)
+        else:
+            self.obj_name_list = self.origin_obj_name_list
+
         self.target_obj = None
         self.stopped = False
         self.random_obj_list = []
@@ -218,8 +232,6 @@ class DataGenerator:
                 file_path = os.path.join(depth_dir, depth_file)
                 new_file_path = os.path.join(depth_dir, f"{name}.png")
                 os.rename(file_path,new_file_path)
-        BlenderUtils.save_blend(scene_dir)
-        exit(0)
         
         return True
     
@@ -245,11 +257,7 @@ class DataGenerator:
 
         msg = self.check_and_adjust_target()
         
-        if msg == "adjusted" and depth < 3:
-            bpy.context.view_layer.update()
-            bpy.context.scene.frame_set(0)
-            return self.simulate_scene(depth = depth + 1, diag=diag)
-        elif msg == "success":
+        if msg == "success":
             print("Scene generation completed.")
             result = self.start_render(diag=diag)
             if not result:
@@ -272,7 +280,7 @@ class DataGenerator:
 
     
     def gen_all_scene_data(self):
-        max_retry_times = 3
+        max_retry_times = 5
         total = len(self.obj_name_list)
         count = 0
         count_success = 0
@@ -285,11 +293,12 @@ class DataGenerator:
             if os.path.exists(scene_info_path):
                 self.add_log(f"Scene for object <{target_obj_name}> already exists, skipping", "warning")
                 count += 1
+                count_success += 1  
                 continue
             
             retry_times = 0
             self.set_status("target_object", target_obj_name)
-            while retry_times < 3 and result == "retry":
+            while retry_times < max_retry_times and result == "retry":
                 self.reset()  
                 try:
                     result = self.gen_scene_data(target_obj_name)

data_renderer.py

@@ -2,11 +2,12 @@ import os
 import bpy
 import sys
 import json
+import time
 import mathutils
 import numpy as np
 sys.path.append(os.path.dirname(os.path.abspath(__file__)))
 from utils.blender_util import BlenderUtils
-
+from utils.material_util import MaterialUtil
 
 class DataRenderer:
     def __init__(self):
@@ -23,8 +24,8 @@ class DataRenderer:
                                 "near_plane": 0.01,
                                 "far_plane": 5,
                                 "fov_vertical": 25,
-                                "resolution": [1280, 800],
-                                "eye_distance": 0.15,
+                                "resolution": [640, 400],
+                                "eye_distance": 0.10,
                                 "eye_angle": 25
                             },
                             "Light": {
@@ -42,7 +43,6 @@ class DataRenderer:
         self.light_and_camera_config = config["renderer"]["generate"]["light_and_camera_config"]
         self.obj_dir = config["renderer"]["generate"]["object_dir"]
         self.binocular_vision = config["renderer"]["generate"]["binocular_vision"]
-        self.obj_name_list = os.listdir(self.obj_dir)
         self.target_obj = None
         self.random_obj_list = []
 
@@ -58,14 +58,31 @@ class DataRenderer:
     def do_render(self, cam_pose, restore_info, temp_dir):
         
         self.reset()
+        start_time = time.time()
         self.restore_scene(restore_info=restore_info)
+        end_time = time.time()
+        print(f"Time taken for restoring scene: {end_time - start_time} seconds")
         object_name = self.target_obj.name
         temp_file_name = f"tmp"
         if "." in object_name:
             object_name = object_name.split(".")[0]
+        start_time = time.time()
         BlenderUtils.set_camera_at(cam_pose)
-        BlenderUtils.render_and_save(temp_dir, temp_file_name, binocular_vision=self.binocular_vision)
+        BlenderUtils.render_mask(temp_dir, temp_file_name, binocular_vision=self.binocular_vision)
+        MaterialUtil.change_object_material(self.target_obj, MaterialUtil.create_normal_material())
+        BlenderUtils.render_normal_and_depth(temp_dir, temp_file_name, binocular_vision=self.binocular_vision)
+        end_time = time.time()
+        print(f"Time taken for rendering: {end_time - start_time} seconds")
         
+        BlenderUtils.save_cam_params(temp_dir, temp_file_name, binocular_vision=self.binocular_vision)
+        
+        depth_dir = os.path.join(temp_dir, "depth")
+        for depth_file in os.listdir(depth_dir):
+            if not depth_file.endswith(".png"):
+                name, _ = os.path.splitext(depth_file)
+                file_path = os.path.join(depth_dir, depth_file)
+                new_file_path = os.path.join(depth_dir, f"{name}.png")
+                os.rename(file_path,new_file_path)
                 
     def restore_scene(self, restore_info):
 
@@ -82,12 +99,7 @@ class DataRenderer:
                     obj.location = mathutils.Vector(obj_info["location"])
                     obj.rotation_euler = mathutils.Vector(obj_info["rotation_euler"])
                     obj.scale = mathutils.Vector(obj_info["scale"])
-                    mat = bpy.data.materials.new(name="GreenMaterial")
-                    mat.diffuse_color = (0.0, 1.0, 0.0, 1.0)  # Green with full alpha (1.0)
-                    if len(obj.data.materials) > 0:
-                        obj.data.materials[0] = mat
-                    else:
-                        obj.data.materials.append(mat)
+                    MaterialUtil.change_object_material(obj, MaterialUtil.create_mask_material(color=(0, 1.0, 0)))
                     self.target_obj = obj
 
     def restore_display_platform(self, platform_info):
@@ -96,16 +108,9 @@ class DataRenderer:
         platform.name = BlenderUtils.DISPLAY_TABLE_NAME
         platform.location = mathutils.Vector(platform_info["location"])
 
-        mat = bpy.data.materials.new(name="RedMaterial")
-        mat.diffuse_color = (1.0, 0.0, 0.0, 1.0)  # Red with full alpha (1.0)
-        if len(platform.data.materials) > 0:
-            platform.data.materials[0] = mat
-        else:
-            platform.data.materials.append(mat)
-
         bpy.ops.rigidbody.object_add()
         bpy.context.object.rigid_body.type = 'PASSIVE'
-        bpy.ops.object.shade_auto_smooth()
+        MaterialUtil.change_object_material(platform, MaterialUtil.create_mask_material(color=(1.0, 0, 0)))
 
 def main(temp_dir):
     params_data_path = os.path.join(temp_dir, "params.json")
@@ -116,9 +121,14 @@ def main(temp_dir):
     scene_info_path = os.path.join(params_data["scene_path"], "scene_info.json")
     with open(scene_info_path, 'r') as f:
         scene_info = json.load(f)
-    
+    start_time = time.time()
     data_renderer = DataRenderer()
+    end_time = time.time()
+    print(f"Time taken for initialization: {end_time - start_time} seconds")
+    start_time = time.time()
     data_renderer.do_render(cam_pose, scene_info, temp_dir)
+    end_time = time.time()
+    print(f"Time taken for rendering: {end_time - start_time} seconds")
     depth_dir = os.path.join(temp_dir, "depth")
     for depth_file in os.listdir(depth_dir):
         if not depth_file.endswith(".png"):

utils/blender_util.py

@@ -279,6 +279,15 @@ class BlenderUtils:
             scene.render.filepath = os.path.join(
                 output_dir, mask_dir, f"{file_name}_{cam_suffix}.png"
             )
+            scene.use_nodes = True
+            tree = scene.node_tree
+
+            for node in tree.nodes:
+                tree.nodes.remove(node)
+
+            rl = tree.nodes.new("CompositorNodeRLayers")
+            scene.render.image_settings.file_format = "PNG"
+            scene.render.image_settings.color_mode = "RGB"
             scene.render.image_settings.color_depth = "8"
             scene.render.resolution_percentage = 100
             scene.render.use_overwrite = False
@@ -385,8 +394,8 @@ class BlenderUtils:
                 and obj.name != BlenderUtils.DISPLAY_TABLE_NAME
             ):
                 obj_info = {
-                    "location": list(obj.location),
-                    "rotation_euler": list(obj.rotation_euler),
+                    "location": list(obj.matrix_world.translation),
+                    "rotation_euler": list(obj.matrix_world.to_euler()),
                     "scale": list(obj.scale),
                 }
                 scene_info[obj.name] = obj_info

utils/view_sample_util.py

@@ -131,7 +131,7 @@ class ViewSampleUtil:
             dot_product = np.dot(forward_vector, up_vector)
             angle = np.degrees(np.arccos(dot_product))
             right_vector = np.cross(up_vector, forward_vector)
-            print(angle)
+           
             if angle > 90 - min_cam_table_included_degree:
                 max_angle = 90 - min_cam_table_included_degree
                 min_angle = max(90 - min_cam_table_included_degree*2, 30)