add training, specify texture size

README.md

@@ -6,7 +6,7 @@
 
 
 
-This is an official python implementation of PRN.  The training code will be released(about two months later).
+This is an official python implementation of PRN. 
 
 PRN is a method to jointly regress dense alignment and 3D face shape in an end-to-end manner. More examples on Multi-PIE and 300VW can be seen in [YouTube](https://youtu.be/tXTgLSyIha8) .
 
@@ -41,7 +41,7 @@ Get the 3D vertices and corresponding colours from a single image.  Save the res
 **New**: 
 
 1. you can choose to output mesh with its original pose(default) or with front view(which means all output meshes are aligned)
-2. obj file can now also written with texture map, and you can set non-visible texture to 0. 
+2. obj file can now also written with texture map(with specified texture size), and you can set non-visible texture to 0. 
 
 
 
@@ -125,8 +125,73 @@ cd PRNet
 
 
 
+## Training
+
+The core idea of the paper is:
+
+Using position map to represent face geometry&alignment information, then learning this with an Encoder-Decoder Network.
+
+So, the training steps:
+
+1. generate position map ground truth.
+
+   the example of generating position map of 300W_LP dataset can be seen in [generate_posmap_300WLP](https://github.com/YadiraF/face3d/blob/master/examples/8_generate_posmap_300WLP.py)
+
+2. an encoder-decoder network to learn mapping from rgb image to position map.
+
+   the weight mask can be found in the folder `Data/uv-data`
+
+What you can custom:
+
+1. the UV space of position map.
+
+   you can change the parameterization method, or change the resolution of UV space.
+
+2. the backbone of encoder-decoder network
+
+   this demo uses residual blocks. VGG, mobile-net are also ok.
+
+3. the weight mask
+
+   you can change the weight to focus more on which part your project need more.
+
+4. the training data
+
+   if you have scanned 3d face, it's better to train PRN with your own data. Before that, you may need use ICP to align your face meshes.
+
+
+
+## Q&A
+
+1. How to **speed up**?
+
+   a. network inference part
+
+   you can train a smaller network or use a smaller position map as input.
+
+   b. render part
+
+   you can refer to  [c++ version](https://github.com/YadiraF/face3d/blob/master/face3d/mesh_cython/render.py). 
+
+   c. other parts like detecting face, writing obj
+
+   the best way is to rewrite them in c++.
+
+2. How to improve the **precision**?
+
+   a. geometry precision.
+
+   Due to the restriction of training data, the precision of reconstructed face from this demo has little detail. You can train the network with your own detailed data or do post-processing like shape-from-shading to add details.
+
+   b. texture precision.
+
+   I just added an option to specify the texture size. When the texture size > face size in original image, and render new facial image with [texture mapping](https://github.com/YadiraF/face3d/blob/04869dcee1455d1fa5b157f165a6878c550cf695/face3d/mesh/render.py#L217), there will be little resample error.
+
+
+
 ## Changelog
 
+* 2018/7/19 add training part. can specify the resolution of the texture map.
 * 2018/5/10 add texture editing examples(for data augmentation, face swapping)
 * 2018/4/28 add visibility of vertices, output obj file with texture map, depth image
 * 2018/4/26 can output mesh with front view
@@ -135,6 +200,14 @@ cd PRNet
 
 
 
+## License
+
+Code: under MIT license.
+
+Trained model file: please see [issue 28](https://github.com/YadiraF/PRNet/issues/28), thank [Kyle McDonald](https://github.com/kylemcdonald) for his answer.
+
+
+
 ## Contacts
 
 Please contact _[email protected]_  or open an issue for any questions or suggestions(like, push me to add more applications).

demo.py

@@ -13,7 +13,7 @@
 from utils.estimate_pose import estimate_pose
 from utils.rotate_vertices import frontalize
 from utils.render_app import get_visibility, get_uv_mask, get_depth_image
-from utils.write import write_obj, write_obj_with_texture
+from utils.write import write_obj_with_colors, write_obj_with_texture
 
 def main(args):
     if args.isShow or args.isTexture:
@@ -42,7 +42,9 @@ def main(args):
 
         # read image
         image = imread(image_path)
-        [h, w, _] = image.shape
+        [h, w, c] = image.shape
+        if c>3:
+            image = image[:,:,:3]
 
         # the core: regress position map
         if args.isDlib:
@@ -80,14 +82,19 @@ def main(args):
             colors = prn.get_colors(image, vertices)
 
             if args.isTexture:
-                texture = cv2.remap(image, pos[:,:,:2].astype(np.float32), None, interpolation=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT,borderValue=(0))
+                if args.texture_size != 256:
+                    pos_interpolated = resize(pos, (args.texture_size, args.texture_size), preserve_range = True)
+                else:
+                    pos_interpolated = pos.copy()
+                texture = cv2.remap(image, pos_interpolated[:,:,:2].astype(np.float32), None, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT,borderValue=(0))
                 if args.isMask:
                     vertices_vis = get_visibility(vertices, prn.triangles, h, w)
                     uv_mask = get_uv_mask(vertices_vis, prn.triangles, prn.uv_coords, h, w, prn.resolution_op)
+                    uv_mask = resize(uv_mask, (args.texture_size, args.texture_size), preserve_range = True)
                     texture = texture*uv_mask[:,:,np.newaxis]
-                write_obj_with_texture(os.path.join(save_folder, name + '.obj'), save_vertices, colors, prn.triangles, texture, prn.uv_coords/prn.resolution_op)#save 3d face with texture(can open with meshlab)
+                write_obj_with_texture(os.path.join(save_folder, name + '.obj'), save_vertices, prn.triangles, texture, prn.uv_coords/prn.resolution_op)#save 3d face with texture(can open with meshlab)
             else:
-                write_obj(os.path.join(save_folder, name + '.obj'), save_vertices, colors, prn.triangles) #save 3d face(can open with meshlab)
+                write_obj_with_colors(os.path.join(save_folder, name + '.obj'), save_vertices, prn.triangles, colors) #save 3d face(can open with meshlab)
 
         if args.isDepth:
             depth_image = get_depth_image(vertices, prn.triangles, h, w, True)
@@ -132,7 +139,7 @@ def main(args):
     parser.add_argument('--isDlib', default=True, type=ast.literal_eval,
                         help='whether to use dlib for detecting face, default is True, if False, the input image should be cropped in advance')
     parser.add_argument('--is3d', default=True, type=ast.literal_eval,
-                        help='whether to output 3D face(.obj)')
+                        help='whether to output 3D face(.obj). default save colors.')
     parser.add_argument('--isMat', default=False, type=ast.literal_eval,
                         help='whether to save vertices,color,triangles as mat for matlab showing')
     parser.add_argument('--isKpt', default=False, type=ast.literal_eval,
@@ -154,5 +161,7 @@ def main(args):
                         help='whether to save texture in obj file')
     parser.add_argument('--isMask', default=False, type=ast.literal_eval,
                         help='whether to set invisible pixels(due to self-occlusion) in texture as 0')
-
+    # update in 2017/7/19
+    parser.add_argument('--texture_size', default=256, type=int,
+                        help='size of texture map, default is 256. need isTexture is True')
     main(parser.parse_args())

run_basics.py

@@ -6,7 +6,7 @@
 from time import time
 
 from api import PRN
-from utils.write import write_obj
+from utils.write import write_obj_with_colors
 
 # ---- init PRN
 os.environ['CUDA_VISIBLE_DEVICES'] = '0' # GPU number, -1 for CPU
@@ -49,6 +49,6 @@
     # -- save
     name = image_path.strip().split('/')[-1][:-4]
     np.savetxt(os.path.join(save_folder, name + '.txt'), kpt) 
-    write_obj(os.path.join(save_folder, name + '.obj'), vertices, colors, prn.triangles) #save 3d face(can open with meshlab)
+    write_obj_with_colors(os.path.join(save_folder, name + '.obj'), vertices, prn.triangles, colors) #save 3d face(can open with meshlab)
 
     sio.savemat(os.path.join(save_folder, name + '_mesh.mat'), {'vertices': vertices, 'colors': colors, 'triangles': prn.triangles})

utils/render.py

@@ -340,54 +340,3 @@ def vis_of_vertices(vertices, triangles, h, w, depth_buffer = None):
             depth_tmp[py, px] = vertex[2]
 
     return vertices_vis
-
-# TODO: To be modified
-def triangle_buffer(vertices, triangles, h, w):
-    '''
-    Args:
-        vertices: 3 x nver
-        triangles: 3 x ntri
-        h: height  # [h, w, _] = image.shape
-        w: width
-    Returns:
-        depth_buffer: height x width
-        triangle_buffer: height x width
-    ToDo:
-        whether to add x, y by 0.5? the center of the pixel?
-        m3. like somewhere is wrong
-    # Each triangle has 3 vertices & Each vertex has 3 coordinates x, y, z.
-    # Here, the bigger the z, the fronter the point.
-    '''
-    depth_buffer = np.zeros([h, w]) + 999999. #+ np.min(vertices[2,:]) - 999999. # set the initial z to the farest position
-    triangle_buffer = np.zeros_like(depth_buffer, dtype = np.int32) - 1 # if -1, the pixel has no triangle correspondance
-
-    ## calculate the depth(z) of each triangle
-    #-m1. z = the center of shpere(through 3 vertices)
-    #center3d = (vertices[:, triangles[0,:]] + vertices[:,triangles[1,:]] + vertices[:, triangles[2,:]])/3.
-    #tri_depth = np.sum(center3d**2, axis = 0)
-    #-m2. z = the center of z(v0, v1, v2)
-    tri_depth = (vertices[2, triangles[0,:]] + vertices[2,triangles[1,:]] + vertices[2, triangles[2,:]])/3.
-
-    for i in range(int(triangles.shape[1])):
-        tri = triangles[:, i] # 3 vertex indices
-
-        # the inner bounding box
-        umin = max(int(np.ceil(np.min(vertices[0,tri]))), 0)
-        umax = min(int(np.floor(np.max(vertices[0,tri]))), w-1)
-
-        vmin = max(int(np.ceil(np.min(vertices[1,tri]))), 0)
-        vmax = min(int(np.floor(np.max(vertices[1,tri]))), h-1)
-
-        if umax<umin or vmax<vmin:
-            continue
-
-        for u in range(umin, umax+1):
-            for v in range(vmin, vmax+1):
-                #-m3. calculate the accurate depth(z) of each pixel by barycentric weights
-                #w0, w1, w2 = weightsOfpoint([u,v], vertices[:2, tri])
-                #tri_depth = w0*vertices[2,tri[0]] + w1*vertices[2,tri[1]] + w2*vertices[2,tri[2]]
-                if tri_depth[i] > depth_buffer[v, u]: # and is_pointIntri([u,v], vertices[:2, tri]): 
-                    depth_buffer[v, u] = tri_depth[i]
-                    triangle_buffer[v, u] = i
-
-    return depth_buffer , triangle_buffer

utils/write.py

@@ -13,8 +13,8 @@ def write_asc(path, vertices):
         np.savetxt(path + '.asc', vertices)
 
 
-def write_obj(obj_name, vertices, colors, triangles):
-    ''' Save 3D face model
+def write_obj_with_colors(obj_name, vertices, triangles, colors):
+    ''' Save 3D face model with texture represented by colors.
     Args:
         obj_name: str
         vertices: shape = (nver, 3)
@@ -44,8 +44,61 @@ def write_obj(obj_name, vertices, colors, triangles):
             f.write(s)
 
 
-def write_obj_with_texture(obj_name, vertices, colors, triangles, texture, uv_coords):
-    ''' Save 3D face model with texture. Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp
+def write_obj_with_texture(obj_name, vertices, triangles, texture, uv_coords):
+    ''' Save 3D face model with texture represented by texture map.
+    Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp
+    Args:
+        obj_name: str
+        vertices: shape = (nver, 3)
+        triangles: shape = (ntri, 3)
+        texture: shape = (256,256,3)
+        uv_coords: shape = (nver, 3) max value<=1
+    '''
+    if obj_name.split('.')[-1] != 'obj':
+        obj_name = obj_name + '.obj'
+    mtl_name = obj_name.replace('.obj', '.mtl')
+    texture_name = obj_name.replace('.obj', '_texture.png')
+
+    triangles = triangles.copy()
+    triangles += 1 # mesh lab start with 1
+
+    # write obj
+    with open(obj_name, 'w') as f:
+        # first line: write mtlib(material library)
+        s = "mtllib {}\n".format(os.path.abspath(mtl_name))
+        f.write(s)
+
+        # write vertices
+        for i in range(vertices.shape[0]):
+            s = 'v {} {} {}\n'.format(vertices[i, 0], vertices[i, 1], vertices[i, 2])
+            f.write(s)
+
+        # write uv coords
+        for i in range(uv_coords.shape[0]):
+            s = 'vt {} {}\n'.format(uv_coords[i,0], 1 - uv_coords[i,1])
+            f.write(s)
+
+        f.write("usemtl FaceTexture\n")
+
+        # write f: ver ind/ uv ind
+        for i in range(triangles.shape[0]):
+            # s = 'f {}/{} {}/{} {}/{}\n'.format(triangles[i,0], triangles[i,0], triangles[i,1], triangles[i,1], triangles[i,2], triangles[i,2])
+            s = 'f {}/{} {}/{} {}/{}\n'.format(triangles[i,2], triangles[i,2], triangles[i,1], triangles[i,1], triangles[i,0], triangles[i,0])
+            f.write(s)
+
+    # write mtl
+    with open(mtl_name, 'w') as f:
+        f.write("newmtl FaceTexture\n")
+        s = 'map_Kd {}\n'.format(os.path.abspath(texture_name)) # map to image
+        f.write(s)
+
+    # write texture as png
+    imsave(texture_name, texture)
+
+
+def write_obj_with_colors_texture(obj_name, vertices, colors, triangles, texture, uv_coords):
+    ''' Save 3D face model with texture. 
+    Ref: https://github.com/patrikhuber/eos/blob/bd00155ebae4b1a13b08bf5a991694d682abbada/include/eos/core/Mesh.hpp
     Args:
         obj_name: str
         vertices: shape = (nver, 3)