You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
307 lines
12 KiB
307 lines
12 KiB
# ------------------------------------------------------------------------------
|
|
# Adapted from https://github.com/akanazawa/hmr
|
|
# Original licence: Copyright (c) 2018 akanazawa, under the MIT License.
|
|
# ------------------------------------------------------------------------------
|
|
|
|
from abc import abstractmethod
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from mmcv.cnn import normal_init, xavier_init
|
|
|
|
from mmpose.models.utils.geometry import batch_rodrigues
|
|
|
|
|
|
class BaseDiscriminator(nn.Module):
|
|
"""Base linear module for SMPL parameter discriminator.
|
|
|
|
Args:
|
|
fc_layers (Tuple): Tuple of neuron count,
|
|
such as (9, 32, 32, 1)
|
|
use_dropout (Tuple): Tuple of bool define use dropout or not
|
|
for each layer, such as (True, True, False)
|
|
drop_prob (Tuple): Tuple of float defined the drop prob,
|
|
such as (0.5, 0.5, 0)
|
|
use_activation(Tuple): Tuple of bool define use active function
|
|
or not, such as (True, True, False)
|
|
"""
|
|
|
|
def __init__(self, fc_layers, use_dropout, drop_prob, use_activation):
|
|
super().__init__()
|
|
self.fc_layers = fc_layers
|
|
self.use_dropout = use_dropout
|
|
self.drop_prob = drop_prob
|
|
self.use_activation = use_activation
|
|
self._check()
|
|
self.create_layers()
|
|
|
|
def _check(self):
|
|
"""Check input to avoid ValueError."""
|
|
if not isinstance(self.fc_layers, tuple):
|
|
raise TypeError(f'fc_layers require tuple, '
|
|
f'get {type(self.fc_layers)}')
|
|
|
|
if not isinstance(self.use_dropout, tuple):
|
|
raise TypeError(f'use_dropout require tuple, '
|
|
f'get {type(self.use_dropout)}')
|
|
|
|
if not isinstance(self.drop_prob, tuple):
|
|
raise TypeError(f'drop_prob require tuple, '
|
|
f'get {type(self.drop_prob)}')
|
|
|
|
if not isinstance(self.use_activation, tuple):
|
|
raise TypeError(f'use_activation require tuple, '
|
|
f'get {type(self.use_activation)}')
|
|
|
|
l_fc_layer = len(self.fc_layers)
|
|
l_use_drop = len(self.use_dropout)
|
|
l_drop_prob = len(self.drop_prob)
|
|
l_use_activation = len(self.use_activation)
|
|
|
|
pass_check = (
|
|
l_fc_layer >= 2 and l_use_drop < l_fc_layer
|
|
and l_drop_prob < l_fc_layer and l_use_activation < l_fc_layer
|
|
and l_drop_prob == l_use_drop)
|
|
|
|
if not pass_check:
|
|
msg = 'Wrong BaseDiscriminator parameters!'
|
|
raise ValueError(msg)
|
|
|
|
def create_layers(self):
|
|
"""Create layers."""
|
|
l_fc_layer = len(self.fc_layers)
|
|
l_use_drop = len(self.use_dropout)
|
|
l_use_activation = len(self.use_activation)
|
|
|
|
self.fc_blocks = nn.Sequential()
|
|
|
|
for i in range(l_fc_layer - 1):
|
|
self.fc_blocks.add_module(
|
|
name=f'regressor_fc_{i}',
|
|
module=nn.Linear(
|
|
in_features=self.fc_layers[i],
|
|
out_features=self.fc_layers[i + 1]))
|
|
|
|
if i < l_use_activation and self.use_activation[i]:
|
|
self.fc_blocks.add_module(
|
|
name=f'regressor_af_{i}', module=nn.ReLU())
|
|
|
|
if i < l_use_drop and self.use_dropout[i]:
|
|
self.fc_blocks.add_module(
|
|
name=f'regressor_fc_dropout_{i}',
|
|
module=nn.Dropout(p=self.drop_prob[i]))
|
|
|
|
@abstractmethod
|
|
def forward(self, inputs):
|
|
"""Forward function."""
|
|
msg = 'the base class [BaseDiscriminator] is not callable!'
|
|
raise NotImplementedError(msg)
|
|
|
|
def init_weights(self):
|
|
"""Initialize model weights."""
|
|
for m in self.fc_blocks.named_modules():
|
|
if isinstance(m, nn.Linear):
|
|
xavier_init(m, gain=0.01)
|
|
|
|
|
|
class ShapeDiscriminator(BaseDiscriminator):
|
|
"""Discriminator for SMPL shape parameters, the inputs is (batch_size x 10)
|
|
|
|
Args:
|
|
fc_layers (Tuple): Tuple of neuron count, such as (10, 5, 1)
|
|
use_dropout (Tuple): Tuple of bool define use dropout or
|
|
not for each layer, such as (True, True, False)
|
|
drop_prob (Tuple): Tuple of float defined the drop prob,
|
|
such as (0.5, 0)
|
|
use_activation(Tuple): Tuple of bool define use active
|
|
function or not, such as (True, False)
|
|
"""
|
|
|
|
def __init__(self, fc_layers, use_dropout, drop_prob, use_activation):
|
|
if fc_layers[-1] != 1:
|
|
msg = f'the neuron count of the last layer ' \
|
|
f'must be 1, but got {fc_layers[-1]}'
|
|
raise ValueError(msg)
|
|
|
|
super().__init__(fc_layers, use_dropout, drop_prob, use_activation)
|
|
|
|
def forward(self, inputs):
|
|
"""Forward function."""
|
|
return self.fc_blocks(inputs)
|
|
|
|
|
|
class PoseDiscriminator(nn.Module):
|
|
"""Discriminator for SMPL pose parameters of each joint. It is composed of
|
|
discriminators for each joints. The inputs is (batch_size x joint_count x
|
|
9)
|
|
|
|
Args:
|
|
channels (Tuple): Tuple of channel number,
|
|
such as (9, 32, 32, 1)
|
|
joint_count (int): Joint number, such as 23
|
|
"""
|
|
|
|
def __init__(self, channels, joint_count):
|
|
super().__init__()
|
|
if channels[-1] != 1:
|
|
msg = f'the neuron count of the last layer ' \
|
|
f'must be 1, but got {channels[-1]}'
|
|
raise ValueError(msg)
|
|
self.joint_count = joint_count
|
|
|
|
self.conv_blocks = nn.Sequential()
|
|
len_channels = len(channels)
|
|
for idx in range(len_channels - 2):
|
|
self.conv_blocks.add_module(
|
|
name=f'conv_{idx}',
|
|
module=nn.Conv2d(
|
|
in_channels=channels[idx],
|
|
out_channels=channels[idx + 1],
|
|
kernel_size=1,
|
|
stride=1))
|
|
|
|
self.fc_layer = nn.ModuleList()
|
|
for idx in range(joint_count):
|
|
self.fc_layer.append(
|
|
nn.Linear(
|
|
in_features=channels[len_channels - 2], out_features=1))
|
|
|
|
def forward(self, inputs):
|
|
"""Forward function.
|
|
|
|
The input is (batch_size x joint_count x 9).
|
|
"""
|
|
# shape: batch_size x 9 x 1 x joint_count
|
|
inputs = inputs.transpose(1, 2).unsqueeze(2).contiguous()
|
|
# shape: batch_size x c x 1 x joint_count
|
|
internal_outputs = self.conv_blocks(inputs)
|
|
outputs = []
|
|
for idx in range(self.joint_count):
|
|
outputs.append(self.fc_layer[idx](internal_outputs[:, :, 0, idx]))
|
|
|
|
return torch.cat(outputs, 1), internal_outputs
|
|
|
|
def init_weights(self):
|
|
"""Initialize model weights."""
|
|
for m in self.conv_blocks:
|
|
if isinstance(m, nn.Conv2d):
|
|
normal_init(m, std=0.001, bias=0)
|
|
for m in self.fc_layer.named_modules():
|
|
if isinstance(m, nn.Linear):
|
|
xavier_init(m, gain=0.01)
|
|
|
|
|
|
class FullPoseDiscriminator(BaseDiscriminator):
|
|
"""Discriminator for SMPL pose parameters of all joints.
|
|
|
|
Args:
|
|
fc_layers (Tuple): Tuple of neuron count,
|
|
such as (736, 1024, 1024, 1)
|
|
use_dropout (Tuple): Tuple of bool define use dropout or not
|
|
for each layer, such as (True, True, False)
|
|
drop_prob (Tuple): Tuple of float defined the drop prob,
|
|
such as (0.5, 0.5, 0)
|
|
use_activation(Tuple): Tuple of bool define use active
|
|
function or not, such as (True, True, False)
|
|
"""
|
|
|
|
def __init__(self, fc_layers, use_dropout, drop_prob, use_activation):
|
|
if fc_layers[-1] != 1:
|
|
msg = f'the neuron count of the last layer must be 1,' \
|
|
f' but got {fc_layers[-1]}'
|
|
raise ValueError(msg)
|
|
|
|
super().__init__(fc_layers, use_dropout, drop_prob, use_activation)
|
|
|
|
def forward(self, inputs):
|
|
"""Forward function."""
|
|
return self.fc_blocks(inputs)
|
|
|
|
|
|
class SMPLDiscriminator(nn.Module):
|
|
"""Discriminator for SMPL pose and shape parameters. It is composed of a
|
|
discriminator for SMPL shape parameters, a discriminator for SMPL pose
|
|
parameters of all joints and a discriminator for SMPL pose parameters of
|
|
each joint.
|
|
|
|
Args:
|
|
beta_channel (tuple of int): Tuple of neuron count of the
|
|
discriminator of shape parameters. Defaults to (10, 5, 1)
|
|
per_joint_channel (tuple of int): Tuple of neuron count of the
|
|
discriminator of each joint. Defaults to (9, 32, 32, 1)
|
|
full_pose_channel (tuple of int): Tuple of neuron count of the
|
|
discriminator of full pose. Defaults to (23*32, 1024, 1024, 1)
|
|
"""
|
|
|
|
def __init__(self,
|
|
beta_channel=(10, 5, 1),
|
|
per_joint_channel=(9, 32, 32, 1),
|
|
full_pose_channel=(23 * 32, 1024, 1024, 1)):
|
|
super().__init__()
|
|
self.joint_count = 23
|
|
# The count of SMPL shape parameter is 10.
|
|
assert beta_channel[0] == 10
|
|
# Use 3 x 3 rotation matrix as the pose parameters
|
|
# of each joint, so the input channel is 9.
|
|
assert per_joint_channel[0] == 9
|
|
assert self.joint_count * per_joint_channel[-2] \
|
|
== full_pose_channel[0]
|
|
|
|
self.beta_channel = beta_channel
|
|
self.per_joint_channel = per_joint_channel
|
|
self.full_pose_channel = full_pose_channel
|
|
self._create_sub_modules()
|
|
|
|
def _create_sub_modules(self):
|
|
"""Create sub discriminators."""
|
|
|
|
# create theta discriminator for each joint
|
|
self.pose_discriminator = PoseDiscriminator(self.per_joint_channel,
|
|
self.joint_count)
|
|
|
|
# create full pose discriminator for total joints
|
|
fc_layers = self.full_pose_channel
|
|
use_dropout = tuple([False] * (len(fc_layers) - 1))
|
|
drop_prob = tuple([0.5] * (len(fc_layers) - 1))
|
|
use_activation = tuple([True] * (len(fc_layers) - 2) + [False])
|
|
|
|
self.full_pose_discriminator = FullPoseDiscriminator(
|
|
fc_layers, use_dropout, drop_prob, use_activation)
|
|
|
|
# create shape discriminator for betas
|
|
fc_layers = self.beta_channel
|
|
use_dropout = tuple([False] * (len(fc_layers) - 1))
|
|
drop_prob = tuple([0.5] * (len(fc_layers) - 1))
|
|
use_activation = tuple([True] * (len(fc_layers) - 2) + [False])
|
|
self.shape_discriminator = ShapeDiscriminator(fc_layers, use_dropout,
|
|
drop_prob,
|
|
use_activation)
|
|
|
|
def forward(self, thetas):
|
|
"""Forward function."""
|
|
_, poses, shapes = thetas
|
|
|
|
batch_size = poses.shape[0]
|
|
shape_disc_value = self.shape_discriminator(shapes)
|
|
|
|
# The first rotation matrix is global rotation
|
|
# and is NOT used in discriminator.
|
|
if poses.dim() == 2:
|
|
rotate_matrixs = \
|
|
batch_rodrigues(poses.contiguous().view(-1, 3)
|
|
).view(batch_size, 24, 9)[:, 1:, :]
|
|
else:
|
|
rotate_matrixs = poses.contiguous().view(batch_size, 24,
|
|
9)[:, 1:, :].contiguous()
|
|
pose_disc_value, pose_inter_disc_value \
|
|
= self.pose_discriminator(rotate_matrixs)
|
|
full_pose_disc_value = self.full_pose_discriminator(
|
|
pose_inter_disc_value.contiguous().view(batch_size, -1))
|
|
return torch.cat(
|
|
(pose_disc_value, full_pose_disc_value, shape_disc_value), 1)
|
|
|
|
def init_weights(self):
|
|
"""Initialize model weights."""
|
|
self.full_pose_discriminator.init_weights()
|
|
self.pose_discriminator.init_weights()
|
|
self.shape_discriminator.init_weights()
|
|
|