| |
|
|
| 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 |
|
|
| 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 |
|
|
