jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def compute_pointcloud_weights_based_on_voxel_density(points, grid_cell_size):
  """Computes pointcloud weights based on voxel density.

  Args:
    points: A tf.float32 tensor of size [num_points, 3].
    grid_cell_size: The size of the grid cells in x, y, z dimensions in the
      voxel grid. It should be either a tf.float32 tensor, a numpy array or a
      list of size [3].

  Returns:
    A tf.float32 tensor of size [num_points, 1] containing weights that are
      inverse proportional to the denisty of the points in voxels.
  """
  num_points = tf.shape(points)[0]
  features = tf.ones([num_points, 1], dtype=tf.float32)
  voxel_features, _, segment_ids, _ = (
      pointcloud_to_sparse_voxel_grid_unbatched(
          points=points,
          features=features,
          grid_cell_size=grid_cell_size,
          segment_func=tf.math.unsorted_segment_sum))
  num_voxels = tf.shape(voxel_features)[0]
  point_features = sparse_voxel_grid_to_pointcloud(
      voxel_features=tf.expand_dims(voxel_features, axis=0),
      segment_ids=tf.expand_dims(segment_ids, axis=0),
      num_valid_voxels=tf.expand_dims(num_voxels, axis=0),
      num_valid_points=tf.expand_dims(num_points, axis=0))
  inverse_point_densities = 1.0 / tf.squeeze(point_features, axis=0)
  total_inverse_density = tf.reduce_sum(inverse_point_densities)
  return (inverse_point_densities * tf.to_float(num_points) /
          total_inverse_density)
//  [python]
def compute_pointcloud_weights_based_on_voxel_density(points, grid_cell_size):
  """Computes pointcloud weights based on voxel density.

  Args:
    points: A tf.float32 tensor of size [num_points, 3].
    grid_cell_size: The size of the grid cells in x, y, z dimensions in the
      voxel grid. It should be either a tf.float32 tensor, a numpy array or a
      list of size [3].

  Returns:
    A tf.float32 tensor of size [num_points, 1] containing weights that are
      inverse proportional to the denisty of the points in voxels.
  """
  num_points = tf.shape(points)[0]
  features = tf.ones([num_points, 1], dtype=tf.float32)
  voxel_features, _, segment_ids, _ = (
      pointcloud_to_sparse_voxel_grid_unbatched(
          points=points,
          features=features,
          grid_cell_size=grid_cell_size,
          segment_func=tf.math.unsorted_segment_sum))
  num_voxels = tf.shape(voxel_features)[0]
  point_features = sparse_voxel_grid_to_pointcloud(
      voxel_features=tf.expand_dims(voxel_features, axis=0),
      segment_ids=tf.expand_dims(segment_ids, axis=0),
      num_valid_voxels=tf.expand_dims(num_voxels, axis=0),
      num_valid_points=tf.expand_dims(num_points, axis=0))
  inverse_point_densities = 1.0 / tf.squeeze(point_features, axis=0)
  total_inverse_density = tf.reduce_sum(inverse_point_densities)
  return (inverse_point_densities * tf.cast(num_points, dtype=tf.float32) /
          total_inverse_density)
//  [python]
def pointcloud_to_voxel_grid(points,
                             features,
                             grid_cell_size,
                             start_location,
                             end_location,
                             segment_func=tf.math.unsorted_segment_mean):
  """Converts a pointcloud into a voxel grid.

  Args:
    points: A tf.float32 tensor of size [N, 3].
    features: A tf.float32 tensor of size [N, F].
    grid_cell_size: A tf.float32 tensor of size [3].
    start_location: A tf.float32 tensor of size [3].
    end_location: A tf.float32 tensor of size [3].
    segment_func: A tensorflow function that operates on segments. Expect one
      of tf.math.unsorted_segment_{min/max/mean/prod/sum}. Defaults to
      tf.math.unsorted_segment_mean

  Returns:
    voxel_features: A tf.float32 tensor of
      size [grid_x_len, grid_y_len, grid_z_len, F].
    segment_ids: A tf.int32 tensor of IDs for each point indicating
      which (flattened) voxel cell its data was mapped to.
    point_indices: A tf.int32 tensor of size [num_points, 3] containing the
      location of each point in the 3d voxel grid.
  """
  grid_cell_size = tf.convert_to_tensor(grid_cell_size, dtype=tf.float32)
  start_location = tf.convert_to_tensor(start_location, dtype=tf.float32)
  end_location = tf.convert_to_tensor(end_location, dtype=tf.float32)
  point_indices = tf.to_int32(
      (points - tf.expand_dims(start_location, axis=0)) /
      tf.expand_dims(grid_cell_size, axis=0))
  grid_size = tf.to_int32(
      tf.math.ceil((end_location - start_location) / grid_cell_size))
  # Note: all points outside the grid are added to the edges
  # Cap index at grid_size - 1 (so a 10x10x10 grid's max cell is (9,9,9))
  point_indices = tf.minimum(point_indices, tf.expand_dims(grid_size - 1,
                                                           axis=0))
  # Don't allow any points below index (0, 0, 0)
  point_indices = tf.maximum(point_indices, 0)
  segment_ids = tf.reduce_sum(
      point_indices * tf.stack(
          [grid_size[1] * grid_size[2], grid_size[2], 1], axis=0),
      axis=1)
  voxel_features = segment_func(
      data=features,
      segment_ids=segment_ids,
      num_segments=(grid_size[0] * grid_size[1] * grid_size[2]))
  return (tf.reshape(voxel_features,
                     [grid_size[0],
                      grid_size[1],
                      grid_size[2],
                      features.get_shape().as_list()[1]]),
          segment_ids,
          point_indices)
