blender_interface.py

import bpy 
import numpy as np 
import os 
import sys
sys.path.append('/Volumes/Data/GitLab/Blender-renderer')
import util 

class BlenderInterface():
    def __init__(self, resolution=128):
        self.resolution = resolution

        # Delete the default cube (default selected)
        bpy.ops.object.delete()
        
        # Set the world background to white 
        world = bpy.context.scene.world
        world.use_nodes = True
        bg_node = world.node_tree.nodes['Background']
        bg_node.inputs[0].default_value = (1, 1, 1, .1) 
        #bpy.context.scene.view_settings.view_transform = 'Standard'
        bpy.context.scene.view_settings.view_transform = 'Raw'

        # Deselect all. All new object added to the scene will automatically selected.
        self.blender_renderer = bpy.context.scene.render
        self.blender_renderer.resolution_x = resolution
        self.blender_renderer.resolution_y = resolution
        self.blender_renderer.resolution_percentage = 100
        self.blender_renderer.image_settings.file_format = 'PNG'  # set output format to .png
        #self.blender_renderer.image_settings.color_mode = 'RGB'

        # set up lighting 
        lamp1 = bpy.data.lights['Light']
        lamp1.type = 'SUN'
        lamp1.use_shadow = False
        lamp1.specular_factor = 0.0
        lamp1.energy = 1.0
        
        bpy.ops.object.light_add(type='SUN')
        lamp2 = bpy.data.lights['Sun']
        lamp2.use_shadow = False
        lamp2.specular_factor = 0.0
        lamp2.energy = 1.0
        bpy.data.objects['Sun'].rotation_euler = bpy.data.objects['Light'].rotation_euler
        bpy.data.objects['Sun'].rotation_euler[0] += 180
        
        bpy.ops.object.light_add(type='SUN')
        lamp2 = bpy.data.lights['Sun.001']
        lamp2.use_shadow = False
        lamp2.specular_factor = 0.0
        lamp2.energy = 0.3
        bpy.data.objects['Sun.001'].rotation_euler = bpy.data.objects['Light'].rotation_euler
        bpy.data.objects['Sun.001'].rotation_euler[0] += 90
        
        # Set up the camera
        self.camera = bpy.context.scene.camera
        self.camera.data.sensor_height = self.camera.data.sensor_width # Square sensor
        util.set_camera_focal_length_in_world_units(self.camera.data, 525./512*resolution) # Set focal length to a common value (kinect)
        
        # set render engine 
        #  'CYCLES', 'BLENDER_EEVEE', 'BLENDER_WORKBENCH'
        #   slow -> fast
        bpy.context.scene.render.engine = 'BLENDER_EEVEE'
        bpy.context.scene.eevee.taa_render_samples = 64
        
        # set up for output depth information 
        bpy.context.view_layer.use_pass_z = True
        bpy.context.scene.use_nodes = True
        tree = bpy.context.scene.node_tree
        for node in tree.nodes:
            tree.nodes.remove(node)
        render_layers = tree.nodes.new('CompositorNodeRLayers')
        composite = tree.nodes.new('CompositorNodeComposite')
        normalize = tree.nodes.new('CompositorNodeNormalize')
        bpy.data.scenes["Scene"].node_tree.nodes["Composite"].use_alpha = False
        self.depth_output = tree.nodes.new('CompositorNodeOutputFile')
        self.depth_output.format.file_format = "PNG"
        self.depth_output.format.color_mode = 'BW'

        tree.links.new(render_layers.outputs['Image'], composite.inputs['Image'])
        tree.links.new(render_layers.outputs['Depth'], normalize.inputs[0])
        tree.links.new(normalize.outputs[0], self.depth_output.inputs[0])
        
        bpy.data.objects['Camera'].visible_shadow = False
        bpy.ops.object.select_all(action='DESELECT')

    def import_mesh(self, fpath, scale=1., object_world_matrix=None):
        ext = os.path.splitext(fpath)[-1] 
        bpy.ops.wm.obj_import(filepath=str(fpath))

        obj = bpy.context.selected_objects[0]
        #bpy.context.object.display.show_shadows = False


        if object_world_matrix is not None:
            obj.matrix_world = object_world_matrix

        bpy.ops.object.origin_set(type='ORIGIN_GEOMETRY', center='BOUNDS')
        obj.location = (0., 0., 0.) # center the bounding box!
        
        # rescale the object 
        """
        max_dim = max(obj.dimensions[i] for i in range(3))
        obj_scaling = 0.5 / max_dim 
        for i in range(3):
            obj.scale[i] *= obj_scaling
        """
            
        # fix the material problem (maybe)
        #bpy.context.object.active_material_index = 0
        #bpy.context.object.active_material.use_nodes = False
        for mat in bpy.data.materials:
            mat.use_backface_culling = True

        if scale != 1.:
            bpy.ops.transform.resize(value=(scale, scale, scale))

        # Disable transparency & specularities
        M = bpy.data.materials
        for i in range(len(M)):
            M[i].blend_method = 'OPAQUE'
            M[i].specular_intensity = 0.0 
            
    def render(self, output_dir, blender_cam2world_matrices, write_cam_params=False):
        
        if write_cam_params: 
            img_dir = os.path.join(output_dir, 'rgb')
            pose_dir = os.path.join(output_dir, 'pose')
            depth_dir = os.path.join(output_dir, 'depth')
            
            util.cond_mkdir(img_dir)
            util.cond_mkdir(pose_dir)
            util.cond_mkdir(depth_dir)
        else:
            img_dir = output_dir
            util.cond_mkdir(img_dir)
            
        if write_cam_params:
            K = util.get_calibration_matrix_K_from_blender(self.camera.data)
            with open(os.path.join(output_dir, 'intrinsics.txt'),'w') as intrinsics_file:
                intrinsics_file.write('%f %f %f 0.\n'%(K[0][0], K[0][2], K[1][2]))
                intrinsics_file.write('0. 0. 0.\n')
                intrinsics_file.write('1.\n')
                intrinsics_file.write('%d %d\n'%(self.resolution, self.resolution))

        for i in range(len(blender_cam2world_matrices)):
            self.camera.matrix_world = blender_cam2world_matrices[i]

            # Render the object
            if os.path.exists(os.path.join(img_dir, '%06d.png' % i)):
                continue
            
            # set z-depth output directory 
            self.depth_output.base_path = depth_dir

            # Render the color image
            self.blender_renderer.filepath = os.path.join(img_dir, '%06d.png'%i)
            bpy.ops.render.render(write_still=True, use_viewport=True)

            if write_cam_params:
                # Write out camera pose
                RT = util.get_world2cam_from_blender_cam(self.camera)
                cam2world = RT.inverted()
                with open(os.path.join(pose_dir, '%06d.txt'%i),'w') as pose_file:
                    matrix_flat = []
                    for j in range(4):
                        for k in range(4):
                            matrix_flat.append(cam2world[j][k])
                    pose_file.write(' '.join(map(str, matrix_flat)) + '\n')
                    
            # rename the depth map since it's given a stupid default name by Blender
            stupid_default_name = os.path.join(depth_dir, 'Image0001.png')
            new_name = os.path.join(depth_dir, f'{i:06d}.png')
            os.system(f"mv {stupid_default_name} {new_name}")

                    
        
        # Remember which meshes were just imported
        meshes_to_remove = []
        for ob in bpy.context.selected_objects:
            meshes_to_remove.append(ob.data)

        bpy.ops.object.delete()

        # Remove the meshes from memory too
        for mesh in meshes_to_remove:
            bpy.data.meshes.remove(mesh)