code: def draw_axis(
      image: np.ndarray,
      rotation: np.ndarray,
      translation: np.ndarray,
      focal_length: Tuple[float, float] = (1.0, 1.0),
      principal_point: Tuple[float, float] = (0.0, 0.0),
      axis_length: float = 0.1,
      axis_drawing_spec: DrawingSpec = DrawingSpec(),
  ):
      """Draws the 3D axis on the image.
  
      Args:
        image: A three channel BGR image represented as numpy ndarray.
        rotation: Rotation matrix from object to camera coordinate frame.
        translation: Translation vector from object to camera coordinate frame.
        focal_length: camera focal length along x and y directions.
        principal_point: camera principal point in x and y.
        axis_length: length of the axis in the drawing.
        axis_drawing_spec: A DrawingSpec object that specifies the xyz axis drawing
          settings such as line thickness.
  
      Raises:
        ValueError: If one of the followings:
          a) If the input image is not three channel BGR.
      """
      if image.shape[2] != _BGR_CHANNELS:
          raise ValueError(
              "Input image must contain three channel bgr data."
          )
      image_rows, image_cols, _ = image.shape
      # Create axis points in camera coordinate frame.
      axis_world = np.float32(
          [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]]
      )
      axis_cam = (
          np.matmul(rotation, axis_length * axis_world.T).T
          + translation
      )
      x = axis_cam[..., 0]
      y = axis_cam[..., 1]
      z = axis_cam[..., 2]
      # Project 3D points to NDC space.
      fx, fy = focal_length
      px, py = principal_point
      x_ndc = np.clip(-fx * x / (z + 1e-5) + px, -1.0, 1.0)
      y_ndc = np.clip(-fy * y / (z + 1e-5) + py, -1.0, 1.0)
      # Convert from NDC space to image space.
      x_im = np.int32((1 + x_ndc) * 0.5 * image_cols)
      y_im = np.int32((1 - y_ndc) * 0.5 * image_rows)
      # Draw xyz axis on the image.
      origin = (x_im[0], y_im[0])
      x_axis = (x_im[1], y_im[1])
      y_axis = (x_im[2], y_im[2])
      z_axis = (x_im[3], y_im[3])
      cv2.arrowedLine(
          image, origin, x_axis, RED_COLOR, axis_drawing_spec.thickness
      )
      cv2.arrowedLine(
          image,
          origin,
          y_axis,
          GREEN_COLOR,
          axis_drawing_spec.thickness,
      )
      cv2.arrowedLine(
          image, origin, z_axis, BLUE_COLOR, axis_drawing_spec.thickness
      )
  
file_path: mediapipe/mediapipe/python/solutions/drawing_utils.py
imports: import dataclasses
  import math
  from typing import List, Mapping, Optional, Tuple, Union
  import cv2
  import matplotlib.pyplot
  import numpy
  from mediapipe.framework.formats import detection_pb2
  from mediapipe.framework.formats import landmark_pb2
  from mediapipe.framework.formats import location_data_pb2
label_desc: Use image node_id 4 for calling a global function with example usage: draw_axis(image, rotation, translation, focal_length, principal_point, axis_length, axis_drawing_spec) without return types