render_textured_images.py

import bpy
import bmesh
import numpy as np
import time
import os
import colorsys
import random
import yaml
import sys

# ++++++++++++++++++++++++++++++++++++++++++++++++++++
# +++++++++++++++++ DEFINE SCENE +++++++++++++++++++++
# ++++++++++++++++++++++++++++++++++++++++++++++++++++

def setupScene(x_dim,y_dim):

    for scene in bpy.data.scenes:
        scene.render.resolution_x = x_dim
        scene.render.resolution_y = y_dim
        scene.render.resolution_percentage = 100
        scene.world.horizon_color = (0.2 , 0.2 , 0.2)
    # bpy.context.scene.world.horizon_color = (0.3, 0.3, 0.3)




# ++++++++++++++++++++++++++++++++++++++++++++++++++++
# +++++++++++++++++ ADD LAMP +++++++++++++++++++++++++
# ++++++++++++++++++++++++++++++++++++++++++++++++++++

def setupLamps(l_energy,l_color,l_type):

    lamp_energy = l_energy
    lamp_color = l_color
    lamp_type = l_type

    lamp_list = []


    # ADD LAMP_1
    lamp_data_1 = bpy.data.lamps.new(name="Lamp_1", type=lamp_type)
    lamp_data_1.energy = lamp_energy
    lamp_data_1.color = lamp_color

    lamp_object_1 = bpy.data.objects.new(name="Lamp_1", object_data=lamp_data_1)
    lamp_object_1.location = (2,-2,2)

    bpy.context.scene.objects.link(lamp_object_1)
    lamp_object_1.select = True
    bpy.context.scene.objects.active = lamp_object_1

    lamp_list.append(lamp_object_1)



    # ADD LAMP_2
    lamp_data_2 = bpy.data.lamps.new(name="Lamp_2", type=lamp_type)
    lamp_data_2.energy = lamp_energy
    lamp_data_2.color = lamp_color

    lamp_object_2 = bpy.data.objects.new(name="Lamp_2", object_data=lamp_data_1)
    lamp_object_2.location = (-2,2,2)

    bpy.context.scene.objects.link(lamp_object_2)
    lamp_object_2.select = True
    bpy.context.scene.objects.active = lamp_object_2

    lamp_list.append(lamp_object_2)



    return lamp_list





# ++++++++++++++++++++++++++++++++++++++++++++++++++++
# +++++++++++++++++ ADD CAMERA +++++++++++++++++++++++
# ++++++++++++++++++++++++++++++++++++++++++++++++++++

def setupCamera():

    bpy.ops.object.camera_add()
    camera = bpy.data.objects["Camera"]
    bpy.context.scene.camera = camera

    return camera



def makeCameraPath(camera,fov_degrees,theta_step_degrees,phi_degrees,dist_factor):

    # Define Camera Properties
    fov = fov_degrees * np.pi / 180
    camera.data.angle = fov
    distance_camera = dist_factor / np.tan(fov / 2)

    print(camera.data.lens , file=sys.stderr)
    print(distance_camera , file=sys.stderr)

    # Define all locations of the camera when doing a sweep
    thetas = np.arange(0, 360, theta_step_degrees) * np.pi / 180.0
    phi = phi_degrees * np.pi / 180.0

    # (I.e. trasnforming spherical coordinates to Cartsian coordinates)
    x_cam = distance_camera * np.cos(thetas) * np.sin(phi)
    y_cam = distance_camera * np.cos(phi)
    z_cam = distance_camera * np.sin(thetas) * np.sin(phi)


    return (x_cam , y_cam , z_cam , thetas)



def setupNodes():

    bpy.context.scene.use_nodes = True
    node_tree = bpy.context.scene.node_tree

    for node in node_tree.nodes:
        node_tree.nodes.remove(node)

    input_node = node_tree.nodes.new(type="CompositorNodeRLayers")

    color_invert_node = node_tree.nodes.new(type="CompositorNodeInvert")
    color_invert_node.invert_rgb = True

    normalize_node = node_tree.nodes.new(type="CompositorNodeNormalize")

    output_node = node_tree.nodes.new(type="CompositorNodeComposite")

    tree_links = node_tree.links

    input_color_link = tree_links.new(input_node.outputs[2]        , color_invert_node.inputs[1])
    color_norm_link =  tree_links.new(color_invert_node.outputs[0] , normalize_node.inputs[0])
    norm_output_link = tree_links.new(normalize_node.outputs[0]    , output_node.inputs[0])

    bpy.context.scene.use_nodes = False



