text: # ------------------------------------------------------------------------------
  # Copyright (c) Microsoft
  # Licensed under the MIT License.
  # Written by Bin Xiao (leoxiaobin@gmail.com)
  # Modified by Bowen Cheng (bcheng9@illinois.edu)
  # ------------------------------------------------------------------------------
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  
  
  class HeatmapGenerator():
      def __init__(self, output_res, num_joints, sigma=-1):
          self.output_res = output_res
          self.num_joints = num_joints
          if sigma < 0:
              sigma = self.output_res/64
          self.sigma = sigma
          size = 6*sigma + 3
          x = np.arange(0, size, 1, float)
          y = x[:, np.newaxis]
          x0, y0 = 3*sigma + 1, 3*sigma + 1
          self.g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
  
      def __call__(self, joints):
          hms = np.zeros((self.num_joints, self.output_res, self.output_res),
                         dtype=np.float32)
          sigma = self.sigma
  
          # joints: ndarray(num_people, num_joints, 4)
          # p: ndarray(num_joints, 4)
          # idx: ndarray(num_joints)
          # pt: ndarray(4)
          for p in joints:
              for idx, pt in enumerate(p):
                  if pt[2] > 0:
                      x, y = int(pt[0]), int(pt[1])
                      if x < 0 or y < 0 or \
                         x >= self.output_res or y >= self.output_res:
                          continue
  
                      ul = int(np.round(x - 3 * sigma - 1)), int(np.round(y - 3 * sigma - 1))
                      br = int(np.round(x + 3 * sigma + 2)), int(np.round(y + 3 * sigma + 2))
  
                      c, d = max(0, -ul[0]), min(br[0], self.output_res) - ul[0]
                      a, b = max(0, -ul[1]), min(br[1], self.output_res) - ul[1]
  
                      cc, dd = max(0, ul[0]), min(br[0], self.output_res)
                      aa, bb = max(0, ul[1]), min(br[1], self.output_res)
                      hms[idx, aa:bb, cc:dd] = np.maximum(
                          hms[idx, aa:bb, cc:dd], self.g[a:b, c:d])
  
          return hms
  
  
  class ScaleAwareHeatmapGenerator():
      def __init__(self, output_res, num_joints):
          self.output_res = output_res
          self.num_joints = num_joints
  
      def get_gaussian_kernel(self, sigma):
          size = 6*sigma + 3
          x = np.arange(0, size, 1, float)
          y = x[:, np.newaxis]
          x0, y0 = 3*sigma + 1, 3*sigma + 1
          g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
          return g
  
      def __call__(self, joints):
          hms = np.zeros((self.num_joints, self.output_res, self.output_res),
                         dtype=np.float32)
          for p in joints:
              # sigma = p[0, 3] 
              # g = self.get_gaussian_kernel(sigma)
              for idx, pt in enumerate(p):
                  ################# 3sigma原则 ###################
                  sigma = pt[3]
                  g = self.get_gaussian_kernel(sigma)
                  ################# 3sigma原则 ###################
                  ################### 制定R ######################
                  sigma = pt[3]
                  thr = 0.01
                  radius = np.sqrt(np.log(1 / thr) * 2 * sigma ** 2)
                  radius = int(np.round(radius))  # 取整
                  size = 2 * radius + 3
                  x = np.arange(0, size, 1, float)
                  y = x[:, np.newaxis]
                  x0, y0 = radius + 1, radius + 1
                  g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
                  ################### 制定R ######################    
                  if pt[2] > 0:
                      x, y = int(pt[0]), int(pt[1])
                      if x < 0 or y < 0 or \
                         x >= self.output_res or y >= self.output_res:
                          continue
                      ################### 制定R ######################
                      ul = int(np.round(x - radius - 1)), int(np.round(y - radius - 1))
                      br = int(np.round(x + radius + 2)), int(np.round(y + radius + 2))
                      ################### 制定R ######################
                      ################# 3sigma原则 ###################
                      ul = int(np.round(x - 3 * sigma - 1)), int(np.round(y - 3 * sigma - 1))
                      br = int(np.round(x + 3 * sigma + 2)), int(np.round(y + 3 * sigma + 2))
                      ################# 3sigma原则 ###################
  
                      c, d = max(0, -ul[0]), min(br[0], self.output_res) - ul[0]
                      a, b = max(0, -ul[1]), min(br[1], self.output_res) - ul[1]
  
                      cc, dd = max(0, ul[0]), min(br[0], self.output_res)
                      aa, bb = max(0, ul[1]), min(br[1], self.output_res)
                      hms[idx, aa:bb, cc:dd] = np.maximum(
                          hms[idx, aa:bb, cc:dd], g[a:b, c:d])
          return hms
  
  
  class JointsGenerator():
      def __init__(self, max_num_people, num_joints, output_res, tag_per_joint):
          self.max_num_people = max_num_people
          self.num_joints = num_joints
          self.output_res = output_res
          self.tag_per_joint = tag_per_joint
  
      def __call__(self, joints):
          visible_nodes = np.zeros((self.max_num_people, self.num_joints, 2))
          output_res = self.output_res
          for i in range(len(joints)):
              tot = 0
              for idx, pt in enumerate(joints[i]):
                  x, y = int(pt[0]), int(pt[1])
                  if pt[2] > 0 and x >= 0 and y >= 0 \
                     and x < self.output_res and y < self.output_res:
                      if self.tag_per_joint:
                          visible_nodes[i][tot] = \
                              (idx * output_res**2 + y * output_res + x, 1)
                      else:
                          visible_nodes[i][tot] = \
                              (y * output_res + x, 1)
                      tot += 1
          return visible_nodes
  