code: def track_points(
      object_coordinates,
      depth,
      depth_range,
      segmentations,
      surface_normals,
      bboxes_3d,
      obj_quat,
      cam_focal_length,
      cam_positions,
      cam_quaternions,
      cam_sensor_width,
      window,
      tracks_to_sample=256,
      sampling_stride=4,
      max_seg_id=25,
      max_sampled_frac=0.1,
  ):
      """Track points in 2D using Kubric data.
  
      Args:
        object_coordinates: Video of coordinates for each pixel in the object's
          local coordinate frame.  Shape [num_frames, height, width, 3]
        depth: uint16 depth video from Kubric.  Shape [num_frames, height, width]
        depth_range: Values needed to normalize Kubric's int16 depth values into
          metric depth.
        segmentations: Integer object id for each pixel.  Shape
          [num_frames, height, width]
        surface_normals: uint16 surface normal map. Shape
          [num_frames, height, width, 3]
        bboxes_3d: The set of all object bounding boxes from Kubric
        obj_quat: Quaternion rotation for each object.  Shape
          [num_objects, num_frames, 4]
        cam_focal_length: Camera focal length
        cam_positions: Camera positions, with shape [num_frames, 3]
        cam_quaternions: Camera orientations, with shape [num_frames, 4]
        cam_sensor_width: Camera sensor width parameter
        window: the window inside which we're sampling points.  Integer valued
          in the format [x_min, y_min, x_max, y_max], where min is inclusive and
          max is exclusive.
        tracks_to_sample: Total number of tracks to sample per video.
        sampling_stride: For efficiency, query points are sampled from a random grid
          of this stride.
        max_seg_id: The maxium segment id in the video.
        max_sampled_frac: The maximum fraction of points to sample from each
          object, out of all points that lie on the sampling grid.
  
      Returns:
        A set of queries, randomly sampled from the video (with a bias toward
          objects), of shape [num_points, 3].  Each point is [t, y, x], where
          t is time.  All points are in pixel/frame coordinates.
        The trajectory for each query point, of shape [num_points, num_frames, 3].
          Each point is [x, y].  Points are in pixel coordinates
        Occlusion flag for each point, of shape [num_points, num_frames].  This is
          a boolean, where True means the point is occluded.
  
      """
      chosen_points = []
      all_reproj = []
      all_occ = []
  
      # Convert to metric depth
  
      depth_range_f32 = tf.cast(depth_range, tf.float32)
      depth_min = depth_range_f32[0]
      depth_max = depth_range_f32[1]
      depth_f32 = tf.cast(depth, tf.float32)
      depth_map = (
          depth_min + depth_f32 * (depth_max - depth_min) / 65535
      )
  
      surface_normal_map = surface_normals / 65535 * 2.0 - 1.0
  
      input_size = object_coordinates.shape.as_list()[1:3]
      num_frames = object_coordinates.shape.as_list()[0]
  
      # We first sample query points within the given window.  That means first
      # extracting the window from the segmentation tensor, because we want to have
      # a bias toward moving objects.
      # Note: for speed we sample points on a grid.  The grid start position is
      # randomized within the window.
      start_vec = [
          tf.random.uniform(
              [], minval=0, maxval=sampling_stride, dtype=tf.int32
          )
          for _ in range(3)
      ]
      start_vec[1] += window[0]
      start_vec[2] += window[1]
      end_vec = [num_frames, window[2], window[3]]
  
      def extract_box(x):
          x = x[
              start_vec[0] :: sampling_stride,
              start_vec[1] : window[2] : sampling_stride,
              start_vec[2] : window[3] : sampling_stride,
          ]
          return x
  
      segmentations_box = extract_box(segmentations)
      object_coordinates_box = extract_box(object_coordinates)
  
      # Next, get the number of points to sample from each object.  First count
      # how many points are available for each object.
  
      cnt = tf.math.bincount(
          tf.cast(tf.reshape(segmentations_box, [-1]), tf.int32)
      )
      num_to_sample = get_num_to_sample(
          cnt,
          max_seg_id,
          max_sampled_frac,
          tracks_to_sample,
      )
      num_to_sample.set_shape([max_seg_id])
      intrinsics, matrix_world = get_camera_matrices(
          cam_focal_length,
          cam_positions,
          cam_quaternions,
          cam_sensor_width,
          input_size,
          num_frames=num_frames,
      )
  
      # If the normal map is very rough, it's often because they come from a normal
      # map rather than the mesh.  These aren't trustworthy, and the normal test
      # may fail (i.e. the normal is pointing away from the camera even though the
      # point is still visible).  So don't use the normal test when inferring
      # occlusion.
      trust_sn = True
      sn_pad = tf.pad(
          surface_normal_map, [(0, 0), (1, 1), (1, 1), (0, 0)]
      )
      shp = surface_normal_map.shape
      sum_thresh = 0
      for i in [0, 2]:
          for j in [0, 2]:
              diff = (
                  sn_pad[:, i : shp[1] + i, j : shp[2] + j, :]
                  - surface_normal_map
              )
              diff = tf.reduce_sum(tf.square(diff), axis=-1)
              sum_thresh += tf.cast(diff > 0.05 * 0.05, tf.int32)
      trust_sn = tf.logical_and(trust_sn, (sum_thresh <= 2))[
          ..., tf.newaxis
      ]
      surface_normals_box = extract_box(surface_normal_map)
      trust_sn_box = extract_box(trust_sn)
  
      def get_camera(fr=None):
          if fr is None:
              return {
                  "intrinsics": intrinsics,
                  "matrix_world": matrix_world,
              }
          return {
              "intrinsics": intrinsics[fr],
              "matrix_world": matrix_world[fr],
          }
  
      # Construct pixel coordinates for each pixel within the window.
      window = tf.cast(window, tf.float32)
      z, y, x = tf.meshgrid(
          *[
              tf.range(st, ed, sampling_stride)
              for st, ed in zip(start_vec, end_vec)
          ],
          indexing="ij",
      )
      pix_coords = tf.reshape(tf.stack([z, y, x], axis=-1), [-1, 3])
  
      for i in range(max_seg_id):
          # sample points on object i in the first frame.  obj_id is the position
          # within the object_coordinates array, which is one lower than the value
          # in the segmentation mask (0 in the segmentation mask is the background
          # object, which has no bounding box).
          obj_id = i - 1
          mask = tf.equal(tf.reshape(segmentations_box, [-1]), i)
          pt = tf.boolean_mask(
              tf.reshape(object_coordinates_box, [-1, 3]), mask
          )
          normals = tf.boolean_mask(
              tf.reshape(surface_normals_box, [-1, 3]), mask
          )
          trust_sn_mask = tf.boolean_mask(
              tf.reshape(trust_sn_box, [-1, 1]), mask
          )
          idx = tf.cond(
              tf.shape(pt)[0] > 0,
              lambda: tf.multinomial(  # pylint: disable=g-long-lambda
                  tf.zeros(tf.shape(pt)[0:1])[tf.newaxis, :],
                  tf.gather(num_to_sample, i),
              )[0],
              lambda: tf.zeros([0], dtype=tf.int64),
          )
          # note: pt_coords is pixel coordinates, not raster coordinates.
          pt_coords = tf.gather(tf.boolean_mask(pix_coords, mask), idx)
          normals = tf.gather(normals, idx)
          trust_sn_gather = tf.gather(trust_sn_mask, idx)
  
          pixel_to_raster = tf.constant([0.0, 0.5, 0.5])[tf.newaxis, :]
  
          if obj_id == -1:
              # For the background object, no bounding box is available.  However,
              # this doesn't move, so we use the depth map to backproject these points
              # into 3D and use those positions throughout the video.
              pt_3d = []
              pt_coords_reorder = []
              for fr in range(num_frames):
                  # We need to loop over frames because we need to use the correct depth
                  # map for each frame.
                  pt_coords_chunk = tf.boolean_mask(
                      pt_coords, tf.equal(pt_coords[:, 0], fr)
                  )
                  pt_coords_reorder.append(pt_coords_chunk)
  
                  pt_3d.append(
                      unproject(
                          pt_coords_chunk[:, 1:],
                          get_camera(fr),
                          depth_map[fr],
                      )
                  )
              pt = tf.concat(pt_3d, axis=0)
              chosen_points.append(
                  tf.cast(
                      tf.concat(pt_coords_reorder, axis=0), tf.float32
                  )
                  + pixel_to_raster
              )
              bbox = None
              quat = None
              frame_for_pt = None
          else:
              # For any other object, we just use the point coordinates supplied by
              # kubric.
              pt = tf.gather(pt, idx)
              pt = pt / np.iinfo(np.uint16).max - 0.5
              chosen_points.append(
                  tf.cast(pt_coords, tf.float32) + pixel_to_raster
              )
              # if obj_id>num_objects, then we won't have a box.  We also won't have
              # points, so just use a dummy to prevent tf from crashing.
              bbox = tf.cond(
                  obj_id >= tf.shape(bboxes_3d)[0],
                  lambda: bboxes_3d[0, :],
                  lambda: bboxes_3d[obj_id, :],
              )
              quat = tf.cond(
                  obj_id >= tf.shape(obj_quat)[0],
                  lambda: obj_quat[0, :],
                  lambda: obj_quat[obj_id, :],
              )
              frame_for_pt = pt_coords[..., 0]
  
          # Finally, compute the reprojections for this particular object.
          obj_reproj, obj_occ = tf.cond(
              tf.shape(pt)[0] > 0,
              functools.partial(
                  single_object_reproject,
                  bbox_3d=bbox,
                  pt=pt,
                  pt_segments=i,
                  camera=get_camera(),
                  cam_positions=cam_positions,
                  num_frames=num_frames,
                  depth_map=depth_map,
                  segments=segmentations,
                  window=window,
                  input_size=input_size,
                  quat=quat,
                  normals=normals,
                  frame_for_pt=frame_for_pt,
                  trust_normals=trust_sn_gather,
              ),
              lambda: (  # pylint: disable=g-long-lambda
                  tf.zeros([0, num_frames, 2], dtype=tf.float32),
                  tf.zeros([0, num_frames], dtype=tf.bool),
              ),
          )
          all_reproj.append(obj_reproj)
          all_occ.append(obj_occ)
  
      # Points are currently in pixel coordinates of the original video.  We now
      # convert them to coordinates within the window frame, and rescale to
      # pixel coordinates.  Note that this produces the pixel coordinates after
      # the window gets cropped and rescaled to the full image size.
      wd = tf.concat(
          [
              np.array([0.0]),
              window[0:2],
              np.array([num_frames]),
              window[2:4],
          ],
          axis=0,
      )
      wd = wd[tf.newaxis, tf.newaxis, :]
      coord_multiplier = [num_frames, input_size[0], input_size[1]]
      all_reproj = tf.concat(all_reproj, axis=0)
      # We need to extract x,y, but the format of the window is [t1,y1,x1,t2,y2,x2]
      window_size = wd[:, :, 5:3:-1] - wd[:, :, 2:0:-1]
      window_top_left = wd[:, :, 2:0:-1]
      all_reproj = (all_reproj - window_top_left) / window_size
      all_reproj = all_reproj * coord_multiplier[2:0:-1]
      all_occ = tf.concat(all_occ, axis=0)
  
      # chosen_points is [num_points, (z,y,x)]
      chosen_points = tf.concat(chosen_points, axis=0)
  
      chosen_points = tf.cast(chosen_points, tf.float32)
  
      # renormalize so the box corners are at [-1,1]
      chosen_points = (chosen_points - wd[:, 0, :3]) / (
          wd[:, 0, 3:] - wd[:, 0, :3]
      )
      chosen_points = chosen_points * coord_multiplier
      # Note: all_reproj is in (x,y) format, but chosen_points is in (z,y,x) format
  
      return (
          tf.cast(chosen_points, tf.float32),
          tf.cast(all_reproj, tf.float32),
          all_occ,
      )
  
imports: import functools
  import itertools
  import matplotlib.pyplot
  import mediapy
  import numpy
  import tensorflow.compat.v1
  import tensorflow_datasets
  from tensorflow_graphics.geometry.transformation import rotation_matrix_3d
label_desc: Use image node_id 10 for calling a global function with example usage: track_points(object_coordinates, depth, depth_range, segmentations, surface_normals, bboxes_3d, obj_quat, cam_focal_length, cam_positions, cam_quaternions, cam_sensor_width, window, tracks_to_sample, sampling_stride, max_seg_id, max_sampled_frac) and returns: tf,  tf,  all_occ, x, dict, dict