code: def append_or_add_object_mask(
      feature_dict,
      num_frames,
      keep_full_frames,
      bbox_num=100,
      mask_radius=-1.0,
      concat_mask=False,
      min_box_size=0,
      add_detections=False,
      detection_stride_spatial=1,
      detection_stride_temporal=1,
      tracked_objects=False,
      add_global_box=False,
  ):
      """Generate object heatmaps.
  
      Args:
        feature_dict: dict of tensors from the dataloader
        num_frames: int; frames of a single model.
        keep_full_frames: number of paired frames that will allways have full masks.
          We use this to provide global features. Note the actual number of full
          frames will be 2 * keep_full_frames, where 2 is the temporal token size
          in ViViT.
        bbox_num: int; the max number of boxes in one frame.
        mask_radius: hyperparameter that controls the size of the heatmap peaks.
        concat_mask: bool: if True, it concatenane the heatmap to RPB pixels so
          that the output channels will be 4. If False, it multiplied the heatmap
          values to the pixels.
        min_box_size: threshold to filter out small boxes.
        add_detections: bool: if True, add mask as an additional key "detections".
        detection_stride_spatial: int: spatial stride of the mask
        detection_stride_temporal: int: temporal stride of the mask
        tracked_objects: bool; if True, return object-specific heatmaps
        add_global_box: if True, add a box for the whole image.
      Returns:
        the updated feature_dict.
      """
      assert not tracked_objects or bbox_num < 100
      bboxes = feature_dict["bboxes"]  # T x N x 4
      images = feature_dict["image"]  # T x H x W x 3
      scores = feature_dict["score"]  # T x N
  
      if add_global_box:
          t = bboxes.shape[0]
          new_box = tf.constant(
              [0, 0, 1.0, 1.0], dtype=tf.float32, shape=[1, 1, 4]
          )
          new_box = tf.broadcast_to(new_box, (t, 1, 4))
          bboxes = tf.concat([new_box, bboxes], axis=1)
          new_score = tf.constant([1.0], dtype=tf.float32, shape=[1, 1])
          new_score = tf.broadcast_to(new_score, (scores.shape[0], 1))
          scores = tf.concat([new_score, scores], axis=1)
  
      if bbox_num < 100:
          n = tf.shape(bboxes)[0]
          scores = scores[:, :bbox_num]
          bboxes = bboxes[:, :bbox_num]
          input_length = tf.shape(bboxes)[1]
          input_length = tf.clip_by_value(input_length, 0, bbox_num)
          padding_length = tf.maximum(0, bbox_num - input_length)
          paddings_boxes = tf.zeros((n, padding_length, 4), tf.float32)
          bboxes = tf.concat([bboxes, paddings_boxes], axis=1)
  
      t, h, w = images.shape[0], images.shape[1], images.shape[2]
  
      if keep_full_frames == 0:  # skip every frame-group
          skip = tf.convert_to_tensor(value=[False] * t, dtype=tf.bool)
          skip = tf.reshape(skip, (t, 1))
          skip = tf.broadcast_to(skip, [t, h * w])
      elif keep_full_frames == 1:  # keep middle frame
          value = (
              [False, False] * (num_frames // 4)
              + [True, True]
              + [False, False] * (num_frames // 4 - 1)
          ) * (t // num_frames)
          skip = tf.convert_to_tensor(value=value, dtype=tf.bool)
          skip = tf.reshape(skip, (t, 1))
          skip = tf.broadcast_to(skip, [t, h * w])
      elif keep_full_frames == 2:  # keep first and last
          value = (
              [True, True]
              + [False, False] * ((num_frames - 4) // 2)
              + [True, True]
          ) * (t // num_frames)
          skip = tf.convert_to_tensor(value=value, dtype=tf.bool)
          skip = tf.reshape(skip, (t, 1))
          skip = tf.broadcast_to(skip, [t, h * w])
      else:
          assert False, keep_full_frames
      # skip: t x h x w
  
      dw, dh = (
          w // detection_stride_spatial,
          h // detection_stride_spatial,
      )
      xs = tf.range(0, dw) * detection_stride_spatial + (
          detection_stride_spatial // 2
      )
      ys = tf.range(0, dh) * detection_stride_spatial + (
          detection_stride_spatial // 2
      )
      grid_x, grid_y = tf.meshgrid(xs, ys)
      grid_x = tf.cast(tf.reshape(grid_x, (dh * dw,)) / w, tf.float32)
      grid_y = tf.cast(tf.reshape(grid_y, (dh * dw,)) / h, tf.float32)
      grid = tf.stack([grid_x, grid_y], axis=1)  # HW x 2
      grid = tf.broadcast_to(grid, [images.shape[0], dh * dw, 2])
      # grid: T x HW x 2
  
      bboxes = tf.clip_by_value(bboxes, 0.0, 1.0)
      x1y1 = tf.minimum(bboxes[:, :, :2], bboxes[:, :, 2:])
      x2y2 = tf.maximum(bboxes[:, :, :2], bboxes[:, :, 2:])
      bboxes = tf.concat(
          [x1y1, x2y2], axis=2
      )  # fix bugs in Epic box format
  
      if mask_radius <= 0.0:
          inside_box = (
              grid[:, None, :, 0] >= bboxes[:, :, None, 0]
          )  # T x N x HW
          inside_box = tf.logical_and(
              inside_box, grid[:, None, :, 0] <= bboxes[:, :, None, 2]
          )
          inside_box = tf.logical_and(
              inside_box, grid[:, None, :, 1] >= bboxes[:, :, None, 1]
          )
          inside_box = tf.logical_and(
              inside_box, grid[:, None, :, 1] <= bboxes[:, :, None, 3]
          )  # T x N x HW
          if tracked_objects:
              mask = tf.cast(
                  tf.transpose(inside_box, (0, 2, 1)), tf.float32
              )  # T x HW x N
              mask_channels = bbox_num
          else:
              mask = tf.cast(
                  tf.reduce_any(inside_box, 1), tf.float32
              )  # T x HW
              mask_channels = 1
      else:
          centers = tf.stack(
              [
                  (bboxes[:, :, 0] + bboxes[:, :, 2]) / 2,
                  (bboxes[:, :, 1] + bboxes[:, :, 3]) / 2,
              ],
              axis=2,
          )  # T x N x 2
          dist2 = tf.reduce_sum(
              (grid[:, None, :] - centers[:, :, None]) ** 2, axis=3
          )  # T x N x HW
          area = (bboxes[:, :, 2] - bboxes[:, :, 0]) * (
              bboxes[:, :, 3] - bboxes[:, :, 1]
          )  # T x N
          filtered = tf.cast(
              area * h * w <= min_box_size**2, tf.float32
          )  # T x N
          dist2 = (
              dist2 * (1.0 - filtered[:, :, None])
              + filtered[:, :, None] * 10000000.0
          )
          delta = (1.0 - mask_radius) / (1.0 + mask_radius)  # scalar
          radius2 = delta**2 * 2 * area + 1e-6  # T x N
          if tracked_objects:
              weighted_dist = dist2 / radius2[:, :, None]  # T x N x HW
              gaussian_mask = tf.exp(-weighted_dist)  # T x N x HW
              gaussian_mask = tf.transpose(
                  gaussian_mask, (0, 2, 1)
              )  # T x HW x N
              mask_channels = bbox_num
          else:
              weighted_dist = tf.reduce_min(
                  dist2 / radius2[:, :, None], axis=1
              )  # T x HW
              gaussian_mask = tf.exp(-weighted_dist)  # T x HW
              mask_channels = 1
          mask = gaussian_mask
  
      if concat_mask:
          mask = tf.reshape(mask, (t, h, w, mask_channels))
          mask = tf.maximum(
              tf.cast(
                  tf.broadcast_to(
                      tf.reshape(skip, (t, h, w, 1)),
                      (t, h, w, mask_channels),
                  ),
                  tf.float32,
              ),
              mask,
          )
          images = tf.concat(
              [images, tf.cast(mask, tf.float32)], axis=3
          )
      elif add_detections:
          mask = tf.reshape(
              mask,
              (
                  t // detection_stride_temporal,
                  detection_stride_temporal,
                  dh,
                  dw,
              ),
          )
          mask = tf.reduce_max(mask, axis=1)  # t' x h x w
          mask = tf.reshape(
              mask, (t // detection_stride_temporal * dh * dw,)
          )
          feature_dict["detections"] = tf.cast(mask, tf.float32)
      feature_dict["image"] = images
      feature_dict["bboxes"] = bboxes
      return feature_dict
  
imports: from typing import Optional, Union
  from absl import logging
  from dmvr import builders
  from dmvr import processors
  import ml_collections
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: append_or_add_object_mask(feature_dict, num_frames, keep_full_frames, bbox_num, mask_radius, concat_mask, min_box_size, add_detections, detection_stride_spatial, detection_stride_temporal, tracked_objects, add_global_box) and returns: feature_dict