Upload 11 files

lib/net/BasePIFuNet.py

@@ -0,0 +1,84 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import torch.nn as nn
+import pytorch_lightning as pl
+
+from .geometry import index, orthogonal, perspective
+
+
+class BasePIFuNet(pl.LightningModule):
+    def __init__(
+            self,
+            projection_mode='orthogonal',
+            error_term=nn.MSELoss(),
+    ):
+        """
+        :param projection_mode:
+        Either orthogonal or perspective.
+        It will call the corresponding function for projection.
+        :param error_term:
+        nn Loss between the predicted [B, Res, N] and the label [B, Res, N]
+        """
+        super(BasePIFuNet, self).__init__()
+        self.name = 'base'
+
+        self.error_term = error_term
+
+        self.index = index
+        self.projection = orthogonal if projection_mode == 'orthogonal' else perspective
+
+    def forward(self, points, images, calibs, transforms=None):
+        '''
+        :param points: [B, 3, N] world space coordinates of points
+        :param images: [B, C, H, W] input images
+        :param calibs: [B, 3, 4] calibration matrices for each image
+        :param transforms: Optional [B, 2, 3] image space coordinate transforms
+        :return: [B, Res, N] predictions for each point
+        '''
+        features = self.filter(images)
+        preds = self.query(features, points, calibs, transforms)
+        return preds
+
+    def filter(self, images):
+        '''
+        Filter the input images
+        store all intermediate features.
+        :param images: [B, C, H, W] input images
+        '''
+        return None
+
+    def query(self, features, points, calibs, transforms=None):
+        '''
+        Given 3D points, query the network predictions for each point.
+        Image features should be pre-computed before this call.
+        store all intermediate features.
+        query() function may behave differently during training/testing.
+        :param points: [B, 3, N] world space coordinates of points
+        :param calibs: [B, 3, 4] calibration matrices for each image
+        :param transforms: Optional [B, 2, 3] image space coordinate transforms
+        :param labels: Optional [B, Res, N] gt labeling
+        :return: [B, Res, N] predictions for each point
+        '''
+        return None
+
+    def get_error(self, preds, labels):
+        '''
+        Get the network loss from the last query
+        :return: loss term
+        '''
+        return self.error_term(preds, labels)

lib/net/FBNet.py

@@ -0,0 +1,388 @@
+'''
+Copyright (C) 2019 NVIDIA Corporation. Ting-Chun Wang, Ming-Yu Liu, Jun-Yan Zhu.
+BSD License. All rights reserved. 
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING ALL 
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR ANY PARTICULAR PURPOSE. 
+IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL 
+DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, 
+WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING 
+OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+'''
+import torch
+import torch.nn as nn
+import functools
+import numpy as np
+import pytorch_lightning as pl
+
+
+###############################################################################
+# Functions
+###############################################################################
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        m.weight.data.normal_(0.0, 0.02)
+    elif classname.find('BatchNorm2d') != -1:
+        m.weight.data.normal_(1.0, 0.02)
+        m.bias.data.fill_(0)
+
+
+def get_norm_layer(norm_type='instance'):
+    if norm_type == 'batch':
+        norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
+    elif norm_type == 'instance':
+        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False)
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' %
+                                  norm_type)
+    return norm_layer
+
+
+def define_G(input_nc,
+             output_nc,
+             ngf,
+             netG,
+             n_downsample_global=3,
+             n_blocks_global=9,
+             n_local_enhancers=1,
+             n_blocks_local=3,
+             norm='instance',
+             gpu_ids=[],
+             last_op=nn.Tanh()):
+    norm_layer = get_norm_layer(norm_type=norm)
+    if netG == 'global':
+        netG = GlobalGenerator(input_nc,
+                               output_nc,
+                               ngf,
+                               n_downsample_global,
+                               n_blocks_global,
+                               norm_layer,
+                               last_op=last_op)
+    elif netG == 'local':
+        netG = LocalEnhancer(input_nc, output_nc, ngf, n_downsample_global,
+                             n_blocks_global, n_local_enhancers,
+                             n_blocks_local, norm_layer)
+    elif netG == 'encoder':
+        netG = Encoder(input_nc, output_nc, ngf, n_downsample_global,
+                       norm_layer)
+    else:
+        raise ('generator not implemented!')
+    # print(netG)
+    if len(gpu_ids) > 0:
+        assert (torch.cuda.is_available())
+        device=torch.device(f"cuda:{gpu_ids[0]}")
+        netG = netG.to(device)
+    netG.apply(weights_init)
+    return netG
+
+
+def print_network(net):
+    if isinstance(net, list):
+        net = net[0]
+    num_params = 0
+    for param in net.parameters():
+        num_params += param.numel()
+    print(net)
+    print('Total number of parameters: %d' % num_params)
+
+
+##############################################################################
+# Generator
+##############################################################################
+class LocalEnhancer(pl.LightningModule):
+    def __init__(self,
+                 input_nc,
+                 output_nc,
+                 ngf=32,
+                 n_downsample_global=3,
+                 n_blocks_global=9,
+                 n_local_enhancers=1,
+                 n_blocks_local=3,
+                 norm_layer=nn.BatchNorm2d,
+                 padding_type='reflect'):
+        super(LocalEnhancer, self).__init__()
+        self.n_local_enhancers = n_local_enhancers
+
+        ###### global generator model #####
+        ngf_global = ngf * (2**n_local_enhancers)
+        model_global = GlobalGenerator(input_nc, output_nc, ngf_global,
+                                       n_downsample_global, n_blocks_global,
+                                       norm_layer).model
+        model_global = [model_global[i] for i in range(len(model_global) - 3)
+                        ]  # get rid of final convolution layers
+        self.model = nn.Sequential(*model_global)
+
+        ###### local enhancer layers #####
+        for n in range(1, n_local_enhancers + 1):
+            # downsample
+            ngf_global = ngf * (2**(n_local_enhancers - n))
+            model_downsample = [
+                nn.ReflectionPad2d(3),
+                nn.Conv2d(input_nc, ngf_global, kernel_size=7, padding=0),
+                norm_layer(ngf_global),
+                nn.ReLU(True),
+                nn.Conv2d(ngf_global,
+                          ngf_global * 2,
+                          kernel_size=3,
+                          stride=2,
+                          padding=1),
+                norm_layer(ngf_global * 2),
+                nn.ReLU(True)
+            ]
+            # residual blocks
+            model_upsample = []
+            for i in range(n_blocks_local):
+                model_upsample += [
+                    ResnetBlock(ngf_global * 2,
+                                padding_type=padding_type,
+                                norm_layer=norm_layer)
+                ]
+
+            # upsample
+            model_upsample += [
+                nn.ConvTranspose2d(ngf_global * 2,
+                                   ngf_global,
+                                   kernel_size=3,
+                                   stride=2,
+                                   padding=1,
+                                   output_padding=1),
+                norm_layer(ngf_global),
+                nn.ReLU(True)
+            ]
+
+            # final convolution
+            if n == n_local_enhancers:
+                model_upsample += [
+                    nn.ReflectionPad2d(3),
+                    nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0),
+                    nn.Tanh()
+                ]
+
+            setattr(self, 'model' + str(n) + '_1',
+                    nn.Sequential(*model_downsample))
+            setattr(self, 'model' + str(n) + '_2',
+                    nn.Sequential(*model_upsample))
+
+        self.downsample = nn.AvgPool2d(3,
+                                       stride=2,
+                                       padding=[1, 1],
+                                       count_include_pad=False)
+
+    def forward(self, input):
+        # create input pyramid
+        input_downsampled = [input]
+        for i in range(self.n_local_enhancers):
+            input_downsampled.append(self.downsample(input_downsampled[-1]))
+
+        # output at coarest level
+        output_prev = self.model(input_downsampled[-1])
+        # build up one layer at a time
+        for n_local_enhancers in range(1, self.n_local_enhancers + 1):
+            model_downsample = getattr(self,
+                                       'model' + str(n_local_enhancers) + '_1')
+            model_upsample = getattr(self,
+                                     'model' + str(n_local_enhancers) + '_2')
+            input_i = input_downsampled[self.n_local_enhancers -
+                                        n_local_enhancers]
+            output_prev = model_upsample(
+                model_downsample(input_i) + output_prev)
+        return output_prev
+
+
+class GlobalGenerator(pl.LightningModule):
+    def __init__(self,
+                 input_nc,
+                 output_nc,
+                 ngf=64,
+                 n_downsampling=3,
+                 n_blocks=9,
+                 norm_layer=nn.BatchNorm2d,
+                 padding_type='reflect',
+                 last_op=nn.Tanh()):
+        assert (n_blocks >= 0)
+        super(GlobalGenerator, self).__init__()
+        activation = nn.ReLU(True)
+
+        model = [
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0),
+            norm_layer(ngf), activation
+        ]
+        # downsample
+        for i in range(n_downsampling):
+            mult = 2**i
+            model += [
+                nn.Conv2d(ngf * mult,
+                          ngf * mult * 2,
+                          kernel_size=3,
+                          stride=2,
+                          padding=1),
+                norm_layer(ngf * mult * 2), activation
+            ]
+
+        # resnet blocks
+        mult = 2**n_downsampling
+        for i in range(n_blocks):
+            model += [
+                ResnetBlock(ngf * mult,
+                            padding_type=padding_type,
+                            activation=activation,
+                            norm_layer=norm_layer)
+            ]
+
+        # upsample
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            model += [
+                nn.ConvTranspose2d(ngf * mult,
+                                   int(ngf * mult / 2),
+                                   kernel_size=3,
+                                   stride=2,
+                                   padding=1,
+                                   output_padding=1),
+                norm_layer(int(ngf * mult / 2)), activation
+            ]
+        model += [
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)
+        ]
+        if last_op is not None:
+            model += [last_op]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input):
+        return self.model(input)
+
+
+# Define a resnet block
+class ResnetBlock(pl.LightningModule):
+    def __init__(self,
+                 dim,
+                 padding_type,
+                 norm_layer,
+                 activation=nn.ReLU(True),
+                 use_dropout=False):
+        super(ResnetBlock, self).__init__()
+        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer,
+                                                activation, use_dropout)
+
+    def build_conv_block(self, dim, padding_type, norm_layer, activation,
+                         use_dropout):
+        conv_block = []
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' %
+                                      padding_type)
+
+        conv_block += [
+            nn.Conv2d(dim, dim, kernel_size=3, padding=p),
+            norm_layer(dim), activation
+        ]
+        if use_dropout:
+            conv_block += [nn.Dropout(0.5)]
+
+        p = 0
+        if padding_type == 'reflect':
+            conv_block += [nn.ReflectionPad2d(1)]
+        elif padding_type == 'replicate':
+            conv_block += [nn.ReplicationPad2d(1)]
+        elif padding_type == 'zero':
+            p = 1
+        else:
+            raise NotImplementedError('padding [%s] is not implemented' %
+                                      padding_type)
+        conv_block += [
+            nn.Conv2d(dim, dim, kernel_size=3, padding=p),
+            norm_layer(dim)
+        ]
+
+        return nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        out = x + self.conv_block(x)
+        return out
+
+
+class Encoder(pl.LightningModule):
+    def __init__(self,
+                 input_nc,
+                 output_nc,
+                 ngf=32,
+                 n_downsampling=4,
+                 norm_layer=nn.BatchNorm2d):
+        super(Encoder, self).__init__()
+        self.output_nc = output_nc
+
+        model = [
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0),
+            norm_layer(ngf),
+            nn.ReLU(True)
+        ]
+        # downsample
+        for i in range(n_downsampling):
+            mult = 2**i
+            model += [
+                nn.Conv2d(ngf * mult,
+                          ngf * mult * 2,
+                          kernel_size=3,
+                          stride=2,
+                          padding=1),
+                norm_layer(ngf * mult * 2),
+                nn.ReLU(True)
+            ]
+
+        # upsample
+        for i in range(n_downsampling):
+            mult = 2**(n_downsampling - i)
+            model += [
+                nn.ConvTranspose2d(ngf * mult,
+                                   int(ngf * mult / 2),
+                                   kernel_size=3,
+                                   stride=2,
+                                   padding=1,
+                                   output_padding=1),
+                norm_layer(int(ngf * mult / 2)),
+                nn.ReLU(True)
+            ]
+
+        model += [
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0),
+            nn.Tanh()
+        ]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, input, inst):
+        outputs = self.model(input)
+
+        # instance-wise average pooling
+        outputs_mean = outputs.clone()
+        inst_list = np.unique(inst.cpu().numpy().astype(int))
+        for i in inst_list:
+            for b in range(input.size()[0]):
+                indices = (inst[b:b + 1] == int(i)).nonzero()  # n x 4
+                for j in range(self.output_nc):
+                    output_ins = outputs[indices[:, 0] + b, indices[:, 1] + j,
+                                         indices[:, 2], indices[:, 3]]
+                    mean_feat = torch.mean(output_ins).expand_as(output_ins)
+                    outputs_mean[indices[:, 0] + b, indices[:, 1] + j,
+                                 indices[:, 2], indices[:, 3]] = mean_feat
+        return outputs_mean

