# MediaPipe graph that performs iris tracking on desktop with TensorFlow Lite
# on CPU.

# max_queue_size limits the number of packets enqueued on any input stream
# by throttling inputs to the graph. This makes the graph only process one
# frame per time.
max_queue_size: 1

# Decodes an input video file into images and a video header.
node {
  calculator: "OpenCvVideoDecoderCalculator"
  input_side_packet: "INPUT_FILE_PATH:input_video_path"
  output_stream: "VIDEO:input_video"
  output_stream: "VIDEO_PRESTREAM:input_video_header"
}

# Defines how many faces to detect. Iris tracking currently only handles one
# face (left and right eye), and therefore this should always be set to 1.
node {
  calculator: "ConstantSidePacketCalculator"
  output_side_packet: "PACKET:0:num_faces"
  node_options: {
    [type.googleapis.com/mediapipe.ConstantSidePacketCalculatorOptions]: {
      packet { int_value: 1 }
    }
  }
}

# Detects faces and corresponding landmarks.
node {
  calculator: "FaceLandmarkFrontCpu"
  input_stream: "IMAGE:input_video"
  input_side_packet: "NUM_FACES:num_faces"
  output_stream: "LANDMARKS:multi_face_landmarks"
  output_stream: "ROIS_FROM_LANDMARKS:face_rects_from_landmarks"
  output_stream: "DETECTIONS:face_detections"
  output_stream: "ROIS_FROM_DETECTIONS:face_rects_from_detections"
}

# Gets the very first and only face from "multi_face_landmarks" vector.
node {
  calculator: "SplitNormalizedLandmarkListVectorCalculator"
  input_stream: "multi_face_landmarks"
  output_stream: "face_landmarks"
  node_options: {
    [type.googleapis.com/mediapipe.SplitVectorCalculatorOptions] {
      ranges: { begin: 0 end: 1 }
      element_only: true
    }
  }
}

# Gets the very first and only face rect from "face_rects_from_landmarks"
# vector.
node {
  calculator: "SplitNormalizedRectVectorCalculator"
  input_stream: "face_rects_from_landmarks"
  output_stream: "face_rect"
  node_options: {
    [type.googleapis.com/mediapipe.SplitVectorCalculatorOptions] {
      ranges: { begin: 0 end: 1 }
      element_only: true
    }
  }
}

# Gets two landmarks which define left eye boundary.
node {
  calculator: "SplitNormalizedLandmarkListCalculator"
  input_stream: "face_landmarks"
  output_stream: "left_eye_boundary_landmarks"
  node_options: {
    [type.googleapis.com/mediapipe.SplitVectorCalculatorOptions] {
      ranges: { begin: 33 end: 34 }
      ranges: { begin: 133 end: 134 }
      combine_outputs: true
    }
  }
}

# Gets two landmarks which define right eye boundary.
node {
  calculator: "SplitNormalizedLandmarkListCalculator"
  input_stream: "face_landmarks"
  output_stream: "right_eye_boundary_landmarks"
  node_options: {
    [type.googleapis.com/mediapipe.SplitVectorCalculatorOptions] {
      ranges: { begin: 362 end: 363 }
      ranges: { begin: 263 end: 264 }
      combine_outputs: true
    }
  }
}

# Detects iris landmarks, eye contour landmarks, and corresponding rect (ROI).
node {
  calculator: "IrisLandmarkLeftAndRightCpu"
  input_stream: "IMAGE:input_video"
  input_stream: "LEFT_EYE_BOUNDARY_LANDMARKS:left_eye_boundary_landmarks"
  input_stream: "RIGHT_EYE_BOUNDARY_LANDMARKS:right_eye_boundary_landmarks"
  output_stream: "LEFT_EYE_CONTOUR_LANDMARKS:left_eye_contour_landmarks"
  output_stream: "LEFT_EYE_IRIS_LANDMARKS:left_iris_landmarks"
  output_stream: "LEFT_EYE_ROI:left_eye_rect_from_landmarks"
  output_stream: "RIGHT_EYE_CONTOUR_LANDMARKS:right_eye_contour_landmarks"
  output_stream: "RIGHT_EYE_IRIS_LANDMARKS:right_iris_landmarks"
  output_stream: "RIGHT_EYE_ROI:right_eye_rect_from_landmarks"
}

node {
  calculator: "ConcatenateNormalizedLandmarkListCalculator"
  input_stream: "left_eye_contour_landmarks"
  input_stream: "right_eye_contour_landmarks"
  output_stream: "refined_eye_landmarks"
}

node {
  calculator: "UpdateFaceLandmarksCalculator"
  input_stream: "NEW_EYE_LANDMARKS:refined_eye_landmarks"
  input_stream: "FACE_LANDMARKS:face_landmarks"
  output_stream: "UPDATED_FACE_LANDMARKS:updated_face_landmarks"
}

# Renders annotations and overlays them on top of the input images.
node {
  calculator: "IrisRendererCpu"
  input_stream: "IMAGE:input_video"
  input_stream: "FACE_LANDMARKS:updated_face_landmarks"
  input_stream: "EYE_LANDMARKS_LEFT:left_eye_contour_landmarks"
  input_stream: "EYE_LANDMARKS_RIGHT:right_eye_contour_landmarks"
  input_stream: "IRIS_LANDMARKS_LEFT:left_iris_landmarks"
  input_stream: "IRIS_LANDMARKS_RIGHT:right_iris_landmarks"
  input_stream: "NORM_RECT:face_rect"
  input_stream: "LEFT_EYE_RECT:left_eye_rect_from_landmarks"
  input_stream: "RIGHT_EYE_RECT:right_eye_rect_from_landmarks"
  input_stream: "DETECTIONS:face_detections"
  output_stream: "IRIS_LANDMARKS:iris_landmarks"
  output_stream: "IMAGE:output_video"
}

# Encodes the annotated images into a video file, adopting properties specified
# in the input video header, e.g., video framerate.
node {
  calculator: "OpenCvVideoEncoderCalculator"
  input_stream: "VIDEO:output_video"
  input_stream: "VIDEO_PRESTREAM:input_video_header"
  input_side_packet: "OUTPUT_FILE_PATH:output_video_path"
  node_options: {
    [type.googleapis.com/mediapipe.OpenCvVideoEncoderCalculatorOptions]: {
      codec: "avc1"
      video_format: "mp4"
    }
  }
}