//  [python]
def pointcloud_to_voxel_grid(points,
                             features,
                             grid_cell_size,
                             start_location,
                             end_location,
                             segment_func=tf.math.unsorted_segment_mean):
  """Converts a pointcloud into a voxel grid.

  Args:
    points: A tf.float32 tensor of size [N, 3].
    features: A tf.float32 tensor of size [N, F].
    grid_cell_size: A tf.float32 tensor of size [3].
    start_location: A tf.float32 tensor of size [3].
    end_location: A tf.float32 tensor of size [3].
    segment_func: A tensorflow function that operates on segments. Expect one
      of tf.math.unsorted_segment_{min/max/mean/prod/sum}. Defaults to
      tf.math.unsorted_segment_mean

  Returns:
    voxel_features: A tf.float32 tensor of
      size [grid_x_len, grid_y_len, grid_z_len, F].
    segment_ids: A tf.int32 tensor of IDs for each point indicating
      which (flattened) voxel cell its data was mapped to.
    point_indices: A tf.int32 tensor of size [num_points, 3] containing the
      location of each point in the 3d voxel grid.
  """
  grid_cell_size = tf.convert_to_tensor(grid_cell_size, dtype=tf.float32)
  start_location = tf.convert_to_tensor(start_location, dtype=tf.float32)
  end_location = tf.convert_to_tensor(end_location, dtype=tf.float32)
  point_indices = tf.cast(
      (points - tf.expand_dims(start_location, axis=0)) /
      tf.expand_dims(grid_cell_size, axis=0),
      dtype=tf.int32)
  grid_size = tf.cast(
      tf.math.ceil((end_location - start_location) / grid_cell_size),
      dtype=tf.int32)
  # Note: all points outside the grid are added to the edges
  # Cap index at grid_size - 1 (so a 10x10x10 grid's max cell is (9,9,9))
  point_indices = tf.minimum(point_indices, tf.expand_dims(grid_size - 1,
                                                           axis=0))
  # Don't allow any points below index (0, 0, 0)
  point_indices = tf.maximum(point_indices, 0)
  segment_ids = tf.reduce_sum(
      point_indices * tf.stack(
          [grid_size[1] * grid_size[2], grid_size[2], 1], axis=0),
      axis=1)
  voxel_features = segment_func(
      data=features,
      segment_ids=segment_ids,
      num_segments=(grid_size[0] * grid_size[1] * grid_size[2]))
  return (tf.reshape(voxel_features,
                     [grid_size[0],
                      grid_size[1],
                      grid_size[2],
                      features.get_shape().as_list()[1]]),
          segment_ids,
          point_indices)
//  [python]
def predictions(logits):
  """Class prediction from logits."""
  inner_dim = logits.get_shape().as_list()[-1]
  with tf.name_scope('predictions'):
    # For binary classification
    if inner_dim == 1:
      pred = tf.to_int64(tf.greater(tf.squeeze(logits, -1), 0.))
    # For multi-class classification
    else:
      pred = tf.argmax(logits, 2)
    return pred
//  [python]
def predictions(logits):
  """Class prediction from logits."""
  inner_dim = logits.get_shape().as_list()[-1]
  with tf.name_scope('predictions'):
    # For binary classification
    if inner_dim == 1:
      pred = tf.cast(tf.greater(tf.squeeze(logits, -1), 0.), tf.int64)
    # For multi-class classification
    else:
      pred = tf.argmax(logits, 2)
    return pred
//  [python]
def accuracy(logits, targets, weights):
  """Computes prediction accuracy.

  Args:
    logits: 2-D classifier logits [timesteps*batch_size, num_classes]
    targets: 1-D [timesteps*batch_size] integer tensor.
    weights: 1-D [timesteps*batch_size] float tensor.

  Returns:
    Accuracy: float scalar.
  """
  with tf.name_scope('accuracy'):
    eq = tf.to_float(tf.equal(predictions(logits), targets))
    return tf.identity(
        tf.reduce_sum(weights * eq) / _num_labels(weights), name='accuracy')
//  [python]
def accuracy(logits, targets, weights):
  """Computes prediction accuracy.

  Args:
    logits: 2-D classifier logits [timesteps*batch_size, num_classes]
    targets: 1-D [timesteps*batch_size] integer tensor.
    weights: 1-D [timesteps*batch_size] float tensor.

  Returns:
    Accuracy: float scalar.
  """
  with tf.name_scope('accuracy'):
    eq = tf.cast(tf.equal(predictions(logits), targets), tf.float32)
    return tf.identity(
        tf.reduce_sum(weights * eq) / _num_labels(weights), name='accuracy')
//  [python]
def __call__(self, score_outputs, labels):
    """Computes total RPN detection loss.

    Computes total RPN detection loss including box and score from all levels.

    Args:
      score_outputs: an OrderDict with keys representing levels and values
        representing scores in [batch_size, height, width, num_anchors].
      labels: the dictionary that returned from dataloader that includes
        groundturth targets.

    Returns:
      rpn_score_loss: a scalar tensor representing total score loss.
    """
    with tf.name_scope('rpn_loss'):
      levels = sorted(score_outputs.keys())

      score_losses = []
      for level in levels:
        score_losses.append(
            self._rpn_score_loss(
                score_outputs[level],
                labels[level],
                normalizer=tf.to_float(
                    tf.shape(score_outputs[level])[0] *
                    self._rpn_batch_size_per_im)))

      # Sums per level losses to total loss.
      return tf.math.add_n(score_losses)
