├── .gitattributes ├── README.md ├── cmesh ├── __init__.py ├── cython │ ├── build │ │ └── temp.win-amd64-3.5 │ │ │ └── Release │ │ │ ├── mesh_core.obj │ │ │ ├── mesh_core_cython.cp35-win_amd64.exp │ │ │ ├── mesh_core_cython.cp35-win_amd64.lib │ │ │ └── mesh_core_cython.obj │ ├── mesh_core.cpp │ ├── mesh_core.h │ ├── mesh_core_cython.cp35-win_amd64.pyd │ ├── mesh_core_cython.cpp │ ├── mesh_core_cython.pyx │ └── setup.py ├── readme.txt └── render.py ├── color_correction.py ├── convert.py ├── face_detect.py ├── generate.py ├── images ├── canonical_vertices.npy ├── face_ind.txt ├── readme.txt ├── test.jpg ├── test_g.obj ├── test_g_kpt.jpg ├── test_g_texture.png ├── test_kpt.jpg ├── test_mask.jpg ├── test_ver.jpg ├── triangles.txt ├── uv_coords.txt ├── uv_face.png ├── uv_face_eyes.png ├── uv_face_mask.png ├── uv_kpt_ind.txt ├── uv_weight_mask.png └── uv_weight_mask2.png ├── model.py ├── predictor.py ├── prepare.py ├── run_one.py ├── run_one_g.py ├── run_swap.py ├── run_two.py ├── run_vedio.py ├── show.py ├── tmodel.py └── train.py /.gitattributes: -------------------------------------------------------------------------------- 1 | # Auto detect text files and perform LF normalization 2 | * text=auto 3 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DeepFaceRestruction 2 | -------------------------------------------------------------------------------- /cmesh/__init__.py: -------------------------------------------------------------------------------- 1 | from .cython import mesh_core_cython 2 | from . import render 3 | 4 | -------------------------------------------------------------------------------- /cmesh/cython/build/temp.win-amd64-3.5/Release/mesh_core.obj: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/cmesh/cython/build/temp.win-amd64-3.5/Release/mesh_core.obj -------------------------------------------------------------------------------- /cmesh/cython/build/temp.win-amd64-3.5/Release/mesh_core_cython.cp35-win_amd64.exp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/cmesh/cython/build/temp.win-amd64-3.5/Release/mesh_core_cython.cp35-win_amd64.exp -------------------------------------------------------------------------------- /cmesh/cython/build/temp.win-amd64-3.5/Release/mesh_core_cython.cp35-win_amd64.lib: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/cmesh/cython/build/temp.win-amd64-3.5/Release/mesh_core_cython.cp35-win_amd64.lib -------------------------------------------------------------------------------- /cmesh/cython/build/temp.win-amd64-3.5/Release/mesh_core_cython.obj: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/cmesh/cython/build/temp.win-amd64-3.5/Release/mesh_core_cython.obj -------------------------------------------------------------------------------- /cmesh/cython/mesh_core.cpp: -------------------------------------------------------------------------------- 1 | #include "mesh_core.h" 2 | 3 | 4 | /* Judge whether the point is in the triangle 5 | Method: 6 | http://blackpawn.com/texts/pointinpoly/ 7 | Args: 8 | point: [x, y] 9 | tri_points: three vertices(2d points) of a triangle. 2 coords x 3 vertices 10 | Returns: 11 | bool: true for in triangle 12 | */ 13 | bool isPointInTri(point p, point p0, point p1, point p2) 14 | { 15 | // vectors 16 | point v0, v1, v2; 17 | v0 = p2 - p0; 18 | v1 = p1 - p0; 19 | v2 = p - p0; 20 | 21 | // dot products 22 | float dot00 = v0.dot(v0); //v0.x * v0.x + v0.y * v0.y //np.dot(v0.T, v0) 23 | float dot01 = v0.dot(v1); //v0.x * v1.x + v0.y * v1.y //np.dot(v0.T, v1) 24 | float dot02 = v0.dot(v2); //v0.x * v2.x + v0.y * v2.y //np.dot(v0.T, v2) 25 | float dot11 = v1.dot(v1); //v1.x * v1.x + v1.y * v1.y //np.dot(v1.T, v1) 26 | float dot12 = v1.dot(v2); //v1.x * v2.x + v1.y * v2.y//np.dot(v1.T, v2) 27 | 28 | // barycentric coordinates 29 | float inverDeno; 30 | if(dot00*dot11 - dot01*dot01 == 0) 31 | inverDeno = 0; 32 | else 33 | inverDeno = 1/(dot00*dot11 - dot01*dot01); 34 | 35 | float u = (dot11*dot02 - dot01*dot12)*inverDeno; 36 | float v = (dot00*dot12 - dot01*dot02)*inverDeno; 37 | 38 | // check if point in triangle 39 | return (u >= 0) && (v >= 0) && (u + v < 1); 40 | } 41 | 42 | void _render_texture_core( 43 | float* image, float* vertices, int* triangles, 44 | float* tri_depth, float* tri_tex, float* depth_buffer, 45 | int ver_len, int tri_len, int tex_len, int h, int w, int c) 46 | { 47 | int i; 48 | int x, y, k; 49 | int tri_p0_ind, tri_p1_ind, tri_p2_ind; 50 | point p, p0, p1, p2; 51 | float p0_depth, p1_depth, p2_depth; 52 | int x_min, x_max, y_min, y_max; 53 | 54 | for(i = 0; i < tri_len; i++) 55 | { 56 | tri_p0_ind = triangles[i]; 57 | tri_p1_ind = triangles[tri_len + i]; 58 | tri_p2_ind = triangles[tri_len * 2 + i]; 59 | 60 | p0.x = vertices[tri_p0_ind]; p0.y = vertices[ver_len + tri_p0_ind]; p0_depth = vertices[ver_len * 2 + tri_p0_ind]; 61 | p1.x = vertices[tri_p1_ind]; p1.y = vertices[ver_len + tri_p1_ind]; p1_depth = vertices[ver_len * 2 + tri_p1_ind]; 62 | p2.x = vertices[tri_p2_ind]; p2.y = vertices[ver_len + tri_p2_ind]; p2_depth = vertices[ver_len * 2 + tri_p2_ind]; 63 | 64 | x_min = max((int)ceil(min(p0.x, min(p1.x, p2.x))), 0); 65 | x_max = min((int)floor(max(p0.x, max(p1.x, p2.x))), w - 1); 66 | 67 | y_min = max((int)ceil(min(p0.y, min(p1.y, p2.y))), 0); 68 | y_max = min((int)floor(max(p0.y, max(p1.y, p2.y))), h - 1); 69 | 70 | if(x_max < x_min || y_max < y_min) 71 | { 72 | continue; 73 | } 74 | 75 | for(y = y_min; y <= y_max; y++) //h 76 | { 77 | for(x = x_min; x <= x_max; x++) //w 78 | { 79 | p.x = x; p.y = y; 80 | if(tri_depth[i] > depth_buffer[y*w+x] && isPointInTri(p, p0, p1, p2)) 81 | { 82 | for(k = 0; k < c; k++) 83 | { 84 | image[y*w*c + x*c + k] = tri_tex[k*tex_len+i]; 85 | } 86 | depth_buffer[y*w + x] = tri_depth[i]; 87 | } 88 | } 89 | } 90 | } 91 | } 92 | -------------------------------------------------------------------------------- /cmesh/cython/mesh_core.h: -------------------------------------------------------------------------------- 1 | #ifndef MESH_CORE_HPP_ 2 | #define MESH_CORE_HPP_ 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | using namespace std; 12 | 13 | class point 14 | { 15 | public: 16 | float x; 17 | float y; 18 | 19 | float dot(point p) 20 | { 21 | return this->x * p.x + this->y * p.y; 22 | } 23 | 24 | point operator-(const point& p) 25 | { 26 | point np; 27 | np.x = this->x - p.x; 28 | np.y = this->y - p.y; 29 | return np; 30 | } 31 | 32 | point operator+(const point& p) 33 | { 34 | point np; 35 | np.x = this->x + p.x; 36 | np.y = this->y + p.y; 37 | return np; 38 | } 39 | 40 | point operator*(float s) 41 | { 42 | point np; 43 | np.x = s * this->x; 44 | np.y = s * this->y; 45 | return np; 46 | } 47 | }; 48 | 49 | 50 | bool isPointInTri(point p, point p0, point p1, point p2, int h, int w); 51 | 52 | void _render_texture_core( 53 | float* image, float* vertices, int* triangles, 54 | float* tri_depth, float* tri_tex, float* depth_buffer, 55 | int ver_len, int tri_len, int tex_len, int h, int w, int c); 56 | 57 | #endif -------------------------------------------------------------------------------- /cmesh/cython/mesh_core_cython.cp35-win_amd64.pyd: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/cmesh/cython/mesh_core_cython.cp35-win_amd64.pyd -------------------------------------------------------------------------------- /cmesh/cython/mesh_core_cython.pyx: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Tue Nov 20 10:33:02 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import numpy as np 9 | cimport numpy as np 10 | from libcpp.string cimport string 11 | 12 | # use the Numpy-C-API from Cython 13 | np.import_array() 14 | 15 | cdef extern from "mesh_core.h": 16 | 17 | void _render_texture_core( 18 | float* image, float* vertices, int* triangles, 19 | float* tri_depth, float* tri_tex, float* depth_buffer, 20 | int ver_len, int tri_len, int tex_len, int h, int w, int c) 21 | 22 | def render_texture_core(np.ndarray[float, ndim=3, mode="c"] image not None, 23 | np.ndarray[float, ndim=2, mode="c"] vertices not None, 24 | np.ndarray[int, ndim=2, mode="c"] triangles not None, 25 | np.ndarray[float, ndim=1, mode="c"] tri_depth not None, 26 | np.ndarray[float, ndim=2, mode="c"] tri_tex not None, 27 | np.ndarray[float, ndim=2, mode="c"] depth_buffer not None, 28 | int ver_len, int tri_len, int tex_len, int h, int w, int c): 29 | _render_texture_core( np.PyArray_DATA(image), np.PyArray_DATA(vertices), 30 | np.PyArray_DATA(triangles), np.PyArray_DATA(tri_depth), 31 | np.PyArray_DATA(tri_tex), np.PyArray_DATA(depth_buffer), 32 | ver_len, tri_len, tex_len, h, w, c) 33 | -------------------------------------------------------------------------------- /cmesh/cython/setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Tue Nov 20 09:50:46 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import numpy as np 9 | from distutils.core import setup 10 | from distutils.extension import Extension 11 | from Cython.Build import cythonize 12 | from Cython.Distutils import build_ext 13 | 14 | setup( 15 | name='mesh_core_cython', 16 | cmdclass={'build_ext': build_ext}, 17 | ext_modules=[Extension('mesh_core_cython', 18 | sources=['mesh_core_cython.pyx', 'mesh_core.cpp'], 19 | language='c++', 20 | include_dirs=[np.get_include()])], 21 | ) -------------------------------------------------------------------------------- /cmesh/readme.txt: -------------------------------------------------------------------------------- 1 | SET VS90COMNTOOLS=%VS120COMNTOOLS% 2 | python setup.py build_ext -i -------------------------------------------------------------------------------- /cmesh/render.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Tue Nov 20 10:04:13 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | from __future__ import absolute_import 9 | from __future__ import division 10 | from __future__ import print_function 11 | 12 | import numpy as np 13 | from .cython import mesh_core_cython 14 | 15 | def render_texture(vertices, colors, triangles, h, w, c=3): 16 | 17 | image = np.zeros((h, w, c), dtype=np.float32, order='C') 18 | depth_buffer = np.zeros([h, w], dtype=np.float32, order='C') - 999999. 19 | tri_depth = (vertices[2, triangles[0,:]] + vertices[2, triangles[1,:]] + vertices[2, triangles[2,:]])/3. 20 | tri_tex = (colors[:, triangles[0,:]] + colors[:, triangles[1,:]] + colors[:, triangles[2,:]])/3. 21 | 22 | vertices = vertices.astype(np.float32).copy() 23 | triangles = triangles.astype(np.int32).copy() 24 | tri_depth = tri_depth.astype(np.float32).copy() 25 | tri_tex = tri_tex.astype(np.float32).copy() 26 | 27 | mesh_core_cython.render_texture_core( 28 | image, vertices, triangles, 29 | tri_depth, tri_tex, depth_buffer, 30 | vertices.shape[1], triangles.shape[1], tri_tex.shape[1], 31 | h, w, c) 32 | 33 | return image 34 | -------------------------------------------------------------------------------- /color_correction.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | """ Color corretion functions""" 4 | def hist_match(source, template): 5 | # Code borrow from: 6 | # https://stackoverflow.com/questions/32655686/histogram-matching-of-two-images-in-python-2-x 7 | oldshape = source.shape 8 | source = source.ravel() 9 | template = template.ravel() 10 | s_values, bin_idx, s_counts = np.unique(source, return_inverse=True, 11 | return_counts=True) 12 | t_values, t_counts = np.unique(template, return_counts=True) 13 | 14 | s_quantiles = np.cumsum(s_counts).astype(np.float64) 15 | s_quantiles /= s_quantiles[-1] 16 | t_quantiles = np.cumsum(t_counts).astype(np.float64) 17 | t_quantiles /= t_quantiles[-1] 18 | interp_t_values = np.interp(s_quantiles, t_quantiles, t_values) 19 | 20 | return interp_t_values[bin_idx].reshape(oldshape) 21 | 22 | def color_hist_match(src_im, tar_im): 23 | 24 | matched_R = hist_match(src_im[:,:,0], tar_im[:,:,0]) 25 | matched_G = hist_match(src_im[:,:,1], tar_im[:,:,1]) 26 | matched_B = hist_match(src_im[:,:,2], tar_im[:,:,2]) 27 | matched = np.stack((matched_R, matched_G, matched_B), axis=2).astype(np.float32) 28 | 29 | return matched 30 | 31 | def adain(src_im, tar_im, eps=1e-7): 32 | # https://github.com/ftokarev/tf-adain/blob/master/adain/norm.py 33 | mt = np.mean(tar_im, axis=(0,1)) 34 | st = np.std(tar_im, axis=(0,1)) 35 | ms = np.mean(src_im, axis=(0,1)) 36 | ss = np.std(src_im, axis=(0,1)) 37 | if ss.any() <= eps: return src_im 38 | result = st * (src_im.astype(np.float32) - ms) / (ss+eps) + mt 39 | if result.min() < 0: 40 | result = result - result.min() 41 | if result.max() > 255: 42 | result = (255.0/result.max()*result).astype(np.float32) 43 | return result -------------------------------------------------------------------------------- /convert.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Mon Jul 23 09:52:19 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import os 9 | import tensorflow as tf 10 | from tensorflow.python.tools import freeze_graph 11 | 12 | from tmodel import PRNet 13 | 14 | if __name__ == "__main__": 15 | 16 | os.environ['CUDA_VISIBLE_DEVICES']='0' 17 | 18 | x, pos = PRNet() 19 | saver = tf.train.Saver() 20 | with tf.Session() as sess: 21 | sess.run(tf.local_variables_initializer()) 22 | sess.run(tf.global_variables_initializer()) 23 | ckpt = tf.train.latest_checkpoint("model") 24 | saver.restore(sess, ckpt) 25 | 26 | print(x.name, pos.name) 27 | ''' 28 | nodes = [node.name for node in sess.graph.as_graph_def().node] 29 | node_name = 'PRNet/Conv2d_transpose_16' 30 | for node in nodes: 31 | if node[:len(node_name)] == node_name: 32 | print(node) 33 | ''' 34 | tf.train.write_graph(sess.graph.as_graph_def(), 'model', 'model_graph.pb') 35 | freeze_graph.freeze_graph('model/model_graph.pb', 36 | '', 37 | False, 38 | ckpt, 39 | pos.name[:-2], 40 | 'save/restore_all', 41 | 'save/Const:0', 42 | 'model/pico_3dFace_model.pb', 43 | False, 44 | "") 45 | -------------------------------------------------------------------------------- /face_detect.py: -------------------------------------------------------------------------------- 1 | """ Tensorflow implementation of the face detection / alignment algorithm found at 2 | https://github.com/kpzhang93/MTCNN_face_detection_alignment 3 | """ 4 | # MIT License 5 | # 6 | # Copyright (c) 2016 David Sandberg 7 | # 8 | # Permission is hereby granted, free of charge, to any person obtaining a copy 9 | # of this software and associated documentation files (the "Software"), to deal 10 | # in the Software without restriction, including without limitation the rights 11 | # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 12 | # copies of the Software, and to permit persons to whom the Software is 13 | # furnished to do so, subject to the following conditions: 14 | # 15 | # The above copyright notice and this permission notice shall be included in all 16 | # copies or substantial portions of the Software. 17 | # 18 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 19 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 20 | # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 21 | # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 22 | # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 23 | # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 24 | # SOFTWARE. 25 | 26 | from __future__ import absolute_import 27 | from __future__ import division 28 | from __future__ import print_function 29 | from six import string_types, iteritems 30 | 31 | import numpy as np 32 | import tensorflow as tf 33 | #from math import floor 34 | import cv2 35 | import os 36 | 37 | def layer(op): 38 | """Decorator for composable network layers.""" 39 | 40 | def layer_decorated(self, *args, **kwargs): 41 | # Automatically set a name if not provided. 42 | name = kwargs.setdefault('name', self.get_unique_name(op.__name__)) 43 | # Figure out the layer inputs. 44 | if len(self.terminals) == 0: 45 | raise RuntimeError('No input variables found for layer %s.' % name) 46 | elif len(self.terminals) == 1: 47 | layer_input = self.terminals[0] 48 | else: 49 | layer_input = list(self.terminals) 50 | # Perform the operation and get the output. 51 | layer_output = op(self, layer_input, *args, **kwargs) 52 | # Add to layer LUT. 53 | self.layers[name] = layer_output 54 | # This output is now the input for the next layer. 55 | self.feed(layer_output) 56 | # Return self for chained calls. 57 | return self 58 | 59 | return layer_decorated 60 | 61 | class Network(object): 62 | 63 | def __init__(self, inputs, trainable=True): 64 | # The input nodes for this network 65 | self.inputs = inputs 66 | # The current list of terminal nodes 67 | self.terminals = [] 68 | # Mapping from layer names to layers 69 | self.layers = dict(inputs) 70 | # If true, the resulting variables are set as trainable 71 | self.trainable = trainable 72 | 73 | self.setup() 74 | 75 | def setup(self): 76 | """Construct the network. """ 77 | raise NotImplementedError('Must be implemented by the subclass.') 78 | 79 | def load(self, data_path, session, ignore_missing=False): 80 | """Load network weights. 81 | data_path: The path to the numpy-serialized network weights 82 | session: The current TensorFlow session 83 | ignore_missing: If true, serialized weights for missing layers are ignored. 84 | """ 85 | data_dict = np.load(data_path, encoding='latin1').item() #pylint: disable=no-member 86 | 87 | for op_name in data_dict: 88 | with tf.variable_scope(op_name, reuse=True): 89 | for param_name, data in iteritems(data_dict[op_name]): 90 | try: 91 | var = tf.get_variable(param_name) 92 | session.run(var.assign(data)) 93 | except ValueError: 94 | if not ignore_missing: 95 | raise 96 | 97 | def feed(self, *args): 98 | """Set the input(s) for the next operation by replacing the terminal nodes. 99 | The arguments can be either layer names or the actual layers. 100 | """ 101 | assert len(args) != 0 102 | self.terminals = [] 103 | for fed_layer in args: 104 | if isinstance(fed_layer, string_types): 105 | try: 106 | fed_layer = self.layers[fed_layer] 107 | except KeyError: 108 | raise KeyError('Unknown layer name fed: %s' % fed_layer) 109 | self.terminals.append(fed_layer) 110 | return self 111 | 112 | def get_output(self): 113 | """Returns the current network output.""" 114 | return self.terminals[-1] 115 | 116 | def get_unique_name(self, prefix): 117 | """Returns an index-suffixed unique name for the given prefix. 118 | This is used for auto-generating layer names based on the type-prefix. 119 | """ 120 | ident = sum(t.startswith(prefix) for t, _ in self.layers.items()) + 1 121 | return '%s_%d' % (prefix, ident) 122 | 123 | def make_var(self, name, shape): 124 | """Creates a new TensorFlow variable.""" 125 | return tf.get_variable(name, shape, trainable=self.trainable) 126 | 127 | def validate_padding(self, padding): 128 | """Verifies that the padding is one of the supported ones.""" 129 | assert padding in ('SAME', 'VALID') 130 | 131 | @layer 132 | def conv(self, 133 | inp, 134 | k_h, 135 | k_w, 136 | c_o, 137 | s_h, 138 | s_w, 139 | name, 140 | relu=True, 141 | padding='SAME', 142 | group=1, 143 | biased=True): 144 | # Verify that the padding is acceptable 145 | self.validate_padding(padding) 146 | # Get the number of channels in the input 147 | c_i = int(inp.get_shape()[-1]) 148 | # Verify that the grouping parameter is valid 149 | assert c_i % group == 0 150 | assert c_o % group == 0 151 | # Convolution for a given input and kernel 152 | convolve = lambda i, k: tf.nn.conv2d(i, k, [1, s_h, s_w, 1], padding=padding) 153 | with tf.variable_scope(name) as scope: 154 | kernel = self.make_var('weights', shape=[k_h, k_w, c_i // group, c_o]) 155 | # This is the common-case. Convolve the input without any further complications. 156 | output = convolve(inp, kernel) 157 | # Add the biases 158 | if biased: 159 | biases = self.make_var('biases', [c_o]) 160 | output = tf.nn.bias_add(output, biases) 161 | if relu: 162 | # ReLU non-linearity 163 | output = tf.nn.relu(output, name=scope.name) 164 | return output 165 | 166 | @layer 167 | def prelu(self, inp, name): 168 | with tf.variable_scope(name): 169 | i = int(inp.get_shape()[-1]) 170 | alpha = self.make_var('alpha', shape=(i,)) 171 | output = tf.nn.relu(inp) + tf.multiply(alpha, -tf.nn.relu(-inp)) 172 | return output 173 | 174 | @layer 175 | def max_pool(self, inp, k_h, k_w, s_h, s_w, name, padding='SAME'): 176 | self.validate_padding(padding) 177 | return tf.nn.max_pool(inp, 178 | ksize=[1, k_h, k_w, 1], 179 | strides=[1, s_h, s_w, 1], 180 | padding=padding, 181 | name=name) 182 | 183 | @layer 184 | def fc(self, inp, num_out, name, relu=True): 185 | with tf.variable_scope(name): 186 | input_shape = inp.get_shape() 187 | if input_shape.ndims == 4: 188 | # The input is spatial. Vectorize it first. 189 | dim = 1 190 | for d in input_shape[1:].as_list(): 191 | dim *= int(d) 192 | feed_in = tf.reshape(inp, [-1, dim]) 193 | else: 194 | feed_in, dim = (inp, input_shape[-1].value) 195 | weights = self.make_var('weights', shape=[dim, num_out]) 196 | biases = self.make_var('biases', [num_out]) 197 | op = tf.nn.relu_layer if relu else tf.nn.xw_plus_b 198 | fc = op(feed_in, weights, biases, name=name) 199 | return fc 200 | 201 | 202 | """ 203 | Multi dimensional softmax, 204 | refer to https://github.com/tensorflow/tensorflow/issues/210 205 | compute softmax along the dimension of target 206 | the native softmax only supports batch_size x dimension 207 | """ 208 | @layer 209 | def softmax(self, target, axis, name=None): 210 | max_axis = tf.reduce_max(target, axis, keepdims=True) 211 | target_exp = tf.exp(target-max_axis) 212 | normalize = tf.reduce_sum(target_exp, axis, keepdims=True) 213 | softmax = tf.div(target_exp, normalize, name) 214 | return softmax 215 | 216 | class PNet(Network): 217 | def setup(self): 218 | (self.feed('data') #pylint: disable=no-value-for-parameter, no-member 219 | .conv(3, 3, 10, 1, 1, padding='VALID', relu=False, name='conv1') 220 | .prelu(name='PReLU1') 221 | .max_pool(2, 2, 2, 2, name='pool1') 222 | .conv(3, 3, 16, 1, 1, padding='VALID', relu=False, name='conv2') 223 | .prelu(name='PReLU2') 224 | .conv(3, 3, 32, 1, 1, padding='VALID', relu=False, name='conv3') 225 | .prelu(name='PReLU3') 226 | .conv(1, 1, 2, 1, 1, relu=False, name='conv4-1') 227 | .softmax(3,name='prob1')) 228 | 229 | (self.feed('PReLU3') #pylint: disable=no-value-for-parameter 230 | .conv(1, 1, 4, 1, 1, relu=False, name='conv4-2')) 231 | 232 | class RNet(Network): 233 | def setup(self): 234 | (self.feed('data') #pylint: disable=no-value-for-parameter, no-member 235 | .conv(3, 3, 28, 1, 1, padding='VALID', relu=False, name='conv1') 236 | .prelu(name='prelu1') 237 | .max_pool(3, 3, 2, 2, name='pool1') 238 | .conv(3, 3, 48, 1, 1, padding='VALID', relu=False, name='conv2') 239 | .prelu(name='prelu2') 240 | .max_pool(3, 3, 2, 2, padding='VALID', name='pool2') 241 | .conv(2, 2, 64, 1, 1, padding='VALID', relu=False, name='conv3') 242 | .prelu(name='prelu3') 243 | .fc(128, relu=False, name='conv4') 244 | .prelu(name='prelu4') 245 | .fc(2, relu=False, name='conv5-1') 246 | .softmax(1,name='prob1')) 247 | 248 | (self.feed('prelu4') #pylint: disable=no-value-for-parameter 249 | .fc(4, relu=False, name='conv5-2')) 250 | 251 | class ONet(Network): 252 | def setup(self): 253 | (self.feed('data') #pylint: disable=no-value-for-parameter, no-member 254 | .conv(3, 3, 32, 1, 1, padding='VALID', relu=False, name='conv1') 255 | .prelu(name='prelu1') 256 | .max_pool(3, 3, 2, 2, name='pool1') 257 | .conv(3, 3, 64, 1, 1, padding='VALID', relu=False, name='conv2') 258 | .prelu(name='prelu2') 259 | .max_pool(3, 3, 2, 2, padding='VALID', name='pool2') 260 | .conv(3, 3, 64, 1, 1, padding='VALID', relu=False, name='conv3') 261 | .prelu(name='prelu3') 262 | .max_pool(2, 2, 2, 2, name='pool3') 263 | .conv(2, 2, 128, 1, 1, padding='VALID', relu=False, name='conv4') 264 | .prelu(name='prelu4') 265 | .fc(256, relu=False, name='conv5') 266 | .prelu(name='prelu5') 267 | .fc(2, relu=False, name='conv6-1') 268 | .softmax(1, name='prob1')) 269 | 270 | (self.feed('prelu5') #pylint: disable=no-value-for-parameter 271 | .fc(4, relu=False, name='conv6-2')) 272 | 273 | (self.feed('prelu5') #pylint: disable=no-value-for-parameter 274 | .fc(10, relu=False, name='conv6-3')) 275 | 276 | def create_mtcnn(sess, model_path): 277 | if not model_path: 278 | model_path,_ = os.path.split(os.path.realpath(__file__)) 279 | 280 | with tf.variable_scope('pnet'): 281 | data = tf.placeholder(tf.float32, (None,None,None,3), 'input') 282 | pnet = PNet({'data':data}) 283 | pnet.load(os.path.join(model_path, 'mtcnn_weights/det1.npy'), sess) 284 | with tf.variable_scope('rnet'): 285 | data = tf.placeholder(tf.float32, (None,24,24,3), 'input') 286 | rnet = RNet({'data':data}) 287 | rnet.load(os.path.join(model_path, 'mtcnn_weights/det2.npy'), sess) 288 | with tf.variable_scope('onet'): 289 | data = tf.placeholder(tf.float32, (None,48,48,3), 'input') 290 | onet = ONet({'data':data}) 291 | onet.load(os.path.join(model_path, 'mtcnn_weights/det3.npy'), sess) 292 | 293 | pnet_fun = lambda img : sess.run(('pnet/conv4-2/BiasAdd:0', 'pnet/prob1:0'), feed_dict={'pnet/input:0':img}) 294 | rnet_fun = lambda img : sess.run(('rnet/conv5-2/conv5-2:0', 'rnet/prob1:0'), feed_dict={'rnet/input:0':img}) 295 | onet_fun = lambda img : sess.run(('onet/conv6-2/conv6-2:0', 'onet/conv6-3/conv6-3:0', 'onet/prob1:0'), feed_dict={'onet/input:0':img}) 296 | return pnet_fun, rnet_fun, onet_fun 297 | 298 | def detect_face(img, minsize, pnet, rnet, onet, threshold, factor): 299 | """Detects faces in an image, and returns bounding boxes and points for them. 300 | img: input image 301 | minsize: minimum faces' size 302 | pnet, rnet, onet: caffemodel 303 | threshold: threshold=[th1, th2, th3], th1-3 are three steps's threshold 304 | factor: the factor used to create a scaling pyramid of face sizes to detect in the image. 305 | """ 306 | factor_count=0 307 | total_boxes=np.empty((0,9)) 308 | points=np.empty(0) 309 | h=img.shape[0] 310 | w=img.shape[1] 311 | minl=np.amin([h, w]) 312 | m=12.0/minsize 313 | minl=minl*m 314 | # create scale pyramid 315 | scales=[] 316 | while minl>=12: 317 | scales += [m*np.power(factor, factor_count)] 318 | minl = minl*factor 319 | factor_count += 1 320 | 321 | # first stage 322 | for scale in scales: 323 | hs=int(np.ceil(h*scale)) 324 | ws=int(np.ceil(w*scale)) 325 | im_data = imresample(img, (hs, ws)) 326 | im_data = (im_data-127.5)*0.0078125 327 | img_x = np.expand_dims(im_data, 0) 328 | img_y = np.transpose(img_x, (0,2,1,3)) 329 | out = pnet(img_y) 330 | out0 = np.transpose(out[0], (0,2,1,3)) 331 | out1 = np.transpose(out[1], (0,2,1,3)) 332 | 333 | boxes, _ = generateBoundingBox(out1[0,:,:,1].copy(), out0[0,:,:,:].copy(), scale, threshold[0]) 334 | 335 | # inter-scale nms 336 | pick = nms(boxes.copy(), 0.5, 'Union') 337 | if boxes.size>0 and pick.size>0: 338 | boxes = boxes[pick,:] 339 | total_boxes = np.append(total_boxes, boxes, axis=0) 340 | 341 | numbox = total_boxes.shape[0] 342 | if numbox>0: 343 | pick = nms(total_boxes.copy(), 0.7, 'Union') 344 | total_boxes = total_boxes[pick,:] 345 | regw = total_boxes[:,2]-total_boxes[:,0] 346 | regh = total_boxes[:,3]-total_boxes[:,1] 347 | qq1 = total_boxes[:,0]+total_boxes[:,5]*regw 348 | qq2 = total_boxes[:,1]+total_boxes[:,6]*regh 349 | qq3 = total_boxes[:,2]+total_boxes[:,7]*regw 350 | qq4 = total_boxes[:,3]+total_boxes[:,8]*regh 351 | total_boxes = np.transpose(np.vstack([qq1, qq2, qq3, qq4, total_boxes[:,4]])) 352 | total_boxes = rerec(total_boxes.copy()) 353 | total_boxes[:,0:4] = np.fix(total_boxes[:,0:4]).astype(np.int32) 354 | dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph = pad(total_boxes.copy(), w, h) 355 | 356 | numbox = total_boxes.shape[0] 357 | if numbox>0: 358 | # second stage 359 | tempimg = np.zeros((24,24,3,numbox)) 360 | for k in range(0,numbox): 361 | tmp = np.zeros((int(tmph[k]),int(tmpw[k]),3)) 362 | tmp[dy[k]-1:edy[k],dx[k]-1:edx[k],:] = img[y[k]-1:ey[k],x[k]-1:ex[k],:] 363 | if tmp.shape[0]>0 and tmp.shape[1]>0 or tmp.shape[0]==0 and tmp.shape[1]==0: 364 | tempimg[:,:,:,k] = imresample(tmp, (24, 24)) 365 | else: 366 | return np.empty() 367 | tempimg = (tempimg-127.5)*0.0078125 368 | tempimg1 = np.transpose(tempimg, (3,1,0,2)) 369 | out = rnet(tempimg1) 370 | out0 = np.transpose(out[0]) 371 | out1 = np.transpose(out[1]) 372 | score = out1[1,:] 373 | ipass = np.where(score>threshold[1]) 374 | total_boxes = np.hstack([total_boxes[ipass[0],0:4].copy(), np.expand_dims(score[ipass].copy(),1)]) 375 | mv = out0[:,ipass[0]] 376 | if total_boxes.shape[0]>0: 377 | pick = nms(total_boxes, 0.7, 'Union') 378 | total_boxes = total_boxes[pick,:] 379 | total_boxes = bbreg(total_boxes.copy(), np.transpose(mv[:,pick])) 380 | total_boxes = rerec(total_boxes.copy()) 381 | 382 | numbox = total_boxes.shape[0] 383 | if numbox>0: 384 | # third stage 385 | total_boxes = np.fix(total_boxes).astype(np.int32) 386 | dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph = pad(total_boxes.copy(), w, h) 387 | tempimg = np.zeros((48,48,3,numbox)) 388 | for k in range(0,numbox): 389 | tmp = np.zeros((int(tmph[k]),int(tmpw[k]),3)) 390 | tmp[dy[k]-1:edy[k],dx[k]-1:edx[k],:] = img[y[k]-1:ey[k],x[k]-1:ex[k],:] 391 | if tmp.shape[0]>0 and tmp.shape[1]>0 or tmp.shape[0]==0 and tmp.shape[1]==0: 392 | tempimg[:,:,:,k] = imresample(tmp, (48, 48)) 393 | else: 394 | return np.empty() 395 | tempimg = (tempimg-127.5)*0.0078125 396 | tempimg1 = np.transpose(tempimg, (3,1,0,2)) 397 | out = onet(tempimg1) 398 | out0 = np.transpose(out[0]) 399 | out1 = np.transpose(out[1]) 400 | out2 = np.transpose(out[2]) 401 | score = out2[1,:] 402 | points = out1 403 | ipass = np.where(score>threshold[2]) 404 | points = points[:,ipass[0]] 405 | total_boxes = np.hstack([total_boxes[ipass[0],0:4].copy(), np.expand_dims(score[ipass].copy(),1)]) 406 | mv = out0[:,ipass[0]] 407 | 408 | w = total_boxes[:,2]-total_boxes[:,0]+1 409 | h = total_boxes[:,3]-total_boxes[:,1]+1 410 | points[0:5,:] = np.tile(w,(5, 1))*points[0:5,:] + np.tile(total_boxes[:,0],(5, 1))-1 411 | points[5:10,:] = np.tile(h,(5, 1))*points[5:10,:] + np.tile(total_boxes[:,1],(5, 1))-1 412 | if total_boxes.shape[0]>0: 413 | total_boxes = bbreg(total_boxes.copy(), np.transpose(mv)) 414 | pick = nms(total_boxes.copy(), 0.7, 'Min') 415 | total_boxes = total_boxes[pick,:] 416 | points = points[:,pick] 417 | 418 | return total_boxes, points 419 | 420 | 421 | def bulk_detect_face(images, detection_window_size_ratio, pnet, rnet, onet, threshold, factor): 422 | """Detects faces in a list of images 423 | images: list containing input images 424 | detection_window_size_ratio: ratio of minimum face size to smallest image dimension 425 | pnet, rnet, onet: caffemodel 426 | threshold: threshold=[th1 th2 th3], th1-3 are three steps's threshold [0-1] 427 | factor: the factor used to create a scaling pyramid of face sizes to detect in the image. 428 | """ 429 | all_scales = [None] * len(images) 430 | images_with_boxes = [None] * len(images) 431 | 432 | for i in range(len(images)): 433 | images_with_boxes[i] = {'total_boxes': np.empty((0, 9))} 434 | 435 | # create scale pyramid 436 | for index, img in enumerate(images): 437 | all_scales[index] = [] 438 | h = img.shape[0] 439 | w = img.shape[1] 440 | minsize = int(detection_window_size_ratio * np.minimum(w, h)) 441 | factor_count = 0 442 | minl = np.amin([h, w]) 443 | if minsize <= 12: 444 | minsize = 12 445 | 446 | m = 12.0 / minsize 447 | minl = minl * m 448 | while minl >= 12: 449 | all_scales[index].append(m * np.power(factor, factor_count)) 450 | minl = minl * factor 451 | factor_count += 1 452 | 453 | # # # # # # # # # # # # # 454 | # first stage - fast proposal network (pnet) to obtain face candidates 455 | # # # # # # # # # # # # # 456 | 457 | images_obj_per_resolution = {} 458 | 459 | # TODO: use some type of rounding to number module 8 to increase probability that pyramid images will have the same resolution across input images 460 | 461 | for index, scales in enumerate(all_scales): 462 | h = images[index].shape[0] 463 | w = images[index].shape[1] 464 | 465 | for scale in scales: 466 | hs = int(np.ceil(h * scale)) 467 | ws = int(np.ceil(w * scale)) 468 | 469 | if (ws, hs) not in images_obj_per_resolution: 470 | images_obj_per_resolution[(ws, hs)] = [] 471 | 472 | im_data = imresample(images[index], (hs, ws)) 473 | im_data = (im_data - 127.5) * 0.0078125 474 | img_y = np.transpose(im_data, (1, 0, 2)) # caffe uses different dimensions ordering 475 | images_obj_per_resolution[(ws, hs)].append({'scale': scale, 'image': img_y, 'index': index}) 476 | 477 | for resolution in images_obj_per_resolution: 478 | images_per_resolution = [i['image'] for i in images_obj_per_resolution[resolution]] 479 | outs = pnet(images_per_resolution) 480 | 481 | for index in range(len(outs[0])): 482 | scale = images_obj_per_resolution[resolution][index]['scale'] 483 | image_index = images_obj_per_resolution[resolution][index]['index'] 484 | out0 = np.transpose(outs[0][index], (1, 0, 2)) 485 | out1 = np.transpose(outs[1][index], (1, 0, 2)) 486 | 487 | boxes, _ = generateBoundingBox(out1[:, :, 1].copy(), out0[:, :, :].copy(), scale, threshold[0]) 488 | 489 | # inter-scale nms 490 | pick = nms(boxes.copy(), 0.5, 'Union') 491 | if boxes.size > 0 and pick.size > 0: 492 | boxes = boxes[pick, :] 493 | images_with_boxes[image_index]['total_boxes'] = np.append(images_with_boxes[image_index]['total_boxes'], 494 | boxes, 495 | axis=0) 496 | 497 | for index, image_obj in enumerate(images_with_boxes): 498 | numbox = image_obj['total_boxes'].shape[0] 499 | if numbox > 0: 500 | h = images[index].shape[0] 501 | w = images[index].shape[1] 502 | pick = nms(image_obj['total_boxes'].copy(), 0.7, 'Union') 503 | image_obj['total_boxes'] = image_obj['total_boxes'][pick, :] 504 | regw = image_obj['total_boxes'][:, 2] - image_obj['total_boxes'][:, 0] 505 | regh = image_obj['total_boxes'][:, 3] - image_obj['total_boxes'][:, 1] 506 | qq1 = image_obj['total_boxes'][:, 0] + image_obj['total_boxes'][:, 5] * regw 507 | qq2 = image_obj['total_boxes'][:, 1] + image_obj['total_boxes'][:, 6] * regh 508 | qq3 = image_obj['total_boxes'][:, 2] + image_obj['total_boxes'][:, 7] * regw 509 | qq4 = image_obj['total_boxes'][:, 3] + image_obj['total_boxes'][:, 8] * regh 510 | image_obj['total_boxes'] = np.transpose(np.vstack([qq1, qq2, qq3, qq4, image_obj['total_boxes'][:, 4]])) 511 | image_obj['total_boxes'] = rerec(image_obj['total_boxes'].copy()) 512 | image_obj['total_boxes'][:, 0:4] = np.fix(image_obj['total_boxes'][:, 0:4]).astype(np.int32) 513 | dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph = pad(image_obj['total_boxes'].copy(), w, h) 514 | 515 | numbox = image_obj['total_boxes'].shape[0] 516 | tempimg = np.zeros((24, 24, 3, numbox)) 517 | 518 | if numbox > 0: 519 | for k in range(0, numbox): 520 | tmp = np.zeros((int(tmph[k]), int(tmpw[k]), 3)) 521 | tmp[dy[k] - 1:edy[k], dx[k] - 1:edx[k], :] = images[index][y[k] - 1:ey[k], x[k] - 1:ex[k], :] 522 | if tmp.shape[0] > 0 and tmp.shape[1] > 0 or tmp.shape[0] == 0 and tmp.shape[1] == 0: 523 | tempimg[:, :, :, k] = imresample(tmp, (24, 24)) 524 | else: 525 | return np.empty() 526 | 527 | tempimg = (tempimg - 127.5) * 0.0078125 528 | image_obj['rnet_input'] = np.transpose(tempimg, (3, 1, 0, 2)) 529 | 530 | # # # # # # # # # # # # # 531 | # second stage - refinement of face candidates with rnet 532 | # # # # # # # # # # # # # 533 | 534 | bulk_rnet_input = np.empty((0, 24, 24, 3)) 535 | for index, image_obj in enumerate(images_with_boxes): 536 | if 'rnet_input' in image_obj: 537 | bulk_rnet_input = np.append(bulk_rnet_input, image_obj['rnet_input'], axis=0) 538 | 539 | out = rnet(bulk_rnet_input) 540 | out0 = np.transpose(out[0]) 541 | out1 = np.transpose(out[1]) 542 | score = out1[1, :] 543 | 544 | i = 0 545 | for index, image_obj in enumerate(images_with_boxes): 546 | if 'rnet_input' not in image_obj: 547 | continue 548 | 549 | rnet_input_count = image_obj['rnet_input'].shape[0] 550 | score_per_image = score[i:i + rnet_input_count] 551 | out0_per_image = out0[:, i:i + rnet_input_count] 552 | 553 | ipass = np.where(score_per_image > threshold[1]) 554 | image_obj['total_boxes'] = np.hstack([image_obj['total_boxes'][ipass[0], 0:4].copy(), 555 | np.expand_dims(score_per_image[ipass].copy(), 1)]) 556 | 557 | mv = out0_per_image[:, ipass[0]] 558 | 559 | if image_obj['total_boxes'].shape[0] > 0: 560 | h = images[index].shape[0] 561 | w = images[index].shape[1] 562 | pick = nms(image_obj['total_boxes'], 0.7, 'Union') 563 | image_obj['total_boxes'] = image_obj['total_boxes'][pick, :] 564 | image_obj['total_boxes'] = bbreg(image_obj['total_boxes'].copy(), np.transpose(mv[:, pick])) 565 | image_obj['total_boxes'] = rerec(image_obj['total_boxes'].copy()) 566 | 567 | numbox = image_obj['total_boxes'].shape[0] 568 | 569 | if numbox > 0: 570 | tempimg = np.zeros((48, 48, 3, numbox)) 571 | image_obj['total_boxes'] = np.fix(image_obj['total_boxes']).astype(np.int32) 572 | dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph = pad(image_obj['total_boxes'].copy(), w, h) 573 | 574 | for k in range(0, numbox): 575 | tmp = np.zeros((int(tmph[k]), int(tmpw[k]), 3)) 576 | tmp[dy[k] - 1:edy[k], dx[k] - 1:edx[k], :] = images[index][y[k] - 1:ey[k], x[k] - 1:ex[k], :] 577 | if tmp.shape[0] > 0 and tmp.shape[1] > 0 or tmp.shape[0] == 0 and tmp.shape[1] == 0: 578 | tempimg[:, :, :, k] = imresample(tmp, (48, 48)) 579 | else: 580 | return np.empty() 581 | tempimg = (tempimg - 127.5) * 0.0078125 582 | image_obj['onet_input'] = np.transpose(tempimg, (3, 1, 0, 2)) 583 | 584 | i += rnet_input_count 585 | 586 | # # # # # # # # # # # # # 587 | # third stage - further refinement and facial landmarks positions with onet 588 | # # # # # # # # # # # # # 589 | 590 | bulk_onet_input = np.empty((0, 48, 48, 3)) 591 | for index, image_obj in enumerate(images_with_boxes): 592 | if 'onet_input' in image_obj: 593 | bulk_onet_input = np.append(bulk_onet_input, image_obj['onet_input'], axis=0) 594 | 595 | out = onet(bulk_onet_input) 596 | 597 | out0 = np.transpose(out[0]) 598 | out1 = np.transpose(out[1]) 599 | out2 = np.transpose(out[2]) 600 | score = out2[1, :] 601 | points = out1 602 | 603 | i = 0 604 | ret = [] 605 | for index, image_obj in enumerate(images_with_boxes): 606 | if 'onet_input' not in image_obj: 607 | ret.append(None) 608 | continue 609 | 610 | onet_input_count = image_obj['onet_input'].shape[0] 611 | 612 | out0_per_image = out0[:, i:i + onet_input_count] 613 | score_per_image = score[i:i + onet_input_count] 614 | points_per_image = points[:, i:i + onet_input_count] 615 | 616 | ipass = np.where(score_per_image > threshold[2]) 617 | points_per_image = points_per_image[:, ipass[0]] 618 | 619 | image_obj['total_boxes'] = np.hstack([image_obj['total_boxes'][ipass[0], 0:4].copy(), 620 | np.expand_dims(score_per_image[ipass].copy(), 1)]) 621 | mv = out0_per_image[:, ipass[0]] 622 | 623 | w = image_obj['total_boxes'][:, 2] - image_obj['total_boxes'][:, 0] + 1 624 | h = image_obj['total_boxes'][:, 3] - image_obj['total_boxes'][:, 1] + 1 625 | points_per_image[0:5, :] = np.tile(w, (5, 1)) * points_per_image[0:5, :] + np.tile( 626 | image_obj['total_boxes'][:, 0], (5, 1)) - 1 627 | points_per_image[5:10, :] = np.tile(h, (5, 1)) * points_per_image[5:10, :] + np.tile( 628 | image_obj['total_boxes'][:, 1], (5, 1)) - 1 629 | 630 | if image_obj['total_boxes'].shape[0] > 0: 631 | image_obj['total_boxes'] = bbreg(image_obj['total_boxes'].copy(), np.transpose(mv)) 632 | pick = nms(image_obj['total_boxes'].copy(), 0.7, 'Min') 633 | image_obj['total_boxes'] = image_obj['total_boxes'][pick, :] 634 | points_per_image = points_per_image[:, pick] 635 | 636 | ret.append((image_obj['total_boxes'], points_per_image)) 637 | else: 638 | ret.append(None) 639 | 640 | i += onet_input_count 641 | 642 | return ret 643 | 644 | 645 | # function [boundingbox] = bbreg(boundingbox,reg) 646 | def bbreg(boundingbox,reg): 647 | """Calibrate bounding boxes""" 648 | if reg.shape[1]==1: 649 | reg = np.reshape(reg, (reg.shape[2], reg.shape[3])) 650 | 651 | w = boundingbox[:,2]-boundingbox[:,0]+1 652 | h = boundingbox[:,3]-boundingbox[:,1]+1 653 | b1 = boundingbox[:,0]+reg[:,0]*w 654 | b2 = boundingbox[:,1]+reg[:,1]*h 655 | b3 = boundingbox[:,2]+reg[:,2]*w 656 | b4 = boundingbox[:,3]+reg[:,3]*h 657 | boundingbox[:,0:4] = np.transpose(np.vstack([b1, b2, b3, b4 ])) 658 | return boundingbox 659 | 660 | def generateBoundingBox(imap, reg, scale, t): 661 | """Use heatmap to generate bounding boxes""" 662 | stride=2 663 | cellsize=12 664 | 665 | imap = np.transpose(imap) 666 | dx1 = np.transpose(reg[:,:,0]) 667 | dy1 = np.transpose(reg[:,:,1]) 668 | dx2 = np.transpose(reg[:,:,2]) 669 | dy2 = np.transpose(reg[:,:,3]) 670 | y, x = np.where(imap >= t) 671 | if y.shape[0]==1: 672 | dx1 = np.flipud(dx1) 673 | dy1 = np.flipud(dy1) 674 | dx2 = np.flipud(dx2) 675 | dy2 = np.flipud(dy2) 676 | score = imap[(y,x)] 677 | reg = np.transpose(np.vstack([ dx1[(y,x)], dy1[(y,x)], dx2[(y,x)], dy2[(y,x)] ])) 678 | if reg.size==0: 679 | reg = np.empty((0,3)) 680 | bb = np.transpose(np.vstack([y,x])) 681 | q1 = np.fix((stride*bb+1)/scale) 682 | q2 = np.fix((stride*bb+cellsize-1+1)/scale) 683 | boundingbox = np.hstack([q1, q2, np.expand_dims(score,1), reg]) 684 | return boundingbox, reg 685 | 686 | # function pick = nms(boxes,threshold,type) 687 | def nms(boxes, threshold, method): 688 | if boxes.size==0: 689 | return np.empty((0,3)) 690 | x1 = boxes[:,0] 691 | y1 = boxes[:,1] 692 | x2 = boxes[:,2] 693 | y2 = boxes[:,3] 694 | s = boxes[:,4] 695 | area = (x2-x1+1) * (y2-y1+1) 696 | I = np.argsort(s) 697 | pick = np.zeros_like(s, dtype=np.int16) 698 | counter = 0 699 | while I.size>0: 700 | i = I[-1] 701 | pick[counter] = i 702 | counter += 1 703 | idx = I[0:-1] 704 | xx1 = np.maximum(x1[i], x1[idx]) 705 | yy1 = np.maximum(y1[i], y1[idx]) 706 | xx2 = np.minimum(x2[i], x2[idx]) 707 | yy2 = np.minimum(y2[i], y2[idx]) 708 | w = np.maximum(0.0, xx2-xx1+1) 709 | h = np.maximum(0.0, yy2-yy1+1) 710 | inter = w * h 711 | if method is 'Min': 712 | o = inter / np.minimum(area[i], area[idx]) 713 | else: 714 | o = inter / (area[i] + area[idx] - inter) 715 | I = I[np.where(o<=threshold)] 716 | pick = pick[0:counter] 717 | return pick 718 | 719 | # function [dy edy dx edx y ey x ex tmpw tmph] = pad(total_boxes,w,h) 720 | def pad(total_boxes, w, h): 721 | """Compute the padding coordinates (pad the bounding boxes to square)""" 722 | tmpw = (total_boxes[:,2]-total_boxes[:,0]+1).astype(np.int32) 723 | tmph = (total_boxes[:,3]-total_boxes[:,1]+1).astype(np.int32) 724 | numbox = total_boxes.shape[0] 725 | 726 | dx = np.ones((numbox), dtype=np.int32) 727 | dy = np.ones((numbox), dtype=np.int32) 728 | edx = tmpw.copy().astype(np.int32) 729 | edy = tmph.copy().astype(np.int32) 730 | 731 | x = total_boxes[:,0].copy().astype(np.int32) 732 | y = total_boxes[:,1].copy().astype(np.int32) 733 | ex = total_boxes[:,2].copy().astype(np.int32) 734 | ey = total_boxes[:,3].copy().astype(np.int32) 735 | 736 | tmp = np.where(ex>w) 737 | edx.flat[tmp] = np.expand_dims(-ex[tmp]+w+tmpw[tmp],1) 738 | ex[tmp] = w 739 | 740 | tmp = np.where(ey>h) 741 | edy.flat[tmp] = np.expand_dims(-ey[tmp]+h+tmph[tmp],1) 742 | ey[tmp] = h 743 | 744 | tmp = np.where(x<1) 745 | dx.flat[tmp] = np.expand_dims(2-x[tmp],1) 746 | x[tmp] = 1 747 | 748 | tmp = np.where(y<1) 749 | dy.flat[tmp] = np.expand_dims(2-y[tmp],1) 750 | y[tmp] = 1 751 | 752 | return dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph 753 | 754 | # function [bboxA] = rerec(bboxA) 755 | def rerec(bboxA): 756 | """Convert bboxA to square.""" 757 | h = bboxA[:,3]-bboxA[:,1] 758 | w = bboxA[:,2]-bboxA[:,0] 759 | l = np.maximum(w, h) 760 | bboxA[:,0] = bboxA[:,0]+w*0.5-l*0.5 761 | bboxA[:,1] = bboxA[:,1]+h*0.5-l*0.5 762 | bboxA[:,2:4] = bboxA[:,0:2] + np.transpose(np.tile(l,(2,1))) 763 | return bboxA 764 | 765 | def imresample(img, sz): 766 | im_data = cv2.resize(img, (sz[1], sz[0]), interpolation=cv2.INTER_AREA) #@UndefinedVariable 767 | return im_data 768 | 769 | # This method is kept for debugging purpose 770 | # h=img.shape[0] 771 | # w=img.shape[1] 772 | # hs, ws = sz 773 | # dx = float(w) / ws 774 | # dy = float(h) / hs 775 | # im_data = np.zeros((hs,ws,3)) 776 | # for a1 in range(0,hs): 777 | # for a2 in range(0,ws): 778 | # for a3 in range(0,3): 779 | # im_data[a1,a2,a3] = img[int(floor(a1*dy)),int(floor(a2*dx)),a3] 780 | # return im_data 781 | 782 | -------------------------------------------------------------------------------- /generate.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Wed Jul 18 14:48:43 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import os 9 | import cv2 10 | import glob 11 | import numpy as np 12 | from skimage.io import imread 13 | 14 | class Generator(object): 15 | 16 | def __init__(self, 17 | train_paths, 18 | mask_path, 19 | valid_path, 20 | image_shape = (256, 256, 3), 21 | batch_size = 16): 22 | 23 | train_images, train_labels = [], [] 24 | for head_path in train_paths: 25 | image_paths = glob.glob(os.path.join(head_path, '*.jpg')) 26 | label_paths = glob.glob(os.path.join(head_path, '*.npy')) 27 | image_paths = [image_path.replace('\\', '/') for image_path in image_paths] 28 | label_paths = [label_path.replace('\\', '/') for label_path in label_paths] 29 | for image_path in image_paths: 30 | image_name = image_path.split('/')[-1][:-4] 31 | npy_path = head_path + '/' + image_name + '.npy' 32 | if npy_path in label_paths: 33 | train_images.append(image_path) 34 | train_labels.append(npy_path) 35 | 36 | valid_images = glob.glob(os.path.join(valid_path, '*.jpg')) 37 | 38 | self.mask = cv2.imread(mask_path, 0).astype('float32') 39 | 40 | self.image_shape = image_shape 41 | self.batch_size = batch_size 42 | self.train_images = train_images 43 | self.train_labels = train_labels 44 | self.valid_images = valid_images 45 | self.train_length = len(self.train_images) 46 | self.valid_length = len(self.valid_images) 47 | 48 | def generate(self, is_training=True): 49 | while(True): 50 | if is_training: 51 | rand_idx = [one for one in range(self.train_length)] 52 | np.random.shuffle(rand_idx) 53 | self.train_images = [self.train_images[one] for one in rand_idx] 54 | self.train_labels = [self.train_labels[one] for one in rand_idx] 55 | 56 | counter = 0 57 | x_array, y_array, m_array = [], [], [] 58 | for index in range(self.train_length): 59 | image = imread(self.train_images[index]) / 255. 60 | label = np.load(self.train_labels[index]) 61 | 62 | x_array.append(image) 63 | y_array.append(label) 64 | m_array.append(self.mask) 65 | counter += 1 66 | 67 | if counter >= self.batch_size: 68 | yield (np.array(x_array), np.array(y_array), np.array(m_array)) 69 | counter = 0 70 | x_array, y_array, m_array = [], [], [] 71 | else: 72 | rand_idx = [one for one in range(self.valid_length)] 73 | np.random.shuffle(rand_idx) 74 | self.valid_images = [self.valid_images[one] for one in rand_idx] 75 | 76 | counter = 0; i_array = [] 77 | for index in range(self.valid_length): 78 | image = imread(self.valid_images[index]) / 255. 79 | i_array.append(image) 80 | 81 | counter += 1 82 | if counter >= self.batch_size: 83 | yield i_array 84 | counter = 0; i_array = [] 85 | -------------------------------------------------------------------------------- /images/canonical_vertices.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/images/canonical_vertices.npy -------------------------------------------------------------------------------- /images/readme.txt: -------------------------------------------------------------------------------- 1 | datasets: http://www.cbsr.ia.ac.cn/users/xiangyuzhu/projects/3DDFA/main.htm -------------------------------------------------------------------------------- /images/test.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/images/test.jpg -------------------------------------------------------------------------------- /images/test_g_kpt.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/images/test_g_kpt.jpg -------------------------------------------------------------------------------- /images/test_g_texture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/images/test_g_texture.png -------------------------------------------------------------------------------- /images/test_kpt.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/images/test_kpt.jpg -------------------------------------------------------------------------------- /images/test_mask.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/images/test_mask.jpg -------------------------------------------------------------------------------- /images/test_ver.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/images/test_ver.jpg -------------------------------------------------------------------------------- /images/uv_face.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/images/uv_face.png -------------------------------------------------------------------------------- /images/uv_face_eyes.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/images/uv_face_eyes.png -------------------------------------------------------------------------------- /images/uv_face_mask.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/images/uv_face_mask.png -------------------------------------------------------------------------------- /images/uv_kpt_ind.txt: -------------------------------------------------------------------------------- 1 | 1.500000000000000000e+01 2.200000000000000000e+01 2.600000000000000000e+01 3.200000000000000000e+01 4.500000000000000000e+01 6.700000000000000000e+01 9.100000000000000000e+01 1.120000000000000000e+02 1.280000000000000000e+02 1.430000000000000000e+02 1.640000000000000000e+02 1.880000000000000000e+02 2.100000000000000000e+02 2.230000000000000000e+02 2.290000000000000000e+02 2.330000000000000000e+02 2.400000000000000000e+02 5.800000000000000000e+01 7.100000000000000000e+01 8.500000000000000000e+01 9.700000000000000000e+01 1.060000000000000000e+02 1.490000000000000000e+02 1.580000000000000000e+02 1.700000000000000000e+02 1.840000000000000000e+02 1.970000000000000000e+02 1.280000000000000000e+02 1.280000000000000000e+02 1.280000000000000000e+02 1.280000000000000000e+02 1.170000000000000000e+02 1.220000000000000000e+02 1.280000000000000000e+02 1.330000000000000000e+02 1.380000000000000000e+02 7.800000000000000000e+01 8.600000000000000000e+01 9.500000000000000000e+01 1.020000000000000000e+02 9.600000000000000000e+01 8.700000000000000000e+01 1.530000000000000000e+02 1.600000000000000000e+02 1.690000000000000000e+02 1.770000000000000000e+02 1.680000000000000000e+02 1.590000000000000000e+02 1.080000000000000000e+02 1.160000000000000000e+02 1.240000000000000000e+02 1.280000000000000000e+02 1.310000000000000000e+02 1.390000000000000000e+02 1.460000000000000000e+02 1.370000000000000000e+02 1.320000000000000000e+02 1.280000000000000000e+02 1.230000000000000000e+02 1.180000000000000000e+02 1.100000000000000000e+02 1.220000000000000000e+02 1.280000000000000000e+02 1.330000000000000000e+02 1.450000000000000000e+02 1.320000000000000000e+02 1.280000000000000000e+02 1.230000000000000000e+02 2 | 9.600000000000000000e+01 1.180000000000000000e+02 1.410000000000000000e+02 1.650000000000000000e+02 1.830000000000000000e+02 1.900000000000000000e+02 1.880000000000000000e+02 1.870000000000000000e+02 1.930000000000000000e+02 1.870000000000000000e+02 1.880000000000000000e+02 1.900000000000000000e+02 1.830000000000000000e+02 1.650000000000000000e+02 1.410000000000000000e+02 1.180000000000000000e+02 9.600000000000000000e+01 4.900000000000000000e+01 4.200000000000000000e+01 3.900000000000000000e+01 4.000000000000000000e+01 4.200000000000000000e+01 4.200000000000000000e+01 4.000000000000000000e+01 3.900000000000000000e+01 4.200000000000000000e+01 4.900000000000000000e+01 5.900000000000000000e+01 7.300000000000000000e+01 8.600000000000000000e+01 9.600000000000000000e+01 1.110000000000000000e+02 1.130000000000000000e+02 1.150000000000000000e+02 1.130000000000000000e+02 1.110000000000000000e+02 6.700000000000000000e+01 6.000000000000000000e+01 6.100000000000000000e+01 6.500000000000000000e+01 6.800000000000000000e+01 6.900000000000000000e+01 6.500000000000000000e+01 6.100000000000000000e+01 6.000000000000000000e+01 6.700000000000000000e+01 6.900000000000000000e+01 6.800000000000000000e+01 1.420000000000000000e+02 1.310000000000000000e+02 1.270000000000000000e+02 1.280000000000000000e+02 1.270000000000000000e+02 1.310000000000000000e+02 1.420000000000000000e+02 1.480000000000000000e+02 1.500000000000000000e+02 1.500000000000000000e+02 1.500000000000000000e+02 1.480000000000000000e+02 1.410000000000000000e+02 1.350000000000000000e+02 1.340000000000000000e+02 1.350000000000000000e+02 1.420000000000000000e+02 1.430000000000000000e+02 1.420000000000000000e+02 1.430000000000000000e+02 3 | -------------------------------------------------------------------------------- /images/uv_weight_mask.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/images/uv_weight_mask.png -------------------------------------------------------------------------------- /images/uv_weight_mask2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/shen1994/DeepFaceRestruction/3c95fb986cc31a82f45ef3333d23136a3761e16a/images/uv_weight_mask2.png -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Thu Nov 8 11:10:07 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import tensorflow as tf 9 | import tensorflow.contrib.layers as tcl 10 | from tensorflow.contrib.framework import arg_scope 11 | 12 | def res_block(x, num_outputs, kernel_size=4, stride=1): 13 | assert num_outputs % 2 == 0 14 | shortcut = x 15 | if stride != 1 or int(x.get_shape()[3]) != num_outputs: 16 | shortcut = tcl.conv2d(shortcut, num_outputs, kernel_size=1, stride=stride, 17 | activation_fn=None, normalizer_fn=None) 18 | x = tcl.conv2d(x, num_outputs/2, kernel_size=1, stride=1, padding='SAME') 19 | x = tcl.conv2d(x, num_outputs/2, kernel_size=kernel_size, stride=stride, padding='SAME') 20 | x = tcl.conv2d(x, num_outputs, kernel_size=1, stride=1, padding='SAME', 21 | activation_fn=None, normalizer_fn=None) 22 | x += shortcut 23 | x = tcl.batch_norm(x) 24 | x = tf.nn.relu(x) 25 | 26 | return x 27 | 28 | def PRNet(image_shape=(256, 256, 3), output_shape=(256, 256,3), size=16, is_training=True): 29 | 30 | with tf.name_scope('3dface') as scope: 31 | x = tf.placeholder(shape=(None, image_shape[0], image_shape[1], image_shape[2]), 32 | dtype=tf.float32, name='x') 33 | m = tf.placeholder(shape=(None, image_shape[0], image_shape[1]), 34 | dtype=tf.float32, name='m') 35 | y = tf.placeholder(shape=(None, output_shape[0], output_shape[1], output_shape[2]), 36 | dtype=tf.float32, name='y') 37 | with tf.name_scope('PRNet') as scope: 38 | with arg_scope([tcl.batch_norm], is_training=is_training, scale=True): 39 | with arg_scope([tcl.conv2d, tcl.conv2d_transpose], activation_fn=tf.nn.relu, 40 | normalizer_fn=tcl.batch_norm, biases_initializer=None, 41 | padding='SAME', weights_regularizer=tcl.l2_regularizer(0.0002)): 42 | # encoder 43 | se = tcl.conv2d(x, num_outputs=size, kernel_size=4, stride=1) # 256 x 256 x 16 44 | se = res_block(se, num_outputs=size * 2, kernel_size=4, stride=2) # 128 x 128 x 32 45 | se = res_block(se, num_outputs=size * 2, kernel_size=4, stride=1) # 128 x 128 x 32 46 | se = res_block(se, num_outputs=size * 4, kernel_size=4, stride=2) # 64 x 64 x 64 47 | se = res_block(se, num_outputs=size * 4, kernel_size=4, stride=1) # 64 x 64 x 64 48 | se = res_block(se, num_outputs=size * 8, kernel_size=4, stride=2) # 32 x 32 x 128 49 | se = res_block(se, num_outputs=size * 8, kernel_size=4, stride=1) # 32 x 32 x 128 50 | se = res_block(se, num_outputs=size * 16, kernel_size=4, stride=2) # 16 x 16 x 256 51 | se = res_block(se, num_outputs=size * 16, kernel_size=4, stride=1) # 16 x 16 x 256 52 | se = res_block(se, num_outputs=size * 32, kernel_size=4, stride=2) # 8 x 8 x 512 53 | se = res_block(se, num_outputs=size * 32, kernel_size=4, stride=1) # 8 x 8 x 512 54 | # decoder 55 | pd = tcl.conv2d_transpose(se, size * 32, 4, stride=1) # 8 x 8 x 512 56 | pd = tcl.conv2d_transpose(pd, size * 16, 4, stride=2) # 16 x 16 x 256 57 | pd = tcl.conv2d_transpose(pd, size * 16, 4, stride=1) # 16 x 16 x 256 58 | pd = tcl.conv2d_transpose(pd, size * 16, 4, stride=1) # 16 x 16 x 256 59 | pd = tcl.conv2d_transpose(pd, size * 8, 4, stride=2) # 32 x 32 x 128 60 | pd = tcl.conv2d_transpose(pd, size * 8, 4, stride=1) # 32 x 32 x 128 61 | pd = tcl.conv2d_transpose(pd, size * 8, 4, stride=1) # 32 x 32 x 128 62 | pd = tcl.conv2d_transpose(pd, size * 4, 4, stride=2) # 64 x 64 x 64 63 | pd = tcl.conv2d_transpose(pd, size * 4, 4, stride=1) # 64 x 64 x 64 64 | pd = tcl.conv2d_transpose(pd, size * 4, 4, stride=1) # 64 x 64 x 64 65 | 66 | pd = tcl.conv2d_transpose(pd, size * 2, 4, stride=2) # 128 x 128 x 32 67 | pd = tcl.conv2d_transpose(pd, size * 2, 4, stride=1) # 128 x 128 x 32 68 | pd = tcl.conv2d_transpose(pd, size, 4, stride=2) # 256 x 256 x 16 69 | pd = tcl.conv2d_transpose(pd, size, 4, stride=1) # 256 x 256 x 16 70 | 71 | pd = tcl.conv2d_transpose(pd, 3, 4, stride=1) # 256 x 256 x 3 72 | pd = tcl.conv2d_transpose(pd, 3, 4, stride=1) # 256 x 256 x 3 73 | pos = tcl.conv2d_transpose(pd, 3, 4, stride=1, activation_fn = tf.nn.sigmoid) 74 | 75 | with tf.name_scope('loss') as scope: 76 | locc = tf.square(y - pos) 77 | loss_x = tf.reduce_mean(tf.multiply(locc[:,:,:,0], m)) 78 | loss_y = tf.reduce_mean(tf.multiply(locc[:,:,:,1], m)) 79 | loss_z = tf.reduce_mean(tf.multiply(locc[:,:,:,2], m)) 80 | loss = loss_x + loss_y + loss_z 81 | 82 | with tf.name_scope('optimizer') as scope: 83 | o_learning_rate = 1e-4 84 | global_step = tf.Variable(0) 85 | learning_rate = tf.train.exponential_decay(o_learning_rate, global_step, 40000, 0.1, staircase=True) 86 | optimizer = tf.train.AdamOptimizer(learning_rate).minimize(loss) 87 | 88 | return x, m, y, pos, loss, optimizer 89 | -------------------------------------------------------------------------------- /predictor.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Wed Nov 14 09:24:57 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import cv2 9 | import numpy as np 10 | from utils.cv_plot import plot_kpt 11 | 12 | class Predictor: 13 | 14 | def __init__(self, 15 | in_shape=(256, 256, 3), 16 | out_shape=(256, 256, 3), 17 | texture_size=256): 18 | self.in_shape = in_shape 19 | self.out_shape = out_shape 20 | self.texture_size = texture_size 21 | self.uv_kpt_ind = np.loadtxt('images/uv_kpt_ind.txt').astype(np.int32) 22 | self.face_ind = np.loadtxt("images/face_ind.txt").astype(np.int32) 23 | self.triangles = np.loadtxt("images/triangles.txt").astype(np.int32) 24 | self.uv_coords = self.generate_uv_coords() 25 | 26 | def generate_uv_coords(self): 27 | resolution = self.out_shape[0] 28 | uv_coords = np.meshgrid(range(resolution),range(resolution)) 29 | uv_coords = np.transpose(np.array(uv_coords), [1,2,0]) 30 | uv_coords = np.reshape(uv_coords, [resolution**2, -1]); 31 | uv_coords = uv_coords[self.face_ind, :] 32 | uv_coords = np.hstack((uv_coords[:,:2], np.zeros([uv_coords.shape[0], 1]))) 33 | 34 | return uv_coords 35 | 36 | def get_vertices(self, pos): 37 | ''' 38 | Args: 39 | pos: the 3D position map. shape = (256, 256, 3). 40 | Returns: 41 | vertices: the vertices(point cloud). shape = (num of points, 3). n is about 40K here. 42 | ''' 43 | all_vertices = np.reshape(pos, [self.out_shape[0]**2, -1]) 44 | vertices = all_vertices[self.face_ind, :] 45 | 46 | return vertices 47 | 48 | def get_colors(self, image, vertices): 49 | ''' 50 | Args: 51 | pos: the 3D position map. shape = (256, 256, 3). 52 | Returns: 53 | colors: the corresponding colors of vertices. shape = (num of points, 3). n is 45128 here. 54 | ''' 55 | [h, w, _] = image.shape 56 | vertices[:,0] = np.minimum(np.maximum(vertices[:,0], 0), w - 1) # x 57 | vertices[:,1] = np.minimum(np.maximum(vertices[:,1], 0), h - 1) # y 58 | ind = np.round(vertices).astype(np.int32) 59 | colors = image[ind[:,1], ind[:,0], :] # n x 3 60 | 61 | return colors 62 | 63 | def get_landmarks(self, pos): 64 | ''' 65 | Args: 66 | pos: the 3D position map. shape = (256, 256, 3). 67 | Returns: 68 | kpt: 68 3D landmarks. shape = (68, 3). 69 | ''' 70 | kpt = pos[self.uv_kpt_ind[1,:], self.uv_kpt_ind[0,:], :] 71 | return kpt 72 | 73 | def predictor(self, sess, x, y, image): 74 | 75 | # run model to get uv_map 76 | pos = sess.run(y, feed_dict={x: image[np.newaxis, :,:,:]}) 77 | pos = np.squeeze(pos) 78 | max_pos = self.in_shape[0] 79 | pos = pos * max_pos 80 | 81 | t_image = (image*255.).astype(np.uint8) 82 | kpt = self.get_landmarks(pos) 83 | kpt_origin = plot_kpt(image, kpt).astype(np.uint8) 84 | kpt_gray = cv2.cvtColor(kpt_origin, cv2.COLOR_RGB2GRAY) 85 | ret, kpt_mask = cv2.threshold(kpt_gray, 127, 255, cv2.THRESH_BINARY) 86 | kpt_mask = cv2.bitwise_not(kpt_mask) 87 | kpt_and = cv2.bitwise_and(t_image, t_image, mask=kpt_mask) 88 | kpt_image = cv2.add(kpt_and, kpt_origin) 89 | 90 | return kpt_image / 255. 91 | 92 | -------------------------------------------------------------------------------- /prepare.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Fri Nov 9 17:13:50 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import os 9 | import glob 10 | import numpy as np 11 | import scipy.io as sio 12 | import skimage.transform 13 | from skimage import io 14 | from skimage.io import imsave 15 | from multiprocessing import Process 16 | from multiprocessing.managers import BaseManager 17 | import sys 18 | sys.path.append('..') 19 | 20 | from mm3d import mesh 21 | from mm3d.morphable_model import MorphabelModel 22 | from mm3d.morphable_model.load import load_uv_coords 23 | 24 | def postmap_from_300WLP(bfm, kpt_ind, triangles, uv_coords, image_path, mat_path, 25 | uv_h = 256, uv_w = 256, image_h = 256, image_w = 256): 26 | 27 | # 1. load image and fitted parameters 28 | image = io.imread(image_path)/255. 29 | [h, w, c] = image.shape 30 | 31 | info = sio.loadmat(mat_path) 32 | pose_para = info['Pose_Para'].T.astype(np.float32) 33 | shape_para = info['Shape_Para'].astype(np.float32) 34 | exp_para = info['Exp_Para'].astype(np.float32) 35 | 36 | # 2. generate mesh 37 | vertices = bfm.generate_vertices(shape_para, exp_para) 38 | s = pose_para[-1, 0] 39 | angles = pose_para[:3, 0] 40 | t = pose_para[3:6, 0] 41 | transformed_vertices = bfm.transform_3ddfa(vertices, s, angles, t) 42 | # using stantard camera & orth projection as in 3DDFA 43 | projected_vertices = transformed_vertices.copy() 44 | image_vertices = projected_vertices.copy() 45 | image_vertices[:,1] = h - image_vertices[:,1] - 1 46 | 47 | # 3. crop image with key points 48 | kpt = image_vertices[kpt_ind, :].astype(np.int32) 49 | left = np.min(kpt[:, 0]) 50 | right = np.max(kpt[:, 0]) 51 | top = np.min(kpt[:, 1]) 52 | bottom = np.max(kpt[:, 1]) 53 | center = np.array([right - (right - left) / 2.0, 54 | bottom - (bottom - top) / 2.0]) 55 | old_size = (right - left + bottom - top)/2 56 | size = int(old_size * 1.5) 57 | # random pertube. you can change the numbers 58 | marg = old_size * 0.1 59 | t_x = np.random.rand() * marg * 2 - marg 60 | t_y = np.random.rand() * marg * 2 - marg 61 | center[0] = center[0] + t_x; center[1] = center[1] + t_y 62 | size = size*(np.random.rand()*0.2 + 0.9) 63 | 64 | # crop and record the transform parameters 65 | src_pts = np.array([[center[0]-size/2, center[1]-size/2], [center[0] - size/2, center[1]+size/2], [center[0]+size/2, center[1]-size/2]]) 66 | DST_PTS = np.array([[0, 0], [0, image_h - 1], [image_w - 1, 0]]) 67 | tform = skimage.transform.estimate_transform('similarity', src_pts, DST_PTS) 68 | cropped_image = skimage.transform.warp(image, tform.inverse, output_shape=(image_h, image_w)) 69 | 70 | # transform face position(image vertices) along with 2d facial image 71 | position = image_vertices.copy() 72 | position[:, 2] = 1 73 | position = np.dot(position, tform.params.T) 74 | position[:, 2] = image_vertices[:, 2]*tform.params[0, 0] # scale z 75 | position[:, 2] = position[:, 2] - np.min(position[:, 2]) # translate z 76 | 77 | # 4. uv position map: render position in uv space 78 | uv_position_map = mesh.render.render_colors(uv_coords, triangles, position, uv_h, uv_w, c = 3) 79 | uv_position_map = uv_position_map / max(image_h, image_w) 80 | 81 | return cropped_image, uv_position_map 82 | 83 | def process_uv(uv_coords, uv_h = 256, uv_w = 256): 84 | uv_coords[:,0] = uv_coords[:, 0]*(uv_w - 1) 85 | uv_coords[:,1] = uv_coords[:, 1]*(uv_h - 1) 86 | uv_coords[:,1] = uv_h - uv_coords[:,1] - 1 87 | uv_coords = np.hstack((uv_coords, np.zeros((uv_coords.shape[0], 1)))) # add z 88 | return uv_coords 89 | 90 | def process_one(number, bfm, kpt_ind, triangles, uv_coords, head_path, save_path, image_paths): 91 | counter = 0 92 | print('N%d--->Start!' %(number)) 93 | for image_path in image_paths: 94 | counter += 1 95 | image_name = image_path.split('/')[-1][:-4] 96 | mat_path = head_path + '/' + image_name + '.mat' 97 | 98 | cropped_image, uv_position_map = postmap_from_300WLP(bfm, kpt_ind, triangles, uv_coords, image_path, mat_path) 99 | imsave('{}/{}'.format(save_path, image_name + '.jpg'), cropped_image) 100 | np.save('{}/{}'.format(save_path, image_name + '.npy'), uv_position_map) 101 | print('N%d--->Full Samples: %d, Current Counter: %d' %(number, len(image_paths), counter)) 102 | 103 | class MyManager(BaseManager): 104 | pass 105 | MyManager.register('MorphabelModel', MorphabelModel) 106 | def into_manager(): 107 | m = MyManager() 108 | m.start() 109 | return m 110 | 111 | if __name__ == "__main__": 112 | 113 | # 1. load bfm model 114 | bfm = MorphabelModel('mm3d/BFM/BFM.mat') 115 | kpt_ind = bfm.kpt_ind 116 | triangles = bfm.full_triangles 117 | uv_coords = load_uv_coords('mm3d/BFM/BFM_UV.mat') 118 | uv_coords = process_uv(uv_coords) 119 | 120 | # 2. use the method 121 | bfm = into_manager().MorphabelModel('mm3d/BFM/BFM.mat') 122 | 123 | # 3. create head path 124 | head_paths = ['images/300W_LP/AFW', 'images/300W_LP/HELEN', 125 | 'images/300W_LP/IBUG', 'images/300W_LP/LFPW', 126 | 'images/300W_LP/AFW_Flip', 'images/300W_LP/HELEN_Flip', 127 | 'images/300W_LP/IBUG_Flip', 'images/300W_LP/LFPW_Flip', 128 | 'images/AFLW2000'] 129 | 130 | save_paths = [] 131 | for head_path in head_paths: 132 | save_path = head_path + '_GEN' 133 | if not os.path.exists(save_path): 134 | os.mkdir(save_path) 135 | save_paths.append(save_path) 136 | 137 | # 3. check and proceed all files 138 | # change here, if larger, speed is faster, but need more time! 139 | base_number = 4000 140 | base_process = [] 141 | for index, head_path in enumerate(head_paths): 142 | # glob images 143 | origin_image_paths = glob.glob(os.path.join(head_path, '*.jpg')) 144 | origin_image_paths = [image_path.replace('\\', '/') for image_path in origin_image_paths] 145 | 146 | cropped_image_paths = glob.glob(os.path.join(head_path+'_GEN', '*.jpg')) 147 | cropped_image_paths = [image_path.replace('\\', '/') for image_path in cropped_image_paths] 148 | 149 | # avoid repeatted file 150 | image_paths = [] 151 | for image_path in origin_image_paths: 152 | if head_path+'_GEN/' + image_path.split('/')[-1] not in cropped_image_paths: 153 | image_paths.append(image_path) 154 | 155 | if image_paths is None: 156 | continue 157 | 158 | # split image numbers 159 | total_base = 0 160 | if len(image_paths) > base_number: 161 | total_base = len(image_paths) // base_number 162 | extra_base = len(image_paths) % base_number 163 | 164 | # add processing 165 | for one_total in range(total_base): 166 | p = Process(target=process_one, args=(index, bfm, kpt_ind, triangles, uv_coords, head_path, save_paths[index], 167 | image_paths[one_total*base_number:(one_total+1)*base_number])) 168 | base_process.append(p) 169 | p = Process(target=process_one, args=(index, bfm, kpt_ind, triangles, uv_coords, head_path, save_paths[index], 170 | image_paths[total_base*base_number:total_base*base_number+extra_base])) 171 | base_process.append(p) 172 | 173 | # start processing 174 | for p in base_process: 175 | p.start() 176 | for p in base_process: 177 | p.join() 178 | -------------------------------------------------------------------------------- /run_one.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Wed Nov 14 16:10:11 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import os 9 | import cv2 10 | import numpy as np 11 | import scipy.io as sio 12 | import skimage.transform 13 | from skimage import io 14 | 15 | from skimage.io import imsave 16 | from skimage.transform import resize 17 | from utils.render_app import get_visibility 18 | from utils.render_app import get_uv_mask 19 | from utils.render_app import get_depth_image 20 | from utils.write import write_obj_with_texture 21 | from utils.write import write_obj_with_colors 22 | from utils.estimate_pose import estimate_pose 23 | from utils.cv_plot import plot_kpt 24 | from utils.cv_plot import plot_vertices 25 | from utils.cv_plot import plot_pose_box 26 | 27 | from mm3d import mesh 28 | from mm3d.morphable_model import MorphabelModel 29 | from mm3d.morphable_model.load import load_uv_coords 30 | 31 | def process_uv(uv_coords, uv_h = 256, uv_w = 256): 32 | uv_coords[:,0] = uv_coords[:,0]*(uv_w - 1) 33 | uv_coords[:,1] = uv_coords[:,1]*(uv_h - 1) 34 | uv_coords[:,1] = uv_h - uv_coords[:,1] - 1 35 | uv_coords = np.hstack((uv_coords, np.zeros((uv_coords.shape[0], 1)))) # add z 36 | return uv_coords 37 | 38 | def get_vertices(pos, face_ind, resolution_op): 39 | 40 | all_vertices = np.reshape(pos, [resolution_op**2, -1]) 41 | vertices = all_vertices[face_ind, :] 42 | 43 | return vertices 44 | 45 | def get_colors(image, vertices): 46 | 47 | [h, w, _] = image.shape 48 | vertices[:,0] = np.minimum(np.maximum(vertices[:,0], 0), w - 1) 49 | vertices[:,1] = np.minimum(np.maximum(vertices[:,1], 0), h - 1) 50 | ind = np.round(vertices).astype(np.int32) 51 | colors = image[ind[:,1], ind[:,0], :] 52 | 53 | return colors 54 | 55 | def get_landmarks(pos, uv_kpt_ind): 56 | 57 | kpt = pos[uv_kpt_ind[1,:], uv_kpt_ind[0,:], :] 58 | return kpt 59 | 60 | def generate_uv_coords(face_ind, resolution_op): 61 | uv_coords = np.meshgrid(range(resolution_op),range(resolution_op)) 62 | uv_coords = np.transpose(np.array(uv_coords), [1,2,0]) 63 | uv_coords = np.reshape(uv_coords, [resolution_op**2, -1]); 64 | uv_coords = uv_coords[face_ind, :] 65 | uv_coords = np.hstack((uv_coords[:,:2], np.zeros([uv_coords.shape[0], 1]))) 66 | 67 | return uv_coords 68 | 69 | def run_one_image(bfm, uv_coords, uv_kpt_ind, face_ind, triangles, s_uv_coords, 70 | image_path, mat_path, save_folder, name, 71 | uv_h = 256, uv_w = 256, image_h = 256, image_w = 256): 72 | # 1. load image and fitted parameters 73 | image = io.imread(image_path)/255. 74 | [h, w, c] = image.shape 75 | 76 | info = sio.loadmat(mat_path) 77 | pose_para = info['Pose_Para'].T.astype(np.float32) 78 | shape_para = info['Shape_Para'].astype(np.float32) 79 | exp_para = info['Exp_Para'].astype(np.float32) 80 | 81 | # 2. generate mesh 82 | vertices = bfm.generate_vertices(shape_para, exp_para) 83 | s = pose_para[-1, 0] 84 | angles = pose_para[:3, 0] 85 | t = pose_para[3:6, 0] 86 | transformed_vertices = bfm.transform_3ddfa(vertices, s, angles, t) 87 | projected_vertices = transformed_vertices.copy() 88 | image_vertices = projected_vertices.copy() 89 | image_vertices[:,1] = h - image_vertices[:,1] - 1 90 | 91 | # 3. crop image with key points 92 | kpt = image_vertices[bfm.kpt_ind, :].astype(np.int32) 93 | left = np.min(kpt[:, 0]) 94 | right = np.max(kpt[:, 0]) 95 | top = np.min(kpt[:, 1]) 96 | bottom = np.max(kpt[:, 1]) 97 | center = np.array([right - (right - left) / 2.0, 98 | bottom - (bottom - top) / 2.0]) 99 | old_size = (right - left + bottom - top)/2 100 | size = int(old_size * 1.5) 101 | # random pertube. you can change the numbers 102 | marg = old_size * 0.1 103 | t_x = np.random.rand() * marg * 2 - marg 104 | t_y = np.random.rand() * marg * 2 - marg 105 | center[0] = center[0] + t_x; center[1] = center[1] + t_y 106 | size = size*(np.random.rand()*0.2 + 0.9) 107 | 108 | # crop and record the transform parameters 109 | src_pts = np.array([[center[0]-size/2, center[1]-size/2], [center[0]-size/2, center[1]+size/2], [center[0]+size/2, center[1]-size/2]]) 110 | DST_PTS = np.array([[0, 0], [0, image_h - 1], [image_w - 1, 0]]) 111 | tform = skimage.transform.estimate_transform('similarity', src_pts, DST_PTS) 112 | cropped_image = skimage.transform.warp(image, tform.inverse, output_shape=(image_h, image_w)) 113 | 114 | print('input image is ok!') 115 | 116 | # transform face position(image vertices) along with 2d facial image 117 | position = image_vertices.copy() 118 | position[:, 2] = 1 119 | position = np.dot(position, tform.params.T) 120 | position[:, 2] = image_vertices[:, 2]*tform.params[0, 0] # scale z 121 | position[:, 2] = position[:, 2] - np.min(position[:, 2]) # translate z 122 | 123 | # 4. uv position map: render position in uv space 124 | uv_position_map = mesh.render.render_colors(uv_coords, bfm.full_triangles, position, uv_h, uv_w, c = 3) 125 | 126 | print('uv map is ok!') 127 | 128 | # 5. deal uv map 129 | # run model to get uv_map 130 | pos = uv_position_map / image_h 131 | max_pos = image_h 132 | pos = pos * max_pos 133 | 134 | # 6. get useful vertices 135 | vertices = get_vertices(pos, face_ind, uv_h) 136 | save_vertices = vertices.copy() 137 | save_vertices[:,1] = image_h - 1 - save_vertices[:,1] 138 | 139 | # 7. get colors 140 | colors = get_colors(cropped_image, vertices) 141 | write_obj_with_colors(os.path.join(save_folder, name + '_c.obj'), 142 | save_vertices, triangles, colors) 143 | 144 | print('color 3d face is ok!') 145 | 146 | # 8. get texture 147 | pos_interpolated = pos.copy() 148 | texture = cv2.remap(cropped_image, pos_interpolated[:,:,:2].astype(np.float32), 149 | None, interpolation=cv2.INTER_LINEAR, 150 | borderMode=cv2.BORDER_CONSTANT,borderValue=(0)) 151 | vertices_vis = get_visibility(vertices, triangles, image_h, image_w) 152 | uv_mask = get_uv_mask(vertices_vis, triangles, s_uv_coords, image_h, image_w, uv_h) 153 | uv_mask = resize(uv_mask, (256, 256), preserve_range = True) 154 | texture = texture * uv_mask[:, :, np.newaxis] 155 | write_obj_with_texture(os.path.join(save_folder, name + '.obj'), 156 | save_vertices, triangles, 157 | texture, s_uv_coords/uv_h) 158 | 159 | print('texture 3d face is ok!') 160 | 161 | # 9. get depth image 162 | depth_image = get_depth_image(vertices, triangles, image_h, image_w, True) 163 | imsave(os.path.join(save_folder, name + '_depth.jpg'), depth_image) 164 | 165 | print('depth image is ok!') 166 | 167 | # get restore size 168 | restore_top = int(center[0]-size/2) 169 | restore_bottom = int(center[0]+size/2) 170 | restore_left = int(center[1]-size/2) 171 | restore_right = int(center[1]+size/2) 172 | restore_w = restore_right - restore_left 173 | restore_h = restore_bottom - restore_top 174 | 175 | # 10. get landmarks 176 | t_image = (cropped_image*255.).astype(np.uint8) 177 | kpt = get_landmarks(pos, uv_kpt_ind) 178 | kpt_origin = plot_kpt(cropped_image, kpt).astype(np.uint8) 179 | kpt_gray = cv2.cvtColor(kpt_origin, cv2.COLOR_RGB2GRAY) 180 | ret, kpt_mask = cv2.threshold(kpt_gray, 127, 255, cv2.THRESH_BINARY) 181 | kpt_mask = cv2.bitwise_not(kpt_mask) 182 | kpt_and = cv2.bitwise_and(t_image, t_image, mask=kpt_mask) 183 | kpt_image = cv2.add(kpt_and, kpt_origin) 184 | imsave(os.path.join(save_folder, name + '_kpt.jpg'), kpt_image/255.) 185 | 186 | print('kpt image is ok!') 187 | 188 | # 10.0 restore kpt image 189 | resize_kpt_image = resize(kpt_image, (restore_w, restore_h)) 190 | rt_kpt_image = image.copy() 191 | rt_kpt_image[restore_left:restore_right, restore_top:restore_bottom] = resize_kpt_image 192 | imsave(os.path.join(save_folder, name + '_r_kpt.jpg'), rt_kpt_image) 193 | 194 | print('kpt fll image is ok!') 195 | 196 | # 11. get mask 197 | t_image = (cropped_image*255.).astype(np.uint8) 198 | ver_origin = plot_vertices(cropped_image, vertices).astype(np.uint8) 199 | ver_gray = cv2.cvtColor(ver_origin, cv2.COLOR_RGB2GRAY) 200 | ret, ver_mask = cv2.threshold(ver_gray, 127, 255, cv2.THRESH_BINARY) 201 | ver_mask = cv2.bitwise_not(ver_mask) 202 | ver_and = cv2.bitwise_and(t_image, t_image, mask=ver_mask) 203 | ver_image = cv2.add(ver_and, ver_origin) 204 | imsave(os.path.join(save_folder, name + '_ver.jpg'), ver_image/255.) 205 | 206 | print('vertices image is ok!') 207 | 208 | # 11.0 restore ver image 209 | resize_ver_image = resize(ver_image, (restore_w, restore_h)) 210 | rt_ver_image = image.copy() 211 | rt_ver_image[restore_left:restore_right, restore_top:restore_bottom] = resize_ver_image 212 | imsave(os.path.join(save_folder, name + '_r_ver.jpg'), rt_ver_image) 213 | 214 | print('vertices full image is ok!') 215 | 216 | # 12. get camera map 217 | camera_matrix, pose = estimate_pose(vertices) 218 | imsave(os.path.join(save_folder, name + '_cam.jpg'), plot_pose_box(cropped_image, camera_matrix, kpt)/255.) 219 | 220 | print('camera image is ok!') 221 | 222 | 223 | if __name__ == "__main__": 224 | 225 | # load bfm model 226 | bfm = MorphabelModel('mm3d/BFM/BFM.mat') 227 | uv_coords = load_uv_coords('mm3d/BFM/BFM_UV.mat') 228 | uv_coords = process_uv(uv_coords) 229 | 230 | print('load bfm ok!') 231 | 232 | uv_kpt_ind = np.loadtxt('images/uv_kpt_ind.txt').astype(np.int32) 233 | face_ind = np.loadtxt("images/face_ind.txt").astype(np.int32) 234 | triangles = np.loadtxt("images/triangles.txt").astype(np.int32) 235 | s_uv_coords = generate_uv_coords(face_ind, 256) 236 | 237 | print('load location ok!') 238 | 239 | image_path = 'images/300W_LP/AFW/AFW_134212_1_1.jpg' 240 | mat_path = 'images/300W_LP/AFW/AFW_134212_1_1.mat' 241 | 242 | run_one_image(bfm, uv_coords, uv_kpt_ind, face_ind, triangles, s_uv_coords, 243 | image_path, mat_path, 'images', 'test') 244 | -------------------------------------------------------------------------------- /run_one_g.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Wed Nov 14 16:10:11 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import os 9 | import cv2 10 | import numpy as np 11 | 12 | from skimage.io import imread 13 | from skimage.io import imsave 14 | from skimage.transform import resize 15 | from utils.render_app import get_visibility 16 | from utils.render_app import get_uv_mask 17 | from utils.write import write_obj_with_texture 18 | from utils.write import write_obj_with_colors 19 | from utils.cv_plot import plot_kpt 20 | 21 | def get_vertices(pos, face_ind, resolution_op): 22 | 23 | all_vertices = np.reshape(pos, [resolution_op**2, -1]) 24 | vertices = all_vertices[face_ind, :] 25 | 26 | return vertices 27 | 28 | def get_colors(image, vertices): 29 | 30 | [h, w, _] = image.shape 31 | vertices[:,0] = np.minimum(np.maximum(vertices[:,0], 0), w - 1) 32 | vertices[:,1] = np.minimum(np.maximum(vertices[:,1], 0), h - 1) 33 | ind = np.round(vertices).astype(np.int32) 34 | colors = image[ind[:,1], ind[:,0], :] 35 | 36 | return colors 37 | 38 | def get_landmarks(pos, uv_kpt_ind): 39 | 40 | kpt = pos[uv_kpt_ind[1,:], uv_kpt_ind[0,:], :] 41 | return kpt 42 | 43 | def generate_uv_coords(face_ind, resolution_op): 44 | uv_coords = np.meshgrid(range(resolution_op),range(resolution_op)) 45 | uv_coords = np.transpose(np.array(uv_coords), [1,2,0]) 46 | uv_coords = np.reshape(uv_coords, [resolution_op**2, -1]); 47 | uv_coords = uv_coords[face_ind, :] 48 | uv_coords = np.hstack((uv_coords[:,:2], np.zeros([uv_coords.shape[0], 1]))) 49 | 50 | return uv_coords 51 | 52 | def run_one_image(uv_kpt_ind, face_ind, triangles, s_uv_coords, 53 | image_path, npy_path, save_folder, name, 54 | uv_h = 256, uv_w = 256, image_h = 256, image_w = 256): 55 | 56 | # 1. load image 57 | cropped_image = imread(image_path) / 255. 58 | 59 | print('input image is ok!') 60 | 61 | # 2. load uv position map 62 | pos = np.load(npy_path) 63 | 64 | print('uv map is ok!') 65 | 66 | # 3. deal uv map 67 | # run model to get uv_map 68 | max_pos = image_h 69 | pos = pos * max_pos 70 | 71 | # 4. get useful vertices 72 | vertices = get_vertices(pos, face_ind, uv_h) 73 | save_vertices = vertices.copy() 74 | save_vertices[:,1] = image_h - 1 - save_vertices[:,1] 75 | 76 | # 5. get colors 77 | colors = get_colors(cropped_image, vertices) 78 | write_obj_with_colors(os.path.join(save_folder, name + '_c.obj'), 79 | save_vertices, triangles, colors) 80 | 81 | print('color 3d face is ok!') 82 | 83 | # 6. get texture 84 | pos_interpolated = pos.copy() 85 | texture = cv2.remap(cropped_image, pos_interpolated[:,:,:2].astype(np.float32), 86 | None, interpolation=cv2.INTER_LINEAR, 87 | borderMode=cv2.BORDER_CONSTANT,borderValue=(0)) 88 | vertices_vis = get_visibility(vertices, triangles, image_h, image_w) 89 | uv_mask = get_uv_mask(vertices_vis, triangles, s_uv_coords, image_h, image_w, uv_h) 90 | uv_mask = resize(uv_mask, (256, 256), preserve_range = True) 91 | texture = texture * uv_mask[:, :, np.newaxis] 92 | write_obj_with_texture(os.path.join(save_folder, name + '.obj'), 93 | save_vertices, triangles, 94 | texture, s_uv_coords/uv_h) 95 | 96 | print('texture 3d face is ok!') 97 | 98 | # 7. get landmarks 99 | t_image = (cropped_image*255.).astype(np.uint8) 100 | kpt = get_landmarks(pos, uv_kpt_ind) 101 | kpt_origin = plot_kpt(cropped_image, kpt).astype(np.uint8) 102 | kpt_gray = cv2.cvtColor(kpt_origin, cv2.COLOR_RGB2GRAY) 103 | ret, kpt_mask = cv2.threshold(kpt_gray, 127, 255, cv2.THRESH_BINARY) 104 | kpt_mask = cv2.bitwise_not(kpt_mask) 105 | kpt_and = cv2.bitwise_and(t_image, t_image, mask=kpt_mask) 106 | kpt_image = cv2.add(kpt_and, kpt_origin) 107 | imsave(os.path.join(save_folder, name + '_kpt.jpg'), kpt_image / 255.) 108 | 109 | print('kpt image is ok!') 110 | 111 | 112 | if __name__ == "__main__": 113 | 114 | uv_kpt_ind = np.loadtxt('images/uv_kpt_ind.txt').astype(np.int32) 115 | face_ind = np.loadtxt("images/face_ind.txt").astype(np.int32) 116 | triangles = np.loadtxt("images/triangles.txt").astype(np.int32) 117 | s_uv_coords = generate_uv_coords(face_ind, 256) 118 | 119 | print('load location ok!') 120 | 121 | image_path = 'images/300W_LP/AFW_GEN/AFW_111076519_1_1.jpg' 122 | mat_path = 'images/300W_LP/AFW_GEN/AFW_111076519_1_1.npy' 123 | 124 | run_one_image(uv_kpt_ind, face_ind, triangles, s_uv_coords, 125 | image_path, mat_path, 'images', 'test_g') 126 | -------------------------------------------------------------------------------- /run_swap.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Wed Nov 14 16:10:11 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import os 9 | import cv2 10 | import cmesh 11 | import face_detect 12 | import numpy as np 13 | import tensorflow as tf 14 | from color_correction import color_hist_match 15 | from color_correction import adain 16 | 17 | # from utils.render import render_texture 18 | 19 | def process_bbox(bboxes, image_shape, prop=0.16): 20 | 21 | for i, bbox in enumerate(bboxes): 22 | y0, x0, y1, x1 = bboxes[i, 0:4] 23 | w, h = int(y1 - y0), int(x1 - x0) 24 | d_w, d_h = int(w * prop), int(h * prop) 25 | y0 = y0 - d_w; y1 = y1 + d_w 26 | x0 = x0 - d_h; x1 = x1 + d_h 27 | length = (w + h) / 2 28 | center = (int((x1+x0)/2), int((y1+y0)/2)) 29 | new_x0 = np.max([0, center[0]-length//2]) 30 | new_x1 = np.min([image_shape[0], center[0]+length//2]) 31 | new_y0 = np.max([0, center[1]-length//2]) 32 | new_y1 = np.min([image_shape[1], center[1]+length//2]) 33 | bboxes[i, 0:4] = new_x0, new_y1, new_x1, new_y0 34 | 35 | return bboxes 36 | 37 | def get_vertices(pos, face_ind, resolution_op): 38 | 39 | all_vertices = np.reshape(pos, [resolution_op**2, -1]) 40 | vertices = all_vertices[face_ind, :] 41 | 42 | return vertices 43 | 44 | def get_colors(image, vertices): 45 | 46 | [h, w, _] = image.shape 47 | vertices[:,0] = np.minimum(np.maximum(vertices[:,0], 0), w - 1) 48 | vertices[:,1] = np.minimum(np.maximum(vertices[:,1], 0), h - 1) 49 | ind = np.round(vertices).astype(np.int32) 50 | colors = image[ind[:,1], ind[:,0], :] 51 | 52 | return colors 53 | 54 | def get_landmarks(pos, uv_kpt_ind): 55 | 56 | kpt = pos[uv_kpt_ind[1,:], uv_kpt_ind[0,:], :] 57 | return kpt 58 | 59 | def generate_texcoord(uv_coords_path, tex_h, tex_w): 60 | 61 | uv_coords = np.loadtxt(uv_coords_path) 62 | texcoord = np.zeros_like(uv_coords) 63 | texcoord[:,0] = uv_coords[:,0]*(tex_h - 1) 64 | texcoord[:,1] = uv_coords[:,1]*(tex_w - 1) 65 | texcoord[:,1] = tex_w - texcoord[:,1] - 1 66 | texcoord = np.hstack((texcoord, np.zeros((texcoord.shape[0], 1)))) 67 | 68 | return texcoord 69 | 70 | def load_detect_model(): 71 | with tf.Graph().as_default(): 72 | sess = tf.Session() 73 | with sess.as_default(): 74 | pnet, rnet, onet = face_detect.create_mtcnn(sess, None) 75 | return pnet, rnet, onet 76 | 77 | def load_3dface_model(): 78 | 79 | tface_graph_def = tf.GraphDef() 80 | tface_graph_def.ParseFromString(open("model/pico_3dFace_model.pb", "rb").read()) 81 | tf.import_graph_def(tface_graph_def, name="") 82 | tface_sess = tf.Session() 83 | tface_sess.graph.get_operations() 84 | tx = tface_sess.graph.get_tensor_by_name("3dface/x:0") 85 | ty = tface_sess.graph.get_tensor_by_name("PRNet/Conv2d_transpose_16/Sigmoid:0") 86 | 87 | return tx, ty, tface_sess 88 | 89 | def get_colors_from_texture(texture, face_ind, resolution_op): 90 | 91 | all_colors = np.reshape(texture, [resolution_op**2, -1]) 92 | colors = all_colors[face_ind, :] 93 | 94 | return colors 95 | 96 | def run_ref_image(image_path, uv_kpt_ind, face_ind, triangles, 97 | pnet, rnet, onet, x, y, Tsess, 98 | minsize=30, threshold=[0.6, 0.7, 0.7], factor=0.709, best_score=0.7, 99 | uv_h=256, uv_w=256, image_h=256, image_w=256): 100 | 101 | input_image = cv2.imread(image_path, 1) 102 | boxes, pnts = face_detect.detect_face(input_image, minsize, 103 | pnet, rnet, onet, threshold, factor) 104 | faces = process_bbox(boxes, input_image.shape) 105 | 106 | for idx, (x0, y1, x1, y0, conf_score) in enumerate(faces): 107 | 108 | if conf_score > best_score: 109 | 110 | det_face = input_image[int(x0):int(x1), int(y0):int(y1), :] 111 | det_face = cv2.resize(det_face, (256,256)) / 255. 112 | 113 | ref_pos = Tsess.run(y, feed_dict={x: det_face[np.newaxis, :,:,:]}) 114 | ref_pos = np.squeeze(ref_pos) 115 | max_pos = image_h 116 | ref_pos = ref_pos * max_pos 117 | 118 | ref_texture = cv2.remap(det_face, ref_pos[:,:,:2].astype(np.float32), 119 | None, interpolation=cv2.INTER_NEAREST, 120 | borderMode=cv2.BORDER_CONSTANT,borderValue=(0)) 121 | 122 | break 123 | 124 | return ref_texture 125 | 126 | def to_roi(image, prop): 127 | 128 | h, w, _ = image.shape 129 | p_h = h * prop 130 | p_w = w * prop 131 | d_h = int((h - p_h) / 2.) 132 | d_w = int((w - p_w) / 2.) 133 | 134 | return image[d_h:h-d_h, d_w:w-d_w], d_h, d_w 135 | 136 | def image_filter(image, template_image, prop=0.98, kernel_size=5, theta=1.5, color_mode=0, is_origin=False): 137 | 138 | ''' 139 | prop < 1. 140 | ''' 141 | 142 | ver_image, d_h, d_w = to_roi(image, prop) 143 | 144 | t_h, t_w = template_image.shape[0], template_image.shape[1] 145 | 146 | if color_mode == 0: 147 | pass 148 | elif color_mode == 1: 149 | ver_image = color_hist_match(ver_image, template_image[d_h:t_h-d_h, d_w:t_w-d_w]) 150 | else: 151 | ver_image = adain(ver_image, template_image[d_h:t_h-d_h, d_w:t_w-d_w]) 152 | 153 | template_image[d_h:t_h-d_h, d_w:t_w-d_w] = ver_image 154 | 155 | if is_origin: 156 | return template_image 157 | 158 | half_size = kernel_size // 2 159 | if t_h-d_h*4-half_size*2 < 0 or t_w-d_w*4-half_size*2 < 0: 160 | return cv2.GaussianBlur(template_image, (kernel_size, kernel_size), theta) 161 | 162 | s_mask = np.zeros_like(template_image, dtype=np.uint8) 163 | ss_mask = np.zeros((t_h, t_w), dtype=np.uint8) 164 | s_mask[d_h*2+half_size:t_h-d_h*2-half_size, d_w*2+half_size:t_w-d_w*2-half_size] = \ 165 | template_image[d_h*2+half_size:t_h-d_h*2-half_size, d_w*2+half_size:t_w-d_w*2-half_size] 166 | ss_mask[d_h*2:t_h-d_h*2, d_w*2:t_w-d_w*2] = np.full((t_h-d_h*4, t_w-d_w*4), 255, dtype=np.uint8) 167 | t_mask = template_image - s_mask 168 | t_mask = cv2.GaussianBlur(t_mask, (kernel_size, kernel_size), theta) 169 | 170 | ss_mask = cv2.bitwise_not(ss_mask) 171 | conbine_image = cv2.bitwise_and(t_mask, t_mask, mask=ss_mask) 172 | conbine_image[d_h*2:t_h-d_h*2, d_w*2:t_w-d_w*2] = template_image[d_h*2:t_h-d_h*2, d_w*2:t_w-d_w*2] 173 | 174 | return conbine_image 175 | 176 | def run_one_image(input_image, uv_kpt_ind, face_ind, triangles, uv_coords, 177 | pnet, rnet, onet, x, y, Tsess, ref_texture, uv_whole_face, blend_factor=0.35, 178 | minsize=30, threshold=[0.6, 0.7, 0.7], factor=0.709, best_score=0.7, 179 | uv_h=256, uv_w=256, image_h=256, image_w=256): 180 | 181 | output_image = input_image.copy() 182 | boxes, pnts = face_detect.detect_face(input_image, minsize, 183 | pnet, rnet, onet, threshold, factor) 184 | faces = process_bbox(boxes, input_image.shape) 185 | is_face = False 186 | for idx, (x0, y1, x1, y0, conf_score) in enumerate(faces): 187 | 188 | if conf_score > best_score: 189 | 190 | is_face = True 191 | 192 | det_face = input_image[int(x0):int(x1), int(y0):int(y1), :] 193 | template_face = det_face.copy() 194 | face_shape = (int(y1)-int(y0), int(x1)-int(x0)) 195 | det_face = cv2.resize(det_face, (256,256)) / 255. 196 | 197 | pos = Tsess.run(y, feed_dict={x: det_face[np.newaxis, :,:,:]}) 198 | pos = np.squeeze(pos) 199 | max_pos = image_h * 1.1 200 | pos = pos * max_pos 201 | 202 | vertices = get_vertices(pos, face_ind, uv_h) 203 | 204 | vis_colors = np.ones((vertices.shape[0], 1)) 205 | face_mask = cmesh.render.render_texture(vertices.T, vis_colors.T, triangles.T, image_h, image_w, c = 1) 206 | face_mask = np.squeeze(face_mask > 0).astype(np.float32) 207 | ''' 208 | texture = cv2.remap(det_face, pos[:,:,:2].astype(np.float32), 209 | None, interpolation=cv2.INTER_NEAREST, 210 | borderMode=cv2.BORDER_CONSTANT,borderValue=(0)) 211 | 212 | new_texture = texture*(1. - uv_whole_face[:,:,np.newaxis]) + ref_texture*uv_whole_face[:,:,np.newaxis] 213 | ''' 214 | new_texture = ref_texture 215 | new_colors = get_colors_from_texture(new_texture, face_ind, uv_h) 216 | new_image = cmesh.render.render_texture(vertices.T, new_colors.T, triangles.T, image_h, image_w, c = 3) 217 | 218 | 219 | new_image = blend_factor * det_face*face_mask[:,:,np.newaxis] + (1-blend_factor) * new_image*face_mask[:,:,np.newaxis] 220 | # new_image = color_hist_match(new_image, det_face*face_mask[:,:,np.newaxis]) 221 | 222 | new_image = det_face*(1.- face_mask[:,:,np.newaxis]) + new_image 223 | 224 | vis_ind = np.argwhere(face_mask>0) 225 | vis_min = np.min(vis_ind, 0) 226 | vis_max = np.max(vis_ind, 0) 227 | center = (int((vis_min[1] + vis_max[1])/2+0.5), int((vis_min[0] + vis_max[0])/2+0.5)) 228 | 229 | output = cv2.seamlessClone((new_image*255.).astype(np.uint8), (det_face*255.).astype(np.uint8), 230 | (face_mask*255.).astype(np.uint8), center, cv2.NORMAL_CLONE) 231 | 232 | temp_ver_image = cv2.resize(output, face_shape) 233 | last_ver_image = image_filter(temp_ver_image, template_face) 234 | output_image[int(x0):int(x1), int(y0):int(y1)] = last_ver_image 235 | 236 | return is_face, output_image 237 | 238 | if __name__ == "__main__": 239 | 240 | os.environ['CUDA_VISIBLE_DEVICES'] = '0' 241 | 242 | # load params 243 | uv_kpt_ind = np.loadtxt("images/uv_kpt_ind.txt").astype(np.int32) 244 | face_ind = np.loadtxt("images/face_ind.txt").astype(np.int32) 245 | triangles = np.loadtxt("images/triangles.txt").astype(np.int32) 246 | uv_coords = generate_texcoord("images/uv_coords.txt", 256, 256) 247 | uv_whole_face = cv2.imread('images/uv_face_mask.png', 0) / 255. 248 | #uv_face_eye = imread('images/uv_face_eyes.png', as_grey=True) / 255. 249 | #uv_face = imread('images/uv_face.png', as_grey=True) / 255. 250 | #uv_whole_face = (abs(uv_face_eye - uv_face) > 0).astype(np.float32) 251 | 252 | # load model 253 | pnet, rnet, onet = load_detect_model() 254 | tx, ty, tsess = load_3dface_model() 255 | 256 | # load ref image 257 | ref_texture = run_ref_image("ref.jpg", uv_kpt_ind, face_ind, triangles, 258 | pnet, rnet, onet, tx, ty, tsess) 259 | 260 | print('load model ok!') 261 | ''' 262 | image_path = 'images/tfd.jpg' 263 | is_face, ver_image = run_one_image(cv2.imread(image_path, 1), uv_kpt_ind, face_ind, triangles, uv_coords, 264 | pnet, rnet, onet, tx, ty, tsess, ref_texture, uv_whole_face) 265 | cv2.imwrite(os.path.join("images", "test" + '_mask.jpg'), ver_image) 266 | ''' 267 | # camera settins 268 | vedio_shape = [1920, 1080] 269 | cap = cv2.VideoCapture(0) 270 | cap.set(cv2.CAP_PROP_FRAME_WIDTH, vedio_shape[0]) 271 | cap.set(cv2.CAP_PROP_FRAME_HEIGHT, vedio_shape[1]) 272 | cv2.namedWindow("Deep3DFace", cv2.WINDOW_NORMAL) 273 | 274 | while(True): 275 | 276 | _, o_image = cap.read() 277 | 278 | is_face, ver_image = run_one_image(o_image, uv_kpt_ind, face_ind, triangles, uv_coords, 279 | pnet, rnet, onet, tx, ty, tsess, ref_texture, uv_whole_face, blend_factor=0.35) 280 | 281 | cv2.imshow("Deep3DFace", np.hstack((o_image, ver_image))) 282 | 283 | if (cv2.waitKey(1) & 0xFF) == ord('q'): 284 | break 285 | 286 | cv2.destroyAllWindows() 287 | 288 | ''' 289 | v_cap = cv2.VideoCapture("data_A.mp4") 290 | 291 | fps = v_cap.get(cv2.CAP_PROP_FPS) 292 | fourcc = int(v_cap.get(cv2.CAP_PROP_FOURCC)) 293 | size = (int(v_cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(v_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))) 294 | videoWriter = cv2.VideoWriter("data_B.mp4", fourcc, fps, size, True) 295 | 296 | counter = 0 297 | success, frame = v_cap.read() 298 | while success : 299 | try: 300 | is_face, ver_image = run_one_image(frame, uv_kpt_ind, face_ind, triangles, uv_coords, 301 | pnet, rnet, onet, tx, ty, tsess, ref_texture, uv_whole_face) 302 | if is_face: 303 | cv2.imwrite('swap/'+str(counter)+'.jpg', ver_image) 304 | videoWriter.write(ver_image) 305 | else: 306 | videoWriter.write(frame) 307 | except Exception: 308 | pass 309 | is_face, ver_image = run_one_image(frame, uv_kpt_ind, face_ind, triangles, uv_coords, 310 | pnet, rnet, onet, tx, ty, tsess, ref_texture, uv_whole_face) 311 | if is_face: 312 | cv2.imwrite('swap/'+str(counter)+'.jpg', ver_image) 313 | videoWriter.write(ver_image) 314 | else: 315 | videoWriter.write(frame) 316 | 317 | counter += 1 318 | print('PROCESSING %d OK' %(counter)) 319 | cv2.waitKey(1000//int(fps)) 320 | success, frame = v_cap.read() 321 | 322 | print('PROCESSING IS OK!') 323 | v_cap.release() 324 | videoWriter.release() 325 | ''' 326 | -------------------------------------------------------------------------------- /run_two.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Wed Nov 14 16:10:11 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import os 9 | import cv2 10 | import numpy as np 11 | import scipy.io as sio 12 | import skimage.transform 13 | from skimage import io 14 | 15 | from skimage.io import imread 16 | from skimage.io import imsave 17 | from utils.render import render_texture 18 | 19 | from mm3d import mesh 20 | from mm3d.morphable_model import MorphabelModel 21 | from mm3d.morphable_model.load import load_uv_coords 22 | 23 | def process_uv(uv_coords, uv_h = 256, uv_w = 256): 24 | uv_coords[:,0] = uv_coords[:,0]*(uv_w - 1) 25 | uv_coords[:,1] = uv_coords[:,1]*(uv_h - 1) 26 | uv_coords[:,1] = uv_h - uv_coords[:,1] - 1 27 | uv_coords = np.hstack((uv_coords, np.zeros((uv_coords.shape[0], 1)))) # add z 28 | return uv_coords 29 | 30 | def get_vertices(pos, face_ind, resolution_op): 31 | 32 | all_vertices = np.reshape(pos, [resolution_op**2, -1]) 33 | vertices = all_vertices[face_ind, :] 34 | 35 | return vertices 36 | 37 | def get_colors_from_texture(texture, face_ind, resolution_op): 38 | 39 | all_colors = np.reshape(texture, [resolution_op**2, -1]) 40 | colors = all_colors[face_ind, :] 41 | 42 | return colors 43 | 44 | def get_colors(image, vertices): 45 | 46 | [h, w, _] = image.shape 47 | vertices[:,0] = np.minimum(np.maximum(vertices[:,0], 0), w - 1) 48 | vertices[:,1] = np.minimum(np.maximum(vertices[:,1], 0), h - 1) 49 | ind = np.round(vertices).astype(np.int32) 50 | colors = image[ind[:,1], ind[:,0], :] 51 | 52 | return colors 53 | 54 | def get_landmarks(pos, uv_kpt_ind): 55 | 56 | kpt = pos[uv_kpt_ind[1,:], uv_kpt_ind[0,:], :] 57 | return kpt 58 | 59 | def generate_uv_coords(face_ind, resolution_op): 60 | uv_coords = np.meshgrid(range(resolution_op),range(resolution_op)) 61 | uv_coords = np.transpose(np.array(uv_coords), [1,2,0]) 62 | uv_coords = np.reshape(uv_coords, [resolution_op**2, -1]); 63 | uv_coords = uv_coords[face_ind, :] 64 | uv_coords = np.hstack((uv_coords[:,:2], np.zeros([uv_coords.shape[0], 1]))) 65 | 66 | return uv_coords 67 | 68 | def run_one_image(bfm, uv_coords, face_ind, 69 | image_path, mat_path, uv_h, uv_w, image_h, image_w): 70 | 71 | # 1. load image and fitted parameters 72 | image = io.imread(image_path)/255. 73 | [h, w, c] = image.shape 74 | 75 | info = sio.loadmat(mat_path) 76 | pose_para = info['Pose_Para'].T.astype(np.float32) 77 | shape_para = info['Shape_Para'].astype(np.float32) 78 | exp_para = info['Exp_Para'].astype(np.float32) 79 | 80 | # 2. generate mesh 81 | vertices = bfm.generate_vertices(shape_para, exp_para) 82 | s = pose_para[-1, 0] 83 | angles = pose_para[:3, 0] 84 | t = pose_para[3:6, 0] 85 | transformed_vertices = bfm.transform_3ddfa(vertices, s, angles, t) 86 | projected_vertices = transformed_vertices.copy() 87 | image_vertices = projected_vertices.copy() 88 | image_vertices[:,1] = h - image_vertices[:,1] - 1 89 | 90 | # 3. crop image with key points 91 | kpt = image_vertices[bfm.kpt_ind, :].astype(np.int32) 92 | left = np.min(kpt[:, 0]) 93 | right = np.max(kpt[:, 0]) 94 | top = np.min(kpt[:, 1]) 95 | bottom = np.max(kpt[:, 1]) 96 | center = np.array([right - (right - left) / 2.0, 97 | bottom - (bottom - top) / 2.0]) 98 | old_size = (right - left + bottom - top)/2 99 | size = int(old_size * 1.5) 100 | # random pertube. you can change the numbers 101 | marg = old_size * 0.1 102 | t_x = np.random.rand() * marg * 2 - marg 103 | t_y = np.random.rand() * marg * 2 - marg 104 | center[0] = center[0] + t_x; center[1] = center[1] + t_y 105 | size = size*(np.random.rand()*0.2 + 0.9) 106 | 107 | # crop and record the transform parameters 108 | src_pts = np.array([[center[0]-size/2, center[1]-size/2], [center[0]-size/2, center[1]+size/2], [center[0]+size/2, center[1]-size/2]]) 109 | DST_PTS = np.array([[0, 0], [0, image_h - 1], [image_w - 1, 0]]) 110 | tform = skimage.transform.estimate_transform('similarity', src_pts, DST_PTS) 111 | cropped_image = skimage.transform.warp(image, tform.inverse, output_shape=(image_h, image_w)) 112 | 113 | print('input image is ok!') 114 | 115 | # transform face position(image vertices) along with 2d facial image 116 | position = image_vertices.copy() 117 | position[:, 2] = 1 118 | position = np.dot(position, tform.params.T) 119 | position[:, 2] = image_vertices[:, 2]*tform.params[0, 0] # scale z 120 | position[:, 2] = position[:, 2] - np.min(position[:, 2]) # translate z 121 | 122 | # 4. uv position map: render position in uv space 123 | uv_position_map = mesh.render.render_colors(uv_coords, bfm.full_triangles, position, uv_h, uv_w, c = 3) 124 | 125 | print('uv map is ok!') 126 | 127 | # 6. get useful vertices 128 | vertices = get_vertices(uv_position_map, face_ind, uv_h) 129 | 130 | return image, cropped_image, (center[0], center[1]), size, uv_position_map, vertices 131 | 132 | def run_two_image(bfm, uv_coords, uv_kpt_ind, face_ind, triangles, s_uv_coords, 133 | image_path_A, mat_path_A, image_path_B, mat_path_B, save_folder, name, mode=1, 134 | uv_h = 256, uv_w = 256, image_h = 256, image_w = 256): 135 | 136 | image, cropped_image, center, size, pos, vertices = \ 137 | run_one_image(bfm, uv_coords, face_ind, image_path_A, mat_path_A, 138 | uv_h, uv_w, image_h, image_w) 139 | 140 | ref_image, ref_cropped_image, ref_center, ref_size, ref_pos, ref_vertices = \ 141 | run_one_image(bfm, uv_coords, face_ind, image_path_B, mat_path_B, 142 | uv_h, uv_w, image_h, image_w) 143 | 144 | texture = cv2.remap(cropped_image, pos[:,:,:2].astype(np.float32), 145 | None, interpolation=cv2.INTER_NEAREST, 146 | borderMode=cv2.BORDER_CONSTANT,borderValue=(0)) 147 | ref_texture = cv2.remap(ref_cropped_image, ref_pos[:,:,:2].astype(np.float32), 148 | None, interpolation=cv2.INTER_NEAREST, 149 | borderMode=cv2.BORDER_CONSTANT,borderValue=(0)) 150 | 151 | if mode == 0: 152 | # load eye mask 153 | uv_face_eye = imread('images/uv_face_eyes.png', as_grey=True) / 255. 154 | uv_face = imread('images/uv_face.png', as_grey=True) / 255. 155 | eye_mask = (abs(uv_face_eye - uv_face) > 0).astype(np.float32) 156 | # modify texture 157 | new_texture = texture*(1 - eye_mask[:,:,np.newaxis]) + ref_texture*eye_mask[:,:,np.newaxis] 158 | else: 159 | uv_whole_face = imread('images/uv_face_mask.png', as_grey=True) / 255. 160 | new_texture = texture*(1 - uv_whole_face[:,:,np.newaxis]) + ref_texture*uv_whole_face[:,:,np.newaxis] 161 | # new_texture = ref_texture 162 | 163 | #-- 3. remap to input image.(render) 164 | vis_colors = np.ones((vertices.shape[0], 1)) 165 | face_mask = render_texture(vertices.T, vis_colors.T, triangles.T, image_h, image_w, c = 1) 166 | face_mask = np.squeeze(face_mask > 0).astype(np.float32) 167 | new_colors = get_colors_from_texture(new_texture, face_ind, uv_h) 168 | new_image = render_texture(vertices.T, new_colors.T, triangles.T, image_h, image_w, c = 3) 169 | new_image = cropped_image*(1 - face_mask[:,:,np.newaxis]) + new_image*face_mask[:,:,np.newaxis] 170 | 171 | # Possion Editing for blending image 172 | vis_ind = np.argwhere(face_mask>0) 173 | vis_min = np.min(vis_ind, 0) 174 | vis_max = np.max(vis_ind, 0) 175 | center = (int((vis_min[1] + vis_max[1])/2+0.5), int((vis_min[0] + vis_max[0])/2+0.5)) 176 | 177 | output = cv2.seamlessClone((new_image*255).astype(np.uint8), (cropped_image*255).astype(np.uint8), 178 | (face_mask*255).astype(np.uint8), center, cv2.NORMAL_CLONE) 179 | 180 | if mode == 0: 181 | imsave(os.path.join(save_folder, name + '_eyes.jpg'), output) 182 | else: 183 | imsave(os.path.join(save_folder, name + '_swap.jpg'), output) 184 | 185 | if __name__ == "__main__": 186 | 187 | # load bfm model 188 | bfm = MorphabelModel('mm3d/BFM/BFM.mat') 189 | uv_coords = load_uv_coords('mm3d/BFM/BFM_UV.mat') 190 | uv_coords = process_uv(uv_coords) 191 | 192 | print('load bfm ok!') 193 | 194 | uv_kpt_ind = np.loadtxt('images/uv_kpt_ind.txt').astype(np.int32) 195 | face_ind = np.loadtxt("images/face_ind.txt").astype(np.int32) 196 | triangles = np.loadtxt("images/triangles.txt").astype(np.int32) 197 | s_uv_coords = generate_uv_coords(face_ind, 256) 198 | 199 | print('load location ok!') 200 | 201 | image_path_A = 'images/300W_LP/AFW/AFW_261068_1_1.jpg' 202 | mat_path_A = 'images/300W_LP/AFW/AFW_261068_1_1.mat' 203 | image_path_B = 'images/300W_LP/AFW/AFW_1634816_1_0.jpg' 204 | mat_path_B = 'images/300W_LP/AFW/AFW_1634816_1_0.mat' 205 | 206 | run_two_image(bfm, uv_coords, uv_kpt_ind, face_ind, triangles, s_uv_coords, 207 | image_path_A, mat_path_A, image_path_B, mat_path_B, 'images', 'test') 208 | -------------------------------------------------------------------------------- /run_vedio.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Wed Nov 14 16:10:11 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import os 9 | import cv2 10 | import glob 11 | import face_detect 12 | import numpy as np 13 | import tensorflow as tf 14 | 15 | from skimage.io import imread 16 | from skimage.io import imsave 17 | from skimage.transform import resize 18 | from utils.cv_plot import plot_kpt 19 | from utils.cv_plot import plot_vertices 20 | 21 | def process_bbox(bboxes, image_shape): 22 | 23 | for i, bbox in enumerate(bboxes): 24 | y0, x0, y1, x1 = bboxes[i, 0:4] 25 | w, h = int(y1 - y0), int(x1 - x0) 26 | length = (w + h) / 2 27 | center = (int((x1+x0)/2), int((y1+y0)/2)) 28 | new_x0 = np.max([0, center[0]-length//2]) 29 | new_x1 = np.min([image_shape[0], center[0]+length//2]) 30 | new_y0 = np.max([0, center[1]-length//2]) 31 | new_y1 = np.min([image_shape[1], center[1]+length//2]) 32 | bboxes[i, 0:4] = new_x0, new_y1, new_x1, new_y0 33 | 34 | return bboxes 35 | 36 | def get_vertices(pos, face_ind, resolution_op): 37 | 38 | all_vertices = np.reshape(pos, [resolution_op**2, -1]) 39 | vertices = all_vertices[face_ind, :] 40 | 41 | return vertices 42 | 43 | def get_colors(image, vertices): 44 | 45 | [h, w, _] = image.shape 46 | vertices[:,0] = np.minimum(np.maximum(vertices[:,0], 0), w - 1) 47 | vertices[:,1] = np.minimum(np.maximum(vertices[:,1], 0), h - 1) 48 | ind = np.round(vertices).astype(np.int32) 49 | colors = image[ind[:,1], ind[:,0], :] 50 | 51 | return colors 52 | 53 | def get_landmarks(pos, uv_kpt_ind): 54 | 55 | kpt = pos[uv_kpt_ind[1,:], uv_kpt_ind[0,:], :] 56 | return kpt 57 | 58 | def generate_uv_coords(face_ind, resolution_op): 59 | uv_coords = np.meshgrid(range(resolution_op),range(resolution_op)) 60 | uv_coords = np.transpose(np.array(uv_coords), [1,2,0]) 61 | uv_coords = np.reshape(uv_coords, [resolution_op**2, -1]); 62 | uv_coords = uv_coords[face_ind, :] 63 | uv_coords = np.hstack((uv_coords[:,:2], np.zeros([uv_coords.shape[0], 1]))) 64 | 65 | return uv_coords 66 | 67 | def load_detect_model(): 68 | with tf.Graph().as_default(): 69 | sess = tf.Session() 70 | with sess.as_default(): 71 | pnet, rnet, onet = face_detect.create_mtcnn(sess, None) 72 | return pnet, rnet, onet 73 | 74 | def load_3dface_model(): 75 | 76 | tface_graph_def = tf.GraphDef() 77 | tface_graph_def.ParseFromString(open("model/pico_3dFace_model.pb", "rb").read()) 78 | tf.import_graph_def(tface_graph_def, name="") 79 | tface_sess = tf.Session() 80 | tface_sess.graph.get_operations() 81 | tx = tface_sess.graph.get_tensor_by_name("3dface/x:0") 82 | ty = tface_sess.graph.get_tensor_by_name("PRNet/Conv2d_transpose_16/Sigmoid:0") 83 | 84 | return tx, ty, tface_sess 85 | 86 | def run_one_image(image_path, uv_kpt_ind, face_ind, triangles, s_uv_coords, 87 | pnet, rnet, onet, x, y, Tsess, 88 | minsize=30, threshold=[0.6, 0.7, 0.7], factor=0.709, best_score=0.7, 89 | uv_h=256, uv_w=256, image_h=256, image_w=256): 90 | 91 | input_image = cv2.imread(image_path, 1) 92 | output_image = input_image.copy() 93 | boxes, pnts = face_detect.detect_face(input_image, minsize, 94 | pnet, rnet, onet, threshold, factor) 95 | faces = process_bbox(boxes, input_image.shape) 96 | 97 | for idx, (x0, y1, x1, y0, conf_score) in enumerate(faces): 98 | 99 | if conf_score > best_score: 100 | 101 | det_face = input_image[int(x0):int(x1), int(y0):int(y1), :] 102 | 103 | face_shape = (int(y1)-int(y0), int(x1)-int(x0)) 104 | det_face = cv2.resize(det_face, (256,256)) / 255. 105 | 106 | pos = Tsess.run(y, feed_dict={x: det_face[np.newaxis, :,:,:]}) 107 | pos = np.squeeze(pos) 108 | max_pos = image_h 109 | pos = pos * max_pos 110 | 111 | vertices = get_vertices(pos, face_ind, uv_h) 112 | 113 | from utils.write import write_obj_with_colors 114 | save_vertices = vertices.copy() 115 | save_vertices[:,1] = image_h - 1 - save_vertices[:,1] 116 | colors = get_colors(det_face, vertices) 117 | write_obj_with_colors(os.path.join('images', 'test' + '_c.obj'), 118 | save_vertices, triangles, colors) 119 | 120 | t_image = (det_face*255.).astype(np.uint8) 121 | kpt = get_landmarks(pos, uv_kpt_ind) 122 | kpt_origin = plot_kpt(det_face, kpt).astype(np.uint8) 123 | kpt_gray = cv2.cvtColor(kpt_origin, cv2.COLOR_RGB2GRAY) 124 | ret, kpt_mask = cv2.threshold(kpt_gray, 127, 255, cv2.THRESH_BINARY) 125 | kpt_mask = cv2.bitwise_not(kpt_mask) 126 | kpt_and = cv2.bitwise_and(t_image, t_image, mask=kpt_mask) 127 | kpt_image = cv2.add(kpt_and, kpt_origin) 128 | imsave(os.path.join('images', 'test' + '_kpt.jpg'), kpt_image/255.) 129 | 130 | t_image = (det_face*255.).astype(np.uint8) 131 | ver_origin = plot_vertices(det_face, vertices).astype(np.uint8) 132 | ver_gray = cv2.cvtColor(ver_origin, cv2.COLOR_RGB2GRAY) 133 | ret, ver_mask = cv2.threshold(ver_gray, 127, 255, cv2.THRESH_BINARY) 134 | ver_mask = cv2.bitwise_not(ver_mask) 135 | ver_and = cv2.bitwise_and(t_image, t_image, mask=ver_mask) 136 | ver_image = cv2.add(ver_and, ver_origin) 137 | imsave(os.path.join('images', 'test' + '_ver.jpg'), ver_image/255.) 138 | 139 | resize_ver_image = cv2.resize(ver_image, face_shape) 140 | 141 | output_image[int(x0):int(x1), int(y0):int(y1)] = resize_ver_image 142 | 143 | return output_image / 255. 144 | 145 | if __name__ == "__main__": 146 | 147 | os.environ['CUDA_VISIBLE_DEVICES'] = '0' 148 | 149 | # load params 150 | uv_kpt_ind = np.loadtxt('images/uv_kpt_ind.txt').astype(np.int32) 151 | face_ind = np.loadtxt("images/face_ind.txt").astype(np.int32) 152 | triangles = np.loadtxt("images/triangles.txt").astype(np.int32) 153 | s_uv_coords = generate_uv_coords(face_ind, 256) 154 | 155 | # load model 156 | pnet, rnet, onet = load_detect_model() 157 | tx, ty, tsess = load_3dface_model() 158 | 159 | print('load model ok!') 160 | ver_image = run_one_image('images/test.jpg', uv_kpt_ind, face_ind, triangles, s_uv_coords, 161 | pnet, rnet, onet, tx, ty, tsess) 162 | ''' 163 | image_paths = glob.glob(os.path.join("TrainingData", '*.jpg')) 164 | counter = 0 165 | for image_path in image_paths: 166 | # image_path = 'images/300W_LP/AFW_GEN/AFW_111076519_1_1.jpg' 167 | counter += 1 168 | ver_image = run_one_image(image_path, uv_kpt_ind, face_ind, triangles, s_uv_coords, 169 | pnet, rnet, onet, tx, ty, tsess) 170 | 171 | imsave(os.path.join("TrainingData", "test" + '_mask_' + str(counter) + '.jpg'), ver_image) 172 | ''' 173 | -------------------------------------------------------------------------------- /show.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Mon Nov 12 16:46:30 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import cv2 9 | import numpy as np 10 | from skimage.transform import resize 11 | 12 | def get_transpose_axes(n): 13 | 14 | if n % 2 == 0: 15 | y_axes = list(range( 1, n-1, 2 )) 16 | x_axes = list(range( 0, n-1, 2 )) 17 | else: 18 | y_axes = list(range( 0, n-1, 2 )) 19 | x_axes = list(range( 1, n-1, 2 )) 20 | 21 | return y_axes, x_axes, [n-1] 22 | 23 | def stack_images(images): 24 | 25 | images_shape = np.array(images.shape) 26 | new_axes = get_transpose_axes(len( images_shape )) 27 | new_shape = [np.prod(images_shape[x]) for x in new_axes] 28 | 29 | return np.transpose(images, axes=np.concatenate(new_axes)).reshape( new_shape) 30 | 31 | def resize_images(images, size): 32 | 33 | new_images = [] 34 | for image in images: 35 | nimage = resize(image, (size, size), preserve_range = True) 36 | new_images.append(nimage) 37 | return np.array(new_images) 38 | 39 | def show_G(images_A, images_B, batch_size, name): 40 | 41 | images_A = resize_images(images_A, 128) 42 | images_B = resize_images(images_B, 128) 43 | figure = np.stack([images_A, images_B], axis=1) 44 | figure = figure.reshape((4, batch_size//4) + figure.shape[1:]) 45 | figure = stack_images(figure) 46 | figure = np.clip(figure * 255.0, 0, 255).astype('uint8') 47 | figure = cv2.cvtColor(figure, cv2.COLOR_BGR2RGB) 48 | 49 | cv2.imshow(name, figure) 50 | -------------------------------------------------------------------------------- /tmodel.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Thu Nov 8 11:10:07 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import tensorflow as tf 9 | import tensorflow.contrib.layers as tcl 10 | from tensorflow.contrib.framework import arg_scope 11 | 12 | def res_block(x, num_outputs, kernel_size=4, stride=1): 13 | assert num_outputs % 2 == 0 14 | shortcut = x 15 | if stride != 1 or int(x.get_shape()[3]) != num_outputs: 16 | shortcut = tcl.conv2d(shortcut, num_outputs, kernel_size=1, stride=stride, 17 | activation_fn=None, normalizer_fn=None) 18 | x = tcl.conv2d(x, num_outputs/2, kernel_size=1, stride=1, padding='SAME') 19 | x = tcl.conv2d(x, num_outputs/2, kernel_size=kernel_size, stride=stride, padding='SAME') 20 | x = tcl.conv2d(x, num_outputs, kernel_size=1, stride=1, padding='SAME', 21 | activation_fn=None, normalizer_fn=None) 22 | x += shortcut 23 | x = tcl.batch_norm(x) 24 | x = tf.nn.relu(x) 25 | 26 | return x 27 | 28 | def PRNet(image_shape=(256, 256, 3), output_shape=(256, 256,3), size=16, is_training=True): 29 | 30 | with tf.name_scope('3dface') as scope: 31 | x = tf.placeholder(shape=(None, image_shape[0], image_shape[1], image_shape[2]), 32 | dtype=tf.float32, name='x') 33 | with tf.name_scope('PRNet') as scope: 34 | with arg_scope([tcl.batch_norm], is_training=is_training, scale=True): 35 | with arg_scope([tcl.conv2d, tcl.conv2d_transpose], activation_fn=tf.nn.relu, 36 | normalizer_fn=tcl.batch_norm, biases_initializer=None, 37 | padding='SAME', weights_regularizer=tcl.l2_regularizer(0.0002)): 38 | # encoder 39 | se = tcl.conv2d(x, num_outputs=size, kernel_size=4, stride=1) # 256 x 256 x 16 40 | se = res_block(se, num_outputs=size * 2, kernel_size=4, stride=2) # 128 x 128 x 32 41 | se = res_block(se, num_outputs=size * 2, kernel_size=4, stride=1) # 128 x 128 x 32 42 | se = res_block(se, num_outputs=size * 4, kernel_size=4, stride=2) # 64 x 64 x 64 43 | se = res_block(se, num_outputs=size * 4, kernel_size=4, stride=1) # 64 x 64 x 64 44 | se = res_block(se, num_outputs=size * 8, kernel_size=4, stride=2) # 32 x 32 x 128 45 | se = res_block(se, num_outputs=size * 8, kernel_size=4, stride=1) # 32 x 32 x 128 46 | se = res_block(se, num_outputs=size * 16, kernel_size=4, stride=2) # 16 x 16 x 256 47 | se = res_block(se, num_outputs=size * 16, kernel_size=4, stride=1) # 16 x 16 x 256 48 | se = res_block(se, num_outputs=size * 32, kernel_size=4, stride=2) # 8 x 8 x 512 49 | se = res_block(se, num_outputs=size * 32, kernel_size=4, stride=1) # 8 x 8 x 512 50 | # decoder 51 | pd = tcl.conv2d_transpose(se, size * 32, 4, stride=1) # 8 x 8 x 512 52 | pd = tcl.conv2d_transpose(pd, size * 16, 4, stride=2) # 16 x 16 x 256 53 | pd = tcl.conv2d_transpose(pd, size * 16, 4, stride=1) # 16 x 16 x 256 54 | pd = tcl.conv2d_transpose(pd, size * 16, 4, stride=1) # 16 x 16 x 256 55 | pd = tcl.conv2d_transpose(pd, size * 8, 4, stride=2) # 32 x 32 x 128 56 | pd = tcl.conv2d_transpose(pd, size * 8, 4, stride=1) # 32 x 32 x 128 57 | pd = tcl.conv2d_transpose(pd, size * 8, 4, stride=1) # 32 x 32 x 128 58 | pd = tcl.conv2d_transpose(pd, size * 4, 4, stride=2) # 64 x 64 x 64 59 | pd = tcl.conv2d_transpose(pd, size * 4, 4, stride=1) # 64 x 64 x 64 60 | pd = tcl.conv2d_transpose(pd, size * 4, 4, stride=1) # 64 x 64 x 64 61 | 62 | pd = tcl.conv2d_transpose(pd, size * 2, 4, stride=2) # 128 x 128 x 32 63 | pd = tcl.conv2d_transpose(pd, size * 2, 4, stride=1) # 128 x 128 x 32 64 | pd = tcl.conv2d_transpose(pd, size, 4, stride=2) # 256 x 256 x 16 65 | pd = tcl.conv2d_transpose(pd, size, 4, stride=1) # 256 x 256 x 16 66 | 67 | pd = tcl.conv2d_transpose(pd, 3, 4, stride=1) # 256 x 256 x 3 68 | pd = tcl.conv2d_transpose(pd, 3, 4, stride=1) # 256 x 256 x 3 69 | pos = tcl.conv2d_transpose(pd, 3, 4, stride=1, activation_fn = tf.nn.sigmoid) 70 | 71 | return x, pos 72 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Created on Mon Nov 12 09:15:27 2018 4 | 5 | @author: shen1994 6 | """ 7 | 8 | import os 9 | import cv2 10 | import numpy as np 11 | import tensorflow as tf 12 | from model import PRNet 13 | from generate import Generator 14 | from predictor import Predictor 15 | from show import show_G 16 | 17 | if __name__ == "__main__": 18 | 19 | os.environ['CUDA_VISIBLE_DEVICES'] = '2' 20 | 21 | model_path = 'model' 22 | if not os.path.exists(model_path): 23 | os.mkdir(model_path) 24 | 25 | # define params 26 | batch_size = 32 27 | epocs = 100001 28 | image_shape = (256, 256, 3) 29 | 30 | # define model 31 | x, m, y, pos, loss, optimizer = PRNet() 32 | 33 | # define predictor 34 | v_predictor = Predictor() 35 | 36 | # define generator 37 | train_paths = ['images/300W_LP/AFW_GEN', 'images/300W_LP/HELEN_GEN', 38 | 'images/300W_LP/IBUG_GEN', 'images/300W_LP/LFPW_GEN', 39 | 'images/300W_LP/AFW_Flip_GEN', 'images/300W_LP/HELEN_Flip_GEN', 40 | 'images/300W_LP/IBUG_Flip_GEN', 'images/300W_LP/LFPW_Flip_GEN'] 41 | valid_path = 'images/AFLW2000_GEN' 42 | mask_path = 'images/uv_weight_mask2.png' 43 | 44 | t_generator = Generator(train_paths=train_paths, 45 | valid_path=valid_path, 46 | mask_path=mask_path, 47 | image_shape=image_shape, 48 | batch_size=batch_size).generate(is_training=True) 49 | v_generator = Generator(train_paths=train_paths, 50 | valid_path=valid_path, 51 | mask_path=mask_path, 52 | image_shape=image_shape, 53 | batch_size=batch_size).generate(is_training=False) 54 | 55 | saver = tf.train.Saver() 56 | with tf.Session() as sess: 57 | 58 | # initial variables 59 | sess.run(tf.local_variables_initializer()) 60 | sess.run(tf.global_variables_initializer()) 61 | 62 | # restore model 63 | try: 64 | ckpt = tf.train.latest_checkpoint(model_path) 65 | saver.restore(sess, ckpt) 66 | except Exception: 67 | print('No existed model to use!') 68 | 69 | # train data 70 | step = 0 71 | total_coss = 0 72 | while step < epocs: 73 | 74 | x_in, y_in, m_in = t_generator.__next__() 75 | _ = sess.run(optimizer, feed_dict={x: x_in, y: y_in, m: m_in}) 76 | total_coss += sess.run(loss, feed_dict={x: x_in, y: y_in, m: m_in}) 77 | 78 | if step % 100 == 0: 79 | 80 | # show total loss 81 | print(str(step) + ": train --->" + "cost:" + str(total_coss)) 82 | print("---------------------------------------->") 83 | total_coss = 0 84 | 85 | # show keypoints 86 | x_ou = v_generator.__next__() 87 | x_pre_ou = [] 88 | for i in range(16): 89 | x_pre_ou.append(v_predictor.predictor(sess, x, pos, x_ou[i])) 90 | show_G(x_ou[:16], np.array(x_pre_ou), 16, "3DFace") 91 | 92 | # save model 93 | if step % 1000 == 0 and step != 0: 94 | saver.save(sess, 'model/model%d.ckpt' % step) 95 | 96 | # exit programs 97 | if cv2.waitKey(1) == ord('q'): 98 | exit() 99 | 100 | step += 1 101 | --------------------------------------------------------------------------------