text: # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this open-source project.
  
  
  import os.path
  import io
  from PIL import Image
  import numpy as np
  import json
  import glob
  from scipy.optimize import curve_fit
  import warnings
  
  def func(x, a, b, c):
      return a * x**2 + b * x + c
  
  def linear(x, a, b):
      return a * x + b
  
  def setColor(im, yy, xx, color):
      if len(im.shape) == 3:
          if (im[yy, xx] == 0).all():            
              im[yy, xx, 0], im[yy, xx, 1], im[yy, xx, 2] = color[0], color[1], color[2]            
          else:            
              im[yy, xx, 0] = ((im[yy, xx, 0].astype(float) + color[0]) / 2).astype(np.uint8)
              im[yy, xx, 1] = ((im[yy, xx, 1].astype(float) + color[1]) / 2).astype(np.uint8)
              im[yy, xx, 2] = ((im[yy, xx, 2].astype(float) + color[2]) / 2).astype(np.uint8)
      else:
          im[yy, xx] = color[0]
  
  def drawEdge(im, x, y, bw=1, color=(255,255,255), draw_end_points=False):
      if x is not None and x.size:
          h, w = im.shape[0], im.shape[1]
          # edge
          for i in range(-bw, bw):
              for j in range(-bw, bw):
                  yy = np.maximum(0, np.minimum(h-1, y+i))
                  xx = np.maximum(0, np.minimum(w-1, x+j))
                  setColor(im, yy, xx, color)
  
          # edge endpoints
          if draw_end_points:
              for i in range(-bw*2, bw*2):
                  for j in range(-bw*2, bw*2):
                      if (i**2) + (j**2) < (4 * bw**2):
                          yy = np.maximum(0, np.minimum(h-1, np.array([y[0], y[-1]])+i))
                          xx = np.maximum(0, np.minimum(w-1, np.array([x[0], x[-1]])+j))
                          setColor(im, yy, xx, color)
  
  def interpPoints(x, y):
      if abs(x[:-1] - x[1:]).max() < abs(y[:-1] - y[1:]).max():
          curve_y, curve_x = interpPoints(y, x)
          if curve_y is None:
              return None, None
      else:        
          with warnings.catch_warnings():
              warnings.simplefilter("ignore")    
              if len(x) < 3:
                  popt, _ = curve_fit(linear, x, y)
              else:
                  popt, _ = curve_fit(func, x, y)                
                  if abs(popt[0]) > 1:
                      return None, None
          if x[0] > x[-1]:
              x = list(reversed(x))
              y = list(reversed(y))
          #curve_x = np.linspace(x[0], x[-1], (x[-1]-x[0]))
          curve_x = np.linspace(x[0], x[-1], int(x[-1]-x[0]))
          if len(x) < 3:
              curve_y = linear(curve_x, *popt)
          else:
              curve_y = func(curve_x, *popt)
      return curve_x.astype(int), curve_y.astype(int)
  
  def read_keypoints(json_input, size, random_drop_prob=0, remove_face_labels=False, basic_point_only=False):
      with io.open(json_input, encoding='utf-8') as f:
          keypoint_dicts = json.loads(f.read())["people"]
  
      edge_lists = define_edge_lists(basic_point_only)
      w, h = size    
      pose_img = np.zeros((h, w, 3), np.uint8)
      for keypoint_dict in keypoint_dicts:    
          pose_pts = np.array(keypoint_dict["pose_keypoints_2d"]).reshape(25, 3)
          # face_pts = np.array(keypoint_dict["face_keypoints_2d"]).reshape(70, 3)
          face_pts = np.zeros(210).reshape(70, 3)
          # hand_pts_l = np.array(keypoint_dict["hand_left_keypoints_2d"]).reshape(21, 3)
          hand_pts_l = np.zeros(63).reshape(21, 3)
          # hand_pts_r = np.array(keypoint_dict["hand_right_keypoints_2d"]).reshape(21, 3)
          hand_pts_r = np.zeros(63).reshape(21, 3)
          pts = [extract_valid_keypoints(pts, edge_lists) for pts in [pose_pts, face_pts, hand_pts_l, hand_pts_r]]
          pose_img += connect_keypoints(pts, edge_lists, size, random_drop_prob, remove_face_labels, basic_point_only)
      return pose_img
  
  def extract_valid_keypoints(pts, edge_lists):
      pose_edge_list, _, hand_edge_list, _, face_list = edge_lists
      p = pts.shape[0]
      thre = 0.1 if p == 70 else 0.01
      output = np.zeros((p, 2))    
  
      if p == 70:   # face
          for edge_list in face_list:
              for edge in edge_list:
                  if (pts[edge, 2] > thre).all():
                      output[edge, :] = pts[edge, :2]        
      elif p == 21: # hand        
          for edge in hand_edge_list:            
              if (pts[edge, 2] > thre).all():
                  output[edge, :] = pts[edge, :2]
      else:         # pose
          valid = (pts[:, 2] > thre)        
          output[valid, :] = pts[valid, :2]
          
      return output
  
  def connect_keypoints(pts, edge_lists, size, random_drop_prob, remove_face_labels, basic_point_only):
      pose_pts, face_pts, hand_pts_l, hand_pts_r = pts
      w, h = size
      output_edges = np.zeros((h, w, 3), np.uint8)
      pose_edge_list, pose_color_list, hand_edge_list, hand_color_list, face_list = edge_lists
      
      if random_drop_prob > 0 and remove_face_labels:
          # add random noise to keypoints
          pose_pts[[0,15,16,17,18], :] += 5 * np.random.randn(5,2)
          face_pts[:,0] += 2 * np.random.randn()
          face_pts[:,1] += 2 * np.random.randn()
  
      ### pose    
      for i, edge in enumerate(pose_edge_list):
          x, y = pose_pts[edge, 0], pose_pts[edge, 1]
          if (np.random.rand() > random_drop_prob) and (0 not in x):
              curve_x, curve_y = interpPoints(x, y)                                        
              drawEdge(output_edges, curve_x, curve_y, bw=3, color=pose_color_list[i], draw_end_points=True)
  
      if not basic_point_only:
          ### hand       
          for hand_pts in [hand_pts_l, hand_pts_r]:     # for left and right hand
              if np.random.rand() > random_drop_prob:
                  for i, edge in enumerate(hand_edge_list): # for each finger
                      for j in range(0, len(edge)-1):       # for each part of the finger
                          sub_edge = edge[j:j+2] 
                          x, y = hand_pts[sub_edge, 0], hand_pts[sub_edge, 1]                    
                          if 0 not in x:
                              line_x, line_y = interpPoints(x, y)                                        
                              drawEdge(output_edges, line_x, line_y, bw=1, color=hand_color_list[i], draw_end_points=True)
  
          ### face
          edge_len = 2
          if (np.random.rand() > random_drop_prob):
              for edge_list in face_list:
                  for edge in edge_list:
                      for i in range(0, max(1, len(edge)-1), edge_len-1):             
                          sub_edge = edge[i:i+edge_len]
                          x, y = face_pts[sub_edge, 0], face_pts[sub_edge, 1]
                          if 0 not in x:
                              curve_x, curve_y = interpPoints(x, y)
                              drawEdge(output_edges, curve_x, curve_y, draw_end_points=True)
  
      return output_edges
  
  def define_edge_lists(basic_point_only):
      ### pose        
      pose_edge_list = []
      pose_color_list = []
      if not basic_point_only:
          pose_edge_list += [[17, 15], [15,  0], [ 0, 16], [16, 18]]       # head
          pose_color_list += [[153,  0,153], [153,  0,102], [102,  0,153], [ 51,  0,153]]
  
      pose_edge_list += [        
          [ 0,  1], [ 1,  8],                                         # body
          [ 1,  2], [ 2,  3], [ 3,  4],                               # right arm
          [ 1,  5], [ 5,  6], [ 6,  7],                               # left arm
          [ 8,  9], [ 9, 10], [10, 11], [11, 24], [11, 22], [22, 23], # right leg
          [ 8, 12], [12, 13], [13, 14], [14, 21], [14, 19], [19, 20]  # left leg
      ]
      pose_color_list += [
          [153,  0, 51], [153,  0,  0],
          [153, 51,  0], [153,102,  0], [153,153,  0],
          [102,153,  0], [ 51,153,  0], [  0,153,  0],
          [  0,153, 51], [  0,153,102], [  0,153,153], [  0,153,153], [  0,153,153], [  0,153,153],
          [  0,102,153], [  0, 51,153], [  0,  0,153], [  0,  0,153], [  0,  0,153], [  0,  0,153]
      ]
  
      ### hand
      hand_edge_list = [
          [0,  1,  2,  3,  4],
          [0,  5,  6,  7,  8],
          [0,  9, 10, 11, 12],
          [0, 13, 14, 15, 16],
          [0, 17, 18, 19, 20]
      ]
      hand_color_list = [
          [204,0,0], [163,204,0], [0,204,82], [0,82,204], [163,0,204]
      ]
  
      ### face        
      face_list = [
                   #[range(0, 17)], # face
                   [range(17, 22)], # left eyebrow
                   [range(22, 27)], # right eyebrow
                   [range(27, 31), range(31, 36)], # nose
                   [[36,37,38,39], [39,40,41,36]], # left eye
                   [[42,43,44,45], [45,46,47,42]], # right eye
                   [range(48, 55), [54,55,56,57,58,59,48]], # mouth
                  ]
      return pose_edge_list, pose_color_list, hand_edge_list, hand_color_list, face_list
  