# MediaPipe graph that performs box tracking with TensorFlow Lite on GPU.

# Images coming into and out of the graph.
input_stream: "input_video"
input_stream: "WIDTH:input_width"
input_stream: "HEIGHT:input_height"
input_side_packet: "LABELS_CSV:allowed_labels"
input_side_packet: "MODEL_SCALE:model_scale"
input_side_packet: "MODEL_TRANSFORMATION:model_transformation"
input_side_packet: "TEXTURE:box_texture"
input_side_packet: "MAX_NUM_OBJECTS:max_num_objects"
input_side_packet: "ANIMATION_ASSET:box_asset_name"
input_side_packet: "MASK_TEXTURE:obj_texture"
input_side_packet: "MASK_ASSET:obj_asset_name"
output_stream: "output_video"

# Throttles the images flowing downstream for flow control. It passes through
# the very first incoming image unaltered, and waits for downstream nodes
# (calculators and subgraphs) in the graph to finish their tasks before it
# passes through another image. All images that come in while waiting are
# dropped, limiting the number of in-flight images in most part of the graph to
# 1. This prevents the downstream nodes from queuing up incoming images and data
# excessively, which leads to increased latency and memory usage, unwanted in
# real-time mobile applications. It also eliminates unnecessarily computation,
# e.g., the output produced by a node may get dropped downstream if the
# subsequent nodes are still busy processing previous inputs.
node {
  calculator: "FlowLimiterCalculator"
  input_stream: "input_video"
  input_stream: "FINISHED:output_video"
  input_stream_info: {
    tag_index: "FINISHED"
    back_edge: true
  }
  output_stream: "throttled_input_video"
}

# Crops the image from the center to the size WIDTHxHEIGHT.
node: {
  calculator: "ImageCroppingCalculator"
  input_stream: "IMAGE_GPU:throttled_input_video"
  output_stream: "IMAGE_GPU:throttled_input_video_3x4"
  input_stream: "WIDTH:input_width"
  input_stream: "HEIGHT:input_height"
  node_options: {
    [type.googleapis.com/mediapipe.ImageCroppingCalculatorOptions] {
      border_mode: BORDER_REPLICATE
    }
  }
}

node {
  calculator: "ObjectronGpuSubgraph"
  input_stream: "IMAGE_GPU:throttled_input_video_3x4"
  input_side_packet: "LABELS_CSV:allowed_labels"
  input_side_packet: "MAX_NUM_OBJECTS:max_num_objects"
  output_stream: "FRAME_ANNOTATION:lifted_objects"
}

# The rendering nodes:
# We are rendering two meshes: 1) a 3D bounding box, which we overlay directly
# on the texture, and 2) a virtual object, which we use as an occlusion mask.
# These models are designed using different tools, so we supply a transformation
# to bring both of them to the Objectron's coordinate system.

# Creates a model matrices for the tracked object given the lifted 3D points.
# This calculator does two things: 1) Estimates object's pose (orientation,
# translation, and scale) from the 3D vertices, and
# 2) bring the object from the objectron's coordinate system to the renderer
# (OpenGL) coordinate system. Since the final goal is to render a mesh file on
# top of the object, we also supply a transformation to bring the mesh to the
# objectron's coordinate system, and rescale mesh to the unit size.
node {
  calculator: "AnnotationsToModelMatricesCalculator"
  input_stream: "ANNOTATIONS:lifted_objects"
  output_stream: "MODEL_MATRICES:model_matrices"
  node_options: {
    [type.googleapis.com/mediapipe.AnnotationsToModelMatricesCalculatorOptions] {
      # Re-scale the CAD model to the size of a unit box
      model_scale: [0.04, 0.04, 0.04]
      # Bring the box CAD model to objectron's coordinate system. This
      # is equivalent of -pi/2 rotation along the y-axis (right-hand rule):
      # Eigen::AngleAxisf(-M_PI / 2., Eigen::Vector3f::UnitY())
      model_transformation: [0.0,  0.0, -1.0,  0.0]
      model_transformation: [0.0,  1.0,  0.0,  0.0]
      model_transformation: [1.0,  0.0,  0.0,  0.0]
      model_transformation: [0.0,  0.0,  0.0,  1.0]
    }
  }
}

# Compute the model matrices for the CAD model of the virtual object, to be used
# as an occlusion mask. The model will be rendered at the exact same location as
# the bounding box.
node {
  calculator: "AnnotationsToModelMatricesCalculator"
  input_stream: "ANNOTATIONS:lifted_objects"
  input_side_packet: "MODEL_SCALE:model_scale"
  input_side_packet: "MODEL_TRANSFORMATION:model_transformation"
  output_stream: "MODEL_MATRICES:mask_model_matrices"
}

# Render everything together. First we render the 3D bounding box animation,
# then we render the occlusion mask.
node: {
  calculator: "GlAnimationOverlayCalculator"
  input_stream: "VIDEO:throttled_input_video_3x4"
  input_stream: "MODEL_MATRICES:model_matrices"
  input_stream: "MASK_MODEL_MATRICES:mask_model_matrices"
  output_stream: "output_video"
  input_side_packet: "TEXTURE:box_texture"
  input_side_packet: "ANIMATION_ASSET:box_asset_name"
  input_side_packet: "MASK_TEXTURE:obj_texture"
  input_side_packet: "MASK_ASSET:obj_asset_name"
  node_options: {
    [type.googleapis.com/mediapipe.GlAnimationOverlayCalculatorOptions] {
      aspect_ratio: 0.75
      vertical_fov_degrees: 70.
      animation_speed_fps: 25
    }
  }
}