def srgb_to_linear(r, g, b):
    def srgb(c):
        a = .055
        if c <= .04045:
            return c / 12.92
        else:
            return ((c+a) / (1+a)) ** 2.4
    return tuple(srgb(c) for c in (r, g, b))



def makeMaterial(color,color_name,diff_int,diff_shader,spec_int,):

    mat = bpy.data.materials.new(color_name)
    mat.diffuse_shader = diff_shader
    mat.diffuse_intensity = diff_int
    mat.specular_intensity = spec_int
    mat.diffuse_color = color

    return mat





def makeMaterialList(num_materials , color_S , color_V , diff_intensity , diff_shader , spec_intensity):

    h_step = 1.0/num_materials
    random_inc = random.random() / float(num_materials)

    material_list = []

    for i in range(num_materials):
        material_list.append(makeMaterial(srgb_to_linear(*colorsys.hsv_to_rgb((h_step*i) + random_inc  , color_S , color_V)) ,
                                          "COLOR_"+str(i) , 
                                          diff_intensity , 
                                          diff_shader , 
                                          spec_intensity)
                            )

    random.shuffle(material_list)

    return material_list



def assignMaterials(label_path,materials_list,shape_obj):

    # shape_obj = bpy.context.object
    # bpy.ops.object.active = shape_obj
    bpy.context.scene.objects.active = shape_obj

    for i,m in enumerate(materials_list):
        bpy.ops.object.material_slot_add()
        shape_obj.material_slots[i].material = m
        

    bpy.ops.object.mode_set(mode = "EDIT")
    shape_mesh = bmesh.from_edit_mesh(shape_obj.data)

    shape_mesh.faces.ensure_lookup_table()
    num_faces = len(shape_mesh.faces)

    with open(label_path , 'r') as fp:
        label_content = fp.readlines()

    num_lines = len(label_content) + 1
    num_labels = int(num_lines/3)

    for i in range(num_labels):

        first_line = label_content[3*i]
        second_line = label_content[(3*i) + 1]

        label_id = ord(first_line[-2]) - ord('A')

        face_ids = list(map(int , second_line.split()))

        for j , f_id in enumerate(face_ids):
            if((f_id-1) < num_faces):
                shape_mesh.faces[f_id-1].material_index = label_id
        
    shape_obj.data.update()
    bpy.ops.object.mode_set(mode = 'OBJECT')

    return num_faces
    

    

def addTrackingConstraints(camera,lamp_list,origin_obj):     

    bpy.context.scene.objects.active = camera
    bpy.ops.object.constraint_add(type='TRACK_TO')
    camera.constraints['Track To'].target = origin_obj
    camera.constraints['Track To'].track_axis = 'TRACK_NEGATIVE_Z'
    camera.constraints['Track To'].up_axis = 'UP_Y'

    for lamp in lamp_list:

        bpy.context.scene.objects.active = lamp
        bpy.ops.object.constraint_add(type='TRACK_TO')
        lamp.constraints['Track To'].target = origin_obj
        lamp.constraints['Track To'].track_axis = 'TRACK_NEGATIVE_Z'
        lamp.constraints['Track To'].up_axis = 'UP_Y'




def scaleShape(shape_object):

    max_dim = max(shape_object.dimensions)
    scale = 1.0 / max_dim
    shape_object.scale /= scale



def orientShape(deg_rot_tuple , shape_obj):

    x_rot = deg_rot_tuple[0] * np.pi/180.0
    y_rot = deg_rot_tuple[1] * np.pi/180.0
    z_rot = deg_rot_tuple[2] * np.pi/180.0
    

    shape_obj.rotation_mode = "XYZ"
    shape_obj.rotation_euler = (x_rot , y_rot , z_rot)







