def __init__(self,
keypoint_head_cfg,
root_head_cfg,
hand_type_head_cfg,
loss_keypoint=None,
loss_root_depth=None,
loss_hand_type=None,
train_cfg=None,
test_cfg=None):
super().__init__()
# build sub-module heads
self.right_hand_head = Heatmap3DHead(**keypoint_head_cfg)
self.left_hand_head = Heatmap3DHead(**keypoint_head_cfg)
self.root_head = Heatmap1DHead(**root_head_cfg)
self.hand_type_head = MultilabelClassificationHead(
**hand_type_head_cfg)
self.neck = GlobalAveragePooling()
# build losses
self.keypoint_loss = build_loss(loss_keypoint)
self.root_depth_loss = build_loss(loss_root_depth)
self.hand_type_loss = build_loss(loss_hand_type)
self.train_cfg = {} if train_cfg is None else train_cfg
self.test_cfg = {} if test_cfg is None else test_cfg
self.target_type = self.test_cfg.get('target_type', 'GaussianHeatMap')
def __init__(self,
in_channels,
num_joints,
max_norm=None,
loss_keypoint=None,
train_cfg=None,
test_cfg=None):
super().__init__()
self.in_channels = in_channels
self.num_joints = num_joints
self.max_norm = max_norm
self.loss = build_loss(loss_keypoint)
self.train_cfg = {} if train_cfg is None else train_cfg
self.test_cfg = {} if test_cfg is None else test_cfg
self.conv = build_conv_layer(
dict(type='Conv1d'), in_channels, num_joints * 3, 1)
if self.max_norm is not None:
# Apply weight norm clip to conv layers
weight_clip = WeightNormClipHook(self.max_norm)
for module in self.modules():
if isinstance(module, nn.modules.conv._ConvNd):
weight_clip.register(module)
def __init__(self,
in_channels,
num_joints,
num_deconv_layers=3,
num_deconv_filters=(256, 256, 256),
num_deconv_kernels=(4, 4, 4),
tag_per_joint=True,
with_ae_loss=None,
extra=None,
loss_keypoint=None):
super().__init__()
self.loss = build_loss(loss_keypoint)
self.in_channels = in_channels
dim_tag = num_joints if tag_per_joint else 1
if with_ae_loss[0]:
out_channels = num_joints + dim_tag
else:
out_channels = num_joints
if extra is not None and not isinstance(extra, dict):
raise TypeError('extra should be dict or None.')
if num_deconv_layers > 0:
self.deconv_layers = self._make_deconv_layer(
num_deconv_layers,
num_deconv_filters,
num_deconv_kernels,
)
elif num_deconv_layers == 0:
self.deconv_layers = nn.Identity()
else:
raise ValueError(
f'num_deconv_layers ({num_deconv_layers}) should >= 0.')
if extra is not None and 'final_conv_kernel' in extra:
assert extra['final_conv_kernel'] in [1, 3]
if extra['final_conv_kernel'] == 3:
padding = 1
else:
padding = 0
kernel_size = extra['final_conv_kernel']
else:
kernel_size = 1
padding = 0
self.final_layer = build_conv_layer(
cfg=dict(type='Conv2d'),
in_channels=num_deconv_filters[-1]
if num_deconv_layers > 0 else in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=1,
padding=padding)
def __init__(self,
in_channels,
num_joints,
tag_per_joint=True,
extra=None,
num_deconv_layers=1,
num_deconv_filters=(32, ),
num_deconv_kernels=(4, ),
num_basic_blocks=4,
cat_output=None,
with_ae_loss=None,
loss_keypoint=None):
super().__init__()
self.loss = build_loss(loss_keypoint)
dim_tag = num_joints if tag_per_joint else 1
self.num_deconvs = num_deconv_layers
self.cat_output = cat_output
final_layer_output_channels = []
if with_ae_loss[0]:
out_channels = num_joints + dim_tag
else:
out_channels = num_joints
final_layer_output_channels.append(out_channels)
for i in range(num_deconv_layers):
if with_ae_loss[i + 1]:
out_channels = num_joints + dim_tag
else:
out_channels = num_joints
final_layer_output_channels.append(out_channels)
deconv_layer_output_channels = []
for i in range(num_deconv_layers):
if with_ae_loss[i]:
out_channels = num_joints + dim_tag
else:
out_channels = num_joints
deconv_layer_output_channels.append(out_channels)
self.final_layers = self._make_final_layers(
in_channels, final_layer_output_channels, extra, num_deconv_layers,
num_deconv_filters)
self.deconv_layers = self._make_deconv_layers(
in_channels, deconv_layer_output_channels, num_deconv_layers,
num_deconv_filters, num_deconv_kernels, num_basic_blocks,
cat_output)
def __init__(self,
in_channels,
num_joints,
loss_keypoint=None,
train_cfg=None,
test_cfg=None):
super().__init__()
self.in_channels = in_channels
self.num_joints = num_joints
self.loss = build_loss(loss_keypoint)
self.train_cfg = {} if train_cfg is None else train_cfg
self.test_cfg = {} if test_cfg is None else test_cfg
self.fc = nn.Linear(self.in_channels, self.num_joints * 2)
def __init__(self,
in_channels,
num_joints,
loss_keypoint=None,
train_cfg=None,
test_cfg=None):
super().__init__()
self.in_channels = in_channels
self.num_joints = num_joints
self.loss = build_loss(loss_keypoint)
self.train_cfg = {} if train_cfg is None else train_cfg
self.test_cfg = {} if test_cfg is None else test_cfg
self.conv = build_conv_layer(dict(type='Conv1d'), in_channels,
num_joints * 3, 1)
def __init__(self,
in_channels=2048,
heatmap_size=64,
hidden_dims=(512, ),
loss_value=None,
train_cfg=None,
test_cfg=None):
super().__init__()
self.loss = build_loss(loss_value)
self.in_channels = in_channels
self.heatmap_size = heatmap_size
self.train_cfg = {} if train_cfg is None else train_cfg
self.test_cfg = {} if test_cfg is None else test_cfg
feature_dims = [in_channels] + \
[dim for dim in hidden_dims] + \
[heatmap_size]
self.fc = self._make_linear_layers(feature_dims)
def __init__(self,
in_channels=2048,
num_labels=2,
hidden_dims=(512, ),
loss_classification=None,
train_cfg=None,
test_cfg=None):
super().__init__()
self.loss = build_loss(loss_classification)
self.in_channels = in_channels
self.num_labesl = num_labels
self.train_cfg = {} if train_cfg is None else train_cfg
self.test_cfg = {} if test_cfg is None else test_cfg
feature_dims = [in_channels] + \
[dim for dim in hidden_dims] + \
[num_labels]
self.fc = self._make_linear_layers(feature_dims)
def __init__(self,
backbone,
keypoint_head=None,
train_cfg=None,
test_cfg=None,
pretrained=None,
loss_pose=None):
super().__init__()
self.backbone = builder.build_backbone(backbone)
if keypoint_head is not None:
self.keypoint_head = builder.build_head(keypoint_head)
self.train_cfg = train_cfg
self.test_cfg = test_cfg
self.parser = HeatmapParser(self.test_cfg)
self.loss = build_loss(loss_pose)
self.init_weights(pretrained=pretrained)
def __init__(self,
out_shape,
unit_channels=256,
out_channels=17,
num_stages=4,
num_units=4,
use_prm=False,
norm_cfg=dict(type='BN'),
loss_keypoint=None,
train_cfg=None,
test_cfg=None):
# Protect mutable default arguments
norm_cfg = cp.deepcopy(norm_cfg)
super().__init__()
self.train_cfg = {} if train_cfg is None else train_cfg
self.test_cfg = {} if test_cfg is None else test_cfg
self.target_type = self.test_cfg.get('target_type', 'GaussianHeatMap')
self.out_shape = out_shape
self.unit_channels = unit_channels
self.out_channels = out_channels
self.num_stages = num_stages
self.num_units = num_units
self.loss = build_loss(loss_keypoint)
self.predict_layers = nn.ModuleList([])
for i in range(self.num_stages):
for j in range(self.num_units):
self.predict_layers.append(
PredictHeatmap(
unit_channels,
out_channels,
out_shape,
use_prm,
norm_cfg=norm_cfg))
def __init__(self,
in_channels,
out_channels,
num_deconv_layers=3,
num_deconv_filters=(256, 256, 256),
num_deconv_kernels=(4, 4, 4),
extra=None,
in_index=0,
input_transform=None,
align_corners=False,
loss_keypoint=None,
train_cfg=None,
test_cfg=None):
super().__init__()
self.in_channels = in_channels
self.loss = build_loss(loss_keypoint)
self.train_cfg = {} if train_cfg is None else train_cfg
self.test_cfg = {} if test_cfg is None else test_cfg
self.target_type = self.test_cfg.get('target_type', 'GaussianHeatMap')
self._init_inputs(in_channels, in_index, input_transform)
self.in_index = in_index
self.align_corners = align_corners
if extra is not None and not isinstance(extra, dict):
raise TypeError('extra should be dict or None.')
if num_deconv_layers > 0:
self.deconv_layers = self._make_deconv_layer(
num_deconv_layers,
num_deconv_filters,
num_deconv_kernels,
)
elif num_deconv_layers == 0:
self.deconv_layers = nn.Identity()
else:
raise ValueError(
f'num_deconv_layers ({num_deconv_layers}) should >= 0.')
identity_final_layer = False
if extra is not None and 'final_conv_kernel' in extra:
assert extra['final_conv_kernel'] in [0, 1, 3]
if extra['final_conv_kernel'] == 3:
padding = 1
elif extra['final_conv_kernel'] == 1:
padding = 0
else:
# 0 for Identity mapping.
identity_final_layer = True
kernel_size = extra['final_conv_kernel']
else:
kernel_size = 1
padding = 0
if identity_final_layer:
self.final_layer = nn.Identity()
else:
conv_channels = num_deconv_filters[
-1] if num_deconv_layers > 0 else self.in_channels
layers = []
if extra is not None:
num_conv_layers = extra.get('num_conv_layers', 0)
num_conv_kernels = extra.get('num_conv_kernels',
[1] * num_conv_layers)
for i in range(num_conv_layers):
layers.append(
build_conv_layer(
dict(type='Conv2d'),
in_channels=conv_channels,
out_channels=conv_channels,
kernel_size=num_conv_kernels[i],
stride=1,
padding=(num_conv_kernels[i] - 1) // 2))
layers.append(
build_norm_layer(dict(type='BN'), conv_channels)[1])
layers.append(nn.ReLU(inplace=True))
layers.append(
build_conv_layer(
cfg=dict(type='Conv2d'),
in_channels=conv_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=1,
padding=padding))
if len(layers) > 1:
self.final_layer = nn.Sequential(*layers)
else:
self.final_layer = layers[0]
def __init__(self,
in_channels,
out_channels,
num_stages=1,
num_deconv_layers=3,
num_deconv_filters=(256, 256, 256),
num_deconv_kernels=(4, 4, 4),
extra=None,
loss_keypoint=None):
super().__init__()
self.loss = build_loss(loss_keypoint)
self.in_channels = in_channels
self.num_stages = num_stages
if extra is not None and not isinstance(extra, dict):
raise TypeError('extra should be dict or None.')
# build multi-stage deconv layers
self.multi_deconv_layers = nn.ModuleList([])
for _ in range(self.num_stages):
if num_deconv_layers > 0:
deconv_layers = self._make_deconv_layer(
num_deconv_layers,
num_deconv_filters,
num_deconv_kernels,
)
elif num_deconv_layers == 0:
deconv_layers = nn.Identity()
else:
raise ValueError(
f'num_deconv_layers ({num_deconv_layers}) should >= 0.')
self.multi_deconv_layers.append(deconv_layers)
identity_final_layer = False
if extra is not None and 'final_conv_kernel' in extra:
assert extra['final_conv_kernel'] in [0, 1, 3]
if extra['final_conv_kernel'] == 3:
padding = 1
elif extra['final_conv_kernel'] == 1:
padding = 0
else:
# 0 for Identity mapping.
identity_final_layer = True
kernel_size = extra['final_conv_kernel']
else:
kernel_size = 1
padding = 0
# build multi-stage final layers
self.multi_final_layers = nn.ModuleList([])
for i in range(self.num_stages):
if identity_final_layer:
final_layer = nn.Identity()
else:
final_layer = build_conv_layer(
cfg=dict(type='Conv2d'),
in_channels=num_deconv_filters[-1]
if num_deconv_layers > 0 else in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=1,
padding=padding)
self.multi_final_layers.append(final_layer)
