def _compute_densepose_part_and_coordinate_losses(self, input_height, input_width, part_predictions, surface_coord_predictions):
    'Computes the individual losses for the DensePose 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      part_predictions: A list of float tensors of shape [batch_size,\n        out_height, out_width, num_parts].\n      surface_coord_predictions: A list of float tensors of shape [batch_size,\n        out_height, out_width, 2 * num_parts].\n\n    Returns:\n      A tuple with two scalar loss tensors: part_prediction_loss and\n      surface_coord_loss.\n    '
    gt_dp_num_points_list = self.groundtruth_lists(fields.BoxListFields.densepose_num_points)
    gt_dp_part_ids_list = self.groundtruth_lists(fields.BoxListFields.densepose_part_ids)
    gt_dp_surface_coords_list = self.groundtruth_lists(fields.BoxListFields.densepose_surface_coords)
    gt_weights_list = self.groundtruth_lists(fields.BoxListFields.weights)
    assigner = self._target_assigner_dict[DENSEPOSE_TASK]
    (batch_indices, batch_part_ids, batch_surface_coords, batch_weights) = assigner.assign_part_and_coordinate_targets(height=input_height, width=input_width, gt_dp_num_points_list=gt_dp_num_points_list, gt_dp_part_ids_list=gt_dp_part_ids_list, gt_dp_surface_coords_list=gt_dp_surface_coords_list, gt_weights_list=gt_weights_list)
    part_prediction_loss = 0
    surface_coord_loss = 0
    classification_loss_fn = self._densepose_params.classification_loss
    localization_loss_fn = self._densepose_params.localization_loss
    num_predictions = float(len(part_predictions))
    num_valid_points = tf.math.count_nonzero(batch_weights)
    num_valid_points = tf.cast(tf.math.maximum(num_valid_points, 1), tf.float32)
    for (part_pred, surface_coord_pred) in zip(part_predictions, surface_coord_predictions):
        if self._densepose_params.upsample_to_input_res:
            part_pred = tf.keras.layers.UpSampling2D(self._stride, interpolation=self._densepose_params.upsample_method)(part_pred)
            surface_coord_pred = tf.keras.layers.UpSampling2D(self._stride, interpolation=self._densepose_params.upsample_method)(surface_coord_pred)
        part_pred = cn_assigner.get_batch_predictions_from_indices(part_pred, batch_indices[:, 0:3])
        part_prediction_loss += classification_loss_fn(part_pred[:, tf.newaxis, :], batch_part_ids[:, tf.newaxis, :], weights=batch_weights[:, tf.newaxis, tf.newaxis])
        (batch_size, out_height, out_width, _) = _get_shape(surface_coord_pred, 4)
        surface_coord_pred = tf.reshape(surface_coord_pred, [batch_size, out_height, out_width, (- 1), 2])
        surface_coord_pred = cn_assigner.get_batch_predictions_from_indices(surface_coord_pred, batch_indices)
        surface_coord_loss += localization_loss_fn(surface_coord_pred, batch_surface_coords, weights=batch_weights[:, tf.newaxis])
    part_prediction_loss = (tf.reduce_sum(part_prediction_loss) / (num_predictions * num_valid_points))
    surface_coord_loss = (tf.reduce_sum(surface_coord_loss) / (num_predictions * num_valid_points))
    return (part_prediction_loss, surface_coord_loss)
Computes the individual losses for the DensePose task.

Args:
  input_height: An integer scalar tensor representing input image height.
  input_width: An integer scalar tensor representing input image width.
  part_predictions: A list of float tensors of shape [batch_size,
    out_height, out_width, num_parts].
  surface_coord_predictions: A list of float tensors of shape [batch_size,
    out_height, out_width, 2 * num_parts].

Returns:
  A tuple with two scalar loss tensors: part_prediction_loss and
  surface_coord_loss.
def _compute_densepose_part_and_coordinate_losses(self, input_height, input_width, part_predictions, surface_coord_predictions):
    'Computes the individual losses for the DensePose 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      part_predictions: A list of float tensors of shape [batch_size,\n        out_height, out_width, num_parts].\n      surface_coord_predictions: A list of float tensors of shape [batch_size,\n        out_height, out_width, 2 * num_parts].\n\n    Returns:\n      A tuple with two scalar loss tensors: part_prediction_loss and\n      surface_coord_loss.\n    '
    gt_dp_num_points_list = self.groundtruth_lists(fields.BoxListFields.densepose_num_points)
    gt_dp_part_ids_list = self.groundtruth_lists(fields.BoxListFields.densepose_part_ids)
    gt_dp_surface_coords_list = self.groundtruth_lists(fields.BoxListFields.densepose_surface_coords)
    gt_weights_list = self.groundtruth_lists(fields.BoxListFields.weights)
    assigner = self._target_assigner_dict[DENSEPOSE_TASK]
    (batch_indices, batch_part_ids, batch_surface_coords, batch_weights) = assigner.assign_part_and_coordinate_targets(height=input_height, width=input_width, gt_dp_num_points_list=gt_dp_num_points_list, gt_dp_part_ids_list=gt_dp_part_ids_list, gt_dp_surface_coords_list=gt_dp_surface_coords_list, gt_weights_list=gt_weights_list)
    part_prediction_loss = 0
    surface_coord_loss = 0
    classification_loss_fn = self._densepose_params.classification_loss
    localization_loss_fn = self._densepose_params.localization_loss
    num_predictions = float(len(part_predictions))
    num_valid_points = tf.math.count_nonzero(batch_weights)
    num_valid_points = tf.cast(tf.math.maximum(num_valid_points, 1), tf.float32)
    for (part_pred, surface_coord_pred) in zip(part_predictions, surface_coord_predictions):
        if self._densepose_params.upsample_to_input_res:
            part_pred = tf.keras.layers.UpSampling2D(self._stride, interpolation=self._densepose_params.upsample_method)(part_pred)
            surface_coord_pred = tf.keras.layers.UpSampling2D(self._stride, interpolation=self._densepose_params.upsample_method)(surface_coord_pred)
        part_pred = cn_assigner.get_batch_predictions_from_indices(part_pred, batch_indices[:, 0:3])
        part_prediction_loss += classification_loss_fn(part_pred[:, tf.newaxis, :], batch_part_ids[:, tf.newaxis, :], weights=batch_weights[:, tf.newaxis, tf.newaxis])
        (batch_size, out_height, out_width, _) = _get_shape(surface_coord_pred, 4)
        surface_coord_pred = tf.reshape(surface_coord_pred, [batch_size, out_height, out_width, (- 1), 2])
        surface_coord_pred = cn_assigner.get_batch_predictions_from_indices(surface_coord_pred, batch_indices)
        surface_coord_loss += localization_loss_fn(surface_coord_pred, batch_surface_coords, weights=batch_weights[:, tf.newaxis])
    part_prediction_loss = (tf.reduce_sum(part_prediction_loss) / (num_predictions * num_valid_points))
    surface_coord_loss = (tf.reduce_sum(surface_coord_loss) / (num_predictions * num_valid_points))
    return (part_prediction_loss, surface_coord_loss)