if __name__ == "__main__":

    render_config_filename = "../data/render_config.yaml"
    shape_data_filename =    "../data/shape_data.yaml"

    with open(render_config_filename,"r") as fp:
        config_dict = yaml.safe_load(fp)

    with open(shape_data_filename,"r") as fp:
        all_shape_dicts = yaml.safe_load(fp)


    setupScene(config_dict['Scene']['x_size'] , 
               config_dict['Scene']['y_size'])
    
    lamp_list = setupLamps(             config_dict['Lamp']['lamp_energy'] ,
                                  tuple(config_dict['Lamp']['lamp_color']) ,
                                        config_dict['Lamp']['lamp_type'])


    camera = setupCamera()

    
    (x_cam , y_cam , z_cam , thetas) = makeCameraPath(  camera , 
                                                        config_dict['Camera']['cam_fov'] , 
                                                        config_dict['Camera']['cam_theta_step'] , 
                                                        config_dict['Camera']['cam_phi'] ,
                                                        config_dict['Camera']['cam_dist_factor'])
    

    setupNodes()

    for shape_dict in all_shape_dicts['shape_data']:

        datatype_start_time = time.time()
        print("\n\n----------------------------------------" , file=sys.stderr)
        print("----------------------------------------" , file=sys.stderr)
        print("PROCESSING DATATYPE:\t" + shape_dict['folder_name'] + "\n\n" , file=sys.stderr)

        materials_list = makeMaterialList(  shape_dict['num_labels'] , 
                                            config_dict['Material']['mat_color_S'] , 
                                            config_dict['Material']['mat_color_V'] ,  
                                            config_dict['Material']['mat_diffuse_intensity'] , 
                                            config_dict['Material']['mat_diffuse_shader'] , 
                                            config_dict['Material']['mat_specular_intensity'])











        # ++++++++++++++++++++++++++++++++++++++++++++++++++++
        # +++++++++++++++++ DEFINE INPUT PATH ++++++++++++++++
        # ++++++++++++++++++++++++++++++++++++++++++++++++++++

        input_folder  = config_dict['Data']['texture_input_root_folder']           + shape_dict['folder_name'] + '/'
        output_folder = config_dict['Data']['texture_output_root_folder']          + shape_dict['folder_name'] + '/'
        # labels_output_folder = config_dict['Data']['labels_output_root_folder'] + shape_dict['folder_name'] + '/'
        # depth_output_folder =  config_dict['Data']['depth_output_root_folder']  + shape_dict['folder_name'] + '/'

        # input_folder = "../Dataset/Temp_Input/Earphone/"
        # output_folder = "../Dataset/Temp_Output/Earphone/"

        if(not (os.path.isdir(output_folder))):
            os.mkdir(output_folder)
        # if(not (os.path.isdir(depth_output_folder))):
        #     os.mkdir(depth_output_folder)

        FOLDER_PATHS = []
        # LABEL_PATHS = []

        for dirname in os.listdir(input_folder):
            FOLDER_PATHS += [os.path.join(input_folder , dirname)]

        # for p in FOLDER_PATHS:
        #     print(p , file=sys.stderr)
        # for root, subdirs, files in os.walk(input_folder):
        #     files = sorted(files)
        #     for filename in files:
        #         file_path = os.path.join(root, filename)
        #         if file_path.endswith(".off"):
        #             SHAPE_PATHS += [file_path]
        #         elif file_path.endswith("labels.txt"):
        #             LABEL_PATHS += [file_path]

        # PATHS = ["../Dataset/SHAPENET_MESHES/Car/1a0c91c02ef35fbe68f60a737d94994a.off" , "../Dataset/SHAPENET_MESHES/Airplane/1d269dbde96f067966cf1b4a8fc3914e.off"]
        # SHAPE_PATHS = ["../Dataset/SHAPENET_MESHES/Car/1a0c91c02ef35fbe68f60a737d94994a.off"]
        # LABEL_PATHS = ["../Dataset/SHAPENET_MESHES/Car/1a0c91c02ef35fbe68f60a737d94994a_labels.txt"]



        # ++++++++++++++++++++++++++++++++++++++++++++++++++++
        # ++++++++++++++ LOOP OVER ALL OBJECTS +++++++++++++++
        # ++++++++++++++++++++++++++++++++++++++++++++++++++++

        shape_counter = 0
        # for shape_path , label_path in zip(SHAPE_PATHS , LABEL_PATHS):
        for obj_folder_path in FOLDER_PATHS:

            shape_time_start = time.time()
            shape_name = obj_folder_path.split("/")[-1]
            # print(obj_folder_path.split("/") , file=sys.stderr)
            bpy.ops.import_scene.obj(filepath=os.path.join(obj_folder_path , "model.obj") , use_smooth_groups=False)
            # print(bpy.context.scene.objects.active , file=sys.stderr)

            flag = False
            for o in bpy.data.objects:
                if(o.type=="MESH"):
                    if(not(flag)):
                        bpy.context.scene.objects.active = o
                        flag = True
                    
                    o.select = True
                else:
                    o.select = False

            # print(bpy.context.selected_objects , file=sys.stderr)
            # bpy.context.scene.objects.active = bpy.context.selected_objects
            temp_name = bpy.ops.object.join()
            shape_object = bpy.context.scene.objects.active
            # print(bpy.context.scene.objects.active , file=sys.stderr)
            # print(bpy.data.objects.keys() , file=sys.stderr)

            # for o in bpy.context.scene.objects:
            #     if(o.select == True):
            #         print(o.name , file=sys.stderr)
            shape_counter += 1
            print("RENDERING " + shape_dict['folder_name'] + " " + str(shape_counter) + "/" + str(len(FOLDER_PATHS)) + "\t" + shape_name + " ...\t" , end= "" , file=sys.stderr)

        #     # Make the imported object active
        #     shape_object = bpy.data.objects[shape_name]
        #     bpy.context.scene.objects.active = shape_object
            shape_object.select = True

            orientShape(tuple(shape_dict['orientation']) , shape_object)
            scaleShape(shape_object)
            addTrackingConstraints(camera,lamp_list,shape_object)

            num_faces = len(shape_object.data.polygons)

        #     num_faces = assignMaterials(label_path,materials_list,shape_object)


            # LOOPING OVER ALL THE CAMERA POSITIONS
            for i, t in enumerate(thetas):

                # Place the camera (Note that it WILL ALWAYS point towards the object, due to the constraint)
                camera.location = (x_cam[i], z_cam[i], y_cam)

                # Render the scene, save to file
                bpy.data.scenes['Scene'].render.filepath = os.path.abspath(output_folder + shape_name + "_" + str(int(i*config_dict['Camera']['cam_theta_step'])).zfill(3) + "_texture" + '.png')
                bpy.ops.render.render(write_still=True)

            
            
        #     bpy.context.scene.use_nodes = True

        #     for i, t in enumerate(thetas):

        #         # Place the camera (Note that it WILL ALWAYS point towards the object, due to the constraint)
        #         camera.location = (x_cam[i], z_cam[i], y_cam)

        #         # Render the scene, save to file
        #         bpy.data.scenes['Scene'].render.filepath = depth_output_folder + shape_name + "_" + str(int(i*config_dict['Camera']['cam_theta_step'])).zfill(3) + "_depth" + '.png'
        #         bpy.ops.render.render(write_still=True)

        #     bpy.context.scene.use_nodes = False
            
            bpy.data.objects.remove(shape_object)

            shape_time_end = time.time()
            print("Done in " + "{0:.3f}".format(shape_time_end-shape_time_start) + " seconds\tRate (/ms): " + "{0:.3f}".format(1000*(shape_time_end-shape_time_start)/float(num_faces)) , file=sys.stderr)
            sys.stderr.flush()

        datatype_end_time = time.time()
        print("\n " + shape_dict['folder_name'] + " done in total " + "{0:.3f}".format(datatype_end_time-datatype_start_time) + " seconds\t Rate: " + "{0:.3f}".format((datatype_end_time-datatype_start_time)/float(shape_counter)) + " / shape" , file=sys.stderr)
        sys.stderr.flush()