Method forward

detrsmpl/models/heads/pare_head.py:509–537

(self, features)

        return nn.Sequential(*layers)

    def forward(self, features):
        batch_size = features.shape[0]

        init_pose = self.init_pose.expand(batch_size, -1)  # N, Jx6
        init_shape = self.init_shape.expand(batch_size, -1)
        init_cam = self.init_cam.expand(batch_size, -1)

        output = {}

        part_feats = self._get_2d_branch_feats(features)

        part_attention = self._get_part_attention_map(part_feats, output)

        smpl_feats = self._get_3d_smpl_feats(features, part_feats)

        point_local_feat, cam_shape_feats = self._get_local_feats(
            smpl_feats, part_attention, output)

        pred_pose, pred_shape, pred_cam = self._get_final_preds(
            point_local_feat, cam_shape_feats, init_pose, init_shape, init_cam)

        pred_rotmat = rot6d_to_rotmat(pred_pose).reshape(batch_size, 24, 3, 3)

        output.update({
            'pred_pose': pred_rotmat,
            'pred_cam': pred_cam,
            'pred_shape': pred_shape,
        })
        return output

    def _get_local_feats(self, smpl_feats, part_attention, output):
        # 1x1 conv