code: def process(example, ann=None, box_key="orvit"):
      """Add bounding boxes as additional fields to TFRecord."""
      height = example.context.feature["image/height"].int64_list.value[
          0
      ]
      width = example.context.feature["image/width"].int64_list.value[0]
      num_frames = example.context.feature[
          "clip/frames"
      ].int64_list.value[0]
      feature_list = example.feature_lists.feature_list
  
      if ann is None:
          annotated_frames = 0
      else:
          xmin = ann[..., 0] / width  # T x O
          ymin = ann[..., 1] / height  # T x O
          xmax = ann[..., 2] / width  # T x O
          ymax = ann[..., 3] / height  # T x O
          boxes = np.stack(
              [xmin, ymin, xmax, ymax], axis=2
          )  # T x O x 4
          scores = ann[..., 4]
          annotated_frames = ann.shape[0]
  
      feature_list[f"{box_key}/bbox/xmin"].Clear()
      feature_list[f"{box_key}/bbox/xmax"].Clear()
      feature_list[f"{box_key}/bbox/ymin"].Clear()
      feature_list[f"{box_key}/bbox/ymax"].Clear()
      feature_list[f"{box_key}/bbox/score"].Clear()
  
      for frame_idx in range(num_frames):
          j = frame_idx - num_frames + annotated_frames
          if j < 0 or j >= annotated_frames:
              boxes_t = np.zeros((0, 4), dtype=np.float32)
              scores_t = np.zeros((0), dtype=np.float32)
          else:
              boxes_t = boxes[j]
              scores_t = scores[j]
          feature_list[
              f"{box_key}/bbox/score"
          ].feature.add().float_list.value.extend(scores_t.tolist())
          feature_list[
              f"{box_key}/bbox/xmin"
          ].feature.add().float_list.value.extend(
              boxes_t[:, 0].tolist()
          )
          feature_list[
              f"{box_key}/bbox/ymin"
          ].feature.add().float_list.value.extend(
              boxes_t[:, 1].tolist()
          )
          feature_list[
              f"{box_key}/bbox/xmax"
          ].feature.add().float_list.value.extend(
              boxes_t[:, 2].tolist()
          )
          feature_list[
              f"{box_key}/bbox/ymax"
          ].feature.add().float_list.value.extend(
              boxes_t[:, 3].tolist()
          )
      return example
  
file_path: scenic/scenic/projects/objectvivit/tools/add_orvit_bbox_to_tfrecord.py
imports: import os
  from absl import app
  from absl import flags
  import numpy
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: process(example, ann, box_key) and returns: example