├── .gitignore ├── LICENSE ├── README.md ├── examples ├── data │ ├── example2_ref.png │ ├── example3_ref.png │ ├── example4_init.png │ ├── example4_ref.png │ └── teapot.obj ├── example1.py ├── example2.py ├── example3.py └── example4.py ├── neural_renderer ├── __init__.py ├── cuda │ ├── __init__.py │ ├── create_texture_image_cuda.cpp │ ├── create_texture_image_cuda_kernel.cu │ ├── load_textures_cuda.cpp │ ├── load_textures_cuda_kernel.cu │ ├── rasterize_cuda.cpp │ └── rasterize_cuda_kernel.cu ├── get_points_from_angles.py ├── lighting.py ├── load_obj.py ├── look.py ├── look_at.py ├── mesh.py ├── perspective.py ├── projection.py ├── rasterize.py ├── renderer.py ├── save_obj.py └── vertices_to_faces.py ├── setup.py └── tests ├── data ├── 1cde62b063e14777c9152a706245d48 │ ├── model.mtl │ └── model.obj ├── 4e49873292196f02574b5684eaec43e9 │ ├── images │ │ ├── texture0.jpg │ │ ├── texture0_.jpg │ │ └── texture1.jpg │ ├── model.mtl │ └── model.obj ├── clean.blend ├── rasterize_silhouettes_case1.png ├── rasterize_silhouettes_case1_v0_x.png ├── rasterize_silhouettes_case1_v0_y.png ├── rasterize_silhouettes_case1_v1_x.png ├── rasterize_silhouettes_case1_v1_y.png ├── rasterize_silhouettes_case1_v2_x.png ├── rasterize_silhouettes_case1_v2_y.png ├── rasterize_silhouettes_case2.png ├── rasterize_silhouettes_case2_v0_x.png ├── rasterize_silhouettes_case2_v0_y.png ├── rasterize_silhouettes_case2_v1_x.png ├── rasterize_silhouettes_case2_v1_y.png ├── rasterize_silhouettes_case2_v2_x.png ├── rasterize_silhouettes_case2_v2_y.png ├── teapot.obj ├── teapot_blender.png ├── test_depth.png └── tetrahedron.obj ├── test_get_points_from_angles.py ├── test_lighting.py ├── test_load_obj.py ├── test_look.py ├── test_look_at.py ├── test_perspective.py ├── test_rasterize.py ├── test_rasterize_depth.py ├── test_rasterize_silhouettes.py ├── test_renderer.py ├── test_save_obj.py ├── test_vertices_to_faces.py └── utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | # Byte-compiled / optimized / DLL files 3 | __pycache__/ 4 | *.py[cod] 5 | *$py.class 6 | 7 | # C extensions 8 | *.so 9 | *.ninja* 10 | *.o 11 | 12 | # Distribution / packaging 13 | .Python 14 | build/ 15 | develop-eggs/ 16 | dist/ 17 | downloads/ 18 | eggs/ 19 | .eggs/ 20 | lib/ 21 | lib64/ 22 | parts/ 23 | sdist/ 24 | var/ 25 | wheels/ 26 | *.egg-info/ 27 | .installed.cfg 28 | *.egg 29 | MANIFEST 30 | 31 | # PyInstaller 32 | # Usually these files are written by a python script from a template 33 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 34 | *.manifest 35 | *.spec 36 | 37 | # Installer logs 38 | pip-log.txt 39 | pip-delete-this-directory.txt 40 | 41 | # Unit test / coverage reports 42 | htmlcov/ 43 | .tox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | 53 | # Translations 54 | *.mo 55 | *.pot 56 | 57 | # Django stuff: 58 | *.log 59 | local_settings.py 60 | db.sqlite3 61 | 62 | # Flask stuff: 63 | instance/ 64 | .webassets-cache 65 | 66 | # Scrapy stuff: 67 | .scrapy 68 | 69 | # Sphinx documentation 70 | docs/_build/ 71 | 72 | # PyBuilder 73 | target/ 74 | 75 | # Jupyter Notebook 76 | .ipynb_checkpoints 77 | 78 | # pyenv 79 | .python-version 80 | 81 | # celery beat schedule file 82 | celerybeat-schedule 83 | 84 | # SageMath parsed files 85 | *.sage.py 86 | 87 | # Environments 88 | .env 89 | .venv 90 | env/ 91 | venv/ 92 | ENV/ 93 | env.bak/ 94 | venv.bak/ 95 | 96 | # Spyder project settings 97 | .spyderproject 98 | .spyproject 99 | 100 | # Rope project settings 101 | .ropeproject 102 | 103 | # mkdocs documentation 104 | /site 105 | 106 | # mypy 107 | .mypy_cache/ 108 | 109 | # vim swap files 110 | *.swp 111 | 112 | # result gifs 113 | *result.gif 114 | tests/data/car.png 115 | tests/data/display.png 116 | tests/data/test_rasterize1.png 117 | tests/data/test_rasterize2.png 118 | examples/data/example1.gif 119 | examples/data/example2_optimization.gif 120 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017 Hiroharu Kato 4 | Copyright (c) 2018 Nikos Kolotouros 5 | A PyTorch implementation of Neural 3D Mesh Renderer (https://github.com/hiroharu-kato/neural_renderer) 6 | 7 | Permission is hereby granted, free of charge, to any person obtaining a copy 8 | of this software and associated documentation files (the "Software"), to deal 9 | in the Software without restriction, including without limitation the rights 10 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 11 | copies of the Software, and to permit persons to whom the Software is 12 | furnished to do so, subject to the following conditions: 13 | 14 | The above copyright notice and this permission notice shall be included in all 15 | copies or substantial portions of the Software. 16 | 17 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 18 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 19 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 20 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 21 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 22 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 23 | SOFTWARE. 24 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # neural_renderer_pytorch-win10 2 | 解决neural_renderer_pytorch在win10上编译安装的问题 3 | 4 | neural_renderer_pytorch使用说明请参照[neural_renderer](https://github.com/daniilidis-group/neural_renderer)，这里主要解决neural_renderer_pytorch在win10上编译安装不通过问题。 5 | 6 | 7 | 8 | ## 问题描述 9 | 10 | 在Linux系统或者是Mac系统上，直接执行命令`pip install neural_renderer_pytorch`是可以直接安装成功的，但是在win10系统上，会报各种编译错误，如何在win10上愉快地玩耍neural_renderer呢？ 11 | 12 | 13 | 14 | ## 解决方案 15 | 16 | 本人经过无数次采坑，尝试，将问题定位为3个，1、vs版本和nvcc版本不匹配导致nvcc编译错误。2、pytorch的源码'cpp_extension.py'第233行，编译信息的解码未设置成"gbk"模式。3、neural_renderer源码中有些地方编译不通过（原因未知）。 17 | 18 | 那么针对上述3个问题，下面给出解决方案 19 | 20 | - vs版本和nvcc版本问题：这里推荐vs2019和CUDA10.1，pytorch和torchvision版本最好是最新的且要匹配 21 | 22 | - pytorch源码问题：在你的python环境中找到pytorch的`cpp_extension.py`文件,路径一般是`/anaconda/Lib/site-packages/torch/utils/cpp_extension.py`，将其第233行的代码 23 | 24 | `match = re.search(r'(\d+)\.(\d+)\.(\d+)', compiler_info.decode().strip())` 25 | 26 | 修改为 27 | 28 | `match = re.search(r'(\d+)\.(\d+)\.(\d+)', compiler_info.decode(' gbk').strip())` 29 | 30 | 注意gbk前面有一个空格，不然会报错。 31 | 32 | - neural_renderer源码问题：我对neural_renderer源码进行了一些修改，请下载修改后的代码进行源码安装`python setup.py install` 33 | 34 | 主要修改内容有: 35 | 36 | - `/cuda/create_texture_image_cuda.cpp`：注释了所有的`AT_CHECK`,`CHECK_INPUT`（不知道为什么win10这个编译的时候会报错） 37 | - `/cuda/load_textures_cuda.cpp`：注释了所有的`AT_CHECK`,`CHECK_INPUT` 38 | - `/cuda/rasterize_cuda.cpp`：注释了所有的`AT_CHECK`,`CHECK_INPUT` 39 | - `/cuda/rasterize_cuda_kernel.cu`：注释了`static __inline__ __device__ double atomicAdd(double* address, double val)`函数，因为高版本pytorch似乎自带double类型的atomicAdd，于是这里再写一遍会冲突。 40 | 41 | 42 | 43 | -------------------------------------------------------------------------------- /examples/data/example2_ref.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/examples/data/example2_ref.png -------------------------------------------------------------------------------- /examples/data/example3_ref.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/examples/data/example3_ref.png -------------------------------------------------------------------------------- /examples/data/example4_init.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/examples/data/example4_init.png -------------------------------------------------------------------------------- /examples/data/example4_ref.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/examples/data/example4_ref.png -------------------------------------------------------------------------------- /examples/example1.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example 1. Drawing a teapot from multiple viewpoints. 3 | """ 4 | import os 5 | import argparse 6 | 7 | import torch 8 | import numpy as np 9 | import tqdm 10 | import imageio 11 | 12 | import neural_renderer as nr 13 | 14 | current_dir = os.path.dirname(os.path.realpath(__file__)) 15 | data_dir = os.path.join(current_dir, 'data') 16 | 17 | 18 | def main(): 19 | parser = argparse.ArgumentParser() 20 | parser.add_argument('-i', '--filename_input', type=str, default=os.path.join(data_dir, 'teapot.obj')) 21 | parser.add_argument('-o', '--filename_output', type=str, default=os.path.join(data_dir, 'example1.gif')) 22 | parser.add_argument('-g', '--gpu', type=int, default=0) 23 | args = parser.parse_args() 24 | 25 | # other settings 26 | camera_distance = 2.732 27 | elevation = 30 28 | texture_size = 2 29 | 30 | # load .obj 31 | vertices, faces = nr.load_obj(args.filename_input) 32 | vertices = vertices[None, :, :] # [num_vertices, XYZ] -> [batch_size=1, num_vertices, XYZ] 33 | faces = faces[None, :, :] # [num_faces, 3] -> [batch_size=1, num_faces, 3] 34 | 35 | # create texture [batch_size=1, num_faces, texture_size, texture_size, texture_size, RGB] 36 | textures = torch.ones(1, faces.shape[1], texture_size, texture_size, texture_size, 3, dtype=torch.float32).cuda() 37 | 38 | # to gpu 39 | 40 | # create renderer 41 | renderer = nr.Renderer(camera_mode='look_at') 42 | 43 | # draw object 44 | loop = tqdm.tqdm(range(0, 360, 4)) 45 | writer = imageio.get_writer(args.filename_output, mode='I') 46 | for num, azimuth in enumerate(loop): 47 | loop.set_description('Drawing') 48 | renderer.eye = nr.get_points_from_angles(camera_distance, elevation, azimuth) 49 | images, _, _ = renderer(vertices, faces, textures) # [batch_size, RGB, image_size, image_size] 50 | image = images.detach().cpu().numpy()[0].transpose((1, 2, 0)) # [image_size, image_size, RGB] 51 | writer.append_data((255*image).astype(np.uint8)) 52 | writer.close() 53 | 54 | if __name__ == '__main__': 55 | main() 56 | -------------------------------------------------------------------------------- /examples/example2.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example 2. Optimizing vertices. 3 | """ 4 | from __future__ import division 5 | import os 6 | import argparse 7 | import glob 8 | 9 | import torch 10 | import torch.nn as nn 11 | import numpy as np 12 | from skimage.io import imread, imsave 13 | import tqdm 14 | import imageio 15 | 16 | import neural_renderer as nr 17 | 18 | current_dir = os.path.dirname(os.path.realpath(__file__)) 19 | data_dir = os.path.join(current_dir, 'data') 20 | 21 | class Model(nn.Module): 22 | def __init__(self, filename_obj, filename_ref): 23 | super(Model, self).__init__() 24 | 25 | # load .obj 26 | vertices, faces = nr.load_obj(filename_obj) 27 | self.vertices = nn.Parameter(vertices[None, :, :]) 28 | self.register_buffer('faces', faces[None, :, :]) 29 | 30 | # create textures 31 | texture_size = 2 32 | textures = torch.ones(1, self.faces.shape[1], texture_size, texture_size, texture_size, 3, dtype=torch.float32) 33 | self.register_buffer('textures', textures) 34 | 35 | # load reference image 36 | image_ref = torch.from_numpy(imread(filename_ref).astype(np.float32).mean(-1) / 255.)[None, ::] 37 | self.register_buffer('image_ref', image_ref) 38 | 39 | # setup renderer 40 | renderer = nr.Renderer(camera_mode='look_at') 41 | self.renderer = renderer 42 | 43 | def forward(self): 44 | self.renderer.eye = nr.get_points_from_angles(2.732, 0, 90) 45 | image = self.renderer(self.vertices, self.faces, mode='silhouettes') 46 | loss = torch.sum((image - self.image_ref[None, :, :])**2) 47 | return loss 48 | 49 | 50 | def make_gif(filename): 51 | with imageio.get_writer(filename, mode='I') as writer: 52 | for filename in sorted(glob.glob('/tmp/_tmp_*.png')): 53 | writer.append_data(imageio.imread(filename)) 54 | os.remove(filename) 55 | writer.close() 56 | 57 | 58 | def main(): 59 | parser = argparse.ArgumentParser() 60 | parser.add_argument('-io', '--filename_obj', type=str, default=os.path.join(data_dir, 'teapot.obj')) 61 | parser.add_argument('-ir', '--filename_ref', type=str, default=os.path.join(data_dir, 'example2_ref.png')) 62 | parser.add_argument( 63 | '-oo', '--filename_output_optimization', type=str, default=os.path.join(data_dir, 'example2_optimization.gif')) 64 | parser.add_argument( 65 | '-or', '--filename_output_result', type=str, default=os.path.join(data_dir, 'example2_result.gif')) 66 | parser.add_argument('-g', '--gpu', type=int, default=0) 67 | args = parser.parse_args() 68 | 69 | model = Model(args.filename_obj, args.filename_ref) 70 | model.cuda() 71 | 72 | optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters())) 73 | # optimizer.setup(model) 74 | loop = tqdm.tqdm(range(300)) 75 | for i in loop: 76 | loop.set_description('Optimizing') 77 | # optimizer.target.cleargrads() 78 | optimizer.zero_grad() 79 | loss = model() 80 | loss.backward() 81 | optimizer.step() 82 | images = model.renderer(model.vertices, model.faces, mode='silhouettes') 83 | image = images.detach().cpu().numpy()[0] 84 | imsave('/tmp/_tmp_%04d.png' % i, image) 85 | make_gif(args.filename_output_optimization) 86 | 87 | # draw object 88 | loop = tqdm.tqdm(range(0, 360, 4)) 89 | for num, azimuth in enumerate(loop): 90 | loop.set_description('Drawing') 91 | model.renderer.eye = nr.get_points_from_angles(2.732, 0, azimuth) 92 | images, _, _ = model.renderer(model.vertices, model.faces, model.textures) 93 | image = images.detach().cpu().numpy()[0].transpose((1, 2, 0)) 94 | imsave('/tmp/_tmp_%04d.png' % num, image) 95 | make_gif(args.filename_output_result) 96 | 97 | 98 | if __name__ == '__main__': 99 | main() 100 | -------------------------------------------------------------------------------- /examples/example3.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example 3. Optimizing textures. 3 | """ 4 | from __future__ import division 5 | import os 6 | import argparse 7 | import glob 8 | 9 | import torch 10 | import torch.nn as nn 11 | import numpy as np 12 | from skimage.io import imread, imsave 13 | import tqdm 14 | import imageio 15 | 16 | import neural_renderer as nr 17 | 18 | current_dir = os.path.dirname(os.path.realpath(__file__)) 19 | data_dir = os.path.join(current_dir, 'data') 20 | 21 | class Model(nn.Module): 22 | def __init__(self, filename_obj, filename_ref): 23 | super(Model, self).__init__() 24 | vertices, faces = nr.load_obj(filename_obj) 25 | self.register_buffer('vertices', vertices[None, :, :]) 26 | self.register_buffer('faces', faces[None, :, :]) 27 | 28 | # create textures 29 | texture_size = 4 30 | textures = torch.zeros(1, self.faces.shape[1], texture_size, texture_size, texture_size, 3, dtype=torch.float32) 31 | self.textures = nn.Parameter(textures) 32 | 33 | # load reference image 34 | image_ref = torch.from_numpy(imread(filename_ref).astype('float32') / 255.).permute(2,0,1)[None, ::] 35 | self.register_buffer('image_ref', image_ref) 36 | 37 | # setup renderer 38 | renderer = nr.Renderer(camera_mode='look_at') 39 | renderer.perspective = False 40 | renderer.light_intensity_directional = 0.0 41 | renderer.light_intensity_ambient = 1.0 42 | self.renderer = renderer 43 | 44 | 45 | def forward(self): 46 | self.renderer.eye = nr.get_points_from_angles(2.732, 0, np.random.uniform(0, 360)) 47 | image, _, _ = self.renderer(self.vertices, self.faces, torch.tanh(self.textures)) 48 | loss = torch.sum((image - self.image_ref) ** 2) 49 | return loss 50 | 51 | 52 | def make_gif(filename): 53 | with imageio.get_writer(filename, mode='I') as writer: 54 | for filename in sorted(glob.glob('/tmp/_tmp_*.png')): 55 | writer.append_data(imageio.imread(filename)) 56 | os.remove(filename) 57 | writer.close() 58 | 59 | 60 | def main(): 61 | parser = argparse.ArgumentParser() 62 | parser.add_argument('-io', '--filename_obj', type=str, default=os.path.join(data_dir, 'teapot.obj')) 63 | parser.add_argument('-ir', '--filename_ref', type=str, default=os.path.join(data_dir, 'example3_ref.png')) 64 | parser.add_argument('-or', '--filename_output', type=str, default=os.path.join(data_dir, 'example3_result.gif')) 65 | parser.add_argument('-g', '--gpu', type=int, default=0) 66 | args = parser.parse_args() 67 | 68 | model = Model(args.filename_obj, args.filename_ref) 69 | model.cuda() 70 | 71 | optimizer = torch.optim.Adam(model.parameters(), lr=0.1, betas=(0.5,0.999)) 72 | loop = tqdm.tqdm(range(300)) 73 | for _ in loop: 74 | loop.set_description('Optimizing') 75 | optimizer.zero_grad() 76 | loss = model() 77 | loss.backward() 78 | optimizer.step() 79 | 80 | # draw object 81 | loop = tqdm.tqdm(range(0, 360, 4)) 82 | for num, azimuth in enumerate(loop): 83 | loop.set_description('Drawing') 84 | model.renderer.eye = nr.get_points_from_angles(2.732, 0, azimuth) 85 | images, _, _ = model.renderer(model.vertices, model.faces, torch.tanh(model.textures)) 86 | image = images.detach().cpu().numpy()[0].transpose((1, 2, 0)) 87 | imsave('/tmp/_tmp_%04d.png' % num, image) 88 | make_gif(args.filename_output) 89 | 90 | 91 | if __name__ == '__main__': 92 | main() 93 | -------------------------------------------------------------------------------- /examples/example4.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example 4. Finding camera parameters. 3 | """ 4 | import os 5 | import argparse 6 | import glob 7 | 8 | import torch 9 | import torch.nn as nn 10 | import numpy as np 11 | from skimage.io import imread, imsave 12 | import tqdm 13 | import imageio 14 | 15 | import neural_renderer as nr 16 | 17 | current_dir = os.path.dirname(os.path.realpath(__file__)) 18 | data_dir = os.path.join(current_dir, 'data') 19 | 20 | class Model(nn.Module): 21 | def __init__(self, filename_obj, filename_ref=None): 22 | super(Model, self).__init__() 23 | # load .obj 24 | vertices, faces = nr.load_obj(filename_obj) 25 | self.register_buffer('vertices', vertices[None, :, :]) 26 | self.register_buffer('faces', faces[None, :, :]) 27 | 28 | # create textures 29 | texture_size = 2 30 | textures = torch.ones(1, self.faces.shape[1], texture_size, texture_size, texture_size, 3, dtype=torch.float32) 31 | self.register_buffer('textures', textures) 32 | 33 | # load reference image 34 | image_ref = torch.from_numpy((imread(filename_ref).max(-1) != 0).astype(np.float32)) 35 | self.register_buffer('image_ref', image_ref) 36 | 37 | # camera parameters 38 | self.camera_position = nn.Parameter(torch.from_numpy(np.array([6, 10, -14], dtype=np.float32))) 39 | 40 | # setup renderer 41 | renderer = nr.Renderer(camera_mode='look_at') 42 | renderer.eye = self.camera_position 43 | self.renderer = renderer 44 | 45 | def forward(self): 46 | image = self.renderer(self.vertices, self.faces, mode='silhouettes') 47 | loss = torch.sum((image - self.image_ref[None, :, :]) ** 2) 48 | return loss 49 | 50 | 51 | def make_gif(filename): 52 | with imageio.get_writer(filename, mode='I') as writer: 53 | for filename in sorted(glob.glob('/tmp/_tmp_*.png')): 54 | writer.append_data(imread(filename)) 55 | os.remove(filename) 56 | writer.close() 57 | 58 | 59 | def make_reference_image(filename_ref, filename_obj): 60 | model = Model(filename_obj) 61 | model.cuda() 62 | 63 | model.renderer.eye = nr.get_points_from_angles(2.732, 30, -15) 64 | images, _, _ = model.renderer.render(model.vertices, model.faces, torch.tanh(model.textures)) 65 | image = images.detach().cpu().numpy()[0] 66 | imsave(filename_ref, image) 67 | 68 | 69 | def main(): 70 | parser = argparse.ArgumentParser() 71 | parser.add_argument('-io', '--filename_obj', type=str, default=os.path.join(data_dir, 'teapot.obj')) 72 | parser.add_argument('-ir', '--filename_ref', type=str, default=os.path.join(data_dir, 'example4_ref.png')) 73 | parser.add_argument('-or', '--filename_output', type=str, default=os.path.join(data_dir, 'example4_result.gif')) 74 | parser.add_argument('-mr', '--make_reference_image', type=int, default=0) 75 | parser.add_argument('-g', '--gpu', type=int, default=0) 76 | args = parser.parse_args() 77 | 78 | if args.make_reference_image: 79 | make_reference_image(args.filename_ref, args.filename_obj) 80 | 81 | model = Model(args.filename_obj, args.filename_ref) 82 | model.cuda() 83 | 84 | # optimizer = chainer.optimizers.Adam(alpha=0.1) 85 | optimizer = torch.optim.Adam(model.parameters(), lr=0.1) 86 | loop = tqdm.tqdm(range(1000)) 87 | for i in loop: 88 | optimizer.zero_grad() 89 | loss = model() 90 | loss.backward() 91 | optimizer.step() 92 | images, _, _ = model.renderer(model.vertices, model.faces, torch.tanh(model.textures)) 93 | image = images.detach().cpu().numpy()[0].transpose(1,2,0) 94 | imsave('/tmp/_tmp_%04d.png' % i, image) 95 | loop.set_description('Optimizing (loss %.4f)' % loss.data) 96 | if loss.item() < 70: 97 | break 98 | make_gif(args.filename_output) 99 | 100 | 101 | if __name__ == '__main__': 102 | main() 103 | -------------------------------------------------------------------------------- /neural_renderer/__init__.py: -------------------------------------------------------------------------------- 1 | from .get_points_from_angles import get_points_from_angles 2 | from .lighting import lighting 3 | from .load_obj import load_obj 4 | from .look import look 5 | from .look_at import look_at 6 | from .mesh import Mesh 7 | from .perspective import perspective 8 | from .projection import projection 9 | from .rasterize import (rasterize_rgbad, rasterize, rasterize_silhouettes, rasterize_depth, Rasterize) 10 | from .renderer import Renderer 11 | from .save_obj import save_obj 12 | from .vertices_to_faces import vertices_to_faces 13 | 14 | __version__ = '1.1.3' 15 | name = 'neural_renderer_pytorch' 16 | -------------------------------------------------------------------------------- /neural_renderer/cuda/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/neural_renderer/cuda/__init__.py -------------------------------------------------------------------------------- /neural_renderer/cuda/create_texture_image_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | // CUDA forward declarations 4 | 5 | at::Tensor create_texture_image_cuda( 6 | at::Tensor vertices_all, 7 | at::Tensor textures, 8 | at::Tensor image, 9 | float eps); 10 | 11 | // C++ interface 12 | 13 | //#define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 14 | //#define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x " must be contiguous") 15 | //#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 16 | 17 | 18 | at::Tensor create_texture_image( 19 | at::Tensor vertices_all, 20 | at::Tensor textures, 21 | at::Tensor image, 22 | float eps) { 23 | 24 | //CHECK_INPUT(vertices_all); 25 | //CHECK_INPUT(textures); 26 | //CHECK_INPUT(image); 27 | 28 | return create_texture_image_cuda(vertices_all, textures, image, eps); 29 | } 30 | 31 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 32 | m.def("create_texture_image", &create_texture_image, "CREATE_TEXTURE_IMAGE (CUDA)"); 33 | } 34 | -------------------------------------------------------------------------------- /neural_renderer/cuda/create_texture_image_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | 5 | #include 6 | #include 7 | 8 | namespace { 9 | template 10 | __global__ void create_texture_image_cuda_kernel( 11 | const scalar_t* __restrict__ vertices_all, 12 | const scalar_t* __restrict__ textures, 13 | scalar_t* __restrict__ image, 14 | size_t image_size, 15 | size_t num_faces, 16 | size_t texture_size_in, 17 | size_t texture_size_out, 18 | size_t tile_width, 19 | scalar_t eps) { 20 | 21 | const int i = blockIdx.x * blockDim.x + threadIdx.x; 22 | if (i >= image_size / 3) { 23 | return; 24 | } 25 | const int x = i % (tile_width * texture_size_out); 26 | const int y = i / (tile_width * texture_size_out); 27 | const int row = x / texture_size_out; 28 | const int column = y / texture_size_out; 29 | const int fn = row + column * tile_width; 30 | const int tsi = texture_size_in; 31 | 32 | const scalar_t* texture = &textures[fn * tsi * tsi * tsi * 3]; 33 | const scalar_t* vertices = &vertices_all[fn * 3 * 2]; 34 | const scalar_t* p0 = &vertices[2 * 0]; 35 | const scalar_t* p1 = &vertices[2 * 1]; 36 | const scalar_t* p2 = &vertices[2 * 2]; 37 | 38 | /* */ 39 | // if ((y % ${texture_size_out}) < (x % ${texture_size_out})) continue; 40 | 41 | /* compute face_inv */ 42 | scalar_t face_inv[9] = { 43 | p1[1] - p2[1], p2[0] - p1[0], p1[0] * p2[1] - p2[0] * p1[1], 44 | p2[1] - p0[1], p0[0] - p2[0], p2[0] * p0[1] - p0[0] * p2[1], 45 | p0[1] - p1[1], p1[0] - p0[0], p0[0] * p1[1] - p1[0] * p0[1]}; 46 | scalar_t face_inv_denominator = ( 47 | p2[0] * (p0[1] - p1[1]) + 48 | p0[0] * (p1[1] - p2[1]) + 49 | p1[0] * (p2[1] - p0[1])); 50 | for (int k = 0; k < 9; k++) face_inv[k] /= face_inv_denominator; 51 | 52 | /* compute w = face_inv * p */ 53 | scalar_t weight[3]; 54 | scalar_t weight_sum = 0; 55 | for (int k = 0; k < 3; k++) { 56 | weight[k] = face_inv[3 * k + 0] * x + face_inv[3 * k + 1] * y + face_inv[3 * k + 2]; 57 | weight_sum += weight[k]; 58 | } 59 | for (int k = 0; k < 3; k++) 60 | weight[k] /= (weight_sum + eps); 61 | 62 | /* get texture index (scalar_t) */ 63 | scalar_t texture_index_scalar_t[3]; 64 | for (int k = 0; k < 3; k++) { 65 | scalar_t tif = weight[k] * (tsi - 1); 66 | tif = max(tif, 0.); 67 | tif = min(tif, tsi - 1 - eps); 68 | texture_index_scalar_t[k] = tif; 69 | } 70 | 71 | /* blend */ 72 | scalar_t new_pixel[3] = {0, 0, 0}; 73 | for (int pn = 0; pn < 8; pn++) { 74 | scalar_t w = 1; // weight 75 | int texture_index_int[3]; // index in source (int) 76 | for (int k = 0; k < 3; k++) { 77 | if ((pn >> k) % 2 == 0) { 78 | w *= 1 - (texture_index_scalar_t[k] - (int)texture_index_scalar_t[k]); 79 | texture_index_int[k] = (int)texture_index_scalar_t[k]; 80 | } 81 | else { 82 | w *= texture_index_scalar_t[k] - (int)texture_index_scalar_t[k]; 83 | texture_index_int[k] = (int)texture_index_scalar_t[k] + 1; 84 | } 85 | } 86 | int isc = texture_index_int[0] * tsi * tsi + texture_index_int[1] * tsi + texture_index_int[2]; 87 | for (int k = 0; k < 3; k++) 88 | new_pixel[k] += w * texture[isc * 3 + k]; 89 | } 90 | for (int k = 0; k < 3; k++) 91 | image[i * 3 + k] = new_pixel[k]; 92 | } 93 | 94 | // didn't really look to see if we fuse the 2 kernels 95 | // probably not because of synchronization issues 96 | template 97 | __global__ void create_texture_image_boundary_cuda_kernel( 98 | scalar_t* image, 99 | size_t image_size, 100 | size_t texture_size_out, 101 | size_t tile_width) { 102 | 103 | const int i = blockIdx.x * blockDim.x + threadIdx.x; 104 | if (i >= image_size / 3) { 105 | return; 106 | } 107 | 108 | const int x = i % (tile_width * texture_size_out); 109 | const int y = i / (tile_width * texture_size_out); 110 | if ((y % texture_size_out + 1) == (x % texture_size_out)) { 111 | for (int k = 0; k < 3; k++) 112 | image[i * 3 + k] = 113 | image[ (y * tile_width * texture_size_out + (x - 1)) * 3 + k]; 114 | } 115 | } 116 | } 117 | 118 | at::Tensor create_texture_image_cuda( 119 | at::Tensor vertices_all, 120 | at::Tensor textures, 121 | at::Tensor image, 122 | float eps) { 123 | 124 | const auto num_faces = textures.size(0); 125 | const auto texture_size_in = textures.size(1); 126 | const auto tile_width = int(sqrt(num_faces - 1)) + 1; 127 | const auto texture_size_out = image.size(1) / tile_width; 128 | 129 | const int threads = 128; 130 | const int image_size = image.numel(); 131 | const dim3 blocks ((image_size / 3 - 1) / threads + 1, 1, 1); 132 | 133 | AT_DISPATCH_FLOATING_TYPES(image.type(), "create_texture_image_cuda", ([&] { 134 | create_texture_image_cuda_kernel<<>>( 135 | vertices_all.data(), 136 | textures.data(), 137 | image.data(), 138 | image_size, 139 | num_faces, 140 | texture_size_in, 141 | texture_size_out, 142 | tile_width, 143 | (scalar_t) eps); 144 | })); 145 | 146 | cudaError_t err = cudaGetLastError(); 147 | if (err != cudaSuccess) 148 | printf("Error in create_texture_image: %s\n", cudaGetErrorString(err)); 149 | 150 | AT_DISPATCH_FLOATING_TYPES(image.type(), "create_texture_image_boundary", ([&] { 151 | create_texture_image_boundary_cuda_kernel<<>>( 152 | image.data(), 153 | image_size, 154 | texture_size_out, 155 | tile_width); 156 | })); 157 | 158 | err = cudaGetLastError(); 159 | if (err != cudaSuccess) 160 | printf("Error in create_texture_image_boundary: %s\n", cudaGetErrorString(err)); 161 | 162 | return image; 163 | } 164 | -------------------------------------------------------------------------------- /neural_renderer/cuda/load_textures_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | // CUDA forward declarations 4 | 5 | at::Tensor load_textures_cuda( 6 | at::Tensor image, 7 | at::Tensor faces, 8 | at::Tensor textures, 9 | at::Tensor is_update, 10 | int texture_wrapping, 11 | int use_bilinear); 12 | 13 | // C++ interface 14 | 15 | //#define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 16 | //#define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x " must be contiguous") 17 | //#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 18 | 19 | 20 | at::Tensor load_textures( 21 | at::Tensor image, 22 | at::Tensor faces, 23 | at::Tensor textures, 24 | at::Tensor is_update, 25 | int texture_wrapping, 26 | int use_bilinear) { 27 | 28 | //CHECK_INPUT(image); 29 | //CHECK_INPUT(faces); 30 | //CHECK_INPUT(is_update); 31 | //CHECK_INPUT(textures); 32 | 33 | return load_textures_cuda(image, faces, textures, is_update, texture_wrapping, use_bilinear); 34 | 35 | } 36 | 37 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 38 | m.def("load_textures", &load_textures, "LOAD_TEXTURES (CUDA)"); 39 | } 40 | -------------------------------------------------------------------------------- /neural_renderer/cuda/load_textures_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | #include 5 | 6 | template 7 | static __inline__ __device__ scalar_t mod(scalar_t x, scalar_t y) { 8 | if (x > 0) { 9 | return fmod(x,y); 10 | } 11 | else { 12 | return y + fmod(x,y); 13 | } 14 | } 15 | 16 | namespace { 17 | 18 | const int REPEAT = 0; 19 | const int MIRRORED_REPEAT = 1; 20 | const int CLAMP_TO_EDGE = 2; 21 | const int CLAMP_TO_BORDER = 3; 22 | 23 | template 24 | __global__ void load_textures_cuda_kernel( 25 | const scalar_t* image, 26 | const int32_t* is_update, 27 | scalar_t* faces, 28 | scalar_t* __restrict__ textures, 29 | int textures_size, 30 | int texture_size, 31 | int image_height, 32 | int image_width, 33 | int texture_wrapping, 34 | bool use_bilinear) { 35 | const int i = blockIdx.x * blockDim.x + threadIdx.x; 36 | if (i >= textures_size / 3) { 37 | return; 38 | } 39 | const int ts = texture_size; 40 | const int fn = i / (ts * ts * ts); 41 | scalar_t dim0 = ((i / (ts * ts)) % ts) / (ts - 1.) ; 42 | scalar_t dim1 = ((i / ts) % ts) / (ts - 1.); 43 | scalar_t dim2 = (i % ts) / (ts - 1.); 44 | if (0 < dim0 + dim1 + dim2) { 45 | float sum = dim0 + dim1 + dim2; 46 | dim0 /= sum; 47 | dim1 /= sum; 48 | dim2 /= sum; 49 | } 50 | scalar_t* face = &faces[fn * 3 * 2]; 51 | scalar_t* texture_ = &textures[i * 3]; 52 | 53 | if (is_update[fn] != 0) { 54 | if (texture_wrapping == REPEAT) { 55 | #pragma unroll 56 | for (int i = 0; i < 6; ++i) { 57 | face[i] = mod(face[i], (scalar_t)1.); 58 | } 59 | } 60 | else if (texture_wrapping == MIRRORED_REPEAT) { 61 | #pragma unroll 62 | for (int i = 0; i < 6; ++i) { 63 | if (mod(face[i], (scalar_t)2) < 1) { 64 | face[i] = mod(face[i], (scalar_t)1.); 65 | } 66 | else { 67 | face[i] = 1 - mod(face[i], (scalar_t)1.); 68 | } 69 | } 70 | } 71 | else if (texture_wrapping == CLAMP_TO_EDGE) { 72 | #pragma unroll 73 | for (int i = 0; i < 6; ++i) { 74 | face[i] = max(min(face[i], (scalar_t) 1), (scalar_t) 0); 75 | } 76 | } 77 | const scalar_t pos_x = ( 78 | (face[2 * 0 + 0] * dim0 + face[2 * 1 + 0] * dim1 + face[2 * 2 + 0] * dim2) * (image_width - 1)); 79 | const scalar_t pos_y = ( 80 | (face[2 * 0 + 1] * dim0 + face[2 * 1 + 1] * dim1 + face[2 * 2 + 1] * dim2) * (image_height - 1)); 81 | if (use_bilinear) { 82 | /* bilinear sampling */ 83 | const scalar_t weight_x1 = pos_x - (int)pos_x; 84 | const scalar_t weight_x0 = 1 - weight_x1; 85 | const scalar_t weight_y1 = pos_y - (int)pos_y; 86 | const scalar_t weight_y0 = 1 - weight_y1; 87 | for (int k = 0; k < 3; k++) { 88 | if (texture_wrapping != CLAMP_TO_BORDER) { 89 | scalar_t c = 0; 90 | c += image[(int)pos_y * image_width * 3 + (int)pos_x * 3 + k] * (weight_x0 * weight_y0); 91 | c += image[min((int)(pos_y + 1), image_height-1) * image_width * 3 + (int)pos_x * 3 + k] * (weight_x0 * weight_y1); 92 | c += image[(int)pos_y * image_width * 3 + min((int)pos_x + 1, image_width-1) * 3 + k] * (weight_x1 * weight_y0); 93 | c += image[min((int)(pos_y + 1), image_height-1) * image_width * 3 + min((int)pos_x + 1, image_width-1) * 3 + k] * (weight_x1 * weight_y1); 94 | texture_[k] = c; 95 | } 96 | else { 97 | texture_[k] = 0; 98 | } 99 | } 100 | } else { 101 | /* nearest neighbor */ 102 | const int pos_xi = round(pos_x); 103 | const int pos_yi = round(pos_y); 104 | for (int k = 0; k < 3; k++) { 105 | if (texture_wrapping != CLAMP_TO_BORDER) { 106 | texture_[k] = image[pos_yi * image_width * 3 + pos_xi * 3 + k]; 107 | } 108 | else { 109 | texture_[k] = 0; 110 | } 111 | } 112 | } 113 | } 114 | } 115 | } 116 | 117 | at::Tensor load_textures_cuda( 118 | at::Tensor image, 119 | at::Tensor faces, 120 | at::Tensor textures, 121 | at::Tensor is_update, 122 | int texture_wrapping, 123 | int use_bilinear) { 124 | // textures_size = size of the textures tensor 125 | const auto textures_size = textures.numel(); 126 | // notice that texture_size != texture_size 127 | const auto texture_size = textures.size(1); 128 | const auto image_height = image.size(0); 129 | const auto image_width = image.size(1); 130 | 131 | const int threads = 1024; 132 | const dim3 blocks ((textures_size / 3 - 1) / threads + 1); 133 | 134 | AT_DISPATCH_FLOATING_TYPES(image.type(), "load_textures_cuda", ([&] { 135 | load_textures_cuda_kernel<<>>( 136 | image.data(), 137 | is_update.data(), 138 | faces.data(), 139 | textures.data(), 140 | textures_size, 141 | texture_size, 142 | image_height, 143 | image_width, 144 | texture_wrapping, 145 | use_bilinear); 146 | })); 147 | 148 | cudaError_t err = cudaGetLastError(); 149 | if (err != cudaSuccess) 150 | printf("Error in load_textures: %s\n", cudaGetErrorString(err)); 151 | return textures; 152 | } 153 | -------------------------------------------------------------------------------- /neural_renderer/cuda/rasterize_cuda.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | 5 | // CUDA forward declarations 6 | 7 | std::vector forward_face_index_map_cuda( 8 | at::Tensor faces, 9 | at::Tensor face_index_map, 10 | at::Tensor weight_map, 11 | at::Tensor depth_map, 12 | at::Tensor face_inv_map, 13 | at::Tensor faces_inv, 14 | int image_size, 15 | float near, 16 | float far, 17 | int return_rgb, 18 | int return_alpha, 19 | int return_depth); 20 | 21 | std::vector forward_texture_sampling_cuda( 22 | at::Tensor faces, 23 | at::Tensor textures, 24 | at::Tensor face_index_map, 25 | at::Tensor weight_map, 26 | at::Tensor depth_map, 27 | at::Tensor rgb_map, 28 | at::Tensor sampling_index_map, 29 | at::Tensor sampling_weight_map, 30 | int image_size, 31 | float eps); 32 | 33 | at::Tensor backward_pixel_map_cuda( 34 | at::Tensor faces, 35 | at::Tensor face_index_map, 36 | at::Tensor rgb_map, 37 | at::Tensor alpha_map, 38 | at::Tensor grad_rgb_map, 39 | at::Tensor grad_alpha_map, 40 | at::Tensor grad_faces, 41 | int image_size, 42 | float eps, 43 | int return_rgb, 44 | int return_alpha); 45 | 46 | at::Tensor backward_textures_cuda( 47 | at::Tensor face_index_map, 48 | at::Tensor sampling_weight_map, 49 | at::Tensor sampling_index_map, 50 | at::Tensor grad_rgb_map, 51 | at::Tensor grad_textures, 52 | int num_faces); 53 | 54 | at::Tensor backward_depth_map_cuda( 55 | at::Tensor faces, 56 | at::Tensor depth_map, 57 | at::Tensor face_index_map, 58 | at::Tensor face_inv_map, 59 | at::Tensor weight_map, 60 | at::Tensor grad_depth_map, 61 | at::Tensor grad_faces, 62 | int image_size); 63 | 64 | // C++ interface 65 | 66 | //#define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor") 67 | //#define CHECK_CONTIGUOUS(x) AT_CHECK(x.is_contiguous(), #x " must be contiguous") 68 | //#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) 69 | 70 | std::vector forward_face_index_map( 71 | at::Tensor faces, 72 | at::Tensor face_index_map, 73 | at::Tensor weight_map, 74 | at::Tensor depth_map, 75 | at::Tensor face_inv_map, 76 | at::Tensor faces_inv, 77 | int image_size, 78 | float near, 79 | float far, 80 | int return_rgb, 81 | int return_alpha, 82 | int return_depth) { 83 | 84 | //CHECK_INPUT(faces); 85 | //CHECK_INPUT(face_index_map); 86 | //CHECK_INPUT(weight_map); 87 | //CHECK_INPUT(depth_map); 88 | //CHECK_INPUT(face_inv_map); 89 | //CHECK_INPUT(faces_inv); 90 | 91 | return forward_face_index_map_cuda(faces, face_index_map, weight_map, 92 | depth_map, face_inv_map, faces_inv, 93 | image_size, near, far, 94 | return_rgb, return_alpha, return_depth); 95 | } 96 | 97 | std::vector forward_texture_sampling( 98 | at::Tensor faces, 99 | at::Tensor textures, 100 | at::Tensor face_index_map, 101 | at::Tensor weight_map, 102 | at::Tensor depth_map, 103 | at::Tensor rgb_map, 104 | at::Tensor sampling_index_map, 105 | at::Tensor sampling_weight_map, 106 | int image_size, 107 | float eps) { 108 | 109 | //CHECK_INPUT(faces); 110 | //CHECK_INPUT(textures); 111 | //CHECK_INPUT(face_index_map); 112 | //CHECK_INPUT(weight_map); 113 | //CHECK_INPUT(depth_map); 114 | //CHECK_INPUT(rgb_map); 115 | //CHECK_INPUT(sampling_index_map); 116 | //CHECK_INPUT(sampling_weight_map); 117 | 118 | return forward_texture_sampling_cuda(faces, textures, face_index_map, 119 | weight_map, depth_map, rgb_map, 120 | sampling_index_map, sampling_weight_map, 121 | image_size, eps); 122 | } 123 | 124 | at::Tensor backward_pixel_map( 125 | at::Tensor faces, 126 | at::Tensor face_index_map, 127 | at::Tensor rgb_map, 128 | at::Tensor alpha_map, 129 | at::Tensor grad_rgb_map, 130 | at::Tensor grad_alpha_map, 131 | at::Tensor grad_faces, 132 | int image_size, 133 | float eps, 134 | int return_rgb, 135 | int return_alpha) { 136 | 137 | //CHECK_INPUT(faces); 138 | //CHECK_INPUT(face_index_map); 139 | //CHECK_INPUT(rgb_map); 140 | //CHECK_INPUT(alpha_map); 141 | //CHECK_INPUT(grad_rgb_map); 142 | //CHECK_INPUT(grad_alpha_map); 143 | //CHECK_INPUT(grad_faces); 144 | 145 | return backward_pixel_map_cuda(faces, face_index_map, rgb_map, alpha_map, 146 | grad_rgb_map, grad_alpha_map, grad_faces, 147 | image_size, eps, return_rgb, return_alpha); 148 | } 149 | 150 | at::Tensor backward_textures( 151 | at::Tensor face_index_map, 152 | at::Tensor sampling_weight_map, 153 | at::Tensor sampling_index_map, 154 | at::Tensor grad_rgb_map, 155 | at::Tensor grad_textures, 156 | int num_faces) { 157 | 158 | //CHECK_INPUT(face_index_map); 159 | //CHECK_INPUT(sampling_weight_map); 160 | //CHECK_INPUT(sampling_index_map); 161 | //CHECK_INPUT(grad_rgb_map); 162 | //CHECK_INPUT(grad_textures); 163 | 164 | return backward_textures_cuda(face_index_map, sampling_weight_map, 165 | sampling_index_map, grad_rgb_map, 166 | grad_textures, num_faces); 167 | } 168 | 169 | at::Tensor backward_depth_map( 170 | at::Tensor faces, 171 | at::Tensor depth_map, 172 | at::Tensor face_index_map, 173 | at::Tensor face_inv_map, 174 | at::Tensor weight_map, 175 | at::Tensor grad_depth_map, 176 | at::Tensor grad_faces, 177 | int image_size) { 178 | 179 | //CHECK_INPUT(faces); 180 | //CHECK_INPUT(depth_map); 181 | //CHECK_INPUT(face_index_map); 182 | //CHECK_INPUT(face_inv_map); 183 | //CHECK_INPUT(weight_map); 184 | //CHECK_INPUT(grad_depth_map); 185 | //CHECK_INPUT(grad_faces); 186 | 187 | return backward_depth_map_cuda(faces, depth_map, face_index_map, 188 | face_inv_map, weight_map, 189 | grad_depth_map, grad_faces, 190 | image_size); 191 | } 192 | 193 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 194 | m.def("forward_face_index_map", &forward_face_index_map, "FORWARD_FACE_INDEX_MAP (CUDA)"); 195 | m.def("forward_texture_sampling", &forward_texture_sampling, "FORWARD_TEXTURE_SAMPLING (CUDA)"); 196 | m.def("backward_pixel_map", &backward_pixel_map, "BACKWARD_PIXEL_MAP (CUDA)"); 197 | m.def("backward_textures", &backward_textures, "BACKWARD_TEXTURES (CUDA)"); 198 | m.def("backward_depth_map", &backward_depth_map, "BACKWARD_DEPTH_MAP (CUDA)"); 199 | } 200 | -------------------------------------------------------------------------------- /neural_renderer/cuda/rasterize_cuda_kernel.cu: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include 5 | #include 6 | 7 | // for the older gpus atomicAdd with double arguments does not exist 8 | /* 9 | #if __CUDA_ARCH__ < 600 and defined(__CUDA_ARCH__) 10 | static __inline__ __device__ double atomicAdd(double* address, double val) { 11 | unsigned long long int* address_as_ull = (unsigned long long int*)address; 12 | unsigned long long int old = *address_as_ull, assumed; 13 | do { 14 | assumed = old; 15 | old = atomicCAS(address_as_ull, assumed, 16 | __double_as_longlong(val + __longlong_as_double(assumed))); 17 | // Note: uses integer comparison to avoid hang in case of NaN (since NaN != NaN) } while (assumed != old); 18 | } while (assumed != old); 19 | return __longlong_as_double(old); 20 | } 21 | #endif 22 | */ 23 | namespace{ 24 | template 25 | __global__ void forward_face_index_map_cuda_kernel_1( 26 | const scalar_t* __restrict__ faces, 27 | scalar_t* __restrict__ faces_inv, 28 | int batch_size, 29 | int num_faces, 30 | int image_size) { 31 | /* batch number, face, number, image size, face[v012][RGB] */ 32 | const int i = blockIdx.x * blockDim.x + threadIdx.x; 33 | if (i >= batch_size * num_faces) { 34 | return; 35 | } 36 | const int is = image_size; 37 | const scalar_t* face = &faces[i * 9]; 38 | scalar_t* face_inv_g = &faces_inv[i * 9]; 39 | 40 | /* return if backside */ 41 | if ((face[7] - face[1]) * (face[3] - face[0]) < (face[4] - face[1]) * (face[6] - face[0])) 42 | return; 43 | 44 | /* p[num][xy]: x, y is normalized from [-1, 1] to [0, is - 1]. */ 45 | scalar_t p[3][2]; 46 | for (int num = 0; num < 3; num++) { 47 | for (int dim = 0; dim < 2; dim++) { 48 | p[num][dim] = 0.5 * (face[3 * num + dim] * is + is - 1); 49 | } 50 | } 51 | 52 | /* compute face_inv */ 53 | scalar_t face_inv[9] = { 54 | p[1][1] - p[2][1], p[2][0] - p[1][0], p[1][0] * p[2][1] - p[2][0] * p[1][1], 55 | p[2][1] - p[0][1], p[0][0] - p[2][0], p[2][0] * p[0][1] - p[0][0] * p[2][1], 56 | p[0][1] - p[1][1], p[1][0] - p[0][0], p[0][0] * p[1][1] - p[1][0] * p[0][1]}; 57 | scalar_t face_inv_denominator = ( 58 | p[2][0] * (p[0][1] - p[1][1]) + 59 | p[0][0] * (p[1][1] - p[2][1]) + 60 | p[1][0] * (p[2][1] - p[0][1])); 61 | for (int k = 0; k < 9; k++) { 62 | face_inv[k] /= face_inv_denominator; 63 | } 64 | /* set to global memory */ 65 | for (int k = 0; k < 9; k++) { 66 | face_inv_g[k] = face_inv[k]; 67 | } 68 | } 69 | 70 | template 71 | __global__ void forward_face_index_map_cuda_kernel_2( 72 | const scalar_t* faces, 73 | scalar_t* faces_inv, 74 | int32_t* __restrict__ face_index_map, 75 | scalar_t* __restrict__ weight_map, 76 | scalar_t* __restrict__ depth_map, 77 | scalar_t* __restrict__ face_inv_map, 78 | int batch_size, 79 | int num_faces, 80 | int image_size, 81 | scalar_t near, 82 | scalar_t far, 83 | int return_rgb, 84 | int return_alpha, 85 | int return_depth) { 86 | 87 | const int i = blockIdx.x * blockDim.x + threadIdx.x; 88 | if (i >= batch_size * image_size * image_size) { 89 | return; 90 | } 91 | const int is = image_size; 92 | const int nf = num_faces; 93 | const int bn = i / (is * is); 94 | const int pn = i % (is * is); 95 | const int yi = pn / is; 96 | const int xi = pn % is; 97 | const scalar_t yp = (2. * yi + 1 - is) / is; 98 | const scalar_t xp = (2. * xi + 1 - is) / is; 99 | 100 | const scalar_t* face = &faces[bn * nf * 9] - 9; 101 | scalar_t* face_inv = &faces_inv[bn * nf * 9] - 9; 102 | scalar_t depth_min = far; 103 | int face_index_min = -1; 104 | scalar_t weight_min[3]; 105 | scalar_t face_inv_min[9]; 106 | for (int fn = 0; fn < nf; fn++) { 107 | /* go to next face */ 108 | face += 9; 109 | face_inv += 9; 110 | 111 | /* return if backside */ 112 | if ((face[7] - face[1]) * (face[3] - face[0]) < (face[4] - face[1]) * (face[6] - face[0])) 113 | continue; 114 | 115 | /* check [py, px] is inside the face */ 116 | if (((yp - face[1]) * (face[3] - face[0]) < (xp - face[0]) * (face[4] - face[1])) || 117 | ((yp - face[4]) * (face[6] - face[3]) < (xp - face[3]) * (face[7] - face[4])) || 118 | ((yp - face[7]) * (face[0] - face[6]) < (xp - face[6]) * (face[1] - face[7]))) 119 | continue; 120 | 121 | /* compute w = face_inv * p */ 122 | scalar_t w[3]; 123 | w[0] = face_inv[3 * 0 + 0] * xi + face_inv[3 * 0 + 1] * yi + face_inv[3 * 0 + 2]; 124 | w[1] = face_inv[3 * 1 + 0] * xi + face_inv[3 * 1 + 1] * yi + face_inv[3 * 1 + 2]; 125 | w[2] = face_inv[3 * 2 + 0] * xi + face_inv[3 * 2 + 1] * yi + face_inv[3 * 2 + 2]; 126 | 127 | /* sum(w) -> 1, 0 < w < 1 */ 128 | scalar_t w_sum = 0; 129 | for (int k = 0; k < 3; k++) { 130 | w[k] = min(max(w[k], 0.), 1.); 131 | w_sum += w[k]; 132 | } 133 | for (int k = 0; k < 3; k++) { 134 | w[k] /= w_sum; 135 | } 136 | /* compute 1 / zp = sum(w / z) */ 137 | const scalar_t zp = 1. / (w[0] / face[2] + w[1] / face[5] + w[2] / face[8]); 138 | if (zp <= near || far <= zp) { 139 | continue; 140 | } 141 | 142 | /* check z-buffer */ 143 | if (zp < depth_min) { 144 | depth_min = zp; 145 | face_index_min = fn; 146 | for (int k = 0; k < 3; k++) { 147 | weight_min[k] = w[k]; 148 | } 149 | if (return_depth) { 150 | for (int k = 0; k < 9; k++) { 151 | face_inv_min[k] = face_inv[k]; 152 | } 153 | } 154 | } 155 | } 156 | 157 | /* set to global memory */ 158 | if (0 <= face_index_min) { 159 | depth_map[i] = depth_min; 160 | face_index_map[i] = face_index_min; 161 | for (int k = 0; k < 3; k++) { 162 | weight_map[3 * i + k] = weight_min[k]; 163 | } 164 | if (return_depth) { 165 | for (int k = 0; k < 9; k++) { 166 | face_inv_map[9 * i + k] = face_inv_min[k]; 167 | } 168 | } 169 | } 170 | } 171 | 172 | template 173 | __global__ void forward_texture_sampling_cuda_kernel( 174 | const scalar_t* faces, 175 | const scalar_t* textures, 176 | const int32_t* face_index_map, 177 | const scalar_t* weight_map, 178 | const scalar_t* depth_map, 179 | scalar_t* rgb_map, 180 | int32_t* sampling_index_map, 181 | scalar_t* sampling_weight_map, 182 | size_t batch_size, 183 | int num_faces, 184 | int image_size, 185 | int texture_size, 186 | scalar_t eps) { 187 | const int i = blockIdx.x * blockDim.x + threadIdx.x; 188 | if (i >= batch_size * image_size * image_size) { 189 | return; 190 | } 191 | const int face_index = face_index_map[i]; 192 | 193 | if (face_index >= 0) { 194 | /* 195 | from global variables: 196 | batch number, num of faces, image_size, face[v012][RGB], pixel[RGB], weight[v012], 197 | texture[ts][ts][ts][RGB], sampling indices[8], sampling_weights[8]; 198 | */ 199 | const int bn = i / (image_size * image_size); 200 | const int nf = num_faces; 201 | const int ts = texture_size; 202 | const scalar_t* face = &faces[(bn * nf + face_index) * 9]; 203 | const scalar_t* texture = &textures[(bn * nf + face_index) * ts * ts * ts * 3]; 204 | scalar_t* pixel = &rgb_map[i * 3]; 205 | const scalar_t* weight = &weight_map[i * 3]; 206 | const scalar_t depth = depth_map[i]; 207 | int32_t* sampling_indices = &sampling_index_map[i * 8]; 208 | scalar_t* sampling_weights = &sampling_weight_map[i * 8]; 209 | 210 | /* get texture index (float) */ 211 | scalar_t texture_index_float[3]; 212 | for (int k = 0; k < 3; k++) { scalar_t tif = weight[k] * (ts - 1) * (depth / (face[3 * k + 2])); 213 | tif = max(tif, 0.); 214 | tif = min(tif, ts - 1 - eps); 215 | texture_index_float[k] = tif; 216 | } 217 | 218 | /* blend */ 219 | scalar_t new_pixel[3] = {0, 0, 0}; 220 | for (int pn = 0; pn < 8; pn++) { 221 | scalar_t w = 1; // weight 222 | int texture_index_int[3]; // index in source (int) 223 | for (int k = 0; k < 3; k++) { 224 | if ((pn >> k) % 2 == 0) { 225 | w *= 1 - (texture_index_float[k] - (int)texture_index_float[k]); 226 | texture_index_int[k] = (int)texture_index_float[k]; 227 | } 228 | else { 229 | w *= texture_index_float[k] - (int)texture_index_float[k]; 230 | texture_index_int[k] = (int)texture_index_float[k] + 1; 231 | } 232 | } 233 | 234 | int isc = texture_index_int[0] * ts * ts + texture_index_int[1] * ts + texture_index_int[2]; 235 | for (int k = 0; k < 3; k++) 236 | new_pixel[k] += w * texture[isc * 3 + k]; 237 | sampling_indices[pn] = isc; 238 | sampling_weights[pn] = w; 239 | } 240 | for (int k = 0; k < 3; k++) 241 | pixel[k] = new_pixel[k]; 242 | } 243 | } 244 | 245 | template 246 | __global__ void backward_pixel_map_cuda_kernel( 247 | const scalar_t* faces, 248 | int32_t* face_index_map, 249 | scalar_t* rgb_map, 250 | scalar_t* alpha_map, 251 | scalar_t* grad_rgb_map, 252 | scalar_t* grad_alpha_map, 253 | scalar_t* grad_faces, 254 | size_t batch_size, 255 | size_t num_faces, 256 | int image_size, 257 | scalar_t eps, 258 | int return_rgb, 259 | int return_alpha) { 260 | const int i = blockIdx.x * blockDim.x + threadIdx.x; 261 | if (i >= batch_size * num_faces) { 262 | return; 263 | } 264 | const int bn = i / num_faces; 265 | const int fn = i % num_faces; 266 | const int is = image_size; 267 | const scalar_t* face = &faces[i * 9]; 268 | scalar_t grad_face[9] = {}; 269 | 270 | /* check backside */ 271 | if ((face[7] - face[1]) * (face[3] - face[0]) < (face[4] - face[1]) * (face[6] - face[0])) 272 | return; 273 | 274 | /* for each edge */ 275 | for (int edge_num = 0; edge_num < 3; edge_num++) { 276 | /* set points of target edge */ 277 | int pi[3]; 278 | scalar_t pp[3][2]; 279 | for (int num = 0; num < 3; num++) 280 | pi[num] = (edge_num + num) % 3; 281 | for (int num = 0; num < 3; num++) { 282 | for (int dim = 0; dim < 2; dim++) { 283 | pp[num][dim] = 0.5 * (face[3 * pi[num] + dim] * is + is - 1); 284 | } 285 | } 286 | 287 | /* for dy, dx */ 288 | for (int axis = 0; axis < 2; axis++) { 289 | /* */ 290 | scalar_t p[3][2]; 291 | for (int num = 0; num < 3; num++) { 292 | for (int dim = 0; dim < 2; dim++) { 293 | p[num][dim] = pp[num][(dim + axis) % 2]; 294 | } 295 | } 296 | 297 | /* set direction */ 298 | int direction; 299 | if (axis == 0) { 300 | if (p[0][0] < p[1][0]) 301 | direction = -1; 302 | else 303 | direction = 1; 304 | } else { 305 | if (p[0][0] < p[1][0]) 306 | direction = 1; 307 | else 308 | direction = -1; 309 | } 310 | 311 | /* along edge */ 312 | int d0_from, d0_to; 313 | d0_from = max(ceil(min(p[0][0], p[1][0])), 0.); 314 | d0_to = min(max(p[0][0], p[1][0]), is - 1.); 315 | for (int d0 = d0_from; d0 <= d0_to; d0++) { 316 | /* get cross point */ 317 | int d1_in, d1_out; 318 | const scalar_t d1_cross = (p[1][1] - p[0][1]) / (p[1][0] - p[0][0]) * (d0 - p[0][0]) + p[0][1]; 319 | if (0 < direction) 320 | d1_in = floor(d1_cross); 321 | else 322 | d1_in = ceil(d1_cross); 323 | d1_out = d1_in + direction; 324 | 325 | /* continue if cross point is not shown */ 326 | if (d1_in < 0 || is <= d1_in) 327 | continue; 328 | if (d1_out < 0 || is <= d1_out) 329 | continue; 330 | 331 | /* get color of in-pixel and out-pixel */ 332 | scalar_t alpha_in; 333 | scalar_t alpha_out; 334 | scalar_t *rgb_in; 335 | scalar_t *rgb_out; 336 | int map_index_in, map_index_out; 337 | if (axis == 0) { 338 | map_index_in = bn * is * is + d1_in * is + d0; 339 | map_index_out = bn * is * is + d1_out * is + d0; 340 | } 341 | else { 342 | map_index_in = bn * is * is + d0 * is + d1_in; 343 | map_index_out = bn * is * is + d0 * is + d1_out; 344 | } 345 | if (return_alpha) { 346 | alpha_in = alpha_map[map_index_in]; 347 | alpha_out = alpha_map[map_index_out]; 348 | } 349 | if (return_rgb) { 350 | rgb_in = &rgb_map[map_index_in * 3]; 351 | rgb_out = &rgb_map[map_index_out * 3]; 352 | } 353 | 354 | /* out */ 355 | bool is_in_fn = (face_index_map[map_index_in] == fn); 356 | if (is_in_fn) { 357 | int d1_limit; 358 | if (0 < direction) 359 | d1_limit = is - 1; 360 | else 361 | d1_limit = 0; 362 | int d1_from = max(min(d1_out, d1_limit), 0); 363 | int d1_to = min(max(d1_out, d1_limit), is - 1); 364 | scalar_t* alpha_map_p; 365 | scalar_t* grad_alpha_map_p; 366 | scalar_t* rgb_map_p; 367 | scalar_t* grad_rgb_map_p; 368 | int map_offset, map_index_from; 369 | if (axis == 0) { 370 | map_offset = is; 371 | map_index_from = bn * is * is + d1_from * is + d0; 372 | } 373 | else { 374 | map_offset = 1; 375 | map_index_from = bn * is * is + d0 * is + d1_from; 376 | } 377 | if (return_alpha) { 378 | alpha_map_p = &alpha_map[map_index_from]; 379 | grad_alpha_map_p = &grad_alpha_map[map_index_from]; 380 | } 381 | if (return_rgb) { 382 | rgb_map_p = &rgb_map[map_index_from * 3]; 383 | grad_rgb_map_p = &grad_rgb_map[map_index_from * 3]; 384 | } 385 | for (int d1 = d1_from; d1 <= d1_to; d1++) { 386 | scalar_t diff_grad = 0; 387 | if (return_alpha) { 388 | diff_grad += (*alpha_map_p - alpha_in) * *grad_alpha_map_p; 389 | } 390 | if (return_rgb) { 391 | for (int k = 0; k < 3; k++) 392 | diff_grad += (rgb_map_p[k] - rgb_in[k]) * grad_rgb_map_p[k]; 393 | } 394 | if (return_alpha) { 395 | alpha_map_p += map_offset; 396 | grad_alpha_map_p += map_offset; 397 | } 398 | if (return_rgb) { 399 | rgb_map_p += 3 * map_offset; 400 | grad_rgb_map_p += 3 * map_offset; 401 | } 402 | if (diff_grad <= 0) 403 | continue; 404 | if (p[1][0] != d0) { 405 | scalar_t dist = (p[1][0] - p[0][0]) / (p[1][0] - d0) * (d1 - d1_cross) * 2. / is; 406 | dist = (0 < dist) ? dist + eps : dist - eps; 407 | grad_face[pi[0] * 3 + (1 - axis)] -= diff_grad / dist; 408 | } 409 | if (p[0][0] != d0) { 410 | scalar_t dist = (p[1][0] - p[0][0]) / (d0 - p[0][0]) * (d1 - d1_cross) * 2. / is; 411 | dist = (0 < dist) ? dist + eps : dist - eps; 412 | grad_face[pi[1] * 3 + (1 - axis)] -= diff_grad / dist; 413 | } 414 | } 415 | } 416 | 417 | /* in */ 418 | { 419 | int d1_limit; 420 | scalar_t d0_cross2; 421 | if ((d0 - p[0][0]) * (d0 - p[2][0]) < 0) { 422 | d0_cross2 = (p[2][1] - p[0][1]) / (p[2][0] - p[0][0]) * (d0 - p[0][0]) + p[0][1]; 423 | } 424 | else { 425 | d0_cross2 = (p[1][1] - p[2][1]) / (p[1][0] - p[2][0]) * (d0 - p[2][0]) + p[2][1]; 426 | } 427 | if (0 < direction) 428 | d1_limit = ceil(d0_cross2); 429 | else 430 | d1_limit = floor(d0_cross2); 431 | int d1_from = max(min(d1_in, d1_limit), 0); 432 | int d1_to = min(max(d1_in, d1_limit), is - 1); 433 | 434 | int* face_index_map_p; 435 | scalar_t* alpha_map_p; 436 | scalar_t* grad_alpha_map_p; 437 | scalar_t* rgb_map_p; 438 | scalar_t* grad_rgb_map_p; 439 | int map_index_from; 440 | int map_offset; 441 | if (axis == 0) 442 | map_offset = is; 443 | else 444 | map_offset = 1; 445 | if (axis == 0) { 446 | map_index_from = bn * is * is + d1_from * is + d0; 447 | } 448 | else { 449 | map_index_from = bn * is * is + d0 * is + d1_from; 450 | } 451 | face_index_map_p = &face_index_map[map_index_from] - map_offset; 452 | if (return_alpha) { 453 | alpha_map_p = &alpha_map[map_index_from] - map_offset; 454 | grad_alpha_map_p = &grad_alpha_map[map_index_from] - map_offset; 455 | } 456 | if (return_rgb) { 457 | rgb_map_p = &rgb_map[map_index_from * 3] - 3 * map_offset; 458 | grad_rgb_map_p = &grad_rgb_map[map_index_from * 3] - 3 * map_offset; 459 | } 460 | 461 | for (int d1 = d1_from; d1 <= d1_to; d1++) { 462 | face_index_map_p += map_offset; 463 | if (return_alpha) { 464 | alpha_map_p += map_offset; 465 | grad_alpha_map_p += map_offset; 466 | } 467 | if (return_rgb) { 468 | rgb_map_p += 3 * map_offset; 469 | grad_rgb_map_p += 3 * map_offset; 470 | } 471 | if (*face_index_map_p != fn) 472 | continue; 473 | 474 | scalar_t diff_grad = 0; 475 | if (return_alpha) { 476 | diff_grad += (*alpha_map_p - alpha_out) * *grad_alpha_map_p; 477 | } 478 | if (return_rgb) { 479 | for (int k = 0; k < 3; k++) 480 | diff_grad += (rgb_map_p[k] - rgb_out[k]) * grad_rgb_map_p[k]; 481 | } 482 | if (diff_grad <= 0) 483 | continue; 484 | 485 | if (p[1][0] != d0) { 486 | scalar_t dist = (p[1][0] - p[0][0]) / (p[1][0] - d0) * (d1 - d1_cross) * 2. / is; 487 | dist = (0 < dist) ? dist + eps : dist - eps; 488 | grad_face[pi[0] * 3 + (1 - axis)] -= diff_grad / dist; 489 | } 490 | if (p[0][0] != d0) { 491 | scalar_t dist = (p[1][0] - p[0][0]) / (d0 - p[0][0]) * (d1 - d1_cross) * 2. / is; 492 | dist = (0 < dist) ? dist + eps : dist - eps; 493 | grad_face[pi[1] * 3 + (1 - axis)] -= diff_grad / dist; 494 | } 495 | } 496 | } 497 | } 498 | } 499 | } 500 | 501 | /* set to global gradient variable */ 502 | for (int k = 0; k < 9; k++) 503 | grad_faces[i * 9 + k] = grad_face[k]; 504 | } 505 | 506 | template 507 | __global__ void backward_textures_cuda_kernel( 508 | const int32_t* face_index_map, 509 | scalar_t* sampling_weight_map, 510 | int32_t* sampling_index_map, 511 | scalar_t* grad_rgb_map, 512 | scalar_t* grad_textures, 513 | size_t batch_size, 514 | size_t num_faces, 515 | int image_size, 516 | size_t texture_size) { 517 | 518 | const int i = blockIdx.x * blockDim.x + threadIdx.x; 519 | if (i >= batch_size * image_size * image_size) { 520 | return; 521 | } 522 | const int face_index = face_index_map[i]; 523 | if (0 <= face_index) { 524 | int is = image_size; 525 | int nf = num_faces; 526 | int ts = texture_size; 527 | int bn = i / (is * is); // batch number [0 -> bs] 528 | 529 | scalar_t* grad_texture = &grad_textures[(bn * nf + face_index) * ts * ts * ts * 3]; 530 | scalar_t* sampling_weight_map_p = &sampling_weight_map[i * 8]; 531 | int* sampling_index_map_p = &sampling_index_map[i * 8]; 532 | for (int pn = 0; pn < 8; pn++) { 533 | scalar_t w = *sampling_weight_map_p++; 534 | int isc = *sampling_index_map_p++; 535 | scalar_t* grad_texture_p = &grad_texture[isc * 3]; 536 | scalar_t* grad_rgb_map_p = &grad_rgb_map[i * 3]; 537 | for (int k = 0; k < 3; k++) 538 | atomicAdd(grad_texture_p++, w * *grad_rgb_map_p++); 539 | } 540 | } 541 | } 542 | 543 | template 544 | __global__ void backward_depth_map_cuda_kernel( 545 | const scalar_t* faces, 546 | const scalar_t* depth_map, 547 | const int32_t* face_index_map, 548 | const scalar_t* face_inv_map, 549 | const scalar_t* weight_map, 550 | scalar_t* grad_depth_map, 551 | scalar_t* grad_faces, 552 | size_t batch_size, 553 | size_t num_faces, 554 | int image_size) { 555 | 556 | const int i = blockIdx.x * blockDim.x + threadIdx.x; 557 | if (i >= batch_size * image_size * image_size) { 558 | return; 559 | } 560 | const int fn = face_index_map[i]; 561 | if (0 <= fn) { 562 | const int nf = num_faces; 563 | const int is = image_size; 564 | const int bn = i / (is * is); 565 | const scalar_t* face = &faces[(bn * nf + fn) * 9]; 566 | const scalar_t depth = depth_map[i]; 567 | const scalar_t depth2 = depth * depth; 568 | const scalar_t* face_inv = &face_inv_map[i * 9]; 569 | const scalar_t* weight = &weight_map[i * 3]; 570 | const scalar_t grad_depth = grad_depth_map[i]; 571 | scalar_t* grad_face = &grad_faces[(bn * nf + fn) * 9]; 572 | 573 | /* derivative wrt z */ 574 | for (int k = 0; k < 3; k++) { 575 | const scalar_t z_k = face[3 * k + 2]; 576 | atomicAdd(&grad_face[3 * k + 2], grad_depth * weight[k] * depth2 / (z_k * z_k)); 577 | } 578 | 579 | /* derivative wrt x, y */ 580 | scalar_t tmp[3] = {}; 581 | for (int k = 0; k < 3; k++) { 582 | for (int l = 0; l < 3; l++) { 583 | tmp[k] += -face_inv[3 * l + k] / face[3 * l + 2]; 584 | } 585 | } 586 | for (int k = 0; k < 3; k++) { 587 | for (int l = 0; l < 2; l++) { 588 | // k: point number, l: dimension 589 | atomicAdd(&grad_face[3 * k + l], -grad_depth * tmp[l] * weight[k] * depth2 * is / 2); 590 | } 591 | } 592 | } 593 | } 594 | } 595 | 596 | std::vector forward_face_index_map_cuda( 597 | at::Tensor faces, 598 | at::Tensor face_index_map, 599 | at::Tensor weight_map, 600 | at::Tensor depth_map, 601 | at::Tensor face_inv_map, 602 | at::Tensor faces_inv, 603 | int image_size, 604 | float near, 605 | float far, 606 | int return_rgb, 607 | int return_alpha, 608 | int return_depth) { 609 | 610 | const auto batch_size = faces.size(0); 611 | const auto num_faces = faces.size(1); 612 | const int threads = 512; 613 | const dim3 blocks_1 ((batch_size * num_faces - 1) / threads +1); 614 | 615 | AT_DISPATCH_FLOATING_TYPES(faces.type(), "forward_face_index_map_cuda_1", ([&] { 616 | forward_face_index_map_cuda_kernel_1<<>>( 617 | faces.data(), 618 | faces_inv.data(), 619 | batch_size, 620 | num_faces, 621 | image_size); 622 | })); 623 | 624 | cudaError_t err = cudaGetLastError(); 625 | if (err != cudaSuccess) 626 | printf("Error in forward_face_index_map_1: %s\n", cudaGetErrorString(err)); 627 | 628 | const dim3 blocks_2 ((batch_size * image_size * image_size - 1) / threads +1); 629 | AT_DISPATCH_FLOATING_TYPES(faces.type(), "forward_face_index_map_cuda_2", ([&] { 630 | forward_face_index_map_cuda_kernel_2<<>>( 631 | faces.data(), 632 | faces_inv.data(), 633 | face_index_map.data(), 634 | weight_map.data(), 635 | depth_map.data(), 636 | face_inv_map.data(), 637 | (int) batch_size, 638 | (int) num_faces, 639 | (int) image_size, 640 | (scalar_t) near, 641 | (scalar_t) far, 642 | return_rgb, 643 | return_alpha, 644 | return_depth); 645 | })); 646 | 647 | err = cudaGetLastError(); 648 | if (err != cudaSuccess) 649 | printf("Error in forward_face_index_map_2: %s\n", cudaGetErrorString(err)); 650 | return {face_index_map, weight_map, depth_map, face_inv_map}; 651 | } 652 | 653 | std::vector forward_texture_sampling_cuda( at::Tensor faces, 654 | at::Tensor textures, 655 | at::Tensor face_index_map, 656 | at::Tensor weight_map, 657 | at::Tensor depth_map, 658 | at::Tensor rgb_map, 659 | at::Tensor sampling_index_map, 660 | at::Tensor sampling_weight_map, 661 | int image_size, 662 | float eps) { 663 | 664 | const auto batch_size = faces.size(0); 665 | const auto num_faces = faces.size(1); 666 | const auto texture_size = textures.size(2); 667 | const int threads = 512; 668 | const dim3 blocks ((batch_size * image_size * image_size - 1) / threads + 1); 669 | 670 | AT_DISPATCH_FLOATING_TYPES(faces.type(), "forward_texture_sampling_cuda", ([&] { 671 | forward_texture_sampling_cuda_kernel<<>>( 672 | faces.data(), 673 | textures.data(), 674 | face_index_map.data(), 675 | weight_map.data(), 676 | depth_map.data(), 677 | rgb_map.data(), 678 | sampling_index_map.data(), 679 | sampling_weight_map.data(), 680 | batch_size, 681 | num_faces, 682 | image_size, 683 | texture_size, 684 | eps); 685 | })); 686 | 687 | cudaError_t err = cudaGetLastError(); 688 | if (err != cudaSuccess) 689 | printf("Error in forward_texture_sampling: %s\n", cudaGetErrorString(err)); 690 | 691 | return {rgb_map, sampling_index_map, sampling_weight_map}; 692 | } 693 | 694 | at::Tensor backward_pixel_map_cuda( 695 | at::Tensor faces, 696 | at::Tensor face_index_map, 697 | at::Tensor rgb_map, 698 | at::Tensor alpha_map, 699 | at::Tensor grad_rgb_map, 700 | at::Tensor grad_alpha_map, 701 | at::Tensor grad_faces, 702 | int image_size, 703 | float eps, 704 | int return_rgb, 705 | int return_alpha) { 706 | 707 | const auto batch_size = faces.size(0); 708 | const auto num_faces = faces.size(1); 709 | const int threads = 512; 710 | const dim3 blocks ((batch_size * num_faces - 1) / threads + 1); 711 | 712 | AT_DISPATCH_FLOATING_TYPES(faces.type(), "backward_pixel_map_cuda", ([&] { 713 | backward_pixel_map_cuda_kernel<<>>( 714 | faces.data(), 715 | face_index_map.data(), 716 | rgb_map.data(), 717 | alpha_map.data(), 718 | grad_rgb_map.data(), 719 | grad_alpha_map.data(), 720 | grad_faces.data(), 721 | batch_size, 722 | num_faces, 723 | image_size, 724 | (scalar_t) eps, 725 | return_rgb, 726 | return_alpha); 727 | })); 728 | 729 | cudaError_t err = cudaGetLastError(); 730 | if (err != cudaSuccess) 731 | printf("Error in backward_pixel_map: %s\n", cudaGetErrorString(err)); 732 | 733 | return grad_faces; 734 | } 735 | 736 | at::Tensor backward_textures_cuda( 737 | at::Tensor face_index_map, 738 | at::Tensor sampling_weight_map, 739 | at::Tensor sampling_index_map, 740 | at::Tensor grad_rgb_map, 741 | at::Tensor grad_textures, 742 | int num_faces) { 743 | 744 | const auto batch_size = face_index_map.size(0); 745 | const auto image_size = face_index_map.size(1); 746 | const auto texture_size = grad_textures.size(2); 747 | const int threads = 512; 748 | const dim3 blocks ((batch_size * image_size * image_size - 1) / threads + 1); 749 | 750 | AT_DISPATCH_FLOATING_TYPES(sampling_weight_map.type(), "backward_textures_cuda", ([&] { 751 | backward_textures_cuda_kernel<<>>( 752 | face_index_map.data(), 753 | sampling_weight_map.data(), 754 | sampling_index_map.data(), 755 | grad_rgb_map.data(), 756 | grad_textures.data(), 757 | batch_size, 758 | num_faces, 759 | image_size, 760 | texture_size); 761 | })); 762 | 763 | cudaError_t err = cudaGetLastError(); 764 | if (err != cudaSuccess) 765 | printf("Error in backward_textures: %s\n", cudaGetErrorString(err)); 766 | 767 | return grad_textures; 768 | } 769 | at::Tensor backward_depth_map_cuda( 770 | at::Tensor faces, 771 | at::Tensor depth_map, 772 | at::Tensor face_index_map, 773 | at::Tensor face_inv_map, 774 | at::Tensor weight_map, 775 | at::Tensor grad_depth_map, 776 | at::Tensor grad_faces, 777 | int image_size) { 778 | 779 | const auto batch_size = faces.size(0); 780 | const auto num_faces = faces.size(1); 781 | const int threads = 512; 782 | const dim3 blocks ((batch_size * image_size * image_size - 1) / threads + 1); 783 | 784 | AT_DISPATCH_FLOATING_TYPES(faces.type(), "backward_depth_map_cuda", ([&] { 785 | backward_depth_map_cuda_kernel<<>>( 786 | faces.data(), 787 | depth_map.data(), 788 | face_index_map.data(), 789 | face_inv_map.data(), 790 | weight_map.data(), 791 | grad_depth_map.data(), 792 | grad_faces.data(), 793 | batch_size, 794 | num_faces, 795 | image_size); 796 | })); 797 | 798 | cudaError_t err = cudaGetLastError(); 799 | if (err != cudaSuccess) 800 | printf("Error in backward_depth_map: %s\n", cudaGetErrorString(err)); 801 | 802 | return grad_faces; 803 | } 804 | -------------------------------------------------------------------------------- /neural_renderer/get_points_from_angles.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | import math 3 | 4 | import torch 5 | 6 | def get_points_from_angles(distance, elevation, azimuth, degrees=True): 7 | if isinstance(distance, float) or isinstance(distance, int): 8 | if degrees: 9 | elevation = math.radians(elevation) 10 | azimuth = math.radians(azimuth) 11 | return ( 12 | distance * math.cos(elevation) * math.sin(azimuth), 13 | distance * math.sin(elevation), 14 | -distance * math.cos(elevation) * math.cos(azimuth)) 15 | else: 16 | if degrees: 17 | elevation = math.pi/180. * elevation 18 | azimuth = math.pi/180. * azimuth 19 | # 20 | return torch.stack([ 21 | distance * torch.cos(elevation) * torch.sin(azimuth), 22 | distance * torch.sin(elevation), 23 | -distance * torch.cos(elevation) * torch.cos(azimuth) 24 | ]).transpose(1,0) 25 | -------------------------------------------------------------------------------- /neural_renderer/lighting.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | import numpy as np 4 | 5 | def lighting(faces, textures, intensity_ambient=0.5, intensity_directional=0.5, 6 | color_ambient=(1, 1, 1), color_directional=(1, 1, 1), direction=(0, 1, 0)): 7 | 8 | bs, nf = faces.shape[:2] 9 | device = faces.device 10 | 11 | # arguments 12 | # make sure to convert all inputs to float tensors 13 | if isinstance(color_ambient, tuple) or isinstance(color_ambient, list): 14 | color_ambient = torch.tensor(color_ambient, dtype=torch.float32, device=device) 15 | elif isinstance(color_ambient, np.ndarray): 16 | color_ambient = torch.from_numpy(color_ambient).float().to(device) 17 | if isinstance(color_directional, tuple) or isinstance(color_directional, list): 18 | color_directional = torch.tensor(color_directional, dtype=torch.float32, device=device) 19 | elif isinstance(color_directional, np.ndarray): 20 | color_directional = torch.from_numpy(color_directional).float().to(device) 21 | if isinstance(direction, tuple) or isinstance(direction, list): 22 | direction = torch.tensor(direction, dtype=torch.float32, device=device) 23 | elif isinstance(direction, np.ndarray): 24 | direction = torch.from_numpy(direction).float().to(device) 25 | if color_ambient.ndimension() == 1: 26 | color_ambient = color_ambient[None, :] 27 | if color_directional.ndimension() == 1: 28 | color_directional = color_directional[None, :] 29 | if direction.ndimension() == 1: 30 | direction = direction[None, :] 31 | 32 | # create light 33 | light = torch.zeros(bs, nf, 3, dtype=torch.float32).to(device) 34 | 35 | # ambient light 36 | if intensity_ambient != 0: 37 | light += intensity_ambient * color_ambient[:, None, :] 38 | 39 | # directional light 40 | if intensity_directional != 0: 41 | faces = faces.reshape((bs * nf, 3, 3)) 42 | v10 = faces[:, 0] - faces[:, 1] 43 | v12 = faces[:, 2] - faces[:, 1] 44 | # pytorch normalize divides by max(norm, eps) instead of (norm+eps) in chainer 45 | normals = F.normalize(torch.cross(v10, v12), eps=1e-5) 46 | normals = normals.reshape((bs, nf, 3)) 47 | 48 | if direction.ndimension() == 2: 49 | direction = direction[:, None, :] 50 | cos = F.relu(torch.sum(normals * direction, dim=2)) 51 | # may have to verify that the next line is correct 52 | light += intensity_directional * (color_directional[:, None, :] * cos[:, :, None]) 53 | 54 | # apply 55 | light = light[:,:,None, None, None, :] 56 | textures *= light 57 | return textures 58 | -------------------------------------------------------------------------------- /neural_renderer/load_obj.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | import os 3 | 4 | import torch 5 | import numpy as np 6 | from skimage.io import imread 7 | 8 | import neural_renderer.cuda.load_textures as load_textures_cuda 9 | 10 | texture_wrapping_dict = {'REPEAT': 0, 'MIRRORED_REPEAT': 1, 11 | 'CLAMP_TO_EDGE': 2, 'CLAMP_TO_BORDER': 3} 12 | 13 | def load_mtl(filename_mtl): 14 | ''' 15 | load color (Kd) and filename of textures from *.mtl 16 | ''' 17 | texture_filenames = {} 18 | colors = {} 19 | material_name = '' 20 | with open(filename_mtl) as f: 21 | for line in f.readlines(): 22 | if len(line.split()) != 0: 23 | if line.split()[0] == 'newmtl': 24 | material_name = line.split()[1] 25 | if line.split()[0] == 'map_Kd': 26 | texture_filenames[material_name] = line.split()[1] 27 | if line.split()[0] == 'Kd': 28 | colors[material_name] = np.array(list(map(float, line.split()[1:4]))) 29 | return colors, texture_filenames 30 | 31 | 32 | def load_textures(filename_obj, filename_mtl, texture_size, texture_wrapping='REPEAT', use_bilinear=True): 33 | # load vertices 34 | vertices = [] 35 | with open(filename_obj) as f: 36 | lines = f.readlines() 37 | for line in lines: 38 | if len(line.split()) == 0: 39 | continue 40 | if line.split()[0] == 'vt': 41 | vertices.append([float(v) for v in line.split()[1:3]]) 42 | vertices = np.vstack(vertices).astype(np.float32) 43 | 44 | # load faces for textures 45 | faces = [] 46 | material_names = [] 47 | material_name = '' 48 | for line in lines: 49 | if len(line.split()) == 0: 50 | continue 51 | if line.split()[0] == 'f': 52 | vs = line.split()[1:] 53 | nv = len(vs) 54 | if '/' in vs[0] and '//' not in vs[0]: 55 | v0 = int(vs[0].split('/')[1]) 56 | else: 57 | v0 = 0 58 | for i in range(nv - 2): 59 | if '/' in vs[i + 1] and '//' not in vs[i + 1]: 60 | v1 = int(vs[i + 1].split('/')[1]) 61 | else: 62 | v1 = 0 63 | if '/' in vs[i + 2] and '//' not in vs[i + 2]: 64 | v2 = int(vs[i + 2].split('/')[1]) 65 | else: 66 | v2 = 0 67 | faces.append((v0, v1, v2)) 68 | material_names.append(material_name) 69 | if line.split()[0] == 'usemtl': 70 | material_name = line.split()[1] 71 | faces = np.vstack(faces).astype(np.int32) - 1 72 | faces = vertices[faces] 73 | faces = torch.from_numpy(faces).cuda() 74 | 75 | colors, texture_filenames = load_mtl(filename_mtl) 76 | 77 | textures = torch.zeros(faces.shape[0], texture_size, texture_size, texture_size, 3, dtype=torch.float32) + 0.5 78 | textures = textures.cuda() 79 | 80 | # 81 | for material_name, color in colors.items(): 82 | color = torch.from_numpy(color).cuda() 83 | for i, material_name_f in enumerate(material_names): 84 | if material_name == material_name_f: 85 | textures[i, :, :, :, :] = color[None, None, None, :] 86 | 87 | for material_name, filename_texture in texture_filenames.items(): 88 | filename_texture = os.path.join(os.path.dirname(filename_obj), filename_texture) 89 | image = imread(filename_texture).astype(np.float32) / 255. 90 | 91 | # texture image may have one channel (grey color) 92 | if len(image.shape) == 2: 93 | image = np.stack((image,)*3,-1) 94 | # or has extral alpha channel shoule ignore for now 95 | if image.shape[2] == 4: 96 | image = image[:,:,:3] 97 | 98 | # pytorch does not support negative slicing for the moment 99 | image = image[::-1, :, :] 100 | image = torch.from_numpy(image.copy()).cuda() 101 | is_update = (np.array(material_names) == material_name).astype(np.int32) 102 | is_update = torch.from_numpy(is_update).cuda() 103 | textures = load_textures_cuda.load_textures(image, faces, textures, is_update, 104 | texture_wrapping_dict[texture_wrapping], 105 | use_bilinear) 106 | return textures 107 | 108 | def load_obj(filename_obj, normalization=True, texture_size=4, load_texture=False, 109 | texture_wrapping='REPEAT', use_bilinear=True): 110 | """ 111 | Load Wavefront .obj file. 112 | This function only supports vertices (v x x x) and faces (f x x x). 113 | """ 114 | 115 | # load vertices 116 | vertices = [] 117 | with open(filename_obj) as f: 118 | lines = f.readlines() 119 | 120 | for line in lines: 121 | if len(line.split()) == 0: 122 | continue 123 | if line.split()[0] == 'v': 124 | vertices.append([float(v) for v in line.split()[1:4]]) 125 | vertices = torch.from_numpy(np.vstack(vertices).astype(np.float32)).cuda() 126 | 127 | # load faces 128 | faces = [] 129 | for line in lines: 130 | if len(line.split()) == 0: 131 | continue 132 | if line.split()[0] == 'f': 133 | vs = line.split()[1:] 134 | nv = len(vs) 135 | v0 = int(vs[0].split('/')[0]) 136 | for i in range(nv - 2): 137 | v1 = int(vs[i + 1].split('/')[0]) 138 | v2 = int(vs[i + 2].split('/')[0]) 139 | faces.append((v0, v1, v2)) 140 | faces = torch.from_numpy(np.vstack(faces).astype(np.int32)).cuda() - 1 141 | 142 | # load textures 143 | textures = None 144 | if load_texture: 145 | for line in lines: 146 | if line.startswith('mtllib'): 147 | filename_mtl = os.path.join(os.path.dirname(filename_obj), line.split()[1]) 148 | textures = load_textures(filename_obj, filename_mtl, texture_size, 149 | texture_wrapping=texture_wrapping, 150 | use_bilinear=use_bilinear) 151 | if textures is None: 152 | raise Exception('Failed to load textures.') 153 | 154 | # normalize into a unit cube centered zero 155 | if normalization: 156 | vertices -= vertices.min(0)[0][None, :] 157 | vertices /= torch.abs(vertices).max() 158 | vertices *= 2 159 | vertices -= vertices.max(0)[0][None, :] / 2 160 | 161 | if load_texture: 162 | return vertices, faces, textures 163 | else: 164 | return vertices, faces 165 | -------------------------------------------------------------------------------- /neural_renderer/look.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn.functional as F 4 | 5 | 6 | def look(vertices, eye, direction=[0, 1, 0], up=None): 7 | """ 8 | "Look" transformation of vertices. 9 | """ 10 | if (vertices.ndimension() != 3): 11 | raise ValueError('vertices Tensor should have 3 dimensions') 12 | 13 | device = vertices.device 14 | 15 | if isinstance(direction, list) or isinstance(direction, tuple): 16 | direction = torch.tensor(direction, dtype=torch.float32, device=device) 17 | elif isinstance(direction, np.ndarray): 18 | direction = torch.from_numpy(direction).to(device) 19 | elif torch.is_tensor(direction): 20 | direction = direction.to(device) 21 | 22 | if isinstance(eye, list) or isinstance(eye, tuple): 23 | eye = torch.tensor(eye, dtype=torch.float32, device=device) 24 | elif isinstance(eye, np.ndarray): 25 | eye = torch.from_numpy(eye).to(device) 26 | elif torch.is_tensor(eye): 27 | eye = eye.to(device) 28 | 29 | if up is None: 30 | up = torch.cuda.FloatTensor([0, 1, 0]) 31 | if eye.ndimension() == 1: 32 | eye = eye[None, :] 33 | if direction.ndimension() == 1: 34 | direction = direction[None, :] 35 | if up.ndimension() == 1: 36 | up = up[None, :] 37 | 38 | # create new axes 39 | z_axis = F.normalize(direction, eps=1e-5) 40 | x_axis = F.normalize(torch.cross(up, z_axis), eps=1e-5) 41 | y_axis = F.normalize(torch.cross(z_axis, x_axis), eps=1e-5) 42 | 43 | # create rotation matrix: [bs, 3, 3] 44 | r = torch.cat((x_axis[:, None, :], y_axis[:, None, :], z_axis[:, None, :]), dim=1) 45 | 46 | # apply 47 | # [bs, nv, 3] -> [bs, nv, 3] -> [bs, nv, 3] 48 | if vertices.shape != eye.shape: 49 | eye = eye[:, None, :] 50 | vertices = vertices - eye 51 | vertices = torch.matmul(vertices, r.transpose(1,2)) 52 | 53 | return vertices 54 | -------------------------------------------------------------------------------- /neural_renderer/look_at.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | import torch.nn.functional as F 4 | 5 | 6 | def look_at(vertices, eye, at=[0, 0, 0], up=[0, 1, 0]): 7 | """ 8 | "Look at" transformation of vertices. 9 | """ 10 | if (vertices.ndimension() != 3): 11 | raise ValueError('vertices Tensor should have 3 dimensions') 12 | 13 | device = vertices.device 14 | 15 | # if list or tuple convert to numpy array 16 | if isinstance(at, list) or isinstance(at, tuple): 17 | at = torch.tensor(at, dtype=torch.float32, device=device) 18 | # if numpy array convert to tensor 19 | elif isinstance(at, np.ndarray): 20 | at = torch.from_numpy(at).to(device) 21 | elif torch.is_tensor(at): 22 | at.to(device) 23 | 24 | if isinstance(up, list) or isinstance(up, tuple): 25 | up = torch.tensor(up, dtype=torch.float32, device=device) 26 | elif isinstance(up, np.ndarray): 27 | up = torch.from_numpy(up).to(device) 28 | elif torch.is_tensor(up): 29 | up.to(device) 30 | 31 | if isinstance(eye, list) or isinstance(eye, tuple): 32 | eye = torch.tensor(eye, dtype=torch.float32, device=device) 33 | elif isinstance(eye, np.ndarray): 34 | eye = torch.from_numpy(eye).to(device) 35 | elif torch.is_tensor(eye): 36 | eye = eye.to(device) 37 | 38 | batch_size = vertices.shape[0] 39 | if eye.ndimension() == 1: 40 | eye = eye[None, :].repeat(batch_size, 1) 41 | if at.ndimension() == 1: 42 | at = at[None, :].repeat(batch_size, 1) 43 | if up.ndimension() == 1: 44 | up = up[None, :].repeat(batch_size, 1) 45 | 46 | # create new axes 47 | # eps is chosen as 0.5 to match the chainer version 48 | z_axis = F.normalize(at - eye, eps=1e-5) 49 | x_axis = F.normalize(torch.cross(up, z_axis), eps=1e-5) 50 | y_axis = F.normalize(torch.cross(z_axis, x_axis), eps=1e-5) 51 | 52 | # create rotation matrix: [bs, 3, 3] 53 | r = torch.cat((x_axis[:, None, :], y_axis[:, None, :], z_axis[:, None, :]), dim=1) 54 | 55 | # apply 56 | # [bs, nv, 3] -> [bs, nv, 3] -> [bs, nv, 3] 57 | if vertices.shape != eye.shape: 58 | eye = eye[:, None, :] 59 | vertices = vertices - eye 60 | vertices = torch.matmul(vertices, r.transpose(1,2)) 61 | 62 | return vertices 63 | -------------------------------------------------------------------------------- /neural_renderer/mesh.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | import neural_renderer as nr 5 | 6 | class Mesh(object): 7 | ''' 8 | A simple class for creating and manipulating trimesh objects 9 | ''' 10 | def __init__(self, vertices, faces, textures=None, texture_size=4): 11 | ''' 12 | vertices, faces and textures(if not None) are expected to be Tensor objects 13 | ''' 14 | self.vertices = vertices 15 | self.faces = faces 16 | self.num_vertices = self.vertices.shape[0] 17 | self.num_faces = self.faces.shape[0] 18 | 19 | # create textures 20 | if textures is None: 21 | shape = (self.num_faces, texture_size, texture_size, texture_size, 3) 22 | self.textures = nn.Parameter(0.05*torch.randn(*shape)) 23 | self.texture_size = texture_size 24 | else: 25 | self.texture_size = textures.shape[0] 26 | 27 | @classmethod 28 | def fromobj(cls, filename_obj, normalization=True, load_texture=False, texture_size=4): 29 | ''' 30 | Create a Mesh object from a .obj file 31 | ''' 32 | if load_texture: 33 | vertices, faces, textures = nr.load_obj(filename_obj, 34 | normalization=normalization, 35 | texture_size=texture_size, 36 | load_texture=True) 37 | else: 38 | vertices, faces = nr.load_obj(filename_obj, 39 | normalization=normalization, 40 | texture_size=texture_size, 41 | load_texture=False) 42 | textures = None 43 | return cls(vertices, faces, textures, texture_size) 44 | -------------------------------------------------------------------------------- /neural_renderer/perspective.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | import math 3 | 4 | import torch 5 | 6 | def perspective(vertices, angle=30.): 7 | ''' 8 | Compute perspective distortion from a given angle 9 | ''' 10 | if (vertices.ndimension() != 3): 11 | raise ValueError('vertices Tensor should have 3 dimensions') 12 | device = vertices.device 13 | angle = torch.tensor(angle / 180 * math.pi, dtype=torch.float32, device=device) 14 | angle = angle[None] 15 | width = torch.tan(angle) 16 | width = width[:, None] 17 | z = vertices[:, :, 2] 18 | x = vertices[:, :, 0] / z / width 19 | y = vertices[:, :, 1] / z / width 20 | vertices = torch.stack((x,y,z), dim=2) 21 | return vertices 22 | -------------------------------------------------------------------------------- /neural_renderer/projection.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | 3 | import torch 4 | 5 | 6 | def projection(vertices, K, R, t, dist_coeffs, orig_size, eps=1e-9): 7 | ''' 8 | Calculate projective transformation of vertices given a projection matrix 9 | Input parameters: 10 | K: batch_size * 3 * 3 intrinsic camera matrix 11 | R, t: batch_size * 3 * 3, batch_size * 1 * 3 extrinsic calibration parameters 12 | dist_coeffs: vector of distortion coefficients 13 | orig_size: original size of image captured by the camera 14 | Returns: For each point [X,Y,Z] in world coordinates [u,v,z] where u,v are the coordinates of the projection in 15 | pixels and z is the depth 16 | ''' 17 | 18 | # instead of P*x we compute x'*P' 19 | vertices = torch.matmul(vertices, R.transpose(2,1)) + t 20 | x, y, z = vertices[:, :, 0], vertices[:, :, 1], vertices[:, :, 2] 21 | x_ = x / (z + eps) 22 | y_ = y / (z + eps) 23 | 24 | # Get distortion coefficients from vector 25 | k1 = dist_coeffs[:, None, 0] 26 | k2 = dist_coeffs[:, None, 1] 27 | p1 = dist_coeffs[:, None, 2] 28 | p2 = dist_coeffs[:, None, 3] 29 | k3 = dist_coeffs[:, None, 4] 30 | 31 | # we use x_ for x' and x__ for x'' etc. 32 | r = torch.sqrt(x_ ** 2 + y_ ** 2) 33 | x__ = x_*(1 + k1*(r**2) + k2*(r**4) + k3*(r**6)) + 2*p1*x_*y_ + p2*(r**2 + 2*x_**2) 34 | y__ = y_*(1 + k1*(r**2) + k2*(r**4) + k3 *(r**6)) + p1*(r**2 + 2*y_**2) + 2*p2*x_*y_ 35 | vertices = torch.stack([x__, y__, torch.ones_like(z)], dim=-1) 36 | vertices = torch.matmul(vertices, K.transpose(1,2)) 37 | u, v = vertices[:, :, 0], vertices[:, :, 1] 38 | v = orig_size - v 39 | # map u,v from [0, img_size] to [-1, 1] to use by the renderer 40 | u = 2 * (u - orig_size / 2.) / orig_size 41 | v = 2 * (v - orig_size / 2.) / orig_size 42 | vertices = torch.stack([u, v, z], dim=-1) 43 | return vertices 44 | -------------------------------------------------------------------------------- /neural_renderer/rasterize.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.autograd import Function 5 | 6 | import neural_renderer.cuda.rasterize as rasterize_cuda 7 | 8 | DEFAULT_IMAGE_SIZE = 256 9 | DEFAULT_ANTI_ALIASING = True 10 | DEFAULT_NEAR = 0.1 11 | DEFAULT_FAR = 100 12 | DEFAULT_EPS = 1e-4 13 | DEFAULT_BACKGROUND_COLOR = (0, 0, 0) 14 | 15 | class RasterizeFunction(Function): 16 | ''' 17 | Definition of differentiable rasterize operation 18 | Some parts of the code are implemented in CUDA 19 | Currently implemented only for cuda Tensors 20 | ''' 21 | @staticmethod 22 | def forward(ctx, faces, textures, image_size, near, far, eps, background_color, 23 | return_rgb=False, return_alpha=False, return_depth=False): 24 | ''' 25 | Forward pass 26 | ''' 27 | ctx.image_size = image_size 28 | ctx.near = near 29 | ctx.far = far 30 | ctx.eps = eps 31 | ctx.background_color = background_color 32 | ctx.return_rgb = return_rgb 33 | ctx.return_alpha = return_alpha 34 | ctx.return_depth = return_depth 35 | 36 | faces = faces.clone() 37 | 38 | ctx.device = faces.device 39 | ctx.batch_size, ctx.num_faces = faces.shape[:2] 40 | 41 | if ctx.return_rgb: 42 | textures = textures.contiguous() 43 | ctx.texture_size = textures.shape[2] 44 | else: 45 | # initializing with dummy values 46 | textures = torch.cuda.FloatTensor(1).fill_(0) 47 | ctx.texture_size = None 48 | 49 | 50 | face_index_map = torch.cuda.IntTensor(ctx.batch_size, ctx.image_size, ctx.image_size).fill_(-1) 51 | weight_map = torch.cuda.FloatTensor(ctx.batch_size, ctx.image_size, ctx.image_size, 3).fill_(0.0) 52 | depth_map = torch.cuda.FloatTensor(ctx.batch_size, ctx.image_size, ctx.image_size).fill_(ctx.far) 53 | 54 | if ctx.return_rgb: 55 | rgb_map = torch.cuda.FloatTensor(ctx.batch_size, ctx.image_size, ctx.image_size, 3).fill_(0) 56 | sampling_index_map = torch.cuda.IntTensor(ctx.batch_size, ctx.image_size, ctx.image_size, 8).fill_(0) 57 | sampling_weight_map = torch.cuda.FloatTensor(ctx.batch_size, ctx.image_size, ctx.image_size, 8).fill_(0) 58 | else: 59 | rgb_map = torch.cuda.FloatTensor(1).fill_(0) 60 | sampling_index_map = torch.cuda.FloatTensor(1).fill_(0) 61 | sampling_weight_map = torch.cuda.FloatTensor(1).fill_(0) 62 | if ctx.return_alpha: 63 | alpha_map = torch.cuda.FloatTensor(ctx.batch_size, ctx.image_size, ctx.image_size).fill_(0) 64 | else: 65 | alpha_map = torch.cuda.FloatTensor(1).fill_(0) 66 | if ctx.return_depth: 67 | face_inv_map = torch.cuda.FloatTensor(ctx.batch_size, ctx.image_size, ctx.image_size, 3, 3).fill_(0) 68 | else: 69 | face_inv_map = torch.cuda.FloatTensor(1).fill_(0) 70 | 71 | face_index_map, weight_map, depth_map, face_inv_map =\ 72 | RasterizeFunction.forward_face_index_map(ctx, faces, face_index_map, 73 | weight_map, depth_map, 74 | face_inv_map) 75 | 76 | rgb_map, sampling_index_map, sampling_weight_map =\ 77 | RasterizeFunction.forward_texture_sampling(ctx, faces, textures, 78 | face_index_map, weight_map, 79 | depth_map, rgb_map, 80 | sampling_index_map, 81 | sampling_weight_map) 82 | 83 | rgb_map = RasterizeFunction.forward_background(ctx, face_index_map, rgb_map) 84 | alpha_map = RasterizeFunction.forward_alpha_map(ctx, alpha_map, face_index_map) 85 | 86 | ctx.save_for_backward(faces, textures, face_index_map, weight_map, 87 | depth_map, rgb_map, alpha_map, face_inv_map, 88 | sampling_index_map, sampling_weight_map) 89 | 90 | 91 | rgb_r, alpha_r, depth_r = torch.tensor([]), torch.tensor([]), torch.tensor([]) 92 | if ctx.return_rgb: 93 | rgb_r = rgb_map 94 | if ctx.return_alpha: 95 | alpha_r = alpha_map.clone() 96 | if ctx.return_depth: 97 | depth_r = depth_map.clone() 98 | return rgb_r, alpha_r, depth_r 99 | 100 | @staticmethod 101 | def backward(ctx, grad_rgb_map, grad_alpha_map, grad_depth_map): 102 | ''' 103 | Backward pass 104 | ''' 105 | faces, textures, face_index_map, weight_map,\ 106 | depth_map, rgb_map, alpha_map, face_inv_map,\ 107 | sampling_index_map, sampling_weight_map = \ 108 | ctx.saved_tensors 109 | # initialize output buffers 110 | # no need for explicit allocation of cuda.FloatTensor because zeros_like does it automatically 111 | grad_faces = torch.zeros_like(faces, dtype=torch.float32) 112 | if ctx.return_rgb: 113 | grad_textures = torch.zeros_like(textures, dtype=torch.float32) 114 | else: 115 | grad_textures = torch.cuda.FloatTensor(1).fill_(0.0) 116 | 117 | # get grad_outputs 118 | if ctx.return_rgb: 119 | if grad_rgb_map is not None: 120 | grad_rgb_map = grad_rgb_map.contiguous() 121 | else: 122 | grad_rgb_map = torch.zeros_like(rgb_map) 123 | else: 124 | grad_rgb_map = torch.cuda.FloatTensor(1).fill_(0.0) 125 | if ctx.return_alpha: 126 | if grad_alpha_map is not None: 127 | grad_alpha_map = grad_alpha_map.contiguous() 128 | else: 129 | grad_alpha_map = torch.zeros_like(alpha_map) 130 | else: 131 | grad_alpha_map = torch.cuda.FloatTensor(1).fill_(0.0) 132 | if ctx.return_depth: 133 | if grad_depth_map is not None: 134 | grad_depth_map = grad_depth_map.contiguous() 135 | else: 136 | grad_depth_map = torch.zeros_like(ctx.depth_map) 137 | else: 138 | grad_depth_map = torch.cuda.FloatTensor(1).fill_(0.0) 139 | 140 | # backward pass 141 | grad_faces = RasterizeFunction.backward_pixel_map( 142 | ctx, faces, face_index_map, rgb_map, 143 | alpha_map, grad_rgb_map, grad_alpha_map, 144 | grad_faces) 145 | grad_textures = RasterizeFunction.backward_textures( 146 | ctx, face_index_map, sampling_weight_map, 147 | sampling_index_map, grad_rgb_map, grad_textures) 148 | grad_faces = RasterizeFunction.backward_depth_map( 149 | ctx, faces, depth_map, face_index_map, 150 | face_inv_map, weight_map, grad_depth_map, 151 | grad_faces) 152 | 153 | if not textures.requires_grad: 154 | grad_textures = None 155 | 156 | return grad_faces, grad_textures, None, None, None, None, None, None, None, None 157 | 158 | @staticmethod 159 | def forward_face_index_map(ctx, faces, face_index_map, weight_map, 160 | depth_map, face_inv_map): 161 | faces_inv = torch.zeros_like(faces) 162 | return rasterize_cuda.forward_face_index_map(faces, face_index_map, weight_map, 163 | depth_map, face_inv_map, faces_inv, 164 | ctx.image_size, ctx.near, ctx.far, 165 | ctx.return_rgb, ctx.return_alpha, 166 | ctx.return_depth) 167 | 168 | @staticmethod 169 | def forward_texture_sampling(ctx, faces, textures, face_index_map, 170 | weight_map, depth_map, rgb_map, 171 | sampling_index_map, sampling_weight_map): 172 | if not ctx.return_rgb: 173 | return rgb_map, sampling_index_map, sampling_weight_map 174 | else: 175 | return rasterize_cuda.forward_texture_sampling(faces, textures, face_index_map, 176 | weight_map, depth_map, rgb_map, 177 | sampling_index_map, sampling_weight_map, 178 | ctx.image_size, ctx.eps) 179 | 180 | @staticmethod 181 | def forward_alpha_map(ctx, alpha_map, face_index_map): 182 | if ctx.return_alpha: 183 | alpha_map[face_index_map >= 0] = 1 184 | return alpha_map 185 | 186 | @staticmethod 187 | def forward_background(ctx, face_index_map, rgb_map): 188 | if ctx.return_rgb: 189 | background_color = torch.cuda.FloatTensor(ctx.background_color) 190 | mask = (face_index_map >= 0).float()[:, :, :, None] 191 | if background_color.ndimension() == 1: 192 | rgb_map = rgb_map * mask + (1-mask) * background_color[None, None, None, :] 193 | elif background_color.ndimension() == 2: 194 | rgb_map = rgb_map * mask + (1-mask) * background_color[:, None, None, :] 195 | return rgb_map 196 | 197 | @staticmethod 198 | def backward_pixel_map(ctx, faces, face_index_map, rgb_map, 199 | alpha_map, grad_rgb_map, grad_alpha_map, grad_faces): 200 | if (not ctx.return_rgb) and (not ctx.return_alpha): 201 | return grad_faces 202 | else: 203 | return rasterize_cuda.backward_pixel_map(faces, face_index_map, rgb_map, 204 | alpha_map, grad_rgb_map, grad_alpha_map, 205 | grad_faces, ctx.image_size, ctx.eps, ctx.return_rgb, 206 | ctx.return_alpha) 207 | 208 | @staticmethod 209 | def backward_textures(ctx, face_index_map, sampling_weight_map, 210 | sampling_index_map, grad_rgb_map, grad_textures): 211 | if not ctx.return_rgb: 212 | return grad_textures 213 | else: 214 | return rasterize_cuda.backward_textures(face_index_map, sampling_weight_map, 215 | sampling_index_map, grad_rgb_map, 216 | grad_textures, ctx.num_faces) 217 | 218 | @staticmethod 219 | def backward_depth_map(ctx, faces, depth_map, face_index_map, 220 | face_inv_map, weight_map, grad_depth_map, grad_faces): 221 | if not ctx.return_depth: 222 | return grad_faces 223 | else: 224 | return rasterize_cuda.backward_depth_map(faces, depth_map, face_index_map, 225 | face_inv_map, weight_map, 226 | grad_depth_map, grad_faces, ctx.image_size) 227 | 228 | class Rasterize(nn.Module): 229 | ''' 230 | Wrapper around the autograd function RasterizeFunction 231 | Currently implemented only for cuda Tensors 232 | ''' 233 | def __init__(self, image_size, near, far, eps, background_color, 234 | return_rgb=False, return_alpha=False, return_depth=False): 235 | super(Rasterize, self).__init__() 236 | self.image_size = image_size 237 | self.image_size = image_size 238 | self.near = near 239 | self.far = far 240 | self.eps = eps 241 | self.background_color = background_color 242 | self.return_rgb = return_rgb 243 | self.return_alpha = return_alpha 244 | self.return_depth = return_depth 245 | 246 | def forward(self, faces, textures): 247 | if faces.device == "cpu" or (textures is not None and textures.device == "cpu"): 248 | raise TypeError('Rasterize module supports only cuda Tensors') 249 | return RasterizeFunction.apply(faces, textures, self.image_size, self.near, self.far, 250 | self.eps, self.background_color, 251 | self.return_rgb, self.return_alpha, self.return_depth) 252 | 253 | def rasterize_rgbad( 254 | faces, 255 | textures=None, 256 | image_size=DEFAULT_IMAGE_SIZE, 257 | anti_aliasing=DEFAULT_ANTI_ALIASING, 258 | near=DEFAULT_NEAR, 259 | far=DEFAULT_FAR, 260 | eps=DEFAULT_EPS, 261 | background_color=DEFAULT_BACKGROUND_COLOR, 262 | return_rgb=True, 263 | return_alpha=True, 264 | return_depth=True, 265 | ): 266 | """ 267 | Generate RGB, alpha channel, and depth images from faces and textures (for RGB). 268 | 269 | Args: 270 | faces (torch.Tensor): Faces. The shape is [batch size, number of faces, 3 (vertices), 3 (XYZ)]. 271 | textures (torch.Tensor): Textures. 272 | The shape is [batch size, number of faces, texture size, texture size, texture size, 3 (RGB)]. 273 | image_size (int): Width and height of rendered images. 274 | anti_aliasing (bool): do anti-aliasing by super-sampling. 275 | near (float): nearest z-coordinate to draw. 276 | far (float): farthest z-coordinate to draw. 277 | eps (float): small epsilon for approximated differentiation. 278 | background_color (tuple): background color of RGB images. 279 | return_rgb (bool): generate RGB images or not. 280 | return_alpha (bool): generate alpha channels or not. 281 | return_depth (bool): generate depth images or not. 282 | 283 | Returns: 284 | dict: 285 | { 286 | 'rgb': RGB images. The shape is [batch size, 3, image_size, image_size]. 287 | 'alpha': Alpha channels. The shape is [batch size, image_size, image_size]. 288 | 'depth': Depth images. The shape is [batch size, image_size, image_size]. 289 | } 290 | 291 | """ 292 | if textures is None: 293 | inputs = [faces, None] 294 | else: 295 | inputs = [faces, textures] 296 | 297 | if anti_aliasing: 298 | # 2x super-sampling 299 | rgb, alpha, depth = Rasterize( 300 | image_size * 2, near, far, eps, background_color, return_rgb, return_alpha, return_depth)(*inputs) 301 | else: 302 | rgb, alpha, depth = Rasterize( 303 | image_size, near, far, eps, background_color, return_rgb, return_alpha, return_depth)(*inputs) 304 | 305 | # transpose & vertical flip 306 | if return_rgb: 307 | rgb = rgb.permute((0, 3, 1, 2)) 308 | # pytorch does not support negative slicing for the moment 309 | # may need to look at this again because it seems to be very slow 310 | # rgb = rgb[:, :, ::-1, :] 311 | rgb = rgb[:, :, list(reversed(range(rgb.shape[2]))), :] 312 | if return_alpha: 313 | # alpha = alpha[:, ::-1, :] 314 | alpha = alpha[:, list(reversed(range(alpha.shape[1]))), :] 315 | if return_depth: 316 | # depth = depth[:, ::-1, :] 317 | depth = depth[:, list(reversed(range(depth.shape[1]))), :] 318 | 319 | if anti_aliasing: 320 | # 0.5x down-sampling 321 | if return_rgb: 322 | rgb = F.avg_pool2d(rgb, kernel_size=(2,2)) 323 | if return_alpha: 324 | alpha = F.avg_pool2d(alpha[:, None, :, :], kernel_size=(2, 2))[:, 0] 325 | if return_depth: 326 | depth = F.avg_pool2d(depth[:, None, :, :], kernel_size=(2, 2))[:, 0] 327 | 328 | ret = { 329 | 'rgb': rgb if return_rgb else None, 330 | 'alpha': alpha if return_alpha else None, 331 | 'depth': depth if return_depth else None, 332 | } 333 | 334 | return ret 335 | 336 | 337 | def rasterize( 338 | faces, 339 | textures, 340 | image_size=DEFAULT_IMAGE_SIZE, 341 | anti_aliasing=DEFAULT_ANTI_ALIASING, 342 | near=DEFAULT_NEAR, 343 | far=DEFAULT_FAR, 344 | eps=DEFAULT_EPS, 345 | background_color=DEFAULT_BACKGROUND_COLOR, 346 | ): 347 | """ 348 | Generate RGB images from faces and textures. 349 | 350 | Args: 351 | faces: see `rasterize_rgbad`. 352 | textures: see `rasterize_rgbad`. 353 | image_size: see `rasterize_rgbad`. 354 | anti_aliasing: see `rasterize_rgbad`. 355 | near: see `rasterize_rgbad`. 356 | far: see `rasterize_rgbad`. 357 | eps: see `rasterize_rgbad`. 358 | background_color: see `rasterize_rgbad`. 359 | 360 | Returns: 361 | ~torch.Tensor: RGB images. The shape is [batch size, 3, image_size, image_size]. 362 | 363 | """ 364 | return rasterize_rgbad( 365 | faces, textures, image_size, anti_aliasing, near, far, eps, background_color, True, False, False)['rgb'] 366 | 367 | def rasterize( 368 | faces, 369 | textures, 370 | image_size=DEFAULT_IMAGE_SIZE, 371 | anti_aliasing=DEFAULT_ANTI_ALIASING, 372 | near=DEFAULT_NEAR, 373 | far=DEFAULT_FAR, 374 | eps=DEFAULT_EPS, 375 | background_color=DEFAULT_BACKGROUND_COLOR): 376 | """ 377 | Generate RGB images from faces and textures. 378 | 379 | Args: 380 | faces: see `rasterize_rgbad`. 381 | textures: see `rasterize_rgbad`. 382 | image_size: see `rasterize_rgbad`. 383 | anti_aliasing: see `rasterize_rgbad`. 384 | near: see `rasterize_rgbad`. 385 | far: see `rasterize_rgbad`. 386 | eps: see `rasterize_rgbad`. 387 | background_color: see `rasterize_rgbad`. 388 | 389 | Returns: 390 | ~torch.Tensor: RGB images. The shape is [batch size, 3, image_size, image_size]. 391 | 392 | """ 393 | return rasterize_rgbad( 394 | faces, textures, image_size, anti_aliasing, near, far, eps, background_color, True, False, False)['rgb'] 395 | 396 | 397 | def rasterize_silhouettes( 398 | faces, 399 | image_size=DEFAULT_IMAGE_SIZE, 400 | anti_aliasing=DEFAULT_ANTI_ALIASING, 401 | near=DEFAULT_NEAR, 402 | far=DEFAULT_FAR, 403 | eps=DEFAULT_EPS, 404 | ): 405 | """ 406 | Generate alpha channels from faces. 407 | 408 | Args: 409 | faces: see `rasterize_rgbad`. 410 | image_size: see `rasterize_rgbad`. 411 | anti_aliasing: see `rasterize_rgbad`. 412 | near: see `rasterize_rgbad`. 413 | far: see `rasterize_rgbad`. 414 | eps: see `rasterize_rgbad`. 415 | 416 | Returns: 417 | ~torch.Tensor: Alpha channels. The shape is [batch size, image_size, image_size]. 418 | 419 | """ 420 | return rasterize_rgbad(faces, None, image_size, anti_aliasing, near, far, eps, None, False, True, False)['alpha'] 421 | 422 | 423 | def rasterize_depth( 424 | faces, 425 | image_size=DEFAULT_IMAGE_SIZE, 426 | anti_aliasing=DEFAULT_ANTI_ALIASING, 427 | near=DEFAULT_NEAR, 428 | far=DEFAULT_FAR, 429 | eps=DEFAULT_EPS, 430 | ): 431 | """ 432 | Generate depth images from faces. 433 | 434 | Args: 435 | faces: see `rasterize_rgbad`. 436 | image_size: see `rasterize_rgbad`. 437 | anti_aliasing: see `rasterize_rgbad`. 438 | near: see `rasterize_rgbad`. 439 | far: see `rasterize_rgbad`. 440 | eps: see `rasterize_rgbad`. 441 | 442 | Returns: 443 | ~torch.Tensor: Depth images. The shape is [batch size, image_size, image_size]. 444 | 445 | """ 446 | return rasterize_rgbad(faces, None, image_size, anti_aliasing, near, far, eps, None, False, False, True)['depth'] 447 | -------------------------------------------------------------------------------- /neural_renderer/renderer.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | import math 3 | 4 | import torch 5 | import torch.nn as nn 6 | import numpy 7 | 8 | import neural_renderer as nr 9 | 10 | 11 | class Renderer(nn.Module): 12 | def __init__(self, image_size=256, anti_aliasing=True, background_color=[0,0,0], 13 | fill_back=True, camera_mode='projection', 14 | K=None, R=None, t=None, dist_coeffs=None, orig_size=1024, 15 | perspective=True, viewing_angle=30, camera_direction=[0,0,1], 16 | near=0.1, far=100, 17 | light_intensity_ambient=0.5, light_intensity_directional=0.5, 18 | light_color_ambient=[1,1,1], light_color_directional=[1,1,1], 19 | light_direction=[0,1,0]): 20 | super(Renderer, self).__init__() 21 | # rendering 22 | self.image_size = image_size 23 | self.anti_aliasing = anti_aliasing 24 | self.background_color = background_color 25 | self.fill_back = fill_back 26 | 27 | # camera 28 | self.camera_mode = camera_mode 29 | if self.camera_mode == 'projection': 30 | self.K = K 31 | self.R = R 32 | self.t = t 33 | if isinstance(self.K, numpy.ndarray): 34 | self.K = torch.cuda.FloatTensor(self.K) 35 | if isinstance(self.R, numpy.ndarray): 36 | self.R = torch.cuda.FloatTensor(self.R) 37 | if isinstance(self.t, numpy.ndarray): 38 | self.t = torch.cuda.FloatTensor(self.t) 39 | self.dist_coeffs = dist_coeffs 40 | if dist_coeffs is None: 41 | self.dist_coeffs = torch.cuda.FloatTensor([[0., 0., 0., 0., 0.]]) 42 | self.orig_size = orig_size 43 | elif self.camera_mode in ['look', 'look_at']: 44 | self.perspective = perspective 45 | self.viewing_angle = viewing_angle 46 | self.eye = [0, 0, -(1. / math.tan(math.radians(self.viewing_angle)) + 1)] 47 | self.camera_direction = [0, 0, 1] 48 | else: 49 | raise ValueError('Camera mode has to be one of projection, look or look_at') 50 | 51 | 52 | self.near = near 53 | self.far = far 54 | 55 | # light 56 | self.light_intensity_ambient = light_intensity_ambient 57 | self.light_intensity_directional = light_intensity_directional 58 | self.light_color_ambient = light_color_ambient 59 | self.light_color_directional = light_color_directional 60 | self.light_direction = light_direction 61 | 62 | # rasterization 63 | self.rasterizer_eps = 1e-3 64 | 65 | def forward(self, vertices, faces, textures=None, mode=None, K=None, R=None, t=None, dist_coeffs=None, orig_size=None): 66 | ''' 67 | Implementation of forward rendering method 68 | The old API is preserved for back-compatibility with the Chainer implementation 69 | ''' 70 | 71 | if mode is None: 72 | return self.render(vertices, faces, textures, K, R, t, dist_coeffs, orig_size) 73 | elif mode is 'rgb': 74 | return self.render_rgb(vertices, faces, textures, K, R, t, dist_coeffs, orig_size) 75 | elif mode == 'silhouettes': 76 | return self.render_silhouettes(vertices, faces, K, R, t, dist_coeffs, orig_size) 77 | elif mode == 'depth': 78 | return self.render_depth(vertices, faces, K, R, t, dist_coeffs, orig_size) 79 | else: 80 | raise ValueError("mode should be one of None, 'silhouettes' or 'depth'") 81 | 82 | def render_silhouettes(self, vertices, faces, K=None, R=None, t=None, dist_coeffs=None, orig_size=None): 83 | 84 | # fill back 85 | if self.fill_back: 86 | faces = torch.cat((faces, faces[:, :, list(reversed(range(faces.shape[-1])))]), dim=1) 87 | 88 | # viewpoint transformation 89 | if self.camera_mode == 'look_at': 90 | vertices = nr.look_at(vertices, self.eye) 91 | # perspective transformation 92 | if self.perspective: 93 | vertices = nr.perspective(vertices, angle=self.viewing_angle) 94 | elif self.camera_mode == 'look': 95 | vertices = nr.look(vertices, self.eye, self.camera_direction) 96 | # perspective transformation 97 | if self.perspective: 98 | vertices = nr.perspective(vertices, angle=self.viewing_angle) 99 | elif self.camera_mode == 'projection': 100 | if K is None: 101 | K = self.K 102 | if R is None: 103 | R = self.R 104 | if t is None: 105 | t = self.t 106 | if dist_coeffs is None: 107 | dist_coeffs = self.dist_coeffs 108 | if orig_size is None: 109 | orig_size = self.orig_size 110 | vertices = nr.projection(vertices, K, R, t, dist_coeffs, orig_size) 111 | 112 | # rasterization 113 | faces = nr.vertices_to_faces(vertices, faces) 114 | images = nr.rasterize_silhouettes(faces, self.image_size, self.anti_aliasing) 115 | return images 116 | 117 | def render_depth(self, vertices, faces, K=None, R=None, t=None, dist_coeffs=None, orig_size=None): 118 | 119 | # fill back 120 | if self.fill_back: 121 | faces = torch.cat((faces, faces[:, :, list(reversed(range(faces.shape[-1])))]), dim=1).detach() 122 | 123 | # viewpoint transformation 124 | if self.camera_mode == 'look_at': 125 | vertices = nr.look_at(vertices, self.eye) 126 | # perspective transformation 127 | if self.perspective: 128 | vertices = nr.perspective(vertices, angle=self.viewing_angle) 129 | elif self.camera_mode == 'look': 130 | vertices = nr.look(vertices, self.eye, self.camera_direction) 131 | # perspective transformation 132 | if self.perspective: 133 | vertices = nr.perspective(vertices, angle=self.viewing_angle) 134 | elif self.camera_mode == 'projection': 135 | if K is None: 136 | K = self.K 137 | if R is None: 138 | R = self.R 139 | if t is None: 140 | t = self.t 141 | if dist_coeffs is None: 142 | dist_coeffs = self.dist_coeffs 143 | if orig_size is None: 144 | orig_size = self.orig_size 145 | vertices = nr.projection(vertices, K, R, t, dist_coeffs, orig_size) 146 | 147 | # rasterization 148 | faces = nr.vertices_to_faces(vertices, faces) 149 | images = nr.rasterize_depth(faces, self.image_size, self.anti_aliasing) 150 | return images 151 | 152 | def render_rgb(self, vertices, faces, textures, K=None, R=None, t=None, dist_coeffs=None, orig_size=None): 153 | # fill back 154 | if self.fill_back: 155 | faces = torch.cat((faces, faces[:, :, list(reversed(range(faces.shape[-1])))]), dim=1).detach() 156 | textures = torch.cat((textures, textures.permute((0, 1, 4, 3, 2, 5))), dim=1) 157 | 158 | # lighting 159 | faces_lighting = nr.vertices_to_faces(vertices, faces) 160 | textures = nr.lighting( 161 | faces_lighting, 162 | textures, 163 | self.light_intensity_ambient, 164 | self.light_intensity_directional, 165 | self.light_color_ambient, 166 | self.light_color_directional, 167 | self.light_direction) 168 | 169 | # viewpoint transformation 170 | if self.camera_mode == 'look_at': 171 | vertices = nr.look_at(vertices, self.eye) 172 | # perspective transformation 173 | if self.perspective: 174 | vertices = nr.perspective(vertices, angle=self.viewing_angle) 175 | elif self.camera_mode == 'look': 176 | vertices = nr.look(vertices, self.eye, self.camera_direction) 177 | # perspective transformation 178 | if self.perspective: 179 | vertices = nr.perspective(vertices, angle=self.viewing_angle) 180 | elif self.camera_mode == 'projection': 181 | if K is None: 182 | K = self.K 183 | if R is None: 184 | R = self.R 185 | if t is None: 186 | t = self.t 187 | if dist_coeffs is None: 188 | dist_coeffs = self.dist_coeffs 189 | if orig_size is None: 190 | orig_size = self.orig_size 191 | vertices = nr.projection(vertices, K, R, t, dist_coeffs, orig_size) 192 | 193 | # rasterization 194 | faces = nr.vertices_to_faces(vertices, faces) 195 | images = nr.rasterize( 196 | faces, textures, self.image_size, self.anti_aliasing, self.near, self.far, self.rasterizer_eps, 197 | self.background_color) 198 | return images 199 | 200 | def render(self, vertices, faces, textures, K=None, R=None, t=None, dist_coeffs=None, orig_size=None): 201 | # fill back 202 | if self.fill_back: 203 | faces = torch.cat((faces, faces[:, :, list(reversed(range(faces.shape[-1])))]), dim=1).detach() 204 | textures = torch.cat((textures, textures.permute((0, 1, 4, 3, 2, 5))), dim=1) 205 | 206 | # lighting 207 | faces_lighting = nr.vertices_to_faces(vertices, faces) 208 | textures = nr.lighting( 209 | faces_lighting, 210 | textures, 211 | self.light_intensity_ambient, 212 | self.light_intensity_directional, 213 | self.light_color_ambient, 214 | self.light_color_directional, 215 | self.light_direction) 216 | 217 | # viewpoint transformation 218 | if self.camera_mode == 'look_at': 219 | vertices = nr.look_at(vertices, self.eye) 220 | # perspective transformation 221 | if self.perspective: 222 | vertices = nr.perspective(vertices, angle=self.viewing_angle) 223 | elif self.camera_mode == 'look': 224 | vertices = nr.look(vertices, self.eye, self.camera_direction) 225 | # perspective transformation 226 | if self.perspective: 227 | vertices = nr.perspective(vertices, angle=self.viewing_angle) 228 | elif self.camera_mode == 'projection': 229 | if K is None: 230 | K = self.K 231 | if R is None: 232 | R = self.R 233 | if t is None: 234 | t = self.t 235 | if dist_coeffs is None: 236 | dist_coeffs = self.dist_coeffs 237 | if orig_size is None: 238 | orig_size = self.orig_size 239 | vertices = nr.projection(vertices, K, R, t, dist_coeffs, orig_size) 240 | 241 | # rasterization 242 | faces = nr.vertices_to_faces(vertices, faces) 243 | out = nr.rasterize_rgbad( 244 | faces, textures, self.image_size, self.anti_aliasing, self.near, self.far, self.rasterizer_eps, 245 | self.background_color) 246 | return out['rgb'], out['depth'], out['alpha'] 247 | -------------------------------------------------------------------------------- /neural_renderer/save_obj.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | import os 3 | 4 | import torch 5 | from skimage.io import imsave 6 | 7 | import neural_renderer.cuda.create_texture_image as create_texture_image_cuda 8 | 9 | 10 | def create_texture_image(textures, texture_size_out=16): 11 | num_faces, texture_size_in = textures.shape[:2] 12 | tile_width = int((num_faces - 1.) ** 0.5) + 1 13 | tile_height = int((num_faces - 1.) / tile_width) + 1 14 | image = torch.zeros(tile_height * texture_size_out, tile_width * texture_size_out, 3, dtype=torch.float32) 15 | vertices = torch.zeros((num_faces, 3, 2), dtype=torch.float32) # [:, :, XY] 16 | face_nums = torch.arange(num_faces) 17 | column = face_nums % tile_width 18 | row = face_nums / tile_width 19 | vertices[:, 0, 0] = column * texture_size_out 20 | vertices[:, 0, 1] = row * texture_size_out 21 | vertices[:, 1, 0] = column * texture_size_out 22 | vertices[:, 1, 1] = (row + 1) * texture_size_out - 1 23 | vertices[:, 2, 0] = (column + 1) * texture_size_out - 1 24 | vertices[:, 2, 1] = (row + 1) * texture_size_out - 1 25 | image = image.cuda() 26 | vertices = vertices.cuda() 27 | textures = textures.cuda() 28 | image = create_texture_image_cuda.create_texture_image(vertices, textures, image, 1e-5) 29 | 30 | vertices[:, :, 0] /= (image.shape[1] - 1) 31 | vertices[:, :, 1] /= (image.shape[0] - 1) 32 | 33 | image = image.detach().cpu().numpy() 34 | vertices = vertices.detach().cpu().numpy() 35 | image = image[::-1, ::1] 36 | 37 | return image, vertices 38 | 39 | 40 | def save_obj(filename, vertices, faces, textures=None): 41 | assert vertices.ndimension() == 2 42 | assert faces.ndimension() == 2 43 | 44 | if textures is not None: 45 | filename_mtl = filename[:-4] + '.mtl' 46 | filename_texture = filename[:-4] + '.png' 47 | material_name = 'material_1' 48 | texture_image, vertices_textures = create_texture_image(textures) 49 | imsave(filename_texture, texture_image) 50 | 51 | faces = faces.detach().cpu().numpy() 52 | 53 | with open(filename, 'w') as f: 54 | f.write('# %s\n' % os.path.basename(filename)) 55 | f.write('#\n') 56 | f.write('\n') 57 | 58 | if textures is not None: 59 | f.write('mtllib %s\n\n' % os.path.basename(filename_mtl)) 60 | 61 | for vertex in vertices: 62 | f.write('v %.8f %.8f %.8f\n' % (vertex[0], vertex[1], vertex[2])) 63 | f.write('\n') 64 | 65 | if textures is not None: 66 | for vertex in vertices_textures.reshape((-1, 2)): 67 | f.write('vt %.8f %.8f\n' % (vertex[0], vertex[1])) 68 | f.write('\n') 69 | 70 | f.write('usemtl %s\n' % material_name) 71 | for i, face in enumerate(faces): 72 | f.write('f %d/%d %d/%d %d/%d\n' % ( 73 | face[0] + 1, 3 * i + 1, face[1] + 1, 3 * i + 2, face[2] + 1, 3 * i + 3)) 74 | f.write('\n') 75 | else: 76 | for face in faces: 77 | f.write('f %d %d %d\n' % (face[0] + 1, face[1] + 1, face[2] + 1)) 78 | 79 | if textures is not None: 80 | with open(filename_mtl, 'w') as f: 81 | f.write('newmtl %s\n' % material_name) 82 | f.write('map_Kd %s\n' % os.path.basename(filename_texture)) 83 | -------------------------------------------------------------------------------- /neural_renderer/vertices_to_faces.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | def vertices_to_faces(vertices, faces): 5 | """ 6 | :param vertices: [batch size, number of vertices, 3] 7 | :param faces: [batch size, number of faces, 3) 8 | :return: [batch size, number of faces, 3, 3] 9 | """ 10 | assert (vertices.ndimension() == 3) 11 | assert (faces.ndimension() == 3) 12 | assert (vertices.shape[0] == faces.shape[0]) 13 | assert (vertices.shape[2] == 3) 14 | assert (faces.shape[2] == 3) 15 | 16 | bs, nv = vertices.shape[:2] 17 | bs, nf = faces.shape[:2] 18 | device = vertices.device 19 | faces = faces + (torch.arange(bs, dtype=torch.int32).to(device) * nv)[:, None, None] 20 | vertices = vertices.reshape((bs * nv, 3)) 21 | # pytorch only supports long and byte tensors for indexing 22 | return vertices[faces.long()] 23 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | import unittest 3 | 4 | from torch.utils.cpp_extension import BuildExtension, CUDAExtension 5 | 6 | CUDA_FLAGS = [] 7 | INSTALL_REQUIREMENTS = [] 8 | 9 | def test_all(): 10 | test_loader = unittest.TestLoader() 11 | test_suite = test_loader.discover('tests', pattern='test_*.py') 12 | return test_suite 13 | 14 | ext_modules=[ 15 | CUDAExtension('neural_renderer.cuda.load_textures', [ 16 | 'neural_renderer/cuda/load_textures_cuda.cpp', 17 | 'neural_renderer/cuda/load_textures_cuda_kernel.cu', 18 | ]), 19 | CUDAExtension('neural_renderer.cuda.rasterize', [ 20 | 'neural_renderer/cuda/rasterize_cuda.cpp', 21 | 'neural_renderer/cuda/rasterize_cuda_kernel.cu', 22 | ]), 23 | CUDAExtension('neural_renderer.cuda.create_texture_image', [ 24 | 'neural_renderer/cuda/create_texture_image_cuda.cpp', 25 | 'neural_renderer/cuda/create_texture_image_cuda_kernel.cu', 26 | ]), 27 | ] 28 | 29 | setup( 30 | description='PyTorch implementation of "A 3D mesh renderer for neural networks"', 31 | author='Nikolaos Kolotouros', 32 | author_email='nkolot@seas.upenn.edu', 33 | license='MIT License', 34 | version='1.1.3', 35 | name='neural_renderer_pytorch', 36 | test_suite='setup.test_all', 37 | packages=['neural_renderer', 'neural_renderer.cuda'], 38 | install_requires=INSTALL_REQUIREMENTS, 39 | ext_modules=ext_modules, 40 | cmdclass = {'build_ext': BuildExtension} 41 | ) 42 | -------------------------------------------------------------------------------- /tests/data/1cde62b063e14777c9152a706245d48/model.mtl: -------------------------------------------------------------------------------- 1 | # 2 | ## Alias OBJ Material File 3 | # Exported from SketchUp, (c) 2000-2012 Trimble Navigation Limited 4 | 5 | newmtl Carrosserie_plas1 6 | Ka 0.000000 0.000000 0.000000 7 | Kd 0.117647 0.117647 0.117647 8 | Ks 0.330000 0.330000 0.330000 9 | 10 | newmtl Carrosserie1 11 | Ka 0.000000 0.000000 0.000000 12 | Kd 0.596078 0.070588 0.086275 13 | Ks 0.330000 0.330000 0.330000 14 | 15 | newmtl Clignotants_avan1 16 | Ka 0.000000 0.000000 0.000000 17 | Kd 0.752941 0.517647 0.372549 18 | Ks 0.330000 0.330000 0.330000 19 | 20 | newmtl Dessous1 21 | Ka 0.000000 0.000000 0.000000 22 | Kd 0.156863 0.156863 0.156863 23 | Ks 0.330000 0.330000 0.330000 24 | 25 | newmtl Feux_arri_re1 26 | Ka 0.000000 0.000000 0.000000 27 | Kd 0.905882 0.196078 0.211765 28 | Ks 0.330000 0.330000 0.330000 29 | 30 | newmtl Clignotants_cot_1 31 | Ka 0.000000 0.000000 0.000000 32 | Kd 0.729412 0.341176 0.227451 33 | Ks 0.330000 0.330000 0.330000 34 | 35 | newmtl chrome1 36 | Ka 0.000000 0.000000 0.000000 37 | Kd 0.384314 0.384314 0.384314 38 | Ks 0.330000 0.330000 0.330000 39 | 40 | newmtl numberPlates1 41 | Ka 0.000000 0.000000 0.000000 42 | Kd 0.811765 0.811765 0.811765 43 | Ks 0.330000 0.330000 0.330000 44 | 45 | newmtl black_plastic1 46 | Ka 0.000000 0.000000 0.000000 47 | Kd 0.000000 0.000000 0.000000 48 | Ks 0.330000 0.330000 0.330000 49 | 50 | newmtl fogLight1 51 | Ka 0.000000 0.000000 0.000000 52 | Kd 1.000000 1.000000 1.000000 53 | Ks 0.330000 0.330000 0.330000 54 | 55 | newmtl floor1 56 | Ka 0.000000 0.000000 0.000000 57 | Kd 0.082353 0.082353 0.082353 58 | Ks 0.330000 0.330000 0.330000 59 | 60 | newmtl Int_rieur_toit1 61 | Ka 0.000000 0.000000 0.000000 62 | Kd 0.850980 0.811765 0.741176 63 | Ks 0.330000 0.330000 0.330000 64 | 65 | newmtl Carrosserie_Toit1 66 | Ka 0.000000 0.000000 0.000000 67 | Kd 0.772549 0.772549 0.772549 68 | Ks 0.330000 0.330000 0.330000 69 | 70 | newmtl Montants_de_vitr1 71 | Ka 0.000000 0.000000 0.000000 72 | Kd 0.117647 0.117647 0.117647 73 | Ks 0.330000 0.330000 0.330000 74 | 75 | newmtl Int_rieur___si_g1 76 | Ka 0.000000 0.000000 0.000000 77 | Kd 0.129412 0.129412 0.129412 78 | Ks 0.330000 0.330000 0.330000 79 | 80 | newmtl Antenne___essuie1 81 | Ka 0.000000 0.000000 0.000000 82 | Kd 0.219608 0.219608 0.219608 83 | Ks 0.330000 0.330000 0.330000 84 | 85 | newmtl Grilles_capot1 86 | Ka 0.000000 0.000000 0.000000 87 | Kd 0.188235 0.188235 0.188235 88 | Ks 0.330000 0.330000 0.330000 89 | 90 | newmtl Phares_vitre1 91 | Ka 0.000000 0.000000 0.000000 92 | Kd 0.607843 0.623529 0.635294 93 | Ks 0.330000 0.330000 0.330000 94 | d 0.500000 95 | 96 | newmtl Mirroir_r_tro1 97 | Ka 0.000000 0.000000 0.000000 98 | Kd 0.356863 0.356863 0.356863 99 | Ks 0.330000 0.330000 0.330000 100 | 101 | newmtl Int_rieur_1 102 | Ka 0.000000 0.000000 0.000000 103 | Kd 0.564706 0.564706 0.564706 104 | Ks 0.330000 0.330000 0.330000 105 | 106 | newmtl speaker_mesh1 107 | Ka 0.000000 0.000000 0.000000 108 | Kd 0.643137 0.643137 0.643137 109 | Ks 0.330000 0.330000 0.330000 110 | 111 | newmtl steeringwheel_lo1 112 | Ka 0.000000 0.000000 0.000000 113 | Kd 0.082353 0.082353 0.082353 114 | Ks 0.330000 0.330000 0.330000 115 | 116 | newmtl radiopanel1 117 | Ka 0.000000 0.000000 0.000000 118 | Kd 0.000000 0.000000 0.000000 119 | Ks 0.330000 0.330000 0.330000 120 | 121 | newmtl speedodisplay1 122 | Ka 0.000000 0.000000 0.000000 123 | Kd 0.874510 0.874510 0.874510 124 | Ks 0.330000 0.330000 0.330000 125 | 126 | newmtl Jantes1 127 | Ka 0.000000 0.000000 0.000000 128 | Kd 0.647059 0.647059 0.647059 129 | Ks 0.330000 0.330000 0.330000 130 | 131 | newmtl Disques_de_frein1 132 | Ka 0.000000 0.000000 0.000000 133 | Kd 0.341176 0.341176 0.341176 134 | Ks 0.330000 0.330000 0.330000 135 | 136 | newmtl Jantes_Logo1 137 | Ka 0.000000 0.000000 0.000000 138 | Kd 0.235294 0.235294 0.235294 139 | Ks 0.330000 0.330000 0.330000 140 | 141 | newmtl Pneus1 142 | Ka 0.000000 0.000000 0.000000 143 | Kd 0.200000 0.200000 0.192157 144 | Ks 0.330000 0.330000 0.330000 145 | 146 | newmtl Vitres1 147 | Ka 0.000000 0.000000 0.000000 148 | Kd 0.196078 0.219608 0.215686 149 | Ks 0.330000 0.330000 0.330000 150 | d 0.510000 151 | 152 | newmtl Phares1 153 | Ka 0.000000 0.000000 0.000000 154 | Kd 0.247059 0.247059 0.247059 155 | Ks 0.330000 0.330000 0.330000 156 | 157 | newmtl mini_logo1 158 | Ka 0.000000 0.000000 0.000000 159 | Kd 0.164706 0.164706 0.164706 160 | Ks 0.330000 0.330000 0.330000 161 | 162 | newmtl door_handle_foam1 163 | Ka 0.000000 0.000000 0.000000 164 | Kd 0.027451 0.027451 0.027451 165 | Ks 0.330000 0.330000 0.330000 166 | 167 | newmtl Gouttiere1 168 | Ka 0.000000 0.000000 0.000000 169 | Kd 0.117647 0.117647 0.117647 170 | Ks 0.330000 0.330000 0.330000 171 | 172 | newmtl brakelight1 173 | Ka 0.000000 0.000000 0.000000 174 | Kd 0.725490 0.000000 0.000000 175 | Ks 0.330000 0.330000 0.330000 176 | 177 | newmtl Si_ges_centre1 178 | Ka 0.000000 0.000000 0.000000 179 | Kd 0.423529 0.423529 0.423529 180 | Ks 0.330000 0.330000 0.330000 181 | 182 | -------------------------------------------------------------------------------- /tests/data/4e49873292196f02574b5684eaec43e9/images/texture0.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/4e49873292196f02574b5684eaec43e9/images/texture0.jpg -------------------------------------------------------------------------------- /tests/data/4e49873292196f02574b5684eaec43e9/images/texture0_.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/4e49873292196f02574b5684eaec43e9/images/texture0_.jpg -------------------------------------------------------------------------------- /tests/data/4e49873292196f02574b5684eaec43e9/images/texture1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/4e49873292196f02574b5684eaec43e9/images/texture1.jpg -------------------------------------------------------------------------------- /tests/data/4e49873292196f02574b5684eaec43e9/model.mtl: -------------------------------------------------------------------------------- 1 | # Blender MTL File: 'blank.blend' 2 | # Material Count: 7 3 | 4 | newmtl material_0_24 5 | Ns 96.078431 6 | Ka 0.000000 0.000000 0.000000 7 | Kd 0.401569 0.401569 0.401569 8 | Ks 0.500000 0.500000 0.500000 9 | Ni -1.000000 10 | d 1.000000 11 | illum 2 12 | 13 | newmtl material_1_24 14 | Ns 96.078431 15 | Ka 0.000000 0.000000 0.000000 16 | Kd 0.109804 0.109804 0.109804 17 | Ks 0.500000 0.500000 0.500000 18 | Ni -1.000000 19 | d 1.000000 20 | illum 2 21 | 22 | newmtl material_2_24 23 | Ns 96.078431 24 | Ka 0.000000 0.000000 0.000000 25 | Kd 0.602353 0.602353 0.602353 26 | Ks 0.500000 0.500000 0.500000 27 | Ni -1.000000 28 | d 1.000000 29 | illum 2 30 | 31 | newmtl material_3_1_0 32 | Ns 96.078431 33 | Ka 0.000000 0.000000 0.000000 34 | Kd 0.640000 0.640000 0.640000 35 | Ks 0.500000 0.500000 0.500000 36 | Ni -1.000000 37 | d 1.000000 38 | illum 2 39 | map_Kd ./images/texture1.jpg 40 | 41 | newmtl material_3_24 42 | Ns 96.078431 43 | Ka 0.000000 0.000000 0.000000 44 | Kd 0.800000 0.800000 0.800000 45 | Ks 0.500000 0.500000 0.500000 46 | Ni -1.000000 47 | d 1.000000 48 | illum 2 49 | 50 | newmtl material_4_24 51 | Ns 96.078431 52 | Ka 0.000000 0.000000 0.000000 53 | Kd 0.000000 0.000000 0.800000 54 | Ks 0.500000 0.500000 0.500000 55 | Ni -1.000000 56 | d 1.000000 57 | illum 2 58 | 59 | newmtl material_5_2_8 60 | Ns 96.078431 61 | Ka 0.000000 0.000000 0.000000 62 | Kd 0.640000 0.640000 0.640000 63 | Ks 0.500000 0.500000 0.500000 64 | Ni -1.000000 65 | d 1.000000 66 | illum 2 67 | map_Kd ./images/texture0.jpg 68 | -------------------------------------------------------------------------------- /tests/data/clean.blend: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/clean.blend -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case1.png -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case1_v0_x.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case1_v0_x.png -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case1_v0_y.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case1_v0_y.png -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case1_v1_x.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case1_v1_x.png -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case1_v1_y.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case1_v1_y.png -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case1_v2_x.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case1_v2_x.png -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case1_v2_y.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case1_v2_y.png -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case2.png -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case2_v0_x.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case2_v0_x.png -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case2_v0_y.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case2_v0_y.png -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case2_v1_x.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case2_v1_x.png -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case2_v1_y.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case2_v1_y.png -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case2_v2_x.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case2_v2_x.png -------------------------------------------------------------------------------- /tests/data/rasterize_silhouettes_case2_v2_y.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/rasterize_silhouettes_case2_v2_y.png -------------------------------------------------------------------------------- /tests/data/teapot_blender.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/teapot_blender.png -------------------------------------------------------------------------------- /tests/data/test_depth.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/lstcutong/neural_renderer_pytorch-win10/206baa196e687510eb684a8519753121c43f8873/tests/data/test_depth.png -------------------------------------------------------------------------------- /tests/data/tetrahedron.obj: -------------------------------------------------------------------------------- 1 | v 1 0 0 2 | v 0 1 0 3 | v 0 0 1 4 | v 0 0 0 5 | f 2 4 3 6 | f 4 2 1 7 | f 3 1 2 8 | f 1 3 4 -------------------------------------------------------------------------------- /tests/test_get_points_from_angles.py: -------------------------------------------------------------------------------- 1 | # TODO 2 | -------------------------------------------------------------------------------- /tests/test_lighting.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | import torch 4 | 5 | import neural_renderer as nr 6 | 7 | class TestLighting(unittest.TestCase): 8 | 9 | def test_case1(self): 10 | """Test whether it is executable.""" 11 | faces = torch.randn(64, 16, 3, 3, dtype=torch.float32) 12 | textures = torch.randn(64, 16, 8, 8, 8, 3, dtype=torch.float32) 13 | nr.lighting(faces, textures) 14 | 15 | if __name__ == '__main__': 16 | unittest.main() 17 | 18 | 19 | 20 | -------------------------------------------------------------------------------- /tests/test_load_obj.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | import os 3 | 4 | import numpy as np 5 | from skimage.io import imsave 6 | 7 | import torch 8 | import neural_renderer as nr 9 | 10 | current_dir = os.path.dirname(os.path.realpath(__file__)) 11 | data_dir = os.path.join(current_dir, 'data') 12 | 13 | class TestCore(unittest.TestCase): 14 | def test_tetrahedron(self): 15 | vertices_ref = np.array( 16 | [ 17 | [1., 0., 0.], 18 | [0., 1., 0.], 19 | [0., 0., 1.], 20 | [0., 0., 0.]], 21 | 'float32') 22 | faces_ref = np.array( 23 | [ 24 | [1, 3, 2], 25 | [3, 1, 0], 26 | [2, 0, 1], 27 | [0, 2, 3]], 28 | 'int32') 29 | 30 | obj_file = os.path.join(data_dir, 'tetrahedron.obj') 31 | vertices, faces = nr.load_obj(obj_file, False) 32 | assert (torch.allclose(torch.from_numpy(vertices_ref).cuda(), vertices)) 33 | assert (torch.allclose(torch.from_numpy(faces_ref).cuda(), faces)) 34 | vertices, faces = nr.load_obj(obj_file, True) 35 | assert (torch.allclose(torch.from_numpy(vertices_ref).cuda() * 2 - 1.0, vertices)) 36 | assert (torch.allclose(torch.from_numpy(faces_ref).cuda(), faces)) 37 | 38 | def test_teapot(self): 39 | obj_file = os.path.join(data_dir, 'teapot.obj') 40 | vertices, faces = nr.load_obj(obj_file) 41 | assert (faces.shape[0] == 2464) 42 | assert (vertices.shape[0] == 1292) 43 | 44 | def test_texture(self): 45 | renderer = nr.Renderer(camera_mode='look_at') 46 | 47 | vertices, faces, textures = nr.load_obj( 48 | os.path.join(data_dir, '1cde62b063e14777c9152a706245d48/model.obj'), load_texture=True) 49 | 50 | renderer.eye = nr.get_points_from_angles(2, 15, 30) 51 | images, _, _ = renderer.render(vertices[None, :, :], faces[None, :, :], textures[None, :, :, :, :, :]) 52 | images = images.permute(0,2,3,1).detach().cpu().numpy() 53 | imsave(os.path.join(data_dir, 'car.png'), images[0]) 54 | 55 | vertices, faces, textures = nr.load_obj( 56 | os.path.join(data_dir, '4e49873292196f02574b5684eaec43e9/model.obj'), load_texture=True, texture_size=16) 57 | renderer.eye = nr.get_points_from_angles(2, 15, -90) 58 | images, _, _ = renderer.render(vertices[None, :, :], faces[None, :, :], textures[None, :, :, :, :, :]) 59 | images = images.permute(0,2,3,1).detach().cpu().numpy() 60 | imsave(os.path.join(data_dir, 'display.png'), images[0]) 61 | 62 | 63 | if __name__ == '__main__': 64 | unittest.main() 65 | -------------------------------------------------------------------------------- /tests/test_look.py: -------------------------------------------------------------------------------- 1 | # TODO 2 | -------------------------------------------------------------------------------- /tests/test_look_at.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | import torch 4 | import numpy as np 5 | 6 | import neural_renderer as nr 7 | 8 | class TestLookAt(unittest.TestCase): 9 | def test_case1(self): 10 | eyes = [ 11 | [1, 0, 1], 12 | [0, 0, -10], 13 | [-1, 1, 0], 14 | ] 15 | answers = [ 16 | [-np.sqrt(2) / 2, 0, np.sqrt(2) / 2], 17 | [1, 0, 10], 18 | [0, np.sqrt(2) / 2, 3. / 2. * np.sqrt(2)], 19 | ] 20 | vertices = torch.from_numpy(np.array([1, 0, 0], np.float32)) 21 | vertices = vertices[None, None, :] 22 | for e, a in zip(eyes, answers): 23 | eye = np.array(e, np.float32) 24 | transformed = nr.look_at(vertices, eye) 25 | assert(np.allclose(transformed.data.squeeze().numpy(), np.array(a))) 26 | 27 | if __name__ == '__main__': 28 | unittest.main() 29 | -------------------------------------------------------------------------------- /tests/test_perspective.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | import torch 4 | import numpy as np 5 | 6 | import neural_renderer as nr 7 | 8 | class TestPerspective(unittest.TestCase): 9 | def test_case1(self): 10 | vertices = torch.from_numpy(np.array([1,2,10], np.float32)) 11 | v_out = np.array([np.sqrt(3) / 10, 2 * np.sqrt(3) / 10, 10], np.float32) 12 | vertices = vertices[None, None, :] 13 | transformer = nr.perspective(vertices) 14 | assert(np.allclose(transformer.data.squeeze().numpy(), v_out)) 15 | 16 | if __name__ == '__main__': 17 | unittest.main() 18 | -------------------------------------------------------------------------------- /tests/test_rasterize.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | import os 3 | 4 | import torch 5 | import numpy as np 6 | from skimage.io import imread, imsave 7 | 8 | import neural_renderer as nr 9 | import utils 10 | 11 | current_dir = os.path.dirname(os.path.realpath(__file__)) 12 | data_dir = os.path.join(current_dir, 'data') 13 | 14 | class TestRasterize(unittest.TestCase): 15 | def test_forward_case1(self): 16 | """Rendering a teapot without anti-aliasing.""" 17 | 18 | # load teapot 19 | vertices, faces, textures = utils.load_teapot_batch() 20 | vertices = vertices.cuda() 21 | faces = faces.cuda() 22 | textures = textures.cuda() 23 | 24 | # create renderer 25 | renderer = nr.Renderer(camera_mode='look_at') 26 | renderer.image_size = 256 27 | renderer.anti_aliasing = False 28 | 29 | # render 30 | images, _, _ = renderer(vertices, faces, textures) 31 | images = images.detach().cpu().numpy() 32 | image = images[2] 33 | image = image.transpose((1, 2, 0)) 34 | 35 | imsave(os.path.join(data_dir, 'test_rasterize1.png'), image) 36 | 37 | def test_forward_case2(self): 38 | """Rendering a teapot with anti-aliasing and another viewpoint.""" 39 | 40 | # load teapot 41 | vertices, faces, textures = utils.load_teapot_batch() 42 | vertices = vertices.cuda() 43 | faces = faces.cuda() 44 | textures = textures.cuda() 45 | 46 | # create renderer 47 | renderer = nr.Renderer(camera_mode='look_at') 48 | renderer.eye = [1, 1, -2.7] 49 | 50 | # render 51 | images, _, _ = renderer(vertices, faces, textures) 52 | images = images.detach().cpu().numpy() 53 | image = images[2] 54 | image = image.transpose((1, 2, 0)) 55 | 56 | imsave(os.path.join(data_dir, 'test_rasterize2.png'), image) 57 | 58 | def test_forward_case3(self): 59 | """Whether a silhouette by neural renderer matches that by Blender.""" 60 | 61 | # load teapot 62 | vertices, faces, textures = utils.load_teapot_batch() 63 | vertices = vertices.cuda() 64 | faces = faces.cuda() 65 | textures = textures.cuda() 66 | 67 | # create renderer 68 | renderer = nr.Renderer(camera_mode='look_at') 69 | renderer.image_size = 256 70 | renderer.anti_aliasing = False 71 | renderer.light_intensity_ambient = 1.0 72 | renderer.light_intensity_directional = 0.0 73 | 74 | images, _, _ = renderer(vertices, faces, textures) 75 | images = images.detach().cpu().numpy() 76 | image = images[2].mean(0) 77 | 78 | # load reference image by blender 79 | ref = imread(os.path.join(data_dir, 'teapot_blender.png')) 80 | ref = (ref.min(axis=-1) != 255).astype(np.float32) 81 | 82 | assert(np.allclose(ref, image)) 83 | 84 | def test_backward_case1(self): 85 | """Backward if non-zero gradient is out of a face.""" 86 | 87 | vertices = [ 88 | [0.8, 0.8, 1.], 89 | [0.0, -0.5, 1.], 90 | [0.2, -0.4, 1.]] 91 | faces = [[0, 1, 2]] 92 | pxi = 35 93 | pyi = 25 94 | grad_ref = [ 95 | [1.6725862, -0.26021874, 0.], 96 | [1.41986704, -1.64284933, 0.], 97 | [0., 0., 0.], 98 | ] 99 | 100 | renderer = nr.Renderer(camera_mode='look_at') 101 | renderer.image_size = 64 102 | renderer.anti_aliasing = False 103 | renderer.perspective = False 104 | renderer.light_intensity_ambient = 1.0 105 | renderer.light_intensity_directional = 0.0 106 | 107 | vertices = torch.from_numpy(np.array(vertices, dtype=np.float32)).cuda() 108 | faces = torch.from_numpy(np.array(faces, dtype=np.int32)).cuda() 109 | textures = torch.ones(faces.shape[0], 4, 4, 4, 3, dtype=torch.float32).cuda() 110 | grad_ref = torch.from_numpy(np.array(grad_ref, dtype=np.float32)).cuda() 111 | vertices, faces, textures, grad_ref = utils.to_minibatch((vertices, faces, textures, grad_ref)) 112 | vertices, faces, textures, grad_ref = vertices.cuda(), faces.cuda(), textures.cuda(), grad_ref.cuda() 113 | vertices.requires_grad = True 114 | images, _, _ = renderer(vertices, faces, textures) 115 | images = torch.mean(images, dim=1) 116 | loss = torch.sum(torch.abs(images[:, pyi, pxi] - 1)) 117 | loss.backward() 118 | 119 | assert(torch.allclose(vertices.grad, grad_ref, rtol=1e-2)) 120 | 121 | def test_backward_case2(self): 122 | """Backward if non-zero gradient is on a face.""" 123 | 124 | vertices = [ 125 | [0.8, 0.8, 1.], 126 | [-0.5, -0.8, 1.], 127 | [0.8, -0.8, 1.]] 128 | faces = [[0, 1, 2]] 129 | pyi = 40 130 | pxi = 50 131 | grad_ref = [ 132 | [0.98646867, 1.04628897, 0.], 133 | [-1.03415668, - 0.10403691, 0.], 134 | [3.00094461, - 1.55173182, 0.], 135 | ] 136 | 137 | renderer = nr.Renderer(camera_mode='look_at') 138 | renderer.image_size = 64 139 | renderer.anti_aliasing = False 140 | renderer.perspective = False 141 | renderer.light_intensity_ambient = 1.0 142 | renderer.light_intensity_directional = 0.0 143 | 144 | vertices = torch.from_numpy(np.array(vertices, dtype=np.float32)).cuda() 145 | faces = torch.from_numpy(np.array(faces, dtype=np.int32)).cuda() 146 | textures = torch.ones(faces.shape[0], 4, 4, 4, 3, dtype=torch.float32).cuda() 147 | grad_ref = torch.from_numpy(np.array(grad_ref, dtype=np.float32)).cuda() 148 | vertices, faces, textures, grad_ref = utils.to_minibatch((vertices, faces, textures, grad_ref)) 149 | vertices.requires_grad=True 150 | 151 | images, _, _ = renderer(vertices, faces, textures) 152 | images = torch.mean(images, dim=1) 153 | loss = torch.sum(torch.abs(images[:, pyi, pxi])) 154 | loss.backward() 155 | 156 | assert(torch.allclose(vertices.grad, grad_ref, rtol=1e-2)) 157 | 158 | 159 | if __name__ == '__main__': 160 | unittest.main() 161 | -------------------------------------------------------------------------------- /tests/test_rasterize_depth.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | import os 3 | 4 | import torch 5 | import numpy as np 6 | from skimage.io import imread 7 | 8 | import neural_renderer as nr 9 | import utils 10 | 11 | current_dir = os.path.dirname(os.path.realpath(__file__)) 12 | data_dir = os.path.join(current_dir, 'data') 13 | 14 | class TestRasterizeDepth(unittest.TestCase): 15 | def test_forward_case1(self): 16 | """Whether a silhouette by neural renderer matches that by Blender.""" 17 | 18 | # load teapot 19 | vertices, faces, _ = utils.load_teapot_batch() 20 | 21 | # create renderer 22 | renderer = nr.Renderer(camera_mode='look_at') 23 | renderer.image_size = 256 24 | renderer.anti_aliasing = False 25 | 26 | images = renderer(vertices, faces, mode='depth') 27 | images = images.detach().cpu().numpy() 28 | image = images[2] 29 | image = image != image.max() 30 | 31 | # load reference image by blender 32 | ref = imread(os.path.join(data_dir, 'teapot_blender.png')) 33 | ref = (ref.min(axis=-1) != 255).astype(np.float32) 34 | 35 | assert(np.allclose(ref, image)) 36 | 37 | def test_forward_case2(self): 38 | # load teapot 39 | vertices, faces, _ = utils.load_teapot_batch() 40 | 41 | # create renderer 42 | renderer = nr.Renderer(camera_mode='look_at') 43 | renderer.image_size = 256 44 | renderer.anti_aliasing = False 45 | 46 | images = renderer(vertices, faces, mode='depth') 47 | images = images.detach().cpu().numpy() 48 | image = images[2] 49 | image[image == image.max()] = image.min() 50 | image = (image - image.min()) / (image.max() - image.min()) 51 | 52 | ref = imread(os.path.join(data_dir, 'test_depth.png')).astype(np.float32) / 255. 53 | 54 | assert(np.allclose(image, ref, atol=1e-2)) 55 | 56 | def test_backward_case1(self): 57 | vertices = [ 58 | [-0.9, -0.9, 2.], 59 | [-0.8, 0.8, 1.], 60 | [0.8, 0.8, 0.5]] 61 | faces = [[0, 1, 2]] 62 | 63 | renderer = nr.Renderer(camera_mode='look_at') 64 | renderer.image_size = 64 65 | renderer.anti_aliasing = False 66 | renderer.perspective = False 67 | renderer.camera_mode = 'none' 68 | 69 | vertices = torch.from_numpy(np.array(vertices, np.float32)).cuda() 70 | faces = torch.from_numpy(np.array(faces, np.int32)).cuda() 71 | vertices, faces = utils.to_minibatch((vertices, faces)) 72 | vertices.requires_grad = True 73 | 74 | images = renderer(vertices, faces, mode='depth') 75 | loss = torch.sum((images[0, 15, 20] - 1)**2) 76 | loss.backward() 77 | grad = vertices.grad.clone() 78 | grad2 = torch.zeros_like(grad) 79 | 80 | for i in range(3): 81 | for j in range(3): 82 | eps = 1e-3 83 | vertices2 = vertices.clone() 84 | vertices2[i, j] += eps 85 | images = renderer.render_depth(vertices2, faces) 86 | loss2 = torch.sum((images[0, 15, 20] - 1)**2) 87 | grad2[i, j] = ((loss2 - loss) / eps).item() 88 | 89 | assert(torch.allclose(grad, grad2, atol=1e-3)) 90 | 91 | 92 | if __name__ == '__main__': 93 | unittest.main() 94 | -------------------------------------------------------------------------------- /tests/test_rasterize_silhouettes.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | import os 3 | 4 | import torch 5 | import numpy as np 6 | from skimage.io import imread 7 | 8 | import neural_renderer as nr 9 | import utils 10 | 11 | current_dir = os.path.dirname(os.path.realpath(__file__)) 12 | data_dir = os.path.join(current_dir, 'data') 13 | 14 | 15 | class TestRasterizeSilhouettes(unittest.TestCase): 16 | def test_case1(self): 17 | """Whether a silhouette by neural renderer matches that by Blender.""" 18 | 19 | # load teapot 20 | vertices, faces, _ = utils.load_teapot_batch() 21 | 22 | # create renderer 23 | renderer = nr.Renderer(camera_mode='look_at') 24 | renderer.image_size = 256 25 | renderer.anti_aliasing = False 26 | 27 | images = renderer(vertices, faces, mode='silhouettes') 28 | images = images.detach().cpu().numpy() 29 | image = images[2] 30 | 31 | # load reference image by blender 32 | ref = imread(os.path.join(data_dir, 'teapot_blender.png')) 33 | ref = (ref.min(-1) != 255).astype(np.float32) 34 | 35 | assert(np.allclose(ref, image)) 36 | 37 | def test_backward_case1(self): 38 | """Backward if non-zero gradient is out of a face.""" 39 | 40 | vertices = [ 41 | [0.8, 0.8, 1.], 42 | [0.0, -0.5, 1.], 43 | [0.2, -0.4, 1.]] 44 | faces = [[0, 1, 2]] 45 | pxi = 35 46 | pyi = 25 47 | grad_ref = [ 48 | [1.6725862, -0.26021874, 0.], 49 | [1.41986704, -1.64284933, 0.], 50 | [0., 0., 0.], 51 | ] 52 | 53 | renderer = nr.Renderer(camera_mode='look_at') 54 | renderer.image_size = 64 55 | renderer.anti_aliasing = False 56 | renderer.perspective = False 57 | 58 | vertices = torch.from_numpy(np.array(vertices, np.float32)).cuda() 59 | faces = torch.from_numpy(np.array(faces, np.int32)).cuda() 60 | grad_ref = torch.from_numpy(np.array(grad_ref, np.float32)).cuda() 61 | vertices, faces, grad_ref = utils.to_minibatch((vertices, faces, grad_ref)) 62 | vertices.requires_grad = True 63 | images = renderer(vertices, faces, mode='silhouettes') 64 | loss = torch.sum(torch.abs(images[:, pyi, pxi] - 1)) 65 | loss.backward() 66 | 67 | assert(torch.allclose(vertices.grad, grad_ref, rtol=1e-2)) 68 | 69 | def test_backward_case2(self): 70 | """Backward if non-zero gradient is on a face.""" 71 | 72 | vertices = [ 73 | [0.8, 0.8, 1.], 74 | [-0.5, -0.8, 1.], 75 | [0.8, -0.8, 1.]] 76 | faces = [[0, 1, 2]] 77 | pyi = 40 78 | pxi = 50 79 | grad_ref = [ 80 | [0.98646867, 1.04628897, 0.], 81 | [-1.03415668, - 0.10403691, 0.], 82 | [3.00094461, - 1.55173182, 0.], 83 | ] 84 | 85 | renderer = nr.Renderer(camera_mode='look_at') 86 | renderer.image_size = 64 87 | renderer.anti_aliasing = False 88 | renderer.perspective = False 89 | 90 | vertices = torch.from_numpy(np.array(vertices, np.float32)).cuda() 91 | faces = torch.from_numpy(np.array(faces, np.int32)).cuda() 92 | grad_ref = torch.from_numpy(np.array(grad_ref, np.float32)).cuda() 93 | vertices, faces, grad_ref = utils.to_minibatch((vertices, faces, grad_ref)) 94 | vertices.requires_grad = True 95 | images = renderer(vertices, faces, mode='silhouettes') 96 | loss = torch.sum(torch.abs(images[:, pyi, pxi])) 97 | loss.backward() 98 | 99 | assert(torch.allclose(vertices.grad, grad_ref, rtol=1e-2)) 100 | 101 | 102 | if __name__ == '__main__': 103 | unittest.main() 104 | -------------------------------------------------------------------------------- /tests/test_renderer.py: -------------------------------------------------------------------------------- 1 | # TODO 2 | ''' 3 | Might have to do some refactoring because the tests for Renderer are included in the test_rasterize* unit tests 4 | ''' 5 | -------------------------------------------------------------------------------- /tests/test_save_obj.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | import os 3 | 4 | import torch 5 | import numpy as np 6 | 7 | import neural_renderer as nr 8 | 9 | current_dir = os.path.dirname(os.path.realpath(__file__)) 10 | data_dir = os.path.join(current_dir, 'data') 11 | 12 | class TestCore(unittest.TestCase): 13 | def test_save_obj(self): 14 | teapot = os.path.join(data_dir, 'teapot.obj') 15 | teapot2 = os.path.join(data_dir, 'teapot2.obj') 16 | vertices, faces = nr.load_obj(teapot) 17 | nr.save_obj(teapot2, vertices, faces) 18 | vertices2, faces2 = nr.load_obj(teapot2) 19 | os.remove(teapot2) 20 | assert torch.allclose(vertices, vertices2) 21 | assert torch.allclose(faces, faces2) 22 | 23 | def test_texture(self): 24 | pass 25 | 26 | if __name__ == '__main__': 27 | unittest.main() 28 | -------------------------------------------------------------------------------- /tests/test_vertices_to_faces.py: -------------------------------------------------------------------------------- 1 | # TODO 2 | -------------------------------------------------------------------------------- /tests/utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import torch 4 | 5 | import neural_renderer as nr 6 | 7 | current_dir = os.path.dirname(os.path.realpath(__file__)) 8 | data_dir = os.path.join(current_dir, 'data') 9 | 10 | 11 | def to_minibatch(data, batch_size=4, target_num=2): 12 | ret = [] 13 | for d in data: 14 | device = d.device 15 | d2 = torch.unsqueeze(torch.zeros_like(d), 0) 16 | r = [1 for _ in d2.shape] 17 | r[0] = batch_size 18 | d2 = torch.unsqueeze(torch.zeros_like(d), 0).repeat(*r).to(device) 19 | d2[target_num] = d 20 | ret.append(d2) 21 | return ret 22 | 23 | def load_teapot_batch(batch_size=4, target_num=2): 24 | vertices, faces = nr.load_obj(os.path.join(data_dir, 'teapot.obj')) 25 | textures = torch.ones((faces.shape[0], 4, 4, 4, 3), dtype=torch.float32) 26 | vertices, faces, textures = to_minibatch((vertices, faces, textures), batch_size, target_num) 27 | return vertices, faces, textures 28 | --------------------------------------------------------------------------------