def forward()

in model.py [0:0]

• 63 lines of code
• 10 McCabe index (conditional complexity)

    def forward(self, kp_loc=None, kp_vis=None,
                class_mask=None, K=None, **kwargs):

        # dictionary with outputs of the fw pass
        preds = {}

        # input sizes ...
        ba, kp_dim, n_kp = kp_loc.shape

        assert kp_dim == 2, 'bad input keypoint dim'
        assert n_kp == self.n_keypoints, 'bad # of keypoints!'

        if self.projection_type == 'perspective':
            assert K is not None
            kp_loc_cal = self.calibrate_keypoints(kp_loc, K)
        else:
            kp_loc_cal = kp_loc

        # normalize keypoints
        kp_loc_norm, kp_mean = \
            self.normalize_keypoints(
                kp_loc_cal, kp_vis, rescale=self.keypoint_rescale)
        # save for later visualisations ...
        preds['kp_loc_norm'] = kp_loc_norm
        preds['kp_mean'] = kp_mean

        # run the shape predictor
        preds['phi'] = self.run_phi(kp_loc_norm, kp_vis, class_mask=class_mask)

        if self.canonicalization['use']:
            preds['l_canonicalization'], preds['psi'] = \
                self.canonicalization_loss(preds['phi'],
                                           class_mask=class_mask)

        # 3D->2D project shape to camera
        kp_reprojected, depth = self.camera_projection(
            preds['phi']['shape_camera_coord'])
        preds['kp_reprojected'] = kp_reprojected

        # compute the repro loss for backpropagation
        if self.reprojection_normalization == 'kp_count_per_image':
            preds['l_reprojection'] = avg_l2_huber(
                kp_reprojected,
                kp_loc_norm,
                mask=kp_vis,
                squared=self.squared_reprojection_loss)

        elif self.reprojection_normalization == 'kp_total_count':
            def flatten_(x): return x.permute(
                1, 2, 0).contiguous().view(1, 2, self.n_keypoints*ba)
            preds['l_reprojection'] = avg_l2_huber(
                flatten_(kp_reprojected),
                flatten_(kp_loc_norm),
                mask=kp_vis.permute(1, 0).contiguous().view(1, -1),
                scaling=self.huber_scaling)

        else:
            raise ValueError('unknown loss normalization %s' %
                             self.loss_normalization)

        # unnormalize the shape projections
        kp_reprojected_image = \
            self.unnormalize_keypoints(kp_reprojected, kp_mean,
                                       rescale=self.keypoint_rescale)

        # projections in the image coordinate frame
        if self.replace_keypoints_with_input and not self.training:
            # use the input points
            kp_reprojected_image = \
                (1-kp_vis[:, None, :]) * kp_reprojected_image + \
                kp_vis[:, None, :] * kp_loc_cal

        preds['kp_reprojected_image'] = kp_reprojected_image

        # projected 3D shape in the image space
        #   = unprojection of kp_reprojected_image
        shape_image_coord = self.camera_unprojection(
            kp_reprojected_image, depth,
            rescale=self.keypoint_rescale)

        if self.projection_type == 'perspective':
            preds['shape_image_coord_cal'] = shape_image_coord
            shape_image_coord = \
                self.uncalibrate_keypoints(shape_image_coord, K)
            preds['kp_reprojected_image_uncal'], _ = \
                self.camera_projection(shape_image_coord)

        preds['shape_image_coord'] = shape_image_coord

        # get the final loss
        preds['objective'] = self.get_objective(preds)
        assert np.isfinite(
            preds['objective'].sum().data.cpu().numpy()), "nans!"

        return preds