add texture map export

config_test.py

@@ -39,7 +39,13 @@
 flags.DEFINE_boolean('neutralize_expression', False,
                      'If true the camera and flame parameters will be saved')
 
-flags.DEFINE_string('flame_model_path', './flame_model/generic_model.pkl', 'path to the neurtral flame model')
+flags.DEFINE_boolean('save_texture', False,
+                     'If true the texture map will be stored')
+
+flags.DEFINE_string('flame_model_path', './flame_model/generic_model.pkl', 'path to the neutral FLAME model')
+
+flags.DEFINE_string('flame_texture_data_path', './flame_model/texture_data_512.npy', 'path to the FLAME texture data')
+
 
 flags.DEFINE_string('load_path', PRETRAINED_MODEL, 'path to trained model')
 

demo.py

@@ -47,6 +47,7 @@
 
 from util import renderer as vis_util
 from util import image as img_util
+from util.project_on_mesh import compute_texture_map
 from config_test import get_config
 from run_RingNet import RingNet_inference
 
@@ -90,6 +91,14 @@ def visualize(img, proc_param, verts, cam, img_name='test_image'):
     # import ipdb
     # ipdb.set_trace()
 
+
+def create_texture(img, proc_param, verts, faces, cam, texture_data):
+    cam_for_render, vert_shifted = vis_util.get_original(proc_param, verts, cam, img_size=img.shape[:2])
+
+    texture_map = compute_texture_map(img, vert_shifted, faces, cam_for_render, texture_data)
+    return texture_map
+
+
 def preprocess_image(img_path):
     img = io.imread(img_path)
     if np.max(img.shape[:2]) != config.img_size:
@@ -117,7 +126,7 @@ def main(config, template_mesh):
     input_img, proc_param, img = preprocess_image(config.img_path)
     vertices, flame_parameters = model.predict(np.expand_dims(input_img, axis=0), get_parameters=True)
     cams = flame_parameters[0][:3]
-    visualize(img, proc_param, vertices[0], cams, img_name=config.out_folder + '/images/' + config.img_path.split('/')[-1][:-4])
+    # visualize(img, proc_param, vertices[0], cams, img_name=config.out_folder + '/images/' + config.img_path.split('/')[-1][:-4])
 
     if config.save_obj_file:
         if not os.path.exists(config.out_folder + '/mesh'):
@@ -139,6 +148,22 @@ def main(config, template_mesh):
         neutral_mesh = make_prdicted_mesh_neutral(config.out_folder + '/params/' + config.img_path.split('/')[-1][:-4] + '.npy', config.flame_model_path)
         neutral_mesh.write_obj(config.out_folder + '/neutral_mesh/' + config.img_path.split('/')[-1][:-4] + '.obj')
 
+    if config.save_texture:
+        if not os.path.exists(config.flame_texture_data_path):
+            print('FLAME texture data not found')
+            return
+        texture_data = np.load(config.flame_texture_data_path, allow_pickle=True)[()]
+        texture = create_texture(img, proc_param, vertices[0], template_mesh.f, cams, texture_data)
+
+        if not os.path.exists(config.out_folder + '/texture'):
+            os.mkdir(config.out_folder + '/texture')
+
+        cv2.imwrite(config.out_folder + '/texture/' + config.img_path.split('/')[-1][:-4] + '.png', texture[:,:,::-1])
+        mesh = Mesh(v=vertices[0], f=template_mesh.f)
+        mesh.vt = texture_data['vt']
+        mesh.ft = texture_data['ft']
+        mesh.set_texture_image(config.out_folder + '/texture/' + config.img_path.split('/')[-1][:-4] + '.png')
+        mesh.write_obj(config.out_folder + '/texture/' + config.img_path.split('/')[-1][:-4] + '.obj')
 
 
 

util/project_on_mesh.py

@@ -0,0 +1,38 @@
+import numpy as np
+from psbody.mesh import Mesh
+from opendr.camera import ProjectPoints
+
+def compute_texture_map(source_img, verts, faces, cam, texture_data):
+    '''
+    Given an image and a mesh aligned with the image (under scale-orthographic projection), project the image onto the
+    mesh and return a texture map.
+    '''
+
+    x_coords = texture_data.get('x_coords')
+    y_coords = texture_data.get('y_coords')
+    valid_pixel_ids = texture_data.get('valid_pixel_ids')
+    valid_pixel_3d_faces = texture_data.get('valid_pixel_3d_faces')
+    valid_pixel_b_coords = texture_data.get('valid_pixel_b_coords')
+    img_size = texture_data.get('img_size')
+
+    pixel_3d_points = verts[valid_pixel_3d_faces[:, 0], :] * valid_pixel_b_coords[:, 0][:, np.newaxis] + \
+                      verts[valid_pixel_3d_faces[:, 1], :] * valid_pixel_b_coords[:, 1][:, np.newaxis] + \
+                      verts[valid_pixel_3d_faces[:, 2], :] * valid_pixel_b_coords[:, 2][:, np.newaxis]
+
+    vertex_normals = Mesh(verts, faces).estimate_vertex_normals()
+    pixel_3d_normals = vertex_normals[valid_pixel_3d_faces[:, 0], :] * valid_pixel_b_coords[:, 0][:, np.newaxis] + \
+                       vertex_normals[valid_pixel_3d_faces[:, 1], :] * valid_pixel_b_coords[:, 1][:, np.newaxis] + \
+                       vertex_normals[valid_pixel_3d_faces[:, 2], :] * valid_pixel_b_coords[:, 2][:, np.newaxis]
+    n_dot_view = pixel_3d_normals[:,2]
+
+    proj_2d_points = ProjectPoints(f=cam[0] * np.ones(2), rt=np.zeros(3), t=np.zeros(3), k=np.zeros(5), c=cam[1:3])
+    proj_2d_points.v = pixel_3d_points
+    proj_2d_points = np.round(proj_2d_points.r).astype(int)
+
+    texture = np.zeros((img_size, img_size, 3))
+    for i, (x, y) in enumerate(proj_2d_points):
+        if n_dot_view[i] > 0.0:
+            continue
+        if x > 0 and x < source_img.shape[1] and y > 0 and y < source_img.shape[0]:
+            texture[y_coords[valid_pixel_ids[i]].astype(int), x_coords[valid_pixel_ids[i]].astype(int), :3] = source_img[y, x]
+    return texture