added video recon and refactored the code.

Author: shentong7

README.md

@@ -1,21 +1,31 @@
 # 3DMM model fitting using Pytorch
 
-This is a fitting framework implemented in Pytorch for reconstructing faces from images using BFM. 
+This is a fitting framework implemented in Pytorch for reconstructing the face in an image or a video using a 3DMM model. 
 
-The frame only uses Pytorch modules and a differentiable renderer from pytorch3d. The whole module is differentiable and can be integrated into other systems for the gradient propagation. 
+The framework only uses Pytorch modules and a differentiable renderer from pytorch3d. The whole module is differentiable and can be integrated into other systems for the gradient propagation. 
+
+## Updates
+- :star2: Refactored the code to make future extension easier (other 3DMM models will be supported). The hyper-parameters are also re-organized.
+- :star2: Added support for monocular video reconstruction. 
+- :star2: Added support for multi-gpu when processing long videos.
 
 <p align="center">
   <img src="gifs/demo.gif" alt="demo" width="512px">
 </p>
 
+<p align="center">
+  <img src="gifs/video_demo.gif" alt="demo" width="830px">
+</p>
+
 ## Installation
 ### Requirements
 - [pytorch3d](https://github.com/facebookresearch/pytorch3d) It might require a specific version of Pytorch to make pytorch3d run succussfully on gpus, please follow the official instructions.
 - Please refer to "requirements.txt" for other dependences.
 - [Basel Face Model 2009 (BFM09)](https://faces.dmi.unibas.ch/bfm/main.php?nav=1-0&id=basel_face_model)
 - [Expression Basis](https://github.com/Juyong/3DFace) extra expression basis.
 
-## Instruction
+## How to use
+### Installation
 1. Clone the repo:
 ```
 git clone https://github.com/ascust/3DMM-Fitting-Pytorch
@@ -26,17 +36,43 @@ cd 3DMM-Fitting-Pytorch
 
 3. Download the Expression Basis. Go to the [repo](https://github.com/Juyong/3DFace), download the "CoarseData" and put "Exp_Pca.bin" into "BFM".
 
-4. Convert the BMF parameters by:
+4. Convert the BFM parameters by:
 ```
-python convert_bfm_data.py 
+python convert_bfm09_data.py
 ```
 
-5. Run the code on specific images by:
+### Single Image Reconstruction
+```
+python fit_single_img.py --img_path data/000002.jpg --res_folder results 
+```
+The results are stored in "results" folder.
+
+
+### Monocular Video Reconstruction
+To fit a video, simply run:
+```
+python fit_video.py --v_path data/sample_video.mp4 --res_folder results 
+```
+The script will extract frames, detact features and fit all frames. 
+
+Fitting a video is a bit different from fitting an image, because frames are not isolated. In this implementation, we first estimate shape and texture of the target face using some of the frames (indicated by --nframes_shape). Then we estimate the other coefficients (expression, rotatation etc.) for each frame and keep the shape and texture coefficients fixed. 
+
+For the first frame, we use much more iterations to get a good starting point. For remaining frames, each is initialized from the previous estimation.Pose regulerizers are also imposed to increase temporal consistency. 
+
+Please check the reconstructed video in the result folder.
+
+### Multi-gpu and -process support
+
+It could take quite a while to process a long video, so multi-gpu and multi-process are also supported to accelerate fitting. To use this feature, simply run:
 ```
-python fit.py --img data/000002.jpg
+python fit_video.py --v_path data/sample_video.mp4 --res_folder results --ngpus 4 --nworkers 4
 ```
-The code will do rigid fitting as well as non-rigid fitting using landmarks as well as the image as supervision. 
+Here we use 4 gpus and 4 workers. We can also use more workers to assign each gpu with multiple workers. The video will be evenly split and each clip will be fit respectively. 
 
+### Hyperparameters
+There are bunch of parameters that might require further tuning. 
+The iteration numbers for non-rigid fitting "--first_nrf_iters" and "--rest_nrf_iters" affect the fitting speed a lot since we have to render the 3d mesh in each iteration. Try to change it to find a trade-off between speed and accuracy.
+If the result is not good, try to play with the parameters.
 
 ## Acknowledgement
 The code is partially borrowed from [Deep3DFaceReconstrution](https://github.com/microsoft/Deep3DFaceReconstruction), which is a Tensorflow-based deep reconstruction method using CNNs. Please note that our framework does not require any pretrained deep models. We estimate the parameters directly using the landmarks and photometric loss as the supervision.

convert_bfm09_data.py

@@ -0,0 +1,102 @@
+from scipy.io import loadmat, savemat
+import numpy as np
+from array import array
+
+# load expression basis
+
+
+def LoadExpBasis():
+    n_vertex = 53215
+    Expbin = open('BFM/Exp_Pca.bin', 'rb')
+    exp_dim = array('i')
+    exp_dim.fromfile(Expbin, 1)
+    expMU = array('f')
+    expPC = array('f')
+    expMU.fromfile(Expbin, 3*n_vertex)
+    expPC.fromfile(Expbin, 3*exp_dim[0]*n_vertex)
+
+    expPC = np.array(expPC)
+    expPC = np.reshape(expPC, [exp_dim[0], -1])
+    expPC = np.transpose(expPC)
+
+    expEV = np.loadtxt('BFM/std_exp.txt')
+
+    return expPC, expEV
+
+# transfer original BFM09 to our face model
+
+
+def transferBFM09():
+    original_BFM = loadmat('BFM/01_MorphableModel.mat')
+    shapePC = original_BFM['shapePC']  # shape basis
+    shapeEV = original_BFM['shapeEV']  # corresponding eigen value
+    shapeMU = original_BFM['shapeMU']  # mean face
+    texPC = original_BFM['texPC']  # texture basis
+    texEV = original_BFM['texEV']  # eigen value
+    texMU = original_BFM['texMU']  # mean texture
+
+    expPC, expEV = LoadExpBasis()
+
+    # transfer BFM09 to our face model
+
+    idBase = shapePC*np.reshape(shapeEV, [-1, 199])
+    idBase = idBase/1e5  # unify the scale to decimeter
+    idBase = idBase[:, :80]  # use only first 80 basis
+
+    exBase = expPC*np.reshape(expEV, [-1, 79])
+    exBase = exBase/1e5  # unify the scale to decimeter
+    exBase = exBase[:, :64]  # use only first 64 basis
+
+    texBase = texPC*np.reshape(texEV, [-1, 199])
+    texBase = texBase[:, :80]  # use only first 80 basis
+
+    # our face model is cropped align face landmarks which contains only 35709 vertex.
+    # original BFM09 contains 53490 vertex, and expression basis provided by JuYong contains 53215 vertex.
+    # thus we select corresponding vertex to get our face model.
+
+    index_exp = loadmat('BFM/BFM_front_idx.mat')
+    index_exp = index_exp['idx'].astype(
+        np.int32) - 1  # starts from 0 (to 53215)
+
+    index_shape = loadmat('BFM/BFM_exp_idx.mat')
+    index_shape = index_shape['trimIndex'].astype(
+        np.int32) - 1  # starts from 0 (to 53490)
+    index_shape = index_shape[index_exp]
+
+    idBase = np.reshape(idBase, [-1, 3, 80])
+    idBase = idBase[index_shape, :, :]
+    idBase = np.reshape(idBase, [-1, 80])
+
+    texBase = np.reshape(texBase, [-1, 3, 80])
+    texBase = texBase[index_shape, :, :]
+    texBase = np.reshape(texBase, [-1, 80])
+
+    exBase = np.reshape(exBase, [-1, 3, 64])
+    exBase = exBase[index_exp, :, :]
+    exBase = np.reshape(exBase, [-1, 64])
+
+    meanshape = np.reshape(shapeMU, [-1, 3])/1e5
+    meanshape = meanshape[index_shape, :]
+    meanshape = np.reshape(meanshape, [1, -1])
+
+    meantex = np.reshape(texMU, [-1, 3])
+    meantex = meantex[index_shape, :]
+    meantex = np.reshape(meantex, [1, -1])
+
+    # other info contains triangles, region used for computing photometric loss,
+    # region used for skin texture regularization, and 68 landmarks index etc.
+    other_info = loadmat('BFM/facemodel_info.mat')
+    frontmask2_idx = other_info['frontmask2_idx']
+    skinmask = other_info['skinmask']
+    keypoints = other_info['keypoints']
+    point_buf = other_info['point_buf']
+    tri = other_info['tri']
+    tri_mask2 = other_info['tri_mask2']
+
+    # save our face model
+    savemat('BFM/BFM09_model_info.mat', {'meanshape': meanshape, 'meantex': meantex, 'idBase': idBase, 'exBase': exBase, 'texBase': texBase,
+                                         'tri': tri, 'point_buf': point_buf, 'tri_mask2': tri_mask2, 'keypoints': keypoints, 'frontmask2_idx': frontmask2_idx, 'skinmask': skinmask})
+
+
+if __name__ == '__main__':
+    transferBFM09()