lib/net/HGFilters.py

@@ -0,0 +1,197 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+from lib.net.net_util import *
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class HourGlass(nn.Module):
+    def __init__(self, num_modules, depth, num_features, opt):
+        super(HourGlass, self).__init__()
+        self.num_modules = num_modules
+        self.depth = depth
+        self.features = num_features
+        self.opt = opt
+
+        self._generate_network(self.depth)
+
+    def _generate_network(self, level):
+        self.add_module('b1_' + str(level),
+                        ConvBlock(self.features, self.features, self.opt))
+
+        self.add_module('b2_' + str(level),
+                        ConvBlock(self.features, self.features, self.opt))
+
+        if level > 1:
+            self._generate_network(level - 1)
+        else:
+            self.add_module('b2_plus_' + str(level),
+                            ConvBlock(self.features, self.features, self.opt))
+
+        self.add_module('b3_' + str(level),
+                        ConvBlock(self.features, self.features, self.opt))
+
+    def _forward(self, level, inp):
+        # Upper branch
+        up1 = inp
+        up1 = self._modules['b1_' + str(level)](up1)
+
+        # Lower branch
+        low1 = F.avg_pool2d(inp, 2, stride=2)
+        low1 = self._modules['b2_' + str(level)](low1)
+
+        if level > 1:
+            low2 = self._forward(level - 1, low1)
+        else:
+            low2 = low1
+            low2 = self._modules['b2_plus_' + str(level)](low2)
+
+        low3 = low2
+        low3 = self._modules['b3_' + str(level)](low3)
+
+        # NOTE: for newer PyTorch (1.3~), it seems that training results are degraded due to implementation diff in F.grid_sample
+        # if the pretrained model behaves weirdly, switch with the commented line.
+        # NOTE: I also found that "bicubic" works better.
+        up2 = F.interpolate(low3,
+                            scale_factor=2,
+                            mode='bicubic',
+                            align_corners=True)
+        # up2 = F.interpolate(low3, scale_factor=2, mode='nearest)
+
+        return up1 + up2
+
+    def forward(self, x):
+        return self._forward(self.depth, x)
+
+
+class HGFilter(nn.Module):
+    def __init__(self, opt, num_modules, in_dim):
+        super(HGFilter, self).__init__()
+        self.num_modules = num_modules
+
+        self.opt = opt
+        [k, s, d, p] = self.opt.conv1
+
+        # self.conv1 = nn.Conv2d(in_dim, 64, kernel_size=7, stride=2, padding=3)
+        self.conv1 = nn.Conv2d(in_dim,
+                               64,
+                               kernel_size=k,
+                               stride=s,
+                               dilation=d,
+                               padding=p)
+
+        if self.opt.norm == 'batch':
+            self.bn1 = nn.BatchNorm2d(64)
+        elif self.opt.norm == 'group':
+            self.bn1 = nn.GroupNorm(32, 64)
+
+        if self.opt.hg_down == 'conv64':
+            self.conv2 = ConvBlock(64, 64, self.opt)
+            self.down_conv2 = nn.Conv2d(64,
+                                        128,
+                                        kernel_size=3,
+                                        stride=2,
+                                        padding=1)
+        elif self.opt.hg_down == 'conv128':
+            self.conv2 = ConvBlock(64, 128, self.opt)
+            self.down_conv2 = nn.Conv2d(128,
+                                        128,
+                                        kernel_size=3,
+                                        stride=2,
+                                        padding=1)
+        elif self.opt.hg_down == 'ave_pool':
+            self.conv2 = ConvBlock(64, 128, self.opt)
+        else:
+            raise NameError('Unknown Fan Filter setting!')
+
+        self.conv3 = ConvBlock(128, 128, self.opt)
+        self.conv4 = ConvBlock(128, 256, self.opt)
+
+        # Stacking part
+        for hg_module in range(self.num_modules):
+            self.add_module('m' + str(hg_module),
+                            HourGlass(1, opt.num_hourglass, 256, self.opt))
+
+            self.add_module('top_m_' + str(hg_module),
+                            ConvBlock(256, 256, self.opt))
+            self.add_module(
+                'conv_last' + str(hg_module),
+                nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
+            if self.opt.norm == 'batch':
+                self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(256))
+            elif self.opt.norm == 'group':
+                self.add_module('bn_end' + str(hg_module),
+                                nn.GroupNorm(32, 256))
+
+            self.add_module(
+                'l' + str(hg_module),
+                nn.Conv2d(256,
+                          opt.hourglass_dim,
+                          kernel_size=1,
+                          stride=1,
+                          padding=0))
+
+            if hg_module < self.num_modules - 1:
+                self.add_module(
+                    'bl' + str(hg_module),
+                    nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
+                self.add_module(
+                    'al' + str(hg_module),
+                    nn.Conv2d(opt.hourglass_dim,
+                              256,
+                              kernel_size=1,
+                              stride=1,
+                              padding=0))
+
+    def forward(self, x):
+        x = F.relu(self.bn1(self.conv1(x)), True)
+        tmpx = x
+        if self.opt.hg_down == 'ave_pool':
+            x = F.avg_pool2d(self.conv2(x), 2, stride=2)
+        elif self.opt.hg_down in ['conv64', 'conv128']:
+            x = self.conv2(x)
+            x = self.down_conv2(x)
+        else:
+            raise NameError('Unknown Fan Filter setting!')
+
+        x = self.conv3(x)
+        x = self.conv4(x)
+
+        previous = x
+
+        outputs = []
+        for i in range(self.num_modules):
+            hg = self._modules['m' + str(i)](previous)
+
+            ll = hg
+            ll = self._modules['top_m_' + str(i)](ll)
+
+            ll = F.relu(
+                self._modules['bn_end' + str(i)](
+                    self._modules['conv_last' + str(i)](ll)), True)
+
+            # Predict heatmaps
+            tmp_out = self._modules['l' + str(i)](ll)
+            outputs.append(tmp_out)
+
+            if i < self.num_modules - 1:
+                ll = self._modules['bl' + str(i)](ll)
+                tmp_out_ = self._modules['al' + str(i)](tmp_out)
+                previous = previous + ll + tmp_out_
+
+        return outputs

lib/net/HGPIFuNet.py

@@ -0,0 +1,403 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+from lib.net.voxelize import Voxelization
+from lib.dataset.mesh_util import cal_sdf_batch, feat_select, read_smpl_constants
+from lib.net.NormalNet import NormalNet
+from lib.net.MLP import MLP
+from lib.dataset.mesh_util import SMPLX
+from lib.net.VE import VolumeEncoder
+from lib.net.HGFilters import *
+from termcolor import colored
+from lib.net.BasePIFuNet import BasePIFuNet
+import torch.nn as nn
+import torch
+
+
+maskout = False
+
+
+class HGPIFuNet(BasePIFuNet):
+    '''
+    HG PIFu network uses Hourglass stacks as the image filter.
+    It does the following:
+        1. Compute image feature stacks and store it in self.im_feat_list
+            self.im_feat_list[-1] is the last stack (output stack)
+        2. Calculate calibration
+        3. If training, it index on every intermediate stacks,
+            If testing, it index on the last stack.
+        4. Classification.
+        5. During training, error is calculated on all stacks.
+    '''
+
+    def __init__(self,
+                 cfg,
+                 projection_mode='orthogonal',
+                 error_term=nn.MSELoss()):
+
+        super(HGPIFuNet, self).__init__(projection_mode=projection_mode,
+                                        error_term=error_term)
+
+        self.l1_loss = nn.SmoothL1Loss()
+        self.opt = cfg.net
+        self.root = cfg.root
+        self.overfit = cfg.overfit
+
+        channels_IF = self.opt.mlp_dim
+
+        self.use_filter = self.opt.use_filter
+        self.prior_type = self.opt.prior_type
+        self.smpl_feats = self.opt.smpl_feats
+
+        self.smpl_dim = self.opt.smpl_dim
+        self.voxel_dim = self.opt.voxel_dim
+        self.hourglass_dim = self.opt.hourglass_dim
+        self.sdf_clip = cfg.sdf_clip / 100.0
+
+        self.in_geo = [item[0] for item in self.opt.in_geo]
+        self.in_nml = [item[0] for item in self.opt.in_nml]
+
+        self.in_geo_dim = sum([item[1] for item in self.opt.in_geo])
+        self.in_nml_dim = sum([item[1] for item in self.opt.in_nml])
+
+        self.in_total = self.in_geo + self.in_nml
+        self.smpl_feat_dict = None
+        self.smplx_data = SMPLX()
+
+        if self.prior_type == 'icon':
+            if 'image' in self.in_geo:
+                self.channels_filter = [[0, 1, 2, 3, 4, 5], [0, 1, 2, 6, 7, 8]]
+            else:
+                self.channels_filter = [[0, 1, 2], [3, 4, 5]]
+
+        else:
+            if 'image' in self.in_geo:
+                self.channels_filter = [[0, 1, 2, 3, 4, 5, 6, 7, 8]]
+            else:
+                self.channels_filter = [[0, 1, 2, 3, 4, 5]]
+
+        channels_IF[0] = self.hourglass_dim if self.use_filter else len(
+            self.channels_filter[0])
+
+        if self.prior_type == 'icon' and 'vis' not in self.smpl_feats:
+            if self.use_filter:
+                channels_IF[0] += self.hourglass_dim
+            else:
+                channels_IF[0] += len(self.channels_filter[0])
+
+        if self.prior_type == 'icon':
+            channels_IF[0] += self.smpl_dim
+        elif self.prior_type == 'pamir':
+            channels_IF[0] += self.voxel_dim
+            smpl_vertex_code, smpl_face_code, smpl_faces, smpl_tetras = read_smpl_constants(
+                self.smplx_data.tedra_dir)
+            self.voxelization = Voxelization(
+                smpl_vertex_code,
+                smpl_face_code,
+                smpl_faces,
+                smpl_tetras,
+                volume_res=128,
+                sigma=0.05,
+                smooth_kernel_size=7,
+                batch_size=cfg.batch_size,
+                device=torch.device(f"cuda:{cfg.gpus[0]}"))
+            self.ve = VolumeEncoder(3, self.voxel_dim, self.opt.num_stack)
+        else:
+            channels_IF[0] += 1
+
+        self.icon_keys = ["smpl_verts", "smpl_faces", "smpl_vis", "smpl_cmap"]
+        self.pamir_keys = [
+            "voxel_verts", "voxel_faces", "pad_v_num", "pad_f_num"
+        ]
+
+        self.if_regressor = MLP(
+            filter_channels=channels_IF,
+            name='if',
+            res_layers=self.opt.res_layers,
+            norm=self.opt.norm_mlp,
+            last_op=nn.Sigmoid() if not cfg.test_mode else None)
+
+        # network
+        if self.use_filter:
+            if self.opt.gtype == "HGPIFuNet":
+                self.F_filter = HGFilter(self.opt, self.opt.num_stack,
+                                         len(self.channels_filter[0]))
+            else:
+                print(
+                    colored(f"Backbone {self.opt.gtype} is unimplemented",
+                            'green'))
+
+        summary_log = f"{self.prior_type.upper()}:\n" + \
+            f"w/ Global Image Encoder: {self.use_filter}\n" + \
+            f"Image Features used by MLP: {self.in_geo}\n"
+
+        if self.prior_type == "icon":
+            summary_log += f"Geometry Features used by MLP: {self.smpl_feats}\n"
+            summary_log += f"Dim of Image Features (local): 6\n"
+            summary_log += f"Dim of Geometry Features (ICON): {self.smpl_dim}\n"
+        elif self.prior_type == "pamir":
+            summary_log += f"Dim of Image Features (global): {self.hourglass_dim}\n"
+            summary_log += f"Dim of Geometry Features (PaMIR): {self.voxel_dim}\n"
+        else:
+            summary_log += f"Dim of Image Features (global): {self.hourglass_dim}\n"
+            summary_log += f"Dim of Geometry Features (PIFu): 1 (z-value)\n"
+
+        summary_log += f"Dim of MLP's first layer: {channels_IF[0]}\n"
+
+        print(colored(summary_log, "yellow"))
+
+        self.normal_filter = NormalNet(cfg)
+        init_net(self)
+
+    def get_normal(self, in_tensor_dict):
+
+        # insert normal features
+        if (not self.training) and (not self.overfit):
+            # print(colored("infer normal","blue"))
+            with torch.no_grad():
+                feat_lst = []
+                if "image" in self.in_geo:
+                    feat_lst.append(
+                        in_tensor_dict['image'])  # [1, 3, 512, 512]
+                if 'normal_F' in self.in_geo and 'normal_B' in self.in_geo:
+                    if 'normal_F' not in in_tensor_dict.keys(
+                    ) or 'normal_B' not in in_tensor_dict.keys():
+                        (nmlF, nmlB) = self.normal_filter(in_tensor_dict)
+                    else:
+                        nmlF = in_tensor_dict['normal_F']
+                        nmlB = in_tensor_dict['normal_B']
+                    feat_lst.append(nmlF)  # [1, 3, 512, 512]
+                    feat_lst.append(nmlB)  # [1, 3, 512, 512]
+            in_filter = torch.cat(feat_lst, dim=1)
+
+        else:
+            in_filter = torch.cat([in_tensor_dict[key] for key in self.in_geo],
+                                  dim=1)
+
+        return in_filter
+
+    def get_mask(self, in_filter, size=128):
+
+        mask = F.interpolate(in_filter[:, self.channels_filter[0]],
+                             size=(size, size),
+                             mode="bilinear",
+                             align_corners=True).abs().sum(dim=1,
+                                                           keepdim=True) != 0.0
+
+        return mask
+
+    def filter(self, in_tensor_dict, return_inter=False):
+        '''
+        Filter the input images
+        store all intermediate features.
+        :param images: [B, C, H, W] input images
+        '''
+
+        in_filter = self.get_normal(in_tensor_dict)
+
+        features_G = []
+
+        if self.prior_type == 'icon':
+            if self.use_filter:
+                features_F = self.F_filter(in_filter[:,
+                                                     self.channels_filter[0]]
+                                           )  # [(B,hg_dim,128,128) * 4]
+                features_B = self.F_filter(in_filter[:,
+                                                     self.channels_filter[1]]
+                                           )  # [(B,hg_dim,128,128) * 4]
+            else:
+                features_F = [in_filter[:, self.channels_filter[0]]]
+                features_B = [in_filter[:, self.channels_filter[1]]]
+            for idx in range(len(features_F)):
+                features_G.append(
+                    torch.cat([features_F[idx], features_B[idx]], dim=1))
+        else:
+            if self.use_filter:
+                features_G = self.F_filter(in_filter[:,
+                                                     self.channels_filter[0]])
+            else:
+                features_G = [in_filter[:, self.channels_filter[0]]]
+
+        if self.prior_type == 'icon':
+            self.smpl_feat_dict = {
+                k: in_tensor_dict[k]
+                for k in self.icon_keys
+            }
+        elif self.prior_type == "pamir":
+            self.smpl_feat_dict = {
+                k: in_tensor_dict[k]
+                for k in self.pamir_keys
+            }
+        else:
+            pass
+            # print(colored("use z rather than icon or pamir", "green"))
+
+        # If it is not in training, only produce the last im_feat
+        if not self.training:
+            features_out = [features_G[-1]]
+        else:
+            features_out = features_G
+
+        if maskout:
+            features_out_mask = []
+            for feat in features_out:
+                features_out_mask.append(
+                    feat * self.get_mask(in_filter, size=feat.shape[2]))
+            features_out = features_out_mask
+
+        if return_inter:
+            return features_out, in_filter
+        else:
+            return features_out
+
+    def query(self, features, points, calibs, transforms=None, regressor=None):
+
+        xyz = self.projection(points, calibs, transforms)
+
+        (xy, z) = xyz.split([2, 1], dim=1)
+
+        in_cube = (xyz > -1.0) & (xyz < 1.0)
+        in_cube = in_cube.all(dim=1, keepdim=True).detach().float()
+
+        preds_list = []
+
+        if self.prior_type == 'icon':
+
+            # smpl_verts [B, N_vert, 3]
+            # smpl_faces [B, N_face, 3]
+            # points [B, 3, N]
+
+            smpl_sdf, smpl_norm, smpl_cmap, smpl_vis = cal_sdf_batch(
+                self.smpl_feat_dict['smpl_verts'],
+                self.smpl_feat_dict['smpl_faces'],
+                self.smpl_feat_dict['smpl_cmap'],
+                self.smpl_feat_dict['smpl_vis'],
+                xyz.permute(0, 2, 1).contiguous())
+
+            # smpl_sdf [B, N, 1]
+            # smpl_norm [B, N, 3]
+            # smpl_cmap [B, N, 3]
+            # smpl_vis [B, N, 1]
+
+            feat_lst = [smpl_sdf]
+            if 'cmap' in self.smpl_feats:
+                feat_lst.append(smpl_cmap)
+            if 'norm' in self.smpl_feats:
+                feat_lst.append(smpl_norm)
+            if 'vis' in self.smpl_feats:
+                feat_lst.append(smpl_vis)
+
+            smpl_feat = torch.cat(feat_lst, dim=2).permute(0, 2, 1)
+            vol_feats = features
+
+        elif self.prior_type == "pamir":
+
+            voxel_verts = self.smpl_feat_dict[
+                'voxel_verts'][:, :-self.smpl_feat_dict['pad_v_num'][0], :]
+            voxel_faces = self.smpl_feat_dict[
+                'voxel_faces'][:, :-self.smpl_feat_dict['pad_f_num'][0], :]
+
+            self.voxelization.update_param(
+                batch_size=voxel_faces.shape[0],
+                smpl_tetra=voxel_faces[0].detach().cpu().numpy())
+            vol = self.voxelization(voxel_verts)  # vol ~ [0,1]
+            vol_feats = self.ve(vol, intermediate_output=self.training)
+        else:
+            vol_feats = features
+
+        for im_feat, vol_feat in zip(features, vol_feats):
+
+            # [B, Feat_i + z, N]
+            # normal feature choice by smpl_vis
+            if self.prior_type == 'icon':
+                if 'vis' in self.smpl_feats:
+                    point_local_feat = feat_select(self.index(im_feat, xy),
+                                                   smpl_feat[:, [-1], :])
+                    if maskout:
+                        normal_mask = torch.tile(
+                            point_local_feat.sum(dim=1, keepdims=True) == 0.0,
+                            (1, smpl_feat.shape[1], 1))
+                        normal_mask[:, 1:, :] = False
+                        smpl_feat[normal_mask] = -1.0
+                    point_feat_list = [point_local_feat, smpl_feat[:, :-1, :]]
+                else:
+                    point_local_feat = self.index(im_feat, xy)
+                    point_feat_list = [point_local_feat, smpl_feat[:, :, :]]
+
+            elif self.prior_type == 'pamir':
+                # im_feat [B, hg_dim, 128, 128]
+                # vol_feat [B, vol_dim, 32, 32, 32]
+                point_feat_list = [
+                    self.index(im_feat, xy),
+                    self.index(vol_feat, xyz)
+                ]
+
+            else:
+                point_feat_list = [self.index(im_feat, xy), z]
+
+            point_feat = torch.cat(point_feat_list, 1)
+
+            # out of image plane is always set to 0
+            preds = regressor(point_feat)
+            preds = in_cube * preds
+
+            preds_list.append(preds)
+
+        return preds_list
+
+    def get_error(self, preds_if_list, labels):
+        """calcaulate error
+
+        Args:
+            preds_list (list): list of torch.tensor(B, 3, N)
+            labels (torch.tensor): (B, N_knn, N)
+
+        Returns:
+            torch.tensor: error
+        """
+        error_if = 0
+
+        for pred_id in range(len(preds_if_list)):
+            pred_if = preds_if_list[pred_id]
+            error_if += self.error_term(pred_if, labels)
+
+        error_if /= len(preds_if_list)
+
+        return error_if
+
+    def forward(self, in_tensor_dict):
+        """
+        sample_tensor [B, 3, N]
+        calib_tensor [B, 4, 4]
+        label_tensor [B, 1, N]
+        smpl_feat_tensor [B, 59, N]
+        """
+
+        sample_tensor = in_tensor_dict['sample']
+        calib_tensor = in_tensor_dict['calib']
+        label_tensor = in_tensor_dict['label']
+
+        in_feat = self.filter(in_tensor_dict)
+
+        preds_if_list = self.query(in_feat,
+                                   sample_tensor,
+                                   calib_tensor,
+                                   regressor=self.if_regressor)
+
+        error = self.get_error(preds_if_list, label_tensor)
+
+        return preds_if_list[-1], error

lib/net/MLP.py

@@ -0,0 +1,72 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+
+import torch
+import torch.nn as nn
+import pytorch_lightning as pl
+
+
+class MLP(pl.LightningModule):
+    def __init__(self,
+                 filter_channels,
+                 name=None,
+                 res_layers=[],
+                 norm='group',
+                 last_op=None):
+
+        super(MLP, self).__init__()
+
+        self.filters = nn.ModuleList()
+        self.norms = nn.ModuleList()
+        self.res_layers = res_layers
+        self.norm = norm
+        self.last_op = last_op
+        self.name = name
+        self.activate = nn.LeakyReLU(inplace=True)
+
+        for l in range(0, len(filter_channels) - 1):
+            if l in self.res_layers:
+                self.filters.append(
+                    nn.Conv1d(filter_channels[l] + filter_channels[0],
+                              filter_channels[l + 1], 1))
+            else:
+                self.filters.append(
+                    nn.Conv1d(filter_channels[l], filter_channels[l + 1], 1))
+
+            if l != len(filter_channels) - 2:
+                if norm == 'group':
+                    self.norms.append(nn.GroupNorm(32, filter_channels[l + 1]))
+                elif norm == 'batch':
+                    self.norms.append(nn.BatchNorm1d(filter_channels[l + 1]))
+                elif norm == 'instance':
+                    self.norms.append(nn.InstanceNorm1d(filter_channels[l +
+                                                                        1]))
+                elif norm == 'weight':
+                    self.filters[l] = nn.utils.weight_norm(self.filters[l],
+                                                           name='weight')
+                    # print(self.filters[l].weight_g.size(),
+                    #       self.filters[l].weight_v.size())
+
+    def forward(self, feature):
+        '''
+        feature may include multiple view inputs
+        args:
+            feature: [B, C_in, N]
+        return:
+            [B, C_out, N] prediction
+        '''
+        y = feature
+        tmpy = feature
+
+        for i, f in enumerate(self.filters):
+
+            y = f(y if i not in self.res_layers else torch.cat([y, tmpy], 1))
+            if i != len(self.filters) - 1:
+                if self.norm not in ['batch', 'group', 'instance']:
+                    y = self.activate(y)
+                else:
+                    y = self.activate(self.norms[i](y))
+
+        if self.last_op is not None:
+            y = self.last_op(y)
+
+        return y

lib/net/NormalNet.py

@@ -0,0 +1,122 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+from lib.net.FBNet import define_G
+from lib.net.net_util import init_net, VGGLoss
+from lib.net.HGFilters import *
+from lib.net.BasePIFuNet import BasePIFuNet
+import torch
+import torch.nn as nn
+
+
+class NormalNet(BasePIFuNet):
+    '''
+    HG PIFu network uses Hourglass stacks as the image filter.
+    It does the following:
+        1. Compute image feature stacks and store it in self.im_feat_list
+            self.im_feat_list[-1] is the last stack (output stack)
+        2. Calculate calibration
+        3. If training, it index on every intermediate stacks,
+            If testing, it index on the last stack.
+        4. Classification.
+        5. During training, error is calculated on all stacks.
+    '''
+
+    def __init__(self, cfg, error_term=nn.SmoothL1Loss()):
+
+        super(NormalNet, self).__init__(error_term=error_term)
+
+        self.l1_loss = nn.SmoothL1Loss()
+
+        self.opt = cfg.net
+
+        if self.training:
+            self.vgg_loss = [VGGLoss()]
+
+        self.in_nmlF = [
+            item[0] for item in self.opt.in_nml
+            if '_F' in item[0] or item[0] == 'image'
+        ]
+        self.in_nmlB = [
+            item[0] for item in self.opt.in_nml
+            if '_B' in item[0] or item[0] == 'image'
+        ]
+        self.in_nmlF_dim = sum([
+            item[1] for item in self.opt.in_nml
+            if '_F' in item[0] or item[0] == 'image'
+        ])
+        self.in_nmlB_dim = sum([
+            item[1] for item in self.opt.in_nml
+            if '_B' in item[0] or item[0] == 'image'
+        ])
+
+        self.netF = define_G(self.in_nmlF_dim, 3, 64, "global", 4, 9, 1, 3,
+                             "instance")
+        self.netB = define_G(self.in_nmlB_dim, 3, 64, "global", 4, 9, 1, 3,
+                             "instance")
+
+        init_net(self)
+
+    def forward(self, in_tensor):
+
+        inF_list = []
+        inB_list = []
+
+        for name in self.in_nmlF:
+            inF_list.append(in_tensor[name])
+        for name in self.in_nmlB:
+            inB_list.append(in_tensor[name])
+
+        nmlF = self.netF(torch.cat(inF_list, dim=1))
+        nmlB = self.netB(torch.cat(inB_list, dim=1))
+        
+        # ||normal|| == 1
+        nmlF /= torch.norm(nmlF, dim=1)
+        nmlB /= torch.norm(nmlB, dim=1)
+        
+        # output: float_arr [-1,1] with [B, C, H, W]
+
+        mask = (in_tensor['image'].abs().sum(dim=1, keepdim=True) !=
+                0.0).detach().float()
+
+        nmlF = nmlF * mask
+        nmlB = nmlB * mask
+
+        return nmlF, nmlB
+
+    def get_norm_error(self, prd_F, prd_B, tgt):
+        """calculate normal loss
+
+        Args:
+            pred (torch.tensor): [B, 6, 512, 512]
+            tagt (torch.tensor): [B, 6, 512, 512]
+        """
+
+        tgt_F, tgt_B = tgt['normal_F'], tgt['normal_B']
+
+        l1_F_loss = self.l1_loss(prd_F, tgt_F)
+        l1_B_loss = self.l1_loss(prd_B, tgt_B)
+
+        with torch.no_grad():
+            vgg_F_loss = self.vgg_loss[0](prd_F, tgt_F)
+            vgg_B_loss = self.vgg_loss[0](prd_B, tgt_B)
+
+        total_loss = [
+            5.0 * l1_F_loss + vgg_F_loss, 5.0 * l1_B_loss + vgg_B_loss
+        ]
+
+        return total_loss

lib/net/VE.py

@@ -0,0 +1,183 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+
+import torch.nn as nn
+import pytorch_lightning as pl
+
+
+class BaseNetwork(pl.LightningModule):
+    def __init__(self):
+        super(BaseNetwork, self).__init__()
+
+    def init_weights(self, init_type='xavier', gain=0.02):
+        '''
+        initializes network's weights
+        init_type: normal | xavier | kaiming | orthogonal
+        https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/9451e70673400885567d08a9e97ade2524c700d0/models/networks.py#L39
+        '''
+        def init_func(m):
+            classname = m.__class__.__name__
+            if hasattr(m, 'weight') and (classname.find('Conv') != -1
+                                         or classname.find('Linear') != -1):
+                if init_type == 'normal':
+                    nn.init.normal_(m.weight.data, 0.0, gain)
+                elif init_type == 'xavier':
+                    nn.init.xavier_normal_(m.weight.data, gain=gain)
+                elif init_type == 'kaiming':
+                    nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+                elif init_type == 'orthogonal':
+                    nn.init.orthogonal_(m.weight.data, gain=gain)
+
+                if hasattr(m, 'bias') and m.bias is not None:
+                    nn.init.constant_(m.bias.data, 0.0)
+
+            elif classname.find('BatchNorm2d') != -1:
+                nn.init.normal_(m.weight.data, 1.0, gain)
+                nn.init.constant_(m.bias.data, 0.0)
+
+        self.apply(init_func)
+
+
+class Residual3D(BaseNetwork):
+    def __init__(self, numIn, numOut):
+        super(Residual3D, self).__init__()
+        self.numIn = numIn
+        self.numOut = numOut
+        self.with_bias = True
+        # self.bn = nn.GroupNorm(4, self.numIn)
+        self.bn = nn.BatchNorm3d(self.numIn)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv1 = nn.Conv3d(self.numIn,
+                               self.numOut,
+                               bias=self.with_bias,
+                               kernel_size=3,
+                               stride=1,
+                               padding=2,
+                               dilation=2)
+        # self.bn1 = nn.GroupNorm(4, self.numOut)
+        self.bn1 = nn.BatchNorm3d(self.numOut)
+        self.conv2 = nn.Conv3d(self.numOut,
+                               self.numOut,
+                               bias=self.with_bias,
+                               kernel_size=3,
+                               stride=1,
+                               padding=1)
+        # self.bn2 = nn.GroupNorm(4, self.numOut)
+        self.bn2 = nn.BatchNorm3d(self.numOut)
+        self.conv3 = nn.Conv3d(self.numOut,
+                               self.numOut,
+                               bias=self.with_bias,
+                               kernel_size=3,
+                               stride=1,
+                               padding=1)
+
+        if self.numIn != self.numOut:
+            self.conv4 = nn.Conv3d(self.numIn,
+                                   self.numOut,
+                                   bias=self.with_bias,
+                                   kernel_size=1)
+        self.init_weights()
+
+    def forward(self, x):
+        residual = x
+        # out = self.bn(x)
+        # out = self.relu(out)
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+        # out = self.conv3(out)
+        # out = self.relu(out)
+
+        if self.numIn != self.numOut:
+            residual = self.conv4(x)
+
+        return out + residual
+
+
+class VolumeEncoder(BaseNetwork):
+    """CycleGan Encoder"""
+
+    def __init__(self, num_in=3, num_out=32, num_stacks=2):
+        super(VolumeEncoder, self).__init__()
+        self.num_in = num_in
+        self.num_out = num_out
+        self.num_inter = 8
+        self.num_stacks = num_stacks
+        self.with_bias = True
+
+        self.relu = nn.ReLU(inplace=True)
+        self.conv1 = nn.Conv3d(self.num_in,
+                               self.num_inter,
+                               bias=self.with_bias,
+                               kernel_size=5,
+                               stride=2,
+                               padding=4,
+                               dilation=2)
+        # self.bn1 = nn.GroupNorm(4, self.num_inter)
+        self.bn1 = nn.BatchNorm3d(self.num_inter)
+        self.conv2 = nn.Conv3d(self.num_inter,
+                               self.num_out,
+                               bias=self.with_bias,
+                               kernel_size=5,
+                               stride=2,
+                               padding=4,
+                               dilation=2)
+        # self.bn2 = nn.GroupNorm(4, self.num_out)
+        self.bn2 = nn.BatchNorm3d(self.num_out)
+
+        self.conv_out1 = nn.Conv3d(self.num_out,
+                                   self.num_out,
+                                   bias=self.with_bias,
+                                   kernel_size=3,
+                                   stride=1,
+                                   padding=1,
+                                   dilation=1)
+        self.conv_out2 = nn.Conv3d(self.num_out,
+                                   self.num_out,
+                                   bias=self.with_bias,
+                                   kernel_size=3,
+                                   stride=1,
+                                   padding=1,
+                                   dilation=1)
+
+        for idx in range(self.num_stacks):
+            self.add_module("res" + str(idx),
+                            Residual3D(self.num_out, self.num_out))
+
+        self.init_weights()
+
+    def forward(self, x, intermediate_output=True):
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out_lst = []
+        for idx in range(self.num_stacks):
+            out = self._modules["res" + str(idx)](out)
+            out_lst.append(out)
+
+        if intermediate_output:
+            return out_lst
+        else:
+            return [out_lst[-1]]

lib/net/__init__.py

@@ -0,0 +1,4 @@
+from .BasePIFuNet import BasePIFuNet
+from .HGPIFuNet import HGPIFuNet
+from .NormalNet import NormalNet
+from .VE import VolumeEncoder

lib/net/geometry.py

@@ -0,0 +1,82 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import torch
+
+
+def index(feat, uv):
+    '''
+    :param feat: [B, C, H, W] image features
+    :param uv: [B, 2, N] uv coordinates in the image plane, range [0, 1]
+    :return: [B, C, N] image features at the uv coordinates
+    '''
+    uv = uv.transpose(1, 2)  # [B, N, 2]
+
+    (B, N, _) = uv.shape
+    C = feat.shape[1]
+
+    if uv.shape[-1] == 3:
+        # uv = uv[:,:,[2,1,0]]
+        # uv = uv * torch.tensor([1.0,-1.0,1.0]).type_as(uv)[None,None,...]
+        uv = uv.unsqueeze(2).unsqueeze(3)  # [B, N, 1, 1, 3]
+    else:
+        uv = uv.unsqueeze(2)  # [B, N, 1, 2]
+
+    # NOTE: for newer PyTorch, it seems that training results are degraded due to implementation diff in F.grid_sample
+    # for old versions, simply remove the aligned_corners argument.
+    samples = torch.nn.functional.grid_sample(
+        feat, uv, align_corners=True)  # [B, C, N, 1]
+    return samples.view(B, C, N)  # [B, C, N]
+
+
+def orthogonal(points, calibrations, transforms=None):
+    '''
+    Compute the orthogonal projections of 3D points into the image plane by given projection matrix
+    :param points: [B, 3, N] Tensor of 3D points
+    :param calibrations: [B, 3, 4] Tensor of projection matrix
+    :param transforms: [B, 2, 3] Tensor of image transform matrix
+    :return: xyz: [B, 3, N] Tensor of xyz coordinates in the image plane
+    '''
+    rot = calibrations[:, :3, :3]
+    trans = calibrations[:, :3, 3:4]
+    pts = torch.baddbmm(trans, rot, points)  # [B, 3, N]
+    if transforms is not None:
+        scale = transforms[:2, :2]
+        shift = transforms[:2, 2:3]
+        pts[:, :2, :] = torch.baddbmm(shift, scale, pts[:, :2, :])
+    return pts
+
+
+def perspective(points, calibrations, transforms=None):
+    '''
+    Compute the perspective projections of 3D points into the image plane by given projection matrix
+    :param points: [Bx3xN] Tensor of 3D points
+    :param calibrations: [Bx3x4] Tensor of projection matrix
+    :param transforms: [Bx2x3] Tensor of image transform matrix
+    :return: xy: [Bx2xN] Tensor of xy coordinates in the image plane
+    '''
+    rot = calibrations[:, :3, :3]
+    trans = calibrations[:, :3, 3:4]
+    homo = torch.baddbmm(trans, rot, points)  # [B, 3, N]
+    xy = homo[:, :2, :] / homo[:, 2:3, :]
+    if transforms is not None:
+        scale = transforms[:2, :2]
+        shift = transforms[:2, 2:3]
+        xy = torch.baddbmm(shift, scale, xy)
+
+    xyz = torch.cat([xy, homo[:, 2:3, :]], 1)
+    return xyz

lib/net/net_util.py

@@ -0,0 +1,329 @@
+
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+from torchvision import models
+import torch
+from torch.nn import init
+import torch.nn as nn
+import torch.nn.functional as F
+import functools
+from torch.autograd import grad
+
+
+def gradient(inputs, outputs):
+    d_points = torch.ones_like(outputs,
+                               requires_grad=False,
+                               device=outputs.device)
+    points_grad = grad(outputs=outputs,
+                       inputs=inputs,
+                       grad_outputs=d_points,
+                       create_graph=True,
+                       retain_graph=True,
+                       only_inputs=True,
+                       allow_unused=True)[0]
+    return points_grad
+
+
+# def conv3x3(in_planes, out_planes, strd=1, padding=1, bias=False):
+#     "3x3 convolution with padding"
+#     return nn.Conv2d(in_planes, out_planes, kernel_size=3,
+#                      stride=strd, padding=padding, bias=bias)
+
+
+def conv3x3(in_planes,
+            out_planes,
+            kernel=3,
+            strd=1,
+            dilation=1,
+            padding=1,
+            bias=False):
+    "3x3 convolution with padding"
+    return nn.Conv2d(in_planes,
+                     out_planes,
+                     kernel_size=kernel,
+                     dilation=dilation,
+                     stride=strd,
+                     padding=padding,
+                     bias=bias)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes,
+                     out_planes,
+                     kernel_size=1,
+                     stride=stride,
+                     bias=False)
+
+
+def init_weights(net, init_type='normal', init_gain=0.02):
+    """Initialize network weights.
+
+    Parameters:
+        net (network)   -- network to be initialized
+        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.
+
+    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
+    work better for some applications. Feel free to try yourself.
+    """
+    def init_func(m):  # define the initialization function
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and (classname.find('Conv') != -1
+                                     or classname.find('Linear') != -1):
+            if init_type == 'normal':
+                init.normal_(m.weight.data, 0.0, init_gain)
+            elif init_type == 'xavier':
+                init.xavier_normal_(m.weight.data, gain=init_gain)
+            elif init_type == 'kaiming':
+                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            elif init_type == 'orthogonal':
+                init.orthogonal_(m.weight.data, gain=init_gain)
+            else:
+                raise NotImplementedError(
+                    'initialization method [%s] is not implemented' %
+                    init_type)
+            if hasattr(m, 'bias') and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find(
+                'BatchNorm2d'
+        ) != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
+            init.normal_(m.weight.data, 1.0, init_gain)
+            init.constant_(m.bias.data, 0.0)
+
+    # print('initialize network with %s' % init_type)
+    net.apply(init_func)  # apply the initialization function <init_func>
+
+
+def init_net(net, init_type='xavier', init_gain=0.02, gpu_ids=[]):
+    """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
+    Parameters:
+        net (network)      -- the network to be initialized
+        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        gain (float)       -- scaling factor for normal, xavier and orthogonal.
+        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
+
+    Return an initialized network.
+    """
+    if len(gpu_ids) > 0:
+        assert (torch.cuda.is_available())
+        net = torch.nn.DataParallel(net)  # multi-GPUs
+    init_weights(net, init_type, init_gain=init_gain)
+    return net
+
+
+def imageSpaceRotation(xy, rot):
+    '''
+    args:
+        xy: (B, 2, N) input
+        rot: (B, 2) x,y axis rotation angles
+
+    rotation center will be always image center (other rotation center can be represented by additional z translation)
+    '''
+    disp = rot.unsqueeze(2).sin().expand_as(xy)
+    return (disp * xy).sum(dim=1)
+
+
+def cal_gradient_penalty(netD,
+                         real_data,
+                         fake_data,
+                         device,
+                         type='mixed',
+                         constant=1.0,
+                         lambda_gp=10.0):
+    """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028
+
+    Arguments:
+        netD (network)              -- discriminator network
+        real_data (tensor array)    -- real images
+        fake_data (tensor array)    -- generated images from the generator
+        device (str)                -- GPU / CPU: from torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')
+        type (str)                  -- if we mix real and fake data or not [real | fake | mixed].
+        constant (float)            -- the constant used in formula ( | |gradient||_2 - constant)^2
+        lambda_gp (float)           -- weight for this loss
+
+    Returns the gradient penalty loss
+    """
+    if lambda_gp > 0.0:
+        # either use real images, fake images, or a linear interpolation of two.
+        if type == 'real':
+            interpolatesv = real_data
+        elif type == 'fake':
+            interpolatesv = fake_data
+        elif type == 'mixed':
+            alpha = torch.rand(real_data.shape[0], 1)
+            alpha = alpha.expand(
+                real_data.shape[0],
+                real_data.nelement() //
+                real_data.shape[0]).contiguous().view(*real_data.shape)
+            alpha = alpha.to(device)
+            interpolatesv = alpha * real_data + ((1 - alpha) * fake_data)
+        else:
+            raise NotImplementedError('{} not implemented'.format(type))
+        interpolatesv.requires_grad_(True)
+        disc_interpolates = netD(interpolatesv)
+        gradients = torch.autograd.grad(
+            outputs=disc_interpolates,
+            inputs=interpolatesv,
+            grad_outputs=torch.ones(disc_interpolates.size()).to(device),
+            create_graph=True,
+            retain_graph=True,
+            only_inputs=True)
+        gradients = gradients[0].view(real_data.size(0), -1)  # flat the data
+        gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) **
+                            2).mean() * lambda_gp  # added eps
+        return gradient_penalty, gradients
+    else:
+        return 0.0, None
+
+
+def get_norm_layer(norm_type='instance'):
+    """Return a normalization layer
+    Parameters:
+        norm_type (str) -- the name of the normalization layer: batch | instance | none
+    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).
+    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.
+    """
+    if norm_type == 'batch':
+        norm_layer = functools.partial(nn.BatchNorm2d,
+                                       affine=True,
+                                       track_running_stats=True)
+    elif norm_type == 'instance':
+        norm_layer = functools.partial(nn.InstanceNorm2d,
+                                       affine=False,
+                                       track_running_stats=False)
+    elif norm_type == 'group':
+        norm_layer = functools.partial(nn.GroupNorm, 32)
+    elif norm_type == 'none':
+        norm_layer = None
+    else:
+        raise NotImplementedError('normalization layer [%s] is not found' %
+                                  norm_type)
+    return norm_layer
+
+
+class Flatten(nn.Module):
+    def forward(self, input):
+        return input.view(input.size(0), -1)
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_planes, out_planes, opt):
+        super(ConvBlock, self).__init__()
+        [k, s, d, p] = opt.conv3x3
+        self.conv1 = conv3x3(in_planes, int(out_planes / 2), k, s, d, p)
+        self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4), k, s, d,
+                             p)
+        self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4), k, s, d,
+                             p)
+
+        if opt.norm == 'batch':
+            self.bn1 = nn.BatchNorm2d(in_planes)
+            self.bn2 = nn.BatchNorm2d(int(out_planes / 2))
+            self.bn3 = nn.BatchNorm2d(int(out_planes / 4))
+            self.bn4 = nn.BatchNorm2d(in_planes)
+        elif opt.norm == 'group':
+            self.bn1 = nn.GroupNorm(32, in_planes)
+            self.bn2 = nn.GroupNorm(32, int(out_planes / 2))
+            self.bn3 = nn.GroupNorm(32, int(out_planes / 4))
+            self.bn4 = nn.GroupNorm(32, in_planes)
+
+        if in_planes != out_planes:
+            self.downsample = nn.Sequential(
+                self.bn4,
+                nn.ReLU(True),
+                nn.Conv2d(in_planes,
+                          out_planes,
+                          kernel_size=1,
+                          stride=1,
+                          bias=False),
+            )
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        residual = x
+
+        out1 = self.bn1(x)
+        out1 = F.relu(out1, True)
+        out1 = self.conv1(out1)
+
+        out2 = self.bn2(out1)
+        out2 = F.relu(out2, True)
+        out2 = self.conv2(out2)
+
+        out3 = self.bn3(out2)
+        out3 = F.relu(out3, True)
+        out3 = self.conv3(out3)
+
+        out3 = torch.cat((out1, out2, out3), 1)
+
+        if self.downsample is not None:
+            residual = self.downsample(residual)
+
+        out3 += residual
+
+        return out3
+
+
+class Vgg19(torch.nn.Module):
+    def __init__(self, requires_grad=False):
+        super(Vgg19, self).__init__()
+        vgg_pretrained_features = models.vgg19(pretrained=True).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        for x in range(2):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(2, 7):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(7, 12):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(12, 21):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(21, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h_relu1 = self.slice1(X)
+        h_relu2 = self.slice2(h_relu1)
+        h_relu3 = self.slice3(h_relu2)
+        h_relu4 = self.slice4(h_relu3)
+        h_relu5 = self.slice5(h_relu4)
+        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
+        return out
+
+
+class VGGLoss(nn.Module):
+    def __init__(self):
+        super(VGGLoss, self).__init__()
+        self.vgg = Vgg19()
+        self.criterion = nn.L1Loss()
+        self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]
+
+    def forward(self, x, y):
+        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
+        loss = 0
+        for i in range(len(x_vgg)):
+            loss += self.weights[i] * self.criterion(x_vgg[i],
+                                                     y_vgg[i].detach())
+        return loss

