Upload 5 files

lib/pymaf/models/__init__.py

@@ -0,0 +1,3 @@
+from .hmr import hmr
+from .pymaf_net import pymaf_net
+from .smpl import SMPL

lib/pymaf/models/hmr.py

@@ -0,0 +1,303 @@
+# This script is borrowed from https://github.com/nkolot/SPIN/blob/master/models/hmr.py
+
+import torch
+import torch.nn as nn
+import torchvision.models.resnet as resnet
+import numpy as np
+import math
+from lib.pymaf.utils.geometry import rot6d_to_rotmat
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+BN_MOMENTUM = 0.1
+
+
+class Bottleneck(nn.Module):
+    """ Redefinition of Bottleneck residual block
+        Adapted from the official PyTorch implementation
+    """
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super().__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes,
+                               planes,
+                               kernel_size=3,
+                               stride=stride,
+                               padding=1,
+                               bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        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 = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet_Backbone(nn.Module):
+    """ Feature Extrator with ResNet backbone
+    """
+
+    def __init__(self, model='res50', pretrained=True):
+        if model == 'res50':
+            block, layers = Bottleneck, [3, 4, 6, 3]
+        else:
+            pass  # TODO
+
+        self.inplanes = 64
+        super().__init__()
+        npose = 24 * 6
+        self.conv1 = nn.Conv2d(3,
+                               64,
+                               kernel_size=7,
+                               stride=2,
+                               padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AvgPool2d(7, stride=1)
+
+        if pretrained:
+            resnet_imagenet = resnet.resnet50(pretrained=True)
+            self.load_state_dict(resnet_imagenet.state_dict(), strict=False)
+            logger.info('loaded resnet50 imagenet pretrained model')
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes,
+                          planes * block.expansion,
+                          kernel_size=1,
+                          stride=stride,
+                          bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def _make_deconv_layer(self, num_layers, num_filters, num_kernels):
+        assert num_layers == len(num_filters), \
+            'ERROR: num_deconv_layers is different len(num_deconv_filters)'
+        assert num_layers == len(num_kernels), \
+            'ERROR: num_deconv_layers is different len(num_deconv_filters)'
+
+        def _get_deconv_cfg(deconv_kernel, index):
+            if deconv_kernel == 4:
+                padding = 1
+                output_padding = 0
+            elif deconv_kernel == 3:
+                padding = 1
+                output_padding = 1
+            elif deconv_kernel == 2:
+                padding = 0
+                output_padding = 0
+
+            return deconv_kernel, padding, output_padding
+
+        layers = []
+        for i in range(num_layers):
+            kernel, padding, output_padding = _get_deconv_cfg(
+                num_kernels[i], i)
+
+            planes = num_filters[i]
+            layers.append(
+                nn.ConvTranspose2d(in_channels=self.inplanes,
+                                   out_channels=planes,
+                                   kernel_size=kernel,
+                                   stride=2,
+                                   padding=padding,
+                                   output_padding=output_padding,
+                                   bias=self.deconv_with_bias))
+            layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))
+            layers.append(nn.ReLU(inplace=True))
+            self.inplanes = planes
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+
+        batch_size = x.shape[0]
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x1 = self.layer1(x)
+        x2 = self.layer2(x1)
+        x3 = self.layer3(x2)
+        x4 = self.layer4(x3)
+
+        xf = self.avgpool(x4)
+        xf = xf.view(xf.size(0), -1)
+
+        x_featmap = x4
+
+        return x_featmap, xf
+
+
+class HMR(nn.Module):
+    """ SMPL Iterative Regressor with ResNet50 backbone
+    """
+
+    def __init__(self, block, layers, smpl_mean_params):
+        self.inplanes = 64
+        super().__init__()
+        npose = 24 * 6
+        self.conv1 = nn.Conv2d(3,
+                               64,
+                               kernel_size=7,
+                               stride=2,
+                               padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AvgPool2d(7, stride=1)
+        self.fc1 = nn.Linear(512 * block.expansion + npose + 13, 1024)
+        self.drop1 = nn.Dropout()
+        self.fc2 = nn.Linear(1024, 1024)
+        self.drop2 = nn.Dropout()
+        self.decpose = nn.Linear(1024, npose)
+        self.decshape = nn.Linear(1024, 10)
+        self.deccam = nn.Linear(1024, 3)
+        nn.init.xavier_uniform_(self.decpose.weight, gain=0.01)
+        nn.init.xavier_uniform_(self.decshape.weight, gain=0.01)
+        nn.init.xavier_uniform_(self.deccam.weight, gain=0.01)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+        mean_params = np.load(smpl_mean_params)
+        init_pose = torch.from_numpy(mean_params['pose'][:]).unsqueeze(0)
+        init_shape = torch.from_numpy(
+            mean_params['shape'][:].astype('float32')).unsqueeze(0)
+        init_cam = torch.from_numpy(mean_params['cam']).unsqueeze(0)
+        self.register_buffer('init_pose', init_pose)
+        self.register_buffer('init_shape', init_shape)
+        self.register_buffer('init_cam', init_cam)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes,
+                          planes * block.expansion,
+                          kernel_size=1,
+                          stride=stride,
+                          bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self,
+                x,
+                init_pose=None,
+                init_shape=None,
+                init_cam=None,
+                n_iter=3):
+
+        batch_size = x.shape[0]
+
+        if init_pose is None:
+            init_pose = self.init_pose.expand(batch_size, -1)
+        if init_shape is None:
+            init_shape = self.init_shape.expand(batch_size, -1)
+        if init_cam is None:
+            init_cam = self.init_cam.expand(batch_size, -1)
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x1 = self.layer1(x)
+        x2 = self.layer2(x1)
+        x3 = self.layer3(x2)
+        x4 = self.layer4(x3)
+
+        xf = self.avgpool(x4)
+        xf = xf.view(xf.size(0), -1)
+
+        pred_pose = init_pose
+        pred_shape = init_shape
+        pred_cam = init_cam
+        for i in range(n_iter):
+            xc = torch.cat([xf, pred_pose, pred_shape, pred_cam], 1)
+            xc = self.fc1(xc)
+            xc = self.drop1(xc)
+            xc = self.fc2(xc)
+            xc = self.drop2(xc)
+            pred_pose = self.decpose(xc) + pred_pose
+            pred_shape = self.decshape(xc) + pred_shape
+            pred_cam = self.deccam(xc) + pred_cam
+
+        pred_rotmat = rot6d_to_rotmat(pred_pose).view(batch_size, 24, 3, 3)
+
+        return pred_rotmat, pred_shape, pred_cam
+
+
+def hmr(smpl_mean_params, pretrained=True, **kwargs):
+    """ Constructs an HMR model with ResNet50 backbone.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = HMR(Bottleneck, [3, 4, 6, 3], smpl_mean_params, **kwargs)
+    if pretrained:
+        resnet_imagenet = resnet.resnet50(pretrained=True)
+        model.load_state_dict(resnet_imagenet.state_dict(), strict=False)
+    return model

lib/pymaf/models/maf_extractor.py

@@ -0,0 +1,137 @@
+# This script is borrowed and extended from https://github.com/shunsukesaito/PIFu/blob/master/lib/model/SurfaceClassifier.py
+
+from packaging import version
+import torch
+import scipy
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+
+from lib.common.config import cfg
+from lib.pymaf.utils.geometry import projection
+from lib.pymaf.core.path_config import MESH_DOWNSAMPLEING
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class MAF_Extractor(nn.Module):
+    ''' Mesh-aligned Feature Extrator
+
+    As discussed in the paper, we extract mesh-aligned features based on 2D projection of the mesh vertices.
+    The features extrated from spatial feature maps will go through a MLP for dimension reduction.
+    '''
+
+    def __init__(self, device=torch.device('cuda')):
+        super().__init__()
+
+        self.device = device
+        self.filters = []
+        self.num_views = 1
+        filter_channels = cfg.MODEL.PyMAF.MLP_DIM
+        self.last_op = nn.ReLU(True)
+
+        for l in range(0, len(filter_channels) - 1):
+            if 0 != l:
+                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))
+
+            self.add_module("conv%d" % l, self.filters[l])
+
+        self.im_feat = None
+        self.cam = None
+
+        # downsample SMPL mesh and assign part labels
+        # from https://github.com/nkolot/GraphCMR/blob/master/data/mesh_downsampling.npz
+        smpl_mesh_graph = np.load(MESH_DOWNSAMPLEING,
+                                  allow_pickle=True,
+                                  encoding='latin1')
+
+        A = smpl_mesh_graph['A']
+        U = smpl_mesh_graph['U']
+        D = smpl_mesh_graph['D']  # shape: (2,)
+
+        # downsampling
+        ptD = []
+        for i in range(len(D)):
+            d = scipy.sparse.coo_matrix(D[i])
+            i = torch.LongTensor(np.array([d.row, d.col]))
+            v = torch.FloatTensor(d.data)
+            ptD.append(torch.sparse.FloatTensor(i, v, d.shape))
+
+        # downsampling mapping from 6890 points to 431 points
+        # ptD[0].to_dense() - Size: [1723, 6890]
+        # ptD[1].to_dense() - Size: [431. 1723]
+        Dmap = torch.matmul(ptD[1].to_dense(),
+                            ptD[0].to_dense())  # 6890 -> 431
+        self.register_buffer('Dmap', Dmap)
+
+    def reduce_dim(self, feature):
+        '''
+        Dimension reduction by multi-layer perceptrons
+        :param feature: list of [B, C_s, N] point-wise features before dimension reduction
+        :return: [B, C_p x N] concatantion of point-wise features after dimension reduction
+        '''
+        y = feature
+        tmpy = feature
+        for i, f in enumerate(self.filters):
+            y = self._modules['conv' +
+                              str(i)](y if i == 0 else torch.cat([y, tmpy], 1))
+            if i != len(self.filters) - 1:
+                y = F.leaky_relu(y)
+            if self.num_views > 1 and i == len(self.filters) // 2:
+                y = y.view(-1, self.num_views, y.shape[1],
+                           y.shape[2]).mean(dim=1)
+                tmpy = feature.view(-1, self.num_views, feature.shape[1],
+                                    feature.shape[2]).mean(dim=1)
+
+        y = self.last_op(y)
+
+        y = y.view(y.shape[0], -1)
+        return y
+
+    def sampling(self, points, im_feat=None, z_feat=None):
+        '''
+        Given 2D points, sample the point-wise features for each point, 
+        the dimension of point-wise features will be reduced from C_s to C_p by MLP.
+        Image features should be pre-computed before this call.
+        :param points: [B, N, 2] image coordinates of points
+        :im_feat: [B, C_s, H_s, W_s] spatial feature maps 
+        :return: [B, C_p x N] concatantion of point-wise features after dimension reduction
+        '''
+        if im_feat is None:
+            im_feat = self.im_feat
+
+        batch_size = im_feat.shape[0]
+
+        if version.parse(torch.__version__) >= version.parse('1.3.0'):
+            # Default grid_sample behavior has changed to align_corners=False since 1.3.0.
+            point_feat = torch.nn.functional.grid_sample(
+                im_feat, points.unsqueeze(2), align_corners=True)[..., 0]
+        else:
+            point_feat = torch.nn.functional.grid_sample(
+                im_feat, points.unsqueeze(2))[..., 0]
+
+        mesh_align_feat = self.reduce_dim(point_feat)
+        return mesh_align_feat
+
+    def forward(self, p, s_feat=None, cam=None, **kwargs):
+        ''' Returns mesh-aligned features for the 3D mesh points.
+
+        Args:
+            p (tensor): [B, N_m, 3] mesh vertices
+            s_feat (tensor): [B, C_s, H_s, W_s] spatial feature maps
+            cam (tensor): [B, 3] camera
+        Return:
+            mesh_align_feat (tensor): [B, C_p x N_m] mesh-aligned features
+        '''
+        if cam is None:
+            cam = self.cam
+        p_proj_2d = projection(p, cam, retain_z=False)
+        mesh_align_feat = self.sampling(p_proj_2d, s_feat)
+        return mesh_align_feat

lib/pymaf/models/res_module.py

@@ -0,0 +1,385 @@
+# code brought in part from https://github.com/microsoft/human-pose-estimation.pytorch/blob/master/lib/models/pose_resnet.py
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from collections import OrderedDict
+import os
+from lib.pymaf.core.cfgs import cfg
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+BN_MOMENTUM = 0.1
+
+
+def conv3x3(in_planes, out_planes, stride=1, bias=False, groups=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes * groups,
+                     out_planes * groups,
+                     kernel_size=3,
+                     stride=stride,
+                     padding=1,
+                     bias=bias,
+                     groups=groups)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1):
+        super().__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride, groups=groups)
+        self.bn1 = nn.BatchNorm2d(planes * groups, momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes, groups=groups)
+        self.bn2 = nn.BatchNorm2d(planes * groups, momentum=BN_MOMENTUM)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1):
+        super().__init__()
+        self.conv1 = nn.Conv2d(inplanes * groups,
+                               planes * groups,
+                               kernel_size=1,
+                               bias=False,
+                               groups=groups)
+        self.bn1 = nn.BatchNorm2d(planes * groups, momentum=BN_MOMENTUM)
+        self.conv2 = nn.Conv2d(planes * groups,
+                               planes * groups,
+                               kernel_size=3,
+                               stride=stride,
+                               padding=1,
+                               bias=False,
+                               groups=groups)
+        self.bn2 = nn.BatchNorm2d(planes * groups, momentum=BN_MOMENTUM)
+        self.conv3 = nn.Conv2d(planes * groups,
+                               planes * self.expansion * groups,
+                               kernel_size=1,
+                               bias=False,
+                               groups=groups)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion * groups,
+                                  momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        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 = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+resnet_spec = {
+    18: (BasicBlock, [2, 2, 2, 2]),
+    34: (BasicBlock, [3, 4, 6, 3]),
+    50: (Bottleneck, [3, 4, 6, 3]),
+    101: (Bottleneck, [3, 4, 23, 3]),
+    152: (Bottleneck, [3, 8, 36, 3])
+}
+
+
+class IUV_predict_layer(nn.Module):
+    def __init__(self,
+                 feat_dim=256,
+                 final_cov_k=3,
+                 part_out_dim=25,
+                 with_uv=True):
+        super().__init__()
+
+        self.with_uv = with_uv
+        if self.with_uv:
+            self.predict_u = nn.Conv2d(in_channels=feat_dim,
+                                       out_channels=25,
+                                       kernel_size=final_cov_k,
+                                       stride=1,
+                                       padding=1 if final_cov_k == 3 else 0)
+
+            self.predict_v = nn.Conv2d(in_channels=feat_dim,
+                                       out_channels=25,
+                                       kernel_size=final_cov_k,
+                                       stride=1,
+                                       padding=1 if final_cov_k == 3 else 0)
+
+        self.predict_ann_index = nn.Conv2d(
+            in_channels=feat_dim,
+            out_channels=15,
+            kernel_size=final_cov_k,
+            stride=1,
+            padding=1 if final_cov_k == 3 else 0)
+
+        self.predict_uv_index = nn.Conv2d(in_channels=feat_dim,
+                                          out_channels=25,
+                                          kernel_size=final_cov_k,
+                                          stride=1,
+                                          padding=1 if final_cov_k == 3 else 0)
+
+        self.inplanes = feat_dim
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes,
+                          planes * block.expansion,
+                          kernel_size=1,
+                          stride=stride,
+                          bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        return_dict = {}
+
+        predict_uv_index = self.predict_uv_index(x)
+        predict_ann_index = self.predict_ann_index(x)
+
+        return_dict['predict_uv_index'] = predict_uv_index
+        return_dict['predict_ann_index'] = predict_ann_index
+
+        if self.with_uv:
+            predict_u = self.predict_u(x)
+            predict_v = self.predict_v(x)
+            return_dict['predict_u'] = predict_u
+            return_dict['predict_v'] = predict_v
+        else:
+            return_dict['predict_u'] = None
+            return_dict['predict_v'] = None
+            # return_dict['predict_u'] = torch.zeros(predict_uv_index.shape).to(predict_uv_index.device)
+            # return_dict['predict_v'] = torch.zeros(predict_uv_index.shape).to(predict_uv_index.device)
+
+        return return_dict
+
+
+class SmplResNet(nn.Module):
+    def __init__(self,
+                 resnet_nums,
+                 in_channels=3,
+                 num_classes=229,
+                 last_stride=2,
+                 n_extra_feat=0,
+                 truncate=0,
+                 **kwargs):
+        super().__init__()
+
+        self.inplanes = 64
+        self.truncate = truncate
+        # extra = cfg.MODEL.EXTRA
+        # self.deconv_with_bias = extra.DECONV_WITH_BIAS
+        block, layers = resnet_spec[resnet_nums]
+
+        self.conv1 = nn.Conv2d(in_channels,
+                               64,
+                               kernel_size=7,
+                               stride=2,
+                               padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2],
+                                       stride=2) if truncate < 2 else None
+        self.layer4 = self._make_layer(
+            block, 512, layers[3],
+            stride=last_stride) if truncate < 1 else None
+
+        self.avg_pooling = nn.AdaptiveAvgPool2d(1)
+
+        self.num_classes = num_classes
+        if num_classes > 0:
+            self.final_layer = nn.Linear(512 * block.expansion, num_classes)
+            nn.init.xavier_uniform_(self.final_layer.weight, gain=0.01)
+
+        self.n_extra_feat = n_extra_feat
+        if n_extra_feat > 0:
+            self.trans_conv = nn.Sequential(
+                nn.Conv2d(n_extra_feat + 512 * block.expansion,
+                          512 * block.expansion,
+                          kernel_size=1,
+                          bias=False),
+                nn.BatchNorm2d(512 * block.expansion, momentum=BN_MOMENTUM),
+                nn.ReLU(True))
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes,
+                          planes * block.expansion,
+                          kernel_size=1,
+                          stride=stride,
+                          bias=False),
+                nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x, infeat=None):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x1 = self.layer1(x)
+        x2 = self.layer2(x1)
+        x3 = self.layer3(x2) if self.truncate < 2 else x2
+        x4 = self.layer4(x3) if self.truncate < 1 else x3
+
+        if infeat is not None:
+            x4 = self.trans_conv(torch.cat([infeat, x4], 1))
+
+        if self.num_classes > 0:
+            xp = self.avg_pooling(x4)
+            cls = self.final_layer(xp.view(xp.size(0), -1))
+            if not cfg.DANET.USE_MEAN_PARA:
+                # for non-negative scale
+                scale = F.relu(cls[:, 0]).unsqueeze(1)
+                cls = torch.cat((scale, cls[:, 1:]), dim=1)
+        else:
+            cls = None
+
+        return cls, {'x4': x4}
+
+    def init_weights(self, pretrained=''):
+        if os.path.isfile(pretrained):
+            logger.info('=> loading pretrained model {}'.format(pretrained))
+            # self.load_state_dict(pretrained_state_dict, strict=False)
+            checkpoint = torch.load(pretrained)
+            if isinstance(checkpoint, OrderedDict):
+                # state_dict = checkpoint
+                state_dict_old = self.state_dict()
+                for key in state_dict_old.keys():
+                    if key in checkpoint.keys():
+                        if state_dict_old[key].shape != checkpoint[key].shape:
+                            del checkpoint[key]
+                state_dict = checkpoint
+            elif isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
+                state_dict_old = checkpoint['state_dict']
+                state_dict = OrderedDict()
+                # delete 'module.' because it is saved from DataParallel module
+                for key in state_dict_old.keys():
+                    if key.startswith('module.'):
+                        # state_dict[key[7:]] = state_dict[key]
+                        # state_dict.pop(key)
+                        state_dict[key[7:]] = state_dict_old[key]
+                    else:
+                        state_dict[key] = state_dict_old[key]
+            else:
+                raise RuntimeError(
+                    'No state_dict found in checkpoint file {}'.format(
+                        pretrained))
+            self.load_state_dict(state_dict, strict=False)
+        else:
+            logger.error('=> imagenet pretrained model dose not exist')
+            logger.error('=> please download it first')
+            raise ValueError('imagenet pretrained model does not exist')
+
+
+class LimbResLayers(nn.Module):
+    def __init__(self,
+                 resnet_nums,
+                 inplanes,
+                 outplanes=None,
+                 groups=1,
+                 **kwargs):
+        super().__init__()
+
+        self.inplanes = inplanes
+        block, layers = resnet_spec[resnet_nums]
+        self.outplanes = 512 if outplanes == None else outplanes
+        self.layer4 = self._make_layer(block,
+                                       self.outplanes,
+                                       layers[3],
+                                       stride=2,
+                                       groups=groups)
+
+        self.avg_pooling = nn.AdaptiveAvgPool2d(1)
+
+    def _make_layer(self, block, planes, blocks, stride=1, groups=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes * groups,
+                          planes * block.expansion * groups,
+                          kernel_size=1,
+                          stride=stride,
+                          bias=False,
+                          groups=groups),
+                nn.BatchNorm2d(planes * block.expansion * groups,
+                               momentum=BN_MOMENTUM),
+            )
+
+        layers = []
+        layers.append(
+            block(self.inplanes, planes, stride, downsample, groups=groups))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=groups))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.layer4(x)
+        x = self.avg_pooling(x)
+
+        return x

lib/pymaf/models/smpl.py

@@ -0,0 +1,92 @@
+# This script is borrowed from https://github.com/nkolot/SPIN/blob/master/models/smpl.py
+
+import torch
+import numpy as np
+from lib.smplx import SMPL as _SMPL
+from lib.smplx.body_models import ModelOutput
+from lib.smplx.lbs import vertices2joints
+from collections import namedtuple
+
+from lib.pymaf.core import path_config, constants
+
+SMPL_MEAN_PARAMS = path_config.SMPL_MEAN_PARAMS
+SMPL_MODEL_DIR = path_config.SMPL_MODEL_DIR
+
+# Indices to get the 14 LSP joints from the 17 H36M joints
+H36M_TO_J17 = [6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9]
+H36M_TO_J14 = H36M_TO_J17[:14]
+
+
+class SMPL(_SMPL):
+    """ Extension of the official SMPL implementation to support more joints """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        joints = [constants.JOINT_MAP[i] for i in constants.JOINT_NAMES]
+        J_regressor_extra = np.load(path_config.JOINT_REGRESSOR_TRAIN_EXTRA)
+        self.register_buffer(
+            'J_regressor_extra',
+            torch.tensor(J_regressor_extra, dtype=torch.float32))
+        self.joint_map = torch.tensor(joints, dtype=torch.long)
+        self.ModelOutput = namedtuple(
+            'ModelOutput_', ModelOutput._fields + (
+                'smpl_joints',
+                'joints_J19',
+            ))
+        self.ModelOutput.__new__.__defaults__ = (None, ) * len(
+            self.ModelOutput._fields)
+
+    def forward(self, *args, **kwargs):
+        kwargs['get_skin'] = True
+        smpl_output = super().forward(*args, **kwargs)
+        extra_joints = vertices2joints(self.J_regressor_extra,
+                                       smpl_output.vertices)
+        # smpl_output.joints: [B, 45, 3]  extra_joints: [B, 9, 3]
+        vertices = smpl_output.vertices
+        joints = torch.cat([smpl_output.joints, extra_joints], dim=1)
+        smpl_joints = smpl_output.joints[:, :24]
+        joints = joints[:, self.joint_map, :]  # [B, 49, 3]
+        joints_J24 = joints[:, -24:, :]
+        joints_J19 = joints_J24[:, constants.J24_TO_J19, :]
+        output = self.ModelOutput(vertices=vertices,
+                                  global_orient=smpl_output.global_orient,
+                                  body_pose=smpl_output.body_pose,
+                                  joints=joints,
+                                  joints_J19=joints_J19,
+                                  smpl_joints=smpl_joints,
+                                  betas=smpl_output.betas,
+                                  full_pose=smpl_output.full_pose)
+        return output
+
+
+def get_smpl_faces():
+    smpl = SMPL(SMPL_MODEL_DIR, batch_size=1, create_transl=False)
+    return smpl.faces
+
+
+def get_part_joints(smpl_joints):
+    batch_size = smpl_joints.shape[0]
+
+    # part_joints = torch.zeros().to(smpl_joints.device)
+
+    one_seg_pairs = [(0, 1), (0, 2), (0, 3), (3, 6), (9, 12), (9, 13), (9, 14),
+                     (12, 15), (13, 16), (14, 17)]
+    two_seg_pairs = [(1, 4), (2, 5), (4, 7), (5, 8), (16, 18), (17, 19),
+                     (18, 20), (19, 21)]
+
+    one_seg_pairs.extend(two_seg_pairs)
+
+    single_joints = [(10), (11), (15), (22), (23)]
+
+    part_joints = []
+
+    for j_p in one_seg_pairs:
+        new_joint = torch.mean(smpl_joints[:, j_p], dim=1, keepdim=True)
+        part_joints.append(new_joint)
+
+    for j_p in single_joints:
+        part_joints.append(smpl_joints[:, j_p:j_p + 1])
+
+    part_joints = torch.cat(part_joints, dim=1)
+
+    return part_joints