code: def assign_center_targets(
      gt_box_batch,
      gt_class_batch,
      gt_num_boxes_batch,
      image_size,
      stride,
      num_classes,
      gt_weights_batch=None,
      min_overlap=0.7,
  ):
      """Returns the center heatmap target for the CenterNet model.
  
      Args:
        gt_box_batch: A list of BoxList objects representing the ground truth
          detection bounding boxes for each sample in the batch.
        gt_class_batch: A tensors representing with the class label for each box in
          `gt_box_batch`.
        gt_num_boxes_batch: A tensors with the number of valid boxes in
          `gt_box_batch`.
        image_size: int tuple (height, width) of input to the CenterNet model. This
          is used to determine the height of the output.
        stride: int, ratio between the input and output size in the network.
        num_classes: int, total number of classes.
        gt_weights_batch: A list of tensors corresponding to the weight of each
          ground truth detection box.
        min_overlap: A float representing the minimum IoU desired.
  
      Returns:
        heatmap: A Tensor of size [batch_size, output_height, output_width,
          num_classes] representing the per class center heatmap. output_height
          and output_width are computed by dividing the input height and width by
          the stride.
      """
      height, width = image_size
  
      if gt_weights_batch:
          print(gt_weights_batch)
  
      x_range = tf.range(width // stride, dtype=tf.float32)
      y_range = tf.range(height // stride, dtype=tf.float32)
      x_grid, y_grid = tf.meshgrid(x_range, y_range, indexing="xy")
  
      x_grid = tf.expand_dims(x_grid, 2)
      y_grid = tf.expand_dims(y_grid, 2)
  
      def _create_heatmap(elems):
          gt_boxes = elems[0]
          gt_num_boxes = elems[1]
          gt_classes = elems[2]
  
          num_boxes = tf.cast(gt_num_boxes, tf.int32)
          boxes = box_list.BoxList(gt_boxes[:num_boxes, ...])
          class_targets = tf.one_hot(
              gt_classes[:num_boxes, ...], num_classes
          )
          boxes = box_list_ops.to_absolute_coordinates(
              boxes, height // stride, width // stride
          )
  
          (
              y_center,
              x_center,
              boxes_height,
              boxes_width,
          ) = boxes.get_center_coordinates_and_sizes()
  
          # The raw center coordinates in the output space.
          x_diff = x_grid - tf.math.floor(x_center)
          y_diff = y_grid - tf.math.floor(y_center)
          squared_distance = x_diff**2 + y_diff**2
  
          # We are dividing by 3 so that points closer than the computed
          # distance have a >99% CDF. In Pytorch, sigma = (2 * sigma + 1) / 6
          sigma = _max_distance_for_overlap(
              boxes_height, boxes_width, min_overlap
          )
          sigma = (
              2 * tf.math.maximum(tf.math.floor(sigma), 0.0) + 1
          ) / 6
  
          gaussian_map = tf.exp(-squared_distance / (2 * sigma * sigma))
  
          valid_boxes_mask = valid_bounding_boxes_mask(
              gt_boxes[:num_boxes, ...],
              height // stride,
              width // stride,
          )
          gaussian_map = gaussian_map * valid_boxes_mask
  
          output_height = tf.shape(gaussian_map)[0]
          output_width = tf.shape(gaussian_map)[1]
          num_boxes = tf.shape(gaussian_map)[2]
  
          reshaped_gaussian_map = tf.reshape(
              gaussian_map, (output_height, output_width, num_boxes, 1)
          )
          reshaped_class_targets = tf.reshape(
              class_targets, (1, 1, num_boxes, num_classes)
          )
          gaussian_per_box_per_class_map = (
              reshaped_gaussian_map * reshaped_class_targets
          )
          heatmap = tf.reduce_max(
              gaussian_per_box_per_class_map, axis=2
          )
          return heatmap
  
      return tf.map_fn(
          _create_heatmap,
          (gt_box_batch, gt_num_boxes_batch, gt_class_batch),
          dtype=tf.float32,
      )
  
file_path: graphics/tensorflow_graphics/projects/points_to_3Dobjects/transforms/targets.py
imports: import tensorflow
  from google3.third_party.tensorflow_models.object_detection.core import box_list
  from google3.third_party.tensorflow_models.object_detection.core import box_list_ops
label_desc: Use image node_id 4 for calling a global function with example usage: assign_center_targets(gt_box_batch, gt_class_batch, gt_num_boxes_batch, image_size, stride, num_classes, gt_weights_batch, min_overlap) and returns: tf, heatmap