//  [python]
def __call__(self, score_outputs, labels):
    """Computes total RPN detection loss.

    Computes total RPN detection loss including box and score from all levels.

    Args:
      score_outputs: an OrderDict with keys representing levels and values
        representing scores in [batch_size, height, width, num_anchors].
      labels: the dictionary that returned from dataloader that includes
        groundturth targets.

    Returns:
      rpn_score_loss: a scalar tensor representing total score loss.
    """
    with tf.name_scope('rpn_loss'):
      levels = sorted(score_outputs.keys())

      score_losses = []
      for level in levels:
        score_losses.append(
            self._rpn_score_loss(
                score_outputs[level],
                labels[level],
                normalizer=tf.cast(
                    tf.shape(score_outputs[level])[0] *
                    self._rpn_batch_size_per_im, dtype=tf.float32)))

      # Sums per level losses to total loss.
      return tf.math.add_n(score_losses)
//  [python]
def _rpn_box_loss(self, box_outputs, box_targets, normalizer=1.0):
    """Computes box regression loss."""
    with tf.name_scope('rpn_box_loss'):
      mask = tf.cast(tf.not_equal(box_targets, 0.0), dtype=tf.float32)
      box_targets = tf.expand_dims(box_targets, axis=-1)
      box_outputs = tf.expand_dims(box_outputs, axis=-1)
      box_loss = self._huber_loss(box_targets, box_outputs, sample_weight=mask)
      # The loss is normalized by the sum of non-zero weights and additional
      # normalizer provided by the function caller. Using + 0.01 here to avoid
      # division by zero.
      box_loss /= normalizer * (tf.reduce_sum(mask) + 0.01)
      return box_loss
//  [python]
def _rpn_center_l1_loss(self, center_outputs, center_targets,
                          normalizer=1.0):
    """Computes centerness regression loss."""
    # for instances, the regression targets of 512x512 input with 6 anchors on
    # P2-P6 pyramid is about [0.1, 0.1, 0.2, 0.2].
    with tf.name_scope('rpn_center_loss'):

      # mask = tf.greater(center_targets, 0.0)  # inside box only.
      mask = tf.greater(center_targets, -1.0)  # in and out of box.
      center_targets = tf.maximum(center_targets, tf.zeros_like(center_targets))
      center_outputs = tf.sigmoid(center_outputs)
      center_targets = tf.expand_dims(center_targets, -1)
      center_outputs = tf.expand_dims(center_outputs, -1)
      mask = tf.to_float(mask)
      center_loss = self._l1_loss(center_targets, center_outputs,
                                  sample_weight=mask)
      center_loss /= normalizer
      return center_loss
//  [python]
def _rpn_center_l1_loss(self, center_outputs, center_targets,
                          normalizer=1.0):
    """Computes centerness regression loss."""
    # for instances, the regression targets of 512x512 input with 6 anchors on
    # P2-P6 pyramid is about [0.1, 0.1, 0.2, 0.2].
    with tf.name_scope('rpn_center_loss'):

      # mask = tf.greater(center_targets, 0.0)  # inside box only.
      mask = tf.greater(center_targets, -1.0)  # in and out of box.
      center_targets = tf.maximum(center_targets, tf.zeros_like(center_targets))
      center_outputs = tf.sigmoid(center_outputs)
      center_targets = tf.expand_dims(center_targets, -1)
      center_outputs = tf.expand_dims(center_outputs, -1)
      mask = tf.cast(mask, dtype=tf.float32)
      center_loss = self._l1_loss(center_targets, center_outputs,
                                  sample_weight=mask)
      center_loss /= normalizer
      return center_loss
//  [python]
def __call__(self, center_outputs, labels):
    """Computes total RPN centerness regression loss.

    Computes total RPN centerness score regression loss from all levels.

    Args:
      center_outputs: an OrderDict with keys representing levels and values
        representing anchor centerness regression targets in
        [batch_size, height, width, num_anchors * 4].
      labels: the dictionary that returned from dataloader that includes
        groundturth targets.

    Returns:
      rpn_center_loss: a scalar tensor representing total centerness regression
        loss.
    """
    with tf.name_scope('rpn_loss'):
      # Normalizer.
      levels = sorted(center_outputs.keys())
      num_valid = 0
      # 0<pos<1, neg=0, ign=-1
      for level in levels:
        num_valid += tf.reduce_sum(tf.to_float(
            tf.greater(labels[level], -1.0)))  # in and out of box
      num_valid += 1e-12

      # Centerness loss over multi levels.
      center_losses = []
      for level in levels:
        center_losses.append(
            self._rpn_center_l1_loss(
                center_outputs[level], labels[level],
                normalizer=num_valid))

      # Sum per level losses to total loss.
      return tf.add_n(center_losses)
//  [python]
def __call__(self, center_outputs, labels):
    """Computes total RPN centerness regression loss.

    Computes total RPN centerness score regression loss from all levels.

    Args:
      center_outputs: an OrderDict with keys representing levels and values
        representing anchor centerness regression targets in
        [batch_size, height, width, num_anchors * 4].
      labels: the dictionary that returned from dataloader that includes
        groundturth targets.

    Returns:
      rpn_center_loss: a scalar tensor representing total centerness regression
        loss.
    """
    with tf.name_scope('rpn_loss'):
      # Normalizer.
      levels = sorted(center_outputs.keys())
      num_valid = 0
      # 0<pos<1, neg=0, ign=-1
      for level in levels:
        num_valid += tf.reduce_sum(tf.cast(
            tf.greater(labels[level], -1.0), tf.float32))  # in and out of box
      num_valid += 1e-12

      # Centerness loss over multi levels.
      center_losses = []
      for level in levels:
        center_losses.append(
            self._rpn_center_l1_loss(
                center_outputs[level], labels[level],
                normalizer=num_valid))

      # Sum per level losses to total loss.
      return tf.add_n(center_losses)
