render normal

data_generator.py

@@ -83,7 +83,7 @@ class DataGenerator:
         
     def generate_display_platform(self):
         config = self.random_config[BlenderUtils.DISPLAY_TABLE_NAME]
-        
+
         height = random.uniform(config["min_height"], config["max_height"])
         radius = random.uniform(config["min_radius"], config["max_radius"])
         while height > 0.5 * radius:
@@ -103,12 +103,29 @@ class DataGenerator:
         bpy.context.object.rigid_body.type = 'PASSIVE'
         bpy.ops.object.shade_auto_smooth()
         
+        # 创建不受光照影响的材质
         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)
+        mat.use_nodes = True
+        
+        # 清除默认节点
+        nodes = mat.node_tree.nodes
+        for node in nodes:
+            nodes.remove(node)
+
+        # 添加 Emission 节点
+        emission_node = nodes.new(type='ShaderNodeEmission')
+        emission_node.inputs['Color'].default_value = (1.0, 0.0, 0.0, 1.0)  # 红色
+
+        # 添加 Material Output 节点
+        output_node = nodes.new(type='ShaderNodeOutputMaterial')
+
+        # 连接节点
+        links = mat.node_tree.links
+        links.new(emission_node.outputs['Emission'], output_node.inputs['Surface'])
+
+        # 将材质赋给对象
+        platform.data.materials.clear()
+        platform.data.materials.append(mat)
 
         self.display_table_config = {
             "height": height,
@@ -126,19 +143,19 @@ class DataGenerator:
         
         if random.random() <= config["random_rotation_ratio"]:
             rotation = (
-                random.uniform(0, 2*np.pi),
-                random.uniform(0, 2*np.pi),
-                random.uniform(0, 2*np.pi)
+                random.uniform(0, 2 * np.pi),
+                random.uniform(0, 2 * np.pi),
+                random.uniform(0, 2 * np.pi)
             )
         else:
             rotation = (0, 0, 0)
-            z=0.05
+            z = 0.05
         
         platform_bbox = self.platform.bound_box
         platform_bbox_world = [self.platform.matrix_world @ mathutils.Vector(corner) for corner in platform_bbox]
         platform_top_z = max([v.z for v in platform_bbox_world])
         
-        obj_mesh_path = BlenderUtils.get_obj_path(self.obj_dir,name)
+        obj_mesh_path = BlenderUtils.get_obj_path(self.obj_dir, name)
         obj = BlenderUtils.load_obj(name, obj_mesh_path)
         
         obj_bottom_z = BlenderUtils.get_object_bottom_z(obj)
@@ -149,12 +166,31 @@ class DataGenerator:
         
         bpy.ops.rigidbody.object_add()
         bpy.context.object.rigid_body.type = 'ACTIVE'
+
+        # 创建不受光照影响的材质
         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)
+        mat.use_nodes = True
+        
+        # 清除默认节点
+        nodes = mat.node_tree.nodes
+        for node in nodes:
+            nodes.remove(node)
+
+        # 添加 Emission 节点
+        emission_node = nodes.new(type='ShaderNodeEmission')
+        emission_node.inputs['Color'].default_value = (0.0, 1.0, 0.0, 1.0)  # 绿色
+
+        # 添加 Material Output 节点
+        output_node = nodes.new(type='ShaderNodeOutputMaterial')
+
+        # 连接节点
+        links = mat.node_tree.links
+        links.new(emission_node.outputs['Emission'], output_node.inputs['Surface'])
+
+        # 将材质赋给对象
+        obj.data.materials.clear()
+        obj.data.materials.append(mat)
+        
         self.target_obj = obj
 
     
@@ -219,6 +255,13 @@ 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)
+        normal_dir = os.path.join(scene_dir, "normal")
+        for normal_file in os.listdir(normal_dir):
+            if not normal_file.endswith(".png"):
+                name, _ = os.path.splitext(normal_file)
+                file_path = os.path.join(normal_dir, normal_file)
+                new_file_path = os.path.join(normal_dir, f"{name}.png")
+                os.rename(file_path,new_file_path)
         return True
     
     def simulate_scene(self, frame_limit=120, depth = 0, diag = 0):
@@ -274,6 +317,7 @@ class DataGenerator:
         count_success = 0
         self.set_progress("generate scene", 0, total)
         result = "retry"
+        print(f"Generating scene for {total} objects")
         for target_obj_name in self.obj_name_list:
             self.add_log(f"Generating scene for object <{target_obj_name}>", "info")
             scene_info_path = os.path.join(self.output_dir, target_obj_name, "scene_info.json")