@HEADS.register_module
class DLCRNetHead(HeatmapHead):
"""A head for DLCRNet models using Part-Affinity Fields to predict individuals."""
def __init__(
self,
predictor: BasePredictor,
target_generator: BaseGenerator,
criterion: dict[str, BaseCriterion],
aggregator: BaseLossAggregator,
heatmap_config: dict,
locref_config: dict,
paf_config: dict,
num_stages: int = 5,
features_dim: int = 128,
weight_init: str | dict | BaseWeightInitializer | None = None,
) -> None:
self.num_stages = num_stages
# FIXME Cleaner __init__ to avoid initializing unused layers
in_channels = heatmap_config["channels"][0]
num_keypoints = heatmap_config["channels"][-1]
num_limbs = paf_config["channels"][-1] # Already has the 2x multiplier
in_refined_channels = features_dim + num_keypoints + num_limbs
if num_stages > 0:
heatmap_config["channels"][0] = paf_config["channels"][0] = in_refined_channels
locref_config["channels"][0] = locref_config["channels"][-1]
super().__init__(
predictor,
target_generator,
criterion,
aggregator,
heatmap_config,
locref_config,
weight_init,
)
if num_stages > 0:
self.stride *= 2 # extra deconv layer where it's multi-stage
self.paf_head = DeconvModule(**paf_config)
self.convt1 = self._make_layer_same_padding(in_channels=in_channels, out_channels=num_keypoints)
self.convt2 = self._make_layer_same_padding(in_channels=in_channels, out_channels=locref_config["channels"][-1])
self.convt3 = self._make_layer_same_padding(in_channels=in_channels, out_channels=num_limbs)
self.convt4 = self._make_layer_same_padding(in_channels=in_channels, out_channels=features_dim)
self.hm_ref_layers = nn.ModuleList()
self.paf_ref_layers = nn.ModuleList()
for _ in range(num_stages):
self.hm_ref_layers.append(
self._make_refinement_layer(in_channels=in_refined_channels, out_channels=num_keypoints)
)
self.paf_ref_layers.append(
self._make_refinement_layer(in_channels=in_refined_channels, out_channels=num_limbs)
)
self._init_weights()
def _make_layer_same_padding(self, in_channels: int, out_channels: int) -> nn.ConvTranspose2d:
# FIXME There is no consensual solution to emulate TF behavior in pytorch
# see https://github.com/pytorch/pytorch/issues/3867
return nn.ConvTranspose2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=3,
stride=2,
padding=1,
output_padding=1,
)
def _make_refinement_layer(self, in_channels: int, out_channels: int) -> nn.Conv2d:
"""Summary:
Helper function to create a refinement layer.
Args:
in_channels: number of input channels
out_channels: number of output channels
Returns:
refinement_layer: the refinement layer.
"""
return nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding="same")
def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
if self.num_stages > 0:
stage1_hm_out = self.convt1(x)
stage1_paf_out = self.convt3(x)
features = self.convt4(x)
stage2_in = torch.cat((stage1_hm_out, stage1_paf_out, features), dim=1)
stage_in = stage2_in
stage_paf_out = stage1_paf_out
stage_hm_out = stage1_hm_out
for i, (hm_ref_layer, paf_ref_layer) in enumerate(
zip(self.hm_ref_layers, self.paf_ref_layers, strict=True)
):
pre_stage_hm_out = stage_hm_out
stage_hm_out = hm_ref_layer(stage_in)
stage_paf_out = paf_ref_layer(stage_in)
if i > 0:
stage_hm_out += pre_stage_hm_out
stage_in = torch.cat((stage_hm_out, stage_paf_out, features), dim=1)
return {
"heatmap": self.heatmap_head(stage_in),
"locref": self.locref_head(self.convt2(x)),
"paf": self.paf_head(stage_in),
}
return {
"heatmap": self.heatmap_head(x),
"locref": self.locref_head(x),
"paf": self.paf_head(x),
}