//  [python]
def __call__(self, class_outputs, class_targets):
    """Computes the class loss (Fast-RCNN branch) of Mask-RCNN.

    This function implements the classification loss of the Fast-RCNN.

    The classification loss is softmax on all RoIs.
    Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py  # pylint: disable=line-too-long

    Args:
      class_outputs: a float tensor representing the class prediction for each box
        with a shape of [batch_size, num_boxes, num_classes].
      class_targets: a float tensor representing the class label for each box
        with a shape of [batch_size, num_boxes].

    Returns:
      a scalar tensor representing total class loss.
    """
    with tf.name_scope('fast_rcnn_loss'):
      batch_size, num_boxes, num_classes = class_outputs.get_shape().as_list()
      class_targets = tf.to_int32(class_targets)
      class_targets_one_hot = tf.one_hot(class_targets, num_classes)
      return self._fast_rcnn_class_loss(class_outputs, class_targets_one_hot,
                                        normalizer=batch_size * num_boxes / 2.0)
//  [python]
def __call__(self, class_outputs, class_targets):
    """Computes the class loss (Fast-RCNN branch) of Mask-RCNN.

    This function implements the classification loss of the Fast-RCNN.

    The classification loss is softmax on all RoIs.
    Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py  # pylint: disable=line-too-long

    Args:
      class_outputs: a float tensor representing the class prediction for each box
        with a shape of [batch_size, num_boxes, num_classes].
      class_targets: a float tensor representing the class label for each box
        with a shape of [batch_size, num_boxes].

    Returns:
      a scalar tensor representing total class loss.
    """
    with tf.name_scope('fast_rcnn_loss'):
      batch_size, num_boxes, num_classes = class_outputs.get_shape().as_list()
      class_targets = tf.cast(class_targets, dtype=tf.int32)
      class_targets_one_hot = tf.one_hot(class_targets, num_classes)
      return self._fast_rcnn_class_loss(class_outputs, class_targets_one_hot,
                                        normalizer=batch_size * num_boxes / 2.0)
//  [python]
def __call__(self, box_outputs, labels, center_targets):
    """Computes total RPN detection loss.

    Computes total RPN box regression loss from all levels.

    Args:
      box_outputs: an OrderDict with keys representing levels and values
        representing box regression targets in
        [batch_size, height, width, num_anchors * 4].
        last channel: (left, right, top, bottom).
      labels: the dictionary that returned from dataloader that includes
        groundturth targets (left, right, top, bottom).
      center_targets: valid_target mask.

    Returns:
      rpn_iou_loss: a scalar tensor representing total box regression loss.
    """
    with tf.name_scope('rpn_loss'):
      # Normalizer.
      levels = sorted(box_outputs.keys())
      normalizer = 0.
      for level in levels:
        # center_targets pos>0, neg=0, ign=-1.
        mask_ = tf.to_float(tf.logical_and(
            tf.greater(center_targets[level][..., 0], 0.0),
            tf.greater(tf.reduce_min(labels[level], -1), 0.0)))
        normalizer += tf.reduce_sum(mask_)
      normalizer += 1e-8
      # iou_loss over multi levels.
      iou_losses = []
      for level in levels:
        iou_losses.append(
            self._rpn_iou_loss(
                box_outputs[level], labels[level],
                center_weight=center_targets[level][..., 0],
                normalizer=normalizer))
      # Sum per level losses to total loss.
      return tf.add_n(iou_losses)
//  [python]
def __call__(self, box_outputs, labels, center_targets):
    """Computes total RPN detection loss.

    Computes total RPN box regression loss from all levels.

    Args:
      box_outputs: an OrderDict with keys representing levels and values
        representing box regression targets in
        [batch_size, height, width, num_anchors * 4].
        last channel: (left, right, top, bottom).
      labels: the dictionary that returned from dataloader that includes
        groundturth targets (left, right, top, bottom).
      center_targets: valid_target mask.

    Returns:
      rpn_iou_loss: a scalar tensor representing total box regression loss.
    """
    with tf.name_scope('rpn_loss'):
      # Normalizer.
      levels = sorted(box_outputs.keys())
      normalizer = 0.
      for level in levels:
        # center_targets pos>0, neg=0, ign=-1.
        mask_ = tf.cast(tf.logical_and(
            tf.greater(center_targets[level][..., 0], 0.0),
            tf.greater(tf.reduce_min(labels[level], -1), 0.0)), tf.float32)
        normalizer += tf.reduce_sum(mask_)
      normalizer += 1e-8
      # iou_loss over multi levels.
      iou_losses = []
      for level in levels:
        iou_losses.append(
            self._rpn_iou_loss(
                box_outputs[level], labels[level],
                center_weight=center_targets[level][..., 0],
                normalizer=normalizer))
      # Sum per level losses to total loss.
      return tf.add_n(iou_losses)
//  [python]
def _fast_rcnn_box_loss(self, box_outputs, box_targets, class_targets,
                          normalizer=1.0):
    """Computes box regression loss."""
    with tf.name_scope('fast_rcnn_box_loss'):
      mask = tf.tile(tf.expand_dims(tf.greater(class_targets, 0), axis=2),
                     [1, 1, 4])
      mask = tf.cast(mask, dtype=tf.float32)
      box_targets = tf.expand_dims(box_targets, axis=-1)
      box_outputs = tf.expand_dims(box_outputs, axis=-1)
      box_loss = self._huber_loss(box_targets, box_outputs, sample_weight=mask)
      # The loss is normalized by the number of ones in mask,
      # additianal normalizer provided by the user and using 0.01 here to avoid
      # division by 0.
      box_loss /= normalizer * (tf.reduce_sum(mask) + 0.01)
      return box_loss
