code: def get_num_to_sample(
      counts, max_seg_id, max_sampled_frac, tracks_to_sample
  ):
      """Computes the number of points to sample for each object.
  
      Args:
        counts: The number of points available per object.  An int array of length
          n, where n is the number of objects.
        max_seg_id: The maximum number of segment id's 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.
        tracks_to_sample: Total number of tracks to sample per video.
  
      Returns:
        The number of points to sample for each object.  An int array of length n.
      """
      seg_order = tf.argsort(counts)
      sorted_counts = tf.gather(counts, seg_order)
      initializer = (0, tracks_to_sample, 0)
  
      def scan_fn(prev_output, count_seg):
          index = prev_output[0]
          remaining_needed = prev_output[1]
          desired_frac = 1 / (tf.shape(seg_order)[0] - index)
          want_to_sample = tf.cast(
              remaining_needed, tf.float32
          ) * tf.cast(desired_frac, tf.float32)
          want_to_sample = tf.cast(tf.round(want_to_sample), tf.int32)
          max_to_sample = tf.cast(count_seg, tf.float32) * tf.cast(
              max_sampled_frac, tf.float32
          )
          max_to_sample = tf.cast(tf.round(max_to_sample), tf.int32)
          num_to_sample = tf.minimum(want_to_sample, max_to_sample)
  
          remaining_needed = remaining_needed - num_to_sample
          return (index + 1, remaining_needed, num_to_sample)
  
      # outputs 0 and 1 are just bookkeeping; output 2 is the actual number of
      # points to sample per object.
      res = tf.scan(scan_fn, sorted_counts, initializer)[2]
      invert = tf.argsort(seg_order)
      num_to_sample = tf.gather(res, invert)
      num_to_sample = tf.concat(
          [
              num_to_sample,
              tf.zeros(
                  [max_seg_id - tf.shape(num_to_sample)[0]],
                  dtype=tf.int32,
              ),
          ],
          axis=0,
      )
      return num_to_sample
  
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 9 for calling a global function with example usage: get_num_to_sample(counts, max_seg_id, max_sampled_frac, tracks_to_sample) and returns: num_to_sample, unknown,  remaining_needed,  num_to_sample