lib/net/voxelize.py

@@ -0,0 +1,184 @@
+from __future__ import division, print_function
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from torch.autograd import Function
+
+import voxelize_cuda
+
+
+class VoxelizationFunction(Function):
+    """
+    Definition of differentiable voxelization function
+    Currently implemented only for cuda Tensors
+    """
+    @staticmethod
+    def forward(ctx, smpl_vertices, smpl_face_center, smpl_face_normal,
+                smpl_vertex_code, smpl_face_code, smpl_tetrahedrons,
+                volume_res, sigma, smooth_kernel_size):
+        """
+        forward pass
+        Output format: (batch_size, z_dims, y_dims, x_dims, channel_num) 
+        """
+        assert (smpl_vertices.size()[1] == smpl_vertex_code.size()[1])
+        assert (smpl_face_center.size()[1] == smpl_face_normal.size()[1])
+        assert (smpl_face_center.size()[1] == smpl_face_code.size()[1])
+        ctx.batch_size = smpl_vertices.size()[0]
+        ctx.volume_res = volume_res
+        ctx.sigma = sigma
+        ctx.smooth_kernel_size = smooth_kernel_size
+        ctx.smpl_vertex_num = smpl_vertices.size()[1]
+        ctx.device = smpl_vertices.device
+
+        smpl_vertices = smpl_vertices.contiguous()
+        smpl_face_center = smpl_face_center.contiguous()
+        smpl_face_normal = smpl_face_normal.contiguous()
+        smpl_vertex_code = smpl_vertex_code.contiguous()
+        smpl_face_code = smpl_face_code.contiguous()
+        smpl_tetrahedrons = smpl_tetrahedrons.contiguous()
+
+        occ_volume = torch.cuda.FloatTensor(ctx.batch_size, ctx.volume_res,
+                                            ctx.volume_res,
+                                            ctx.volume_res).fill_(0.0)
+        semantic_volume = torch.cuda.FloatTensor(ctx.batch_size,
+                                                 ctx.volume_res,
+                                                 ctx.volume_res,
+                                                 ctx.volume_res, 3).fill_(0.0)
+        weight_sum_volume = torch.cuda.FloatTensor(ctx.batch_size,
+                                                   ctx.volume_res,
+                                                   ctx.volume_res,
+                                                   ctx.volume_res).fill_(1e-3)
+
+        # occ_volume [B, volume_res, volume_res, volume_res]
+        # semantic_volume [B, volume_res, volume_res, volume_res, 3]
+        # weight_sum_volume [B, volume_res, volume_res, volume_res]
+
+        occ_volume, semantic_volume, weight_sum_volume = voxelize_cuda.forward_semantic_voxelization(
+            smpl_vertices, smpl_vertex_code, smpl_tetrahedrons, occ_volume,
+            semantic_volume, weight_sum_volume, sigma)
+
+        return semantic_volume
+
+
+class Voxelization(nn.Module):
+    """
+    Wrapper around the autograd function VoxelizationFunction
+    """
+
+    def __init__(self, smpl_vertex_code, smpl_face_code, smpl_face_indices,
+                 smpl_tetraderon_indices, volume_res, sigma,
+                 smooth_kernel_size, batch_size, device):
+        super(Voxelization, self).__init__()
+        assert (len(smpl_face_indices.shape) == 2)
+        assert (len(smpl_tetraderon_indices.shape) == 2)
+        assert (smpl_face_indices.shape[1] == 3)
+        assert (smpl_tetraderon_indices.shape[1] == 4)
+
+        self.volume_res = volume_res
+        self.sigma = sigma
+        self.smooth_kernel_size = smooth_kernel_size
+        self.batch_size = batch_size
+        self.device = device
+
+        self.smpl_vertex_code = smpl_vertex_code
+        self.smpl_face_code = smpl_face_code
+        self.smpl_face_indices = smpl_face_indices
+        self.smpl_tetraderon_indices = smpl_tetraderon_indices
+
+    def update_param(self, batch_size, smpl_tetra):
+
+        self.batch_size = batch_size
+        self.smpl_tetraderon_indices = smpl_tetra
+
+        smpl_vertex_code_batch = np.tile(self.smpl_vertex_code,
+                                         (self.batch_size, 1, 1))
+        smpl_face_code_batch = np.tile(self.smpl_face_code,
+                                       (self.batch_size, 1, 1))
+        smpl_face_indices_batch = np.tile(self.smpl_face_indices,
+                                          (self.batch_size, 1, 1))
+        smpl_tetraderon_indices_batch = np.tile(self.smpl_tetraderon_indices,
+                                                (self.batch_size, 1, 1))
+
+        smpl_vertex_code_batch = torch.from_numpy(
+            smpl_vertex_code_batch).contiguous().to(self.device)
+        smpl_face_code_batch = torch.from_numpy(
+            smpl_face_code_batch).contiguous().to(self.device)
+        smpl_face_indices_batch = torch.from_numpy(
+            smpl_face_indices_batch).contiguous().to(self.device)
+        smpl_tetraderon_indices_batch = torch.from_numpy(
+            smpl_tetraderon_indices_batch).contiguous().to(self.device)
+
+        self.register_buffer('smpl_vertex_code_batch', smpl_vertex_code_batch)
+        self.register_buffer('smpl_face_code_batch', smpl_face_code_batch)
+        self.register_buffer('smpl_face_indices_batch',
+                             smpl_face_indices_batch)
+        self.register_buffer('smpl_tetraderon_indices_batch',
+                             smpl_tetraderon_indices_batch)
+
+    def forward(self, smpl_vertices):
+        """
+        Generate semantic volumes from SMPL vertices
+        """
+        assert (smpl_vertices.size()[0] == self.batch_size)
+        self.check_input(smpl_vertices)
+        smpl_faces = self.vertices_to_faces(smpl_vertices)
+        smpl_tetrahedrons = self.vertices_to_tetrahedrons(smpl_vertices)
+        smpl_face_center = self.calc_face_centers(smpl_faces)
+        smpl_face_normal = self.calc_face_normals(smpl_faces)
+        smpl_surface_vertex_num = self.smpl_vertex_code_batch.size()[1]
+        smpl_vertices_surface = smpl_vertices[:, :smpl_surface_vertex_num, :]
+        vol = VoxelizationFunction.apply(smpl_vertices_surface,
+                                         smpl_face_center, smpl_face_normal,
+                                         self.smpl_vertex_code_batch,
+                                         self.smpl_face_code_batch,
+                                         smpl_tetrahedrons, self.volume_res,
+                                         self.sigma, self.smooth_kernel_size)
+        return vol.permute((0, 4, 1, 2, 3))  # (bzyxc --> bcdhw)
+
+    def vertices_to_faces(self, vertices):
+        assert (vertices.ndimension() == 3)
+        bs, nv = vertices.shape[:2]
+        device = vertices.device
+        face = self.smpl_face_indices_batch + (
+            torch.arange(bs, dtype=torch.int32).to(device) * nv)[:, None, None]
+        vertices_ = vertices.reshape((bs * nv, 3))
+        return vertices_[face.long()]
+
+    def vertices_to_tetrahedrons(self, vertices):
+        assert (vertices.ndimension() == 3)
+        bs, nv = vertices.shape[:2]
+        device = vertices.device
+        tets = self.smpl_tetraderon_indices_batch + (
+            torch.arange(bs, dtype=torch.int32).to(device) * nv)[:, None, None]
+        vertices_ = vertices.reshape((bs * nv, 3))
+        return vertices_[tets.long()]
+
+    def calc_face_centers(self, face_verts):
+        assert len(face_verts.shape) == 4
+        assert face_verts.shape[2] == 3
+        assert face_verts.shape[3] == 3
+        bs, nf = face_verts.shape[:2]
+        face_centers = (face_verts[:, :, 0, :] + face_verts[:, :, 1, :] +
+                        face_verts[:, :, 2, :]) / 3.0
+        face_centers = face_centers.reshape((bs, nf, 3))
+        return face_centers
+
+    def calc_face_normals(self, face_verts):
+        assert len(face_verts.shape) == 4
+        assert face_verts.shape[2] == 3
+        assert face_verts.shape[3] == 3
+        bs, nf = face_verts.shape[:2]
+        face_verts = face_verts.reshape((bs * nf, 3, 3))
+        v10 = face_verts[:, 0] - face_verts[:, 1]
+        v12 = face_verts[:, 2] - face_verts[:, 1]
+        normals = F.normalize(torch.cross(v10, v12), eps=1e-5)
+        normals = normals.reshape((bs, nf, 3))
+        return normals
+
+    def check_input(self, x):
+        if x.device == 'cpu':
+            raise TypeError('Voxelization module supports only cuda tensors')
+        if x.type() != 'torch.cuda.FloatTensor':
+            raise TypeError(
+                'Voxelization module supports only float32 tensors')