//  [python]
def _rpn_iou_loss(self, box_outputs, box_targets,
                    center_weight=None, normalizer=1.0):
    """Computes box regression loss."""
    # for instances, the regression targets of 512x512 input with 6 anchors on
    # P2-P6 pyramid is about [0.1, 0.1, 0.2, 0.2].
    with tf.name_scope('rpn_iou_loss'):
      mask = tf.logical_and(
          tf.greater(center_weight, 0.0),
          tf.greater(tf.reduce_min(box_targets, -1), 0.0))

      pred_left = box_outputs[..., 0]
      pred_right = box_outputs[..., 1]
      pred_top = box_outputs[..., 2]
      pred_bottom = box_outputs[..., 3]

      gt_left = box_targets[..., 0]
      gt_right = box_targets[..., 1]
      gt_top = box_targets[..., 2]
      gt_bottom = box_targets[..., 3]

      inter_width = (tf.minimum(pred_left, gt_left) +
                     tf.minimum(pred_right, gt_right))
      inter_height = (tf.minimum(pred_top, gt_top) +
                      tf.minimum(pred_bottom, gt_bottom))
      inter_area = inter_width * inter_height
      union_area = ((pred_left + pred_right) * (pred_top + pred_bottom) +
                    (gt_left + gt_right) * (gt_top + gt_bottom) -
                    inter_area)
      iou = inter_area / (union_area + 1e-8)
      mask_ = tf.to_float(mask)
      iou = tf.clip_by_value(iou, clip_value_min=1e-8, clip_value_max=1.0)
      neg_log_iou = -tf.math.log(iou)
      iou_loss = tf.reduce_sum(neg_log_iou * mask_)
      iou_loss /= normalizer
      return iou_loss
//  [python]
def _rpn_iou_loss(self, box_outputs, box_targets,
                    center_weight=None, normalizer=1.0):
    """Computes box regression loss."""
    # for instances, the regression targets of 512x512 input with 6 anchors on
    # P2-P6 pyramid is about [0.1, 0.1, 0.2, 0.2].
    with tf.name_scope('rpn_iou_loss'):
      mask = tf.logical_and(
          tf.greater(center_weight, 0.0),
          tf.greater(tf.reduce_min(box_targets, -1), 0.0))

      pred_left = box_outputs[..., 0]
      pred_right = box_outputs[..., 1]
      pred_top = box_outputs[..., 2]
      pred_bottom = box_outputs[..., 3]

      gt_left = box_targets[..., 0]
      gt_right = box_targets[..., 1]
      gt_top = box_targets[..., 2]
      gt_bottom = box_targets[..., 3]

      inter_width = (tf.minimum(pred_left, gt_left) +
                     tf.minimum(pred_right, gt_right))
      inter_height = (tf.minimum(pred_top, gt_top) +
                      tf.minimum(pred_bottom, gt_bottom))
      inter_area = inter_width * inter_height
      union_area = ((pred_left + pred_right) * (pred_top + pred_bottom) +
                    (gt_left + gt_right) * (gt_top + gt_bottom) -
                    inter_area)
      iou = inter_area / (union_area + 1e-8)
      mask_ = tf.cast(mask, tf.float32)
      iou = tf.clip_by_value(iou, clip_value_min=1e-8, clip_value_max=1.0)
      neg_log_iou = -tf.math.log(iou)
      iou_loss = tf.reduce_sum(neg_log_iou * mask_)
      iou_loss /= normalizer
      return iou_loss
//  [python]
def box_loss(self, box_outputs, box_targets, num_positives):
    """Computes RetinaNet box regression loss."""
    # The delta is typically around the mean value of regression target.
    # for instances, the regression targets of 512x512 input with 6 anchors on
    # P3-P7 pyramid is about [0.1, 0.1, 0.2, 0.2].
    normalizer = num_positives * 4.0
    mask = tf.to_float(tf.not_equal(box_targets, 0.0))
    box_targets = tf.expand_dims(box_targets, axis=-1)
    box_outputs = tf.expand_dims(box_outputs, axis=-1)
    box_loss = self._huber_loss(box_targets, box_outputs, sample_weight=mask)
    box_loss /= normalizer
    return box_loss
//  [python]
def box_loss(self, box_outputs, box_targets, num_positives):
    """Computes RetinaNet box regression loss."""
    # The delta is typically around the mean value of regression target.
    # for instances, the regression targets of 512x512 input with 6 anchors on
    # P3-P7 pyramid is about [0.1, 0.1, 0.2, 0.2].
    normalizer = num_positives * 4.0
    mask = tf.cast(tf.not_equal(box_targets, 0.0), dtype=tf.float32)
    box_targets = tf.expand_dims(box_targets, axis=-1)
    box_outputs = tf.expand_dims(box_outputs, axis=-1)
    box_loss = self._huber_loss(box_targets, box_outputs, sample_weight=mask)
    box_loss /= normalizer
    return box_loss
//  [python]
def __call__(self, score_outputs, score_targets):
    """Computes the class loss (Fast-RCNN branch) of Mask-RCNN.

    This function implements the classification loss of the Fast-RCNN.

    The classification loss is softmax on all RoIs.
    Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py  # pylint: disable=line-too-long

    Args:
      score_outputs: a float tensor representing the class prediction for each box
        with a shape of [batch_size, num_boxes, num_classes].
      score_targets: a float tensor representing the class label for each box
        with a shape of [batch_size, num_boxes].

    Returns:
      a scalar tensor representing total score loss.
    """
    with tf.name_scope('fast_rcnn_loss'):
      score_outputs = tf.squeeze(score_outputs, -1)

      mask = tf.greater(score_targets, self._ignore_threshold)
      num_valid = tf.reduce_sum(tf.to_float(mask))
      score_targets = tf.maximum(score_targets, tf.zeros_like(score_targets))
      score_outputs = tf.sigmoid(score_outputs)
      score_targets = tf.expand_dims(score_targets, -1)
      score_outputs = tf.expand_dims(score_outputs, -1)
      mask = tf.to_float(mask)
      score_loss = self._l1_loss(score_targets, score_outputs,
                                 sample_weight=mask)
      score_loss /= (num_valid + 1e-10)
      return score_loss
