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