add standard rasterizer

.gitignore

@@ -41,9 +41,8 @@ Thumbs.db
 __pycache__
 
 ## deca data
-results
+*results*
 # *_vis.jpg
-
 ## internal use
 cluster_scripts
 internal

README.md

@@ -33,9 +33,9 @@ Clone the repo:
  ```
  Or use virtual environment by runing 
  ```bash
- bash install_pip.sh
+ bash install_conda.sh
  ```
- Then follow the instruction to install [pytorch3d](https://github.com/facebookresearch/pytorch3d/blob/master/INSTALL.md).
+ For visualization, we use our rasterizer that uses pytorch JIT Compiling Extensions. If there occurs a compiling error, you can install [pytorch3d](https://github.com/facebookresearch/pytorch3d/blob/master/INSTALL.md) instead and set --rasterizer_type=pytorch3d when running the demos.
 
 ### Usage
 1. Prepare data 
@@ -50,10 +50,10 @@ Clone the repo:
  ``` 
  to visualize the predicted 2D landmanks, 3D landmarks (red means non-visible points), coarse geometry, detailed geometry, and depth. 
  <p align="center"> 
- <img src="TestSamples/examples/results/id04657-PPHljWCZ53c-000565_inputs_inputs_vis.jpg">
+ <img src="Doc/images/id04657-PPHljWCZ53c-000565_inputs_inputs_vis.jpg">
  </p> 
  <p align="center"> 
- <img src="TestSamples/examples/results/IMG_0392_inputs_vis.jpg">
+ <img src="Doc/images/IMG_0392_inputs_vis.jpg">
  </p> 
  You can also generate an obj file (which can be opened with Meshlab) that includes extracted texture from the input image. 
 

decalib/deca.py

@@ -24,7 +24,7 @@
 from skimage.io import imread
 import cv2
 import pickle
-from .utils.renderer import SRenderY
+from .utils.renderer import SRenderY, set_rasterizer
 from .models.encoders import ResnetEncoder
 from .models.FLAME import FLAME, FLAMETex
 from .models.decoders import Generator
@@ -49,7 +49,8 @@ def __init__(self, config=None, device='cuda'):
  self._setup_renderer(self.cfg.model)
 
  def _setup_renderer(self, model_cfg):
- self.render = SRenderY(self.image_size, obj_filename=model_cfg.topology_path, uv_size=model_cfg.uv_size).to(self.device)
+ set_rasterizer(self.cfg.rasterizer_type)
+ self.render = SRenderY(self.image_size, obj_filename=model_cfg.topology_path, uv_size=model_cfg.uv_size, rasterizer_type=self.cfg.rasterizer_type).to(self.device)
  # face mask for rendering details
  mask = imread(model_cfg.face_eye_mask_path).astype(np.float32)/255.; mask = torch.from_numpy(mask[:,:,0])[None,None,:,:].contiguous()
  self.uv_face_eye_mask = F.interpolate(mask, [model_cfg.uv_size, model_cfg.uv_size]).to(self.device)

decalib/utils/config.py

@@ -15,6 +15,7 @@
 
 cfg.pretrained_modelpath = os.path.join(cfg.deca_dir, 'data', 'deca_model.tar')
 cfg.output_dir = ''
+cfg.rasterizer_type = 'pytorch3d'
 # ---------------------------------------------------------------------------- #
 # Options for Face model
 # ---------------------------------------------------------------------------- #

decalib/utils/rasterizer/INSTALL.md

@@ -0,0 +1,29 @@
+## Install
+from standard_rasterize_cuda import standard_rasterize
+ # from .rasterizer.standard_rasterize_cuda import standard_rasterize
+
+in this folder, run
+```python setup.py build_ext -i ```
+
+then remember to set --rasterizer_type=standard when runing demos :) 
+
+## Alg
+https://www.scratchapixel.com/lessons/3d-basic-rendering/rasterization-practical-implementation
+
+## Speed Comparison
+runtime for raterization only 
+In PIXIE, number of faces in SMPLX: 20908 
+
+for image size = 1024 
+pytorch3d: 0.031s 
+standard: 0.01s 
+
+for image size = 224 
+pytorch3d: 0.0035s 
+standard: 0.0014s 
+
+why standard rasterizer is faster than pytorch3d? 
+Ref: https://github.com/facebookresearch/pytorch3d/blob/master/pytorch3d/csrc/rasterize_meshes/rasterize_meshes.cu 
+pytorch3d: for each pixel in image space (each pixel is parallel in cuda), loop through the faces, check if this pixel is in the projection bounding box of the face, then sorting faces according to z, record the face id of closest K faces. 
+standard rasterization: for each face in mesh (each face is parallel in cuda), loop through pixels in the projection bounding box (normally a very samll number), compare z, record face id of that pixel 
+

decalib/utils/rasterizer/setup.py

@@ -0,0 +1,24 @@
+# To install, run 
+# python setup.py build_ext -i
+# Ref: https://github.com/pytorch/pytorch/blob/11a40410e755b1fe74efe9eaa635e7ba5712846b/test/cpp_extensions/setup.py#L62
+
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+import os
+
+# USE_NINJA = os.getenv('USE_NINJA') == '1'
+os.environ["CC"] = "gcc-7"
+os.environ["CXX"] = "gcc-7"
+
+USE_NINJA = os.getenv('USE_NINJA') == '1'
+
+setup(
+ name='standard_rasterize_cuda',
+ ext_modules=[
+ CUDAExtension('standard_rasterize_cuda', [
+ 'standard_rasterize_cuda.cpp',
+ 'standard_rasterize_cuda_kernel.cu',
+ ])
+ ],
+ cmdclass={'build_ext': BuildExtension.with_options(use_ninja=USE_NINJA)}
+)

decalib/utils/rasterizer/standard_rasterize_cuda.cpp

@@ -0,0 +1,48 @@
+#include <torch/torch.h>
+#include <vector>
+#include <iostream>
+
+std::vector<at::Tensor> forward_rasterize_cuda(
+ at::Tensor face_vertices,
+ at::Tensor depth_buffer,
+ at::Tensor triangle_buffer,
+ at::Tensor baryw_buffer,
+ int h,
+ int w);
+
+std::vector<at::Tensor> standard_rasterize(
+ at::Tensor face_vertices,
+ at::Tensor depth_buffer,
+ at::Tensor triangle_buffer,
+ at::Tensor baryw_buffer,
+ int height, int width
+ ) {
+ return forward_rasterize_cuda(face_vertices, depth_buffer, triangle_buffer, baryw_buffer, height, width);
+}
+
+std::vector<at::Tensor> forward_rasterize_colors_cuda(
+ at::Tensor face_vertices,
+ at::Tensor face_colors,
+ at::Tensor depth_buffer,
+ at::Tensor triangle_buffer,
+ at::Tensor images,
+ int h,
+ int w);
+
+std::vector<at::Tensor> standard_rasterize_colors(
+ at::Tensor face_vertices,
+ at::Tensor face_colors,
+ at::Tensor depth_buffer,
+ at::Tensor triangle_buffer,
+ at::Tensor images,
+ int height, int width
+ ) {
+ return forward_rasterize_colors_cuda(face_vertices, face_colors, depth_buffer, triangle_buffer, images, height, width);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("standard_rasterize", &standard_rasterize, "RASTERIZE (CUDA)");
+ m.def("standard_rasterize_colors", &standard_rasterize_colors, "RASTERIZE COLORS (CUDA)");
+}
+
+// TODO: backward