def _compute_kp_heatmap_loss(self, input_height, input_width, task_name, heatmap_predictions, classification_loss_fn, per_pixel_weights):
    'Computes the heatmap loss of the keypoint estimation task.\n\n    Args:\n      input_height: An integer scalar tensor representing input image height.\n      input_width: An integer scalar tensor representing input image width.\n      task_name: A string representing the name of the keypoint task.\n      heatmap_predictions: A list of float tensors of shape [batch_size,\n        out_height, out_width, num_keypoints] representing the prediction heads\n        of the model for keypoint heatmap.\n      classification_loss_fn: An object_detection.core.losses.Loss object to\n        compute the loss for the class predictions in CenterNet.\n      per_pixel_weights: A float tensor of shape [batch_size,\n        out_height * out_width, 1] with 1s in locations where the spatial\n        coordinates fall within the height and width in true_image_shapes.\n\n    Returns:\n      loss: A float scalar tensor representing the object keypoint heatmap loss\n        normalized by number of instances.\n    '
    gt_keypoints_list = self.groundtruth_lists(fields.BoxListFields.keypoints)
    gt_classes_list = self.groundtruth_lists(fields.BoxListFields.classes)
    gt_weights_list = self.groundtruth_lists(fields.BoxListFields.weights)
    gt_boxes_list = self.groundtruth_lists(fields.BoxListFields.boxes)
    assigner = self._target_assigner_dict[task_name]
    (keypoint_heatmap, num_instances_per_kp_type, valid_mask_batch) = assigner.assign_keypoint_heatmap_targets(height=input_height, width=input_width, gt_keypoints_list=gt_keypoints_list, gt_weights_list=gt_weights_list, gt_classes_list=gt_classes_list, gt_boxes_list=gt_boxes_list)
    flattened_valid_mask = _flatten_spatial_dimensions(tf.expand_dims(valid_mask_batch, axis=(- 1)))
    flattened_heapmap_targets = _flatten_spatial_dimensions(keypoint_heatmap)
    num_instances = tf.maximum(tf.cast(tf.reduce_sum(num_instances_per_kp_type), dtype=tf.float32), 1.0)
    loss = 0.0
    for pred in heatmap_predictions:
        pred = _flatten_spatial_dimensions(pred)
        unweighted_loss = classification_loss_fn(pred, flattened_heapmap_targets, weights=tf.ones_like(per_pixel_weights))
        loss += ((unweighted_loss * per_pixel_weights) * flattened_valid_mask)
    loss = (tf.reduce_sum(loss) / (float(len(heatmap_predictions)) * num_instances))
    return loss
Computes the heatmap loss of the keypoint estimation task.

Args:
  input_height: An integer scalar tensor representing input image height.
  input_width: An integer scalar tensor representing input image width.
  task_name: A string representing the name of the keypoint task.
  heatmap_predictions: A list of float tensors of shape [batch_size,
    out_height, out_width, num_keypoints] representing the prediction heads
    of the model for keypoint heatmap.
  classification_loss_fn: An object_detection.core.losses.Loss object to
    compute the loss for the class predictions in CenterNet.
  per_pixel_weights: A float tensor of shape [batch_size,
    out_height * out_width, 1] with 1s in locations where the spatial
    coordinates fall within the height and width in true_image_shapes.

Returns:
  loss: A float scalar tensor representing the object keypoint heatmap loss
    normalized by number of instances.
def _compute_kp_heatmap_loss(self, input_height, input_width, task_name, heatmap_predictions, classification_loss_fn, per_pixel_weights):
    'Computes the heatmap loss of the keypoint estimation task.\n\n    Args:\n      input_height: An integer scalar tensor representing input image height.\n      input_width: An integer scalar tensor representing input image width.\n      task_name: A string representing the name of the keypoint task.\n      heatmap_predictions: A list of float tensors of shape [batch_size,\n        out_height, out_width, num_keypoints] representing the prediction heads\n        of the model for keypoint heatmap.\n      classification_loss_fn: An object_detection.core.losses.Loss object to\n        compute the loss for the class predictions in CenterNet.\n      per_pixel_weights: A float tensor of shape [batch_size,\n        out_height * out_width, 1] with 1s in locations where the spatial\n        coordinates fall within the height and width in true_image_shapes.\n\n    Returns:\n      loss: A float scalar tensor representing the object keypoint heatmap loss\n        normalized by number of instances.\n    '
    gt_keypoints_list = self.groundtruth_lists(fields.BoxListFields.keypoints)
    gt_classes_list = self.groundtruth_lists(fields.BoxListFields.classes)
    gt_weights_list = self.groundtruth_lists(fields.BoxListFields.weights)
    gt_boxes_list = self.groundtruth_lists(fields.BoxListFields.boxes)
    assigner = self._target_assigner_dict[task_name]
    (keypoint_heatmap, num_instances_per_kp_type, valid_mask_batch) = assigner.assign_keypoint_heatmap_targets(height=input_height, width=input_width, gt_keypoints_list=gt_keypoints_list, gt_weights_list=gt_weights_list, gt_classes_list=gt_classes_list, gt_boxes_list=gt_boxes_list)
    flattened_valid_mask = _flatten_spatial_dimensions(tf.expand_dims(valid_mask_batch, axis=(- 1)))
    flattened_heapmap_targets = _flatten_spatial_dimensions(keypoint_heatmap)
    num_instances = tf.maximum(tf.cast(tf.reduce_sum(num_instances_per_kp_type), dtype=tf.float32), 1.0)
    loss = 0.0
    for pred in heatmap_predictions:
        pred = _flatten_spatial_dimensions(pred)
        unweighted_loss = classification_loss_fn(pred, flattened_heapmap_targets, weights=tf.ones_like(per_pixel_weights))
        loss += ((unweighted_loss * per_pixel_weights) * flattened_valid_mask)
    loss = (tf.reduce_sum(loss) / (float(len(heatmap_predictions)) * num_instances))
    return loss