//  [python]
def __call__(self, score_outputs, score_targets):
    """Computes the class loss (Fast-RCNN branch) of Mask-RCNN.

    This function implements the classification loss of the Fast-RCNN.

    The classification loss is softmax on all RoIs.
    Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py  # pylint: disable=line-too-long

    Args:
      score_outputs: a float tensor representing the class prediction for each box
        with a shape of [batch_size, num_boxes, num_classes].
      score_targets: a float tensor representing the class label for each box
        with a shape of [batch_size, num_boxes].

    Returns:
      a scalar tensor representing total score loss.
    """
    with tf.name_scope('fast_rcnn_loss'):
      score_outputs = tf.squeeze(score_outputs, -1)

      mask = tf.greater(score_targets, self._ignore_threshold)
      num_valid = tf.reduce_sum(tf.cast(mask, tf.float32))
      score_targets = tf.maximum(score_targets, tf.zeros_like(score_targets))
      score_outputs = tf.sigmoid(score_outputs)
      score_targets = tf.expand_dims(score_targets, -1)
      score_outputs = tf.expand_dims(score_outputs, -1)
      mask = tf.cast(mask, dtype=tf.float32)
      score_loss = self._l1_loss(score_targets, score_outputs,
                                 sample_weight=mask)
      score_loss /= (num_valid + 1e-10)
      return score_loss
//  [python]
def class_loss(self, cls_outputs, cls_targets, num_positives,
                 ignore_label=-2):
    """Computes RetinaNet classification loss."""
    # Onehot encoding for classification labels.
    cls_targets_one_hot = tf.one_hot(cls_targets, self._num_classes)
    bs, height, width, _, _ = cls_targets_one_hot.get_shape().as_list()
    cls_targets_one_hot = tf.reshape(cls_targets_one_hot,
                                     [bs, height, width, -1])
    loss = focal_loss(t.cast(cls_outputs, dtype=tf.float32),
                      t.cast(cls_targets_one_hot, dtype=tf.float32),
                      self._focal_loss_alpha,
                      self._focal_loss_gamma,
                      num_positives)

    ignore_loss = tf.where(
        tf.equal(cls_targets, ignore_label),
        tf.zeros_like(cls_targets, dtype=tf.float32),
        tf.ones_like(cls_targets, dtype=tf.float32),
    )
    ignore_loss = tf.expand_dims(ignore_loss, -1)
    ignore_loss = tf.tile(ignore_loss, [1, 1, 1, 1, self._num_classes])
    ignore_loss = tf.reshape(ignore_loss, tf.shape(input=loss))
    return tf.reduce_sum(input_tensor=ignore_loss * loss)
//  [python]
def class_loss(self, cls_outputs, cls_targets, num_positives,
                 ignore_label=-2):
    """Computes RetinaNet classification loss."""
    # Onehot encoding for classification labels.
    cls_targets_one_hot = tf.one_hot(cls_targets, self._num_classes)
    bs, height, width, _, _ = cls_targets_one_hot.get_shape().as_list()
    cls_targets_one_hot = tf.reshape(cls_targets_one_hot,
                                     [bs, height, width, -1])
    loss = focal_loss(tf.cast(cls_outputs, dtype=tf.float32),
                      tf.cast(cls_targets_one_hot, dtype=tf.float32),
                      self._focal_loss_alpha,
                      self._focal_loss_gamma,
                      num_positives)

    ignore_loss = tf.where(
        tf.equal(cls_targets, ignore_label),
        tf.zeros_like(cls_targets, dtype=tf.float32),
        tf.ones_like(cls_targets, dtype=tf.float32),
    )
    ignore_loss = tf.expand_dims(ignore_loss, -1)
    ignore_loss = tf.tile(ignore_loss, [1, 1, 1, 1, self._num_classes])
    ignore_loss = tf.reshape(ignore_loss, tf.shape(input=loss))
    return tf.reduce_sum(input_tensor=ignore_loss * loss)
//  [python]
def __call__(self, box_outputs, class_targets, box_targets):
    """Computes the box loss (Fast-RCNN branch) of Mask-RCNN.

    This function implements the box regression loss of the Fast-RCNN. As the
    `box_outputs` produces `num_classes` boxes for each RoI, the reference model
    expands `box_targets` to match the shape of `box_outputs` and selects only
    the target that the RoI has a maximum overlap. (Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/roi_data/fast_rcnn.py)  # pylint: disable=line-too-long
    Instead, this function selects the `box_outputs` by the `class_targets` so
    that it doesn't expand `box_targets`.

    The box loss is smooth L1-loss on only positive samples of RoIs.
    Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py  # pylint: disable=line-too-long

    Args:
      box_outputs: a float tensor representing the box prediction for each box
        with a shape of [batch_size, num_boxes, num_classes * 4].
      class_targets: a float tensor representing the class label for each box
        with a shape of [batch_size, num_boxes].
      box_targets: a float tensor representing the box label for each box
        with a shape of [batch_size, num_boxes, 4].

    Returns:
      box_loss: a scalar tensor representing total box regression loss.
    """
    with tf.name_scope('fast_rcnn_loss'):
      class_targets = tf.to_int32(class_targets)

      # Selects the box from `box_outputs` based on `class_targets`, with which
      # the box has the maximum overlap.
      (batch_size, num_rois,
       num_class_specific_boxes) = box_outputs.get_shape().as_list()
      num_classes = num_class_specific_boxes // 4
      box_outputs = tf.reshape(box_outputs,
                               [batch_size, num_rois, num_classes, 4])

      box_indices = tf.reshape(
          class_targets + tf.tile(
              tf.expand_dims(
                  tf.range(batch_size) * num_rois * num_classes, 1),
              [1, num_rois]) + tf.tile(
                  tf.expand_dims(tf.range(num_rois) * num_classes, 0),
                  [batch_size, 1]), [-1])

      box_outputs = tf.matmul(
          tf.one_hot(
              box_indices,
              batch_size * num_rois * num_classes,
              dtype=box_outputs.dtype), tf.reshape(box_outputs, [-1, 4]))
      box_outputs = tf.reshape(box_outputs, [batch_size, -1, 4])

      return self._fast_rcnn_box_loss(box_outputs, box_targets, class_targets)
