text: import cv2
  import numpy as np
  from math import cos, sin
  
  
  def draw_axis(img, yaw, pitch, roll, tdx=None, tdy=None, size=100):
      """draw head pose axis."""
  
      pitch = pitch * np.pi / 180
      yaw = -yaw * np.pi / 180
      roll = roll * np.pi / 180
  
      if tdx is None or tdy is None:
          height, width = img.shape[:2]
          tdx = width / 2
          tdy = height / 2
  
      # X axis pointing to right, drawn in red
      x1 = size * (cos(yaw) * cos(roll)) + tdx
      y1 = size * (cos(pitch) * sin(roll) + cos(roll) * sin(pitch) * sin(yaw)) + tdy
      # Y axis poiting downside, drawn in green
      x2 = size * (-cos(yaw) * sin(roll)) + tdx
      y2 = size * (cos(pitch) * cos(roll) - sin(pitch) * sin(yaw) * sin(roll)) + tdy
      # Z axis, out of the screen, drawn in blue
      x3 = size * (sin(yaw)) + tdx
      y3 = size * (-cos(yaw) * sin(pitch)) + tdy
  
      cv2.line(img, (int(tdx), int(tdy)), (int(x1), int(y1)), (0, 0, 255), 3)
      cv2.line(img, (int(tdx), int(tdy)), (int(x2), int(y2)), (0, 255, 0), 3)
      cv2.line(img, (int(tdx), int(tdy)), (int(x3), int(y3)), (255, 0, 0), 2)
  
      return img
  
  
  def draw_pose_cube(img, yaw, pitch, roll, tdx=None, tdy=None, size=150.):
      """draw head pose cube.
      Where (tdx, tdy) is the translation of the face.
      For pose we have [pitch yaw roll tdx tdy tdz scale_factor]
      """
  
      p = pitch * np.pi / 180
      y = -yaw * np.pi / 180
      r = roll * np.pi / 180
      if tdx is not None and tdy is not None:
          face_x = tdx - 0.50 * size
          face_y = tdy - 0.50 * size
      else:
          height, width = img.shape[:2]
          face_x = width / 2 - 0.5 * size
          face_y = height / 2 - 0.5 * size
  
      x1 = size * (cos(y) * cos(r)) + face_x
      y1 = size * (cos(p) * sin(r) + cos(r) * sin(p) * sin(y)) + face_y
      x2 = size * (-cos(y) * sin(r)) + face_x
      y2 = size * (cos(p) * cos(r) - sin(p) * sin(y) * sin(r)) + face_y
      x3 = size * (sin(y)) + face_x
      y3 = size * (-cos(y) * sin(p)) + face_y
  
      # Draw base in red
      cv2.line(img, (int(face_x), int(face_y)), (int(x1), int(y1)), (0, 0, 255), 3)
      cv2.line(img, (int(face_x), int(face_y)), (int(x2), int(y2)), (0, 0, 255), 3)
      cv2.line(img, (int(x2), int(y2)), (int(x2 + x1 - face_x), int(y2 + y1 - face_y)), (0, 0, 255), 3)
      cv2.line(img, (int(x1), int(y1)), (int(x1 + x2 - face_x), int(y1 + y2 - face_y)), (0, 0, 255), 3)
      # Draw pillars in blue
      cv2.line(img, (int(face_x), int(face_y)), (int(x3), int(y3)), (255, 0, 0), 2)
      cv2.line(img, (int(x1), int(y1)), (int(x1 + x3 - face_x), int(y1 + y3 - face_y)), (255, 0, 0), 2)
      cv2.line(img, (int(x2), int(y2)), (int(x2 + x3 - face_x), int(y2 + y3 - face_y)), (255, 0, 0), 2)
      cv2.line(img, (int(x2 + x1 - face_x), int(y2 + y1 - face_y)),
               (int(x3 + x1 + x2 - 2 * face_x), int(y3 + y2 + y1 - 2 * face_y)), (255, 0, 0), 2)
      # Draw top in green
      cv2.line(img, (int(x3 + x1 - face_x), int(y3 + y1 - face_y)),
               (int(x3 + x1 + x2 - 2 * face_x), int(y3 + y2 + y1 - 2 * face_y)), (0, 255, 0), 2)
      cv2.line(img, (int(x2 + x3 - face_x), int(y2 + y3 - face_y)),
               (int(x3 + x1 + x2 - 2 * face_x), int(y3 + y2 + y1 - 2 * face_y)), (0, 255, 0), 2)
      cv2.line(img, (int(x3), int(y3)), (int(x3 + x1 - face_x), int(y3 + y1 - face_y)), (0, 255, 0), 2)
      cv2.line(img, (int(x3), int(y3)), (int(x3 + x2 - face_x), int(y3 + y2 - face_y)), (0, 255, 0), 2)
  
      return img
  
  
  def visualize_headpose(img, yaw, pitch, roll, save_path=None, to_bgr=False):
      img = np.copy(img)
      if to_bgr:
          img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
      show_string = (f'y {yaw[0].item():.2f}, p {pitch[0].item():.2f}, ' + f'r {roll[0].item():.2f}')
      cv2.putText(img, show_string, (30, img.shape[0] - 30), fontFace=1, fontScale=1, color=(0, 0, 255), thickness=2)
      draw_pose_cube(img, yaw[0], pitch[0], roll[0], size=100)
      draw_axis(img, yaw[0], pitch[0], roll[0], tdx=50, tdy=50, size=100)
      # save img
      if save_path is not None:
          cv2.imwrite(save_path, img)
  