//  [python]
def __call__(self, box_outputs, class_targets, box_targets):
    """Computes the box loss (Fast-RCNN branch) of Mask-RCNN.

    This function implements the box regression loss of the Fast-RCNN. As the
    `box_outputs` produces `num_classes` boxes for each RoI, the reference model
    expands `box_targets` to match the shape of `box_outputs` and selects only
    the target that the RoI has a maximum overlap. (Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/roi_data/fast_rcnn.py)  # pylint: disable=line-too-long
    Instead, this function selects the `box_outputs` by the `class_targets` so
    that it doesn't expand `box_targets`.

    The box loss is smooth L1-loss on only positive samples of RoIs.
    Reference: https://github.com/facebookresearch/Detectron/blob/master/detectron/modeling/fast_rcnn_heads.py  # pylint: disable=line-too-long

    Args:
      box_outputs: a float tensor representing the box prediction for each box
        with a shape of [batch_size, num_boxes, num_classes * 4].
      class_targets: a float tensor representing the class label for each box
        with a shape of [batch_size, num_boxes].
      box_targets: a float tensor representing the box label for each box
        with a shape of [batch_size, num_boxes, 4].

    Returns:
      box_loss: a scalar tensor representing total box regression loss.
    """
    with tf.name_scope('fast_rcnn_loss'):
      class_targets = tf.cast(class_targets, dtype=tf.int32)

      # Selects the box from `box_outputs` based on `class_targets`, with which
      # the box has the maximum overlap.
      (batch_size, num_rois,
       num_class_specific_boxes) = box_outputs.get_shape().as_list()
      num_classes = num_class_specific_boxes // 4
      box_outputs = tf.reshape(box_outputs,
                               [batch_size, num_rois, num_classes, 4])

      box_indices = tf.reshape(
          class_targets + tf.tile(
              tf.expand_dims(
                  tf.range(batch_size) * num_rois * num_classes, 1),
              [1, num_rois]) + tf.tile(
                  tf.expand_dims(tf.range(num_rois) * num_classes, 0),
                  [batch_size, 1]), [-1])

      box_outputs = tf.matmul(
          tf.one_hot(
              box_indices,
              batch_size * num_rois * num_classes,
              dtype=box_outputs.dtype), tf.reshape(box_outputs, [-1, 4]))
      box_outputs = tf.reshape(box_outputs, [batch_size, -1, 4])

      return self._fast_rcnn_box_loss(box_outputs, box_targets, class_targets)
//  [python]
def __call__(self, probs, labels, valid_mask):
    """Compute instance segmentation loss.

    Args:
      probs: A Tensor of shape [batch_size * num_points, height, width,
        num_classes]. The logits are not necessarily between 0 and 1.
      labels: A float32/float16 Tensor of shape [batch_size, num_instances,
          mask_size, mask_size], where mask_size =
          mask_crop_size * gt_upsample_scale for fine mask, or mask_crop_size
          for coarse masks and shape priors.
      valid_mask: a binary mask indicating valid training masks.

    Returns:
      loss: an float tensor representing total mask classification loss.
    """
    with tf.name_scope('shapemask_prior_loss'):
      batch_size, num_instances = valid_mask.get_shape().as_list()[:2]
      diff = (tf.to_float(labels) -
              tf.to_float(probs))
      diff *= tf.to_float(
          tf.reshape(valid_mask, [batch_size, num_instances, 1, 1]))
      # Adding 0.001 in the denominator to avoid division by zero.
      loss = tf.nn.l2_loss(diff) / (tf.reduce_sum(labels) + 0.001)
    return loss
//  [python]
def __call__(self, probs, labels, valid_mask):
    """Compute instance segmentation loss.

    Args:
      probs: A Tensor of shape [batch_size * num_points, height, width,
        num_classes]. The logits are not necessarily between 0 and 1.
      labels: A float32/float16 Tensor of shape [batch_size, num_instances,
          mask_size, mask_size], where mask_size =
          mask_crop_size * gt_upsample_scale for fine mask, or mask_crop_size
          for coarse masks and shape priors.
      valid_mask: a binary mask indicating valid training masks.

    Returns:
      loss: an float tensor representing total mask classification loss.
    """
    with tf.name_scope('shapemask_prior_loss'):
      batch_size, num_instances = valid_mask.get_shape().as_list()[:2]
      diff = (tf.cast(labels, dtype=tf.float32) -
              tf.cast(probs, dtype=tf.float32))
      diff *= tf.cast(
          tf.reshape(valid_mask, [batch_size, num_instances, 1, 1]),
          tf.float32)
      # Adding 0.001 in the denominator to avoid division by zero.
      loss = tf.nn.l2_loss(diff) / (tf.reduce_sum(labels) + 0.001)
    return loss
//  [python]
def __call__(self, logits, labels, valid_mask):
    """ShapeMask mask cross entropy loss function wrapper.

    Args:
      logits: A Tensor of shape [batch_size * num_instances, height, width,
        num_classes]. The logits are not necessarily between 0 and 1.
      labels: A float16/float32 Tensor of shape [batch_size, num_instances,
        mask_size, mask_size], where mask_size =
        mask_crop_size * gt_upsample_scale for fine mask, or mask_crop_size
        for coarse masks and shape priors.
      valid_mask: a binary mask of shape [batch_size, num_instances]
        indicating valid training masks.
    Returns:
      loss: an float tensor representing total mask classification loss.
    """
    with tf.name_scope('shapemask_loss'):
      batch_size, num_instances = valid_mask.get_shape().as_list()[:2]
      labels = tf.to_float(labels)
      logits = tf.to_float(logits)
      loss = self._binary_crossentropy(labels, logits)
      loss *= tf.to_float(tf.reshape(
          valid_mask, [batch_size, num_instances, 1, 1]))
      # Adding 0.001 in the denominator to avoid division by zero.
      loss = tf.reduce_sum(loss) / (tf.reduce_sum(labels) + 0.001)
    return loss
//  [python]
def __call__(self, logits, labels, valid_mask):
    """ShapeMask mask cross entropy loss function wrapper.

    Args:
      logits: A Tensor of shape [batch_size * num_instances, height, width,
        num_classes]. The logits are not necessarily between 0 and 1.
      labels: A float16/float32 Tensor of shape [batch_size, num_instances,
        mask_size, mask_size], where mask_size =
        mask_crop_size * gt_upsample_scale for fine mask, or mask_crop_size
        for coarse masks and shape priors.
      valid_mask: a binary mask of shape [batch_size, num_instances]
        indicating valid training masks.
    Returns:
      loss: an float tensor representing total mask classification loss.
    """
    with tf.name_scope('shapemask_loss'):
      batch_size, num_instances = valid_mask.get_shape().as_list()[:2]
      labels = tf.cast(labels, tf.float32)
      logits = tf.cast(logits, tf.float32)
      loss = self._binary_crossentropy(labels, logits)
      loss *= tf.cast(tf.reshape(
          valid_mask, [batch_size, num_instances, 1, 1]), loss.dtype)
      # Adding 0.001 in the denominator to avoid division by zero.
      loss = tf.reduce_sum(loss) / (tf.reduce_sum(labels) + 0.001)
    return loss
//  [python]
def test_horovod_allreduce_cpu_process_sets(self):
        """ Test on CPU that allreduce correctly sums if restricted to non-global process sets"""
        rank = hvd.rank()

        if hvd.ccl_built():
            self.skipTest("Multiple process sets currently do not support CCL.")

        dtypes = self.filter_supported_types([tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64])
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/cpu:0"):
                even_rank_tensor = self.random_uniform([17] * dim, -100, 100)
                even_rank_tensor = tf.cast(even_rank_tensor, dtype=dtype)
                odd_rank_tensor = self.random_uniform([17] * dim, -100, 100)
                odd_rank_tensor = tf.cast(odd_rank_tensor, dtype=dtype)
                if rank in self.even_ranks:
                    summed = hvd.allreduce(even_rank_tensor, average=False, process_set=self.even_set)
                    multiplied = even_rank_tensor * len(self.even_ranks)
                if rank in self.odd_ranks:
                    summed = hvd.allreduce(odd_rank_tensor, average=False, process_set=self.odd_set)
                    multiplied = odd_rank_tensor * len(self.odd_ranks)
                difference = summed - multiplied
                difference = tf.cast(difference, tf.int32) if dtype == tf.uint8 else difference
                max_difference = tf.reduce_max(tf.abs(difference))

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            max_process_set_size = max(len(self.even_ranks), len(self.odd_ranks))
            if max_process_set_size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif max_process_set_size < 10:
                threshold = 1e-4
            elif max_process_set_size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results")
//  [python]
def test_horovod_allreduce_gpu_process_sets(self):
        """ Test on GPU that allreduce correctly sums if restricted to non-global process sets"""
        # Only do this test if there are GPUs available.
        if not tf.test.is_gpu_available(cuda_only=True):
            self.skipTest("No GPUs available")

        if int(os.environ.get('HOROVOD_MIXED_INSTALL', 0)):
            # Skip if compiled with CUDA but without HOROVOD_GPU_OPERATIONS.
            self.skipTest("Not compiled with HOROVOD_GPU_OPERATIONS")

        local_rank = hvd.local_rank()
        rank = hvd.rank()

        dtypes = [tf.uint8, tf.int8, tf.int32, tf.int64, tf.float16, tf.float32, tf.float64]
        dims = [1, 2, 3]
        for dtype, dim in itertools.product(dtypes, dims):
            with tf.device("/gpu:%d" % local_rank):
                even_rank_tensor = self.random_uniform([17] * dim, -100, 100)
                even_rank_tensor = t.cast(even_rank_tensor, dtype=dtype)
                odd_rank_tensor = self.random_uniform([17] * dim, -100, 100)
                odd_rank_tensor = t.cast(odd_rank_tensor, dtype=dtype)
                if rank in self.even_ranks:
                    summed = hvd.allreduce(even_rank_tensor, average=False, process_set=self.even_set)
                    multiplied = even_rank_tensor * len(self.even_ranks)
                if rank in self.odd_ranks:
                    summed = hvd.allreduce(odd_rank_tensor, average=False, process_set=self.odd_set)
                    multiplied = odd_rank_tensor * len(self.odd_ranks)
                difference = summed - multiplied
                difference = t.cast(difference, tf.int32) if dtype == tf.uint8 else difference
                max_difference = tf.reduce_max(tf.abs(difference))

            # Threshold for floating point equality depends on number of
            # ranks, since we're comparing against precise multiplication.
            max_process_set_size = max(len(self.even_ranks), len(self.odd_ranks))
            if max_process_set_size <= 3 or dtype in [tf.uint8, tf.int8, tf.int32, tf.int64]:
                threshold = 0
            elif max_process_set_size < 10:
                threshold = 1e-4
            elif max_process_set_size < 15:
                threshold = 5e-4
            else:
                self.skipTest("Horovod cluster too large for precise multiplication comparison")

            diff = self.evaluate(max_difference)
            self.assertTrue(diff <= threshold, "hvd.allreduce produces incorrect results")