angles.py

import numpy as np
import json
import math
import matplotlib.pyplot as plt


def showAnns(jts):
    # c = (np.random.random((1, 3)) * 0.6 + 0.4).tolist()[0]
    c = [0.7, 0.9, 0.8]

    sks = np.array([[16,14],[14,12],[17,15],[15,13],[12,13],[6,12],[7,13],[6,7],
                    [6,8],[7,9],[8,10],[9,11],[2,3],[1,2],[1,3],[2,4],[3,5],[4,6],[5,7]]) - 1
    kp = np.array(jts)
    x = kp[0::3]
    y = kp[1::3]
    v = kp[2::3]
    for sk in sks:
        if np.all(v[sk] > 0):
            plt.plot(x[sk], y[sk], linewidth=3, color=c)
    plt.plot(x[v > 0], y[v > 0], 'o', markersize=8, markerfacecolor=c, markeredgecolor='k',
             markeredgewidth=2)
    plt.plot(x[v > 1], y[v > 1], 'o', markersize=8, markerfacecolor=c, markeredgecolor=c, markeredgewidth=2)

class Angles():
    def __init__(self, joints):
        joints = np.array(joints)

        joints = np.vstack([joints[0::3], joints[1::3]])

        keys = ["Nose","LEye","REye","LEar","REar","LShoulder","RShoulder","LElbow",
                "RElbow","LWrist","RWrist","LHip","RHip","LKnee","RKnee","LAnkle","RAnkle"]

        self.joints = {}
        for i in range(len(keys)):
            self.joints[keys[i]] = joints[:,i]

        self.normal = np.array([1,0])

        self.values = {}
        self.values['feet_width'] = self.get_feet_width()
        self.values['shoulder_plane'] = self.get_shoulder_plane()
        self.values['hip_height'] = self.get_hip_height()
        self.values['left_elbow_angle'] = self.get_left_elbow_angle()
        self.values['right_elbow_angle'] = self.get_right_elbow_angle()
        self.values['left_ground_angle'] = self.get_left_ground_angle()
        self.values['right_ground_angle'] = self.get_right_ground_angle()
        self.values['left_knee_angle'] = self.get_left_knee_angle()
        self.values['right_knee_angle'] = self.get_right_knee_angle()

    def unit_vector(self, vector):
        """ Returns the unit vector of the vector.  """
        return vector / np.linalg.norm(vector)


    def angle(self, v1, v2):
        v1 = np.array(v1)
        v2 = np.array(v2)
        v1_u = self.unit_vector(v1)
        v2_u = self.unit_vector(v2)
        return np.degrees(np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0)))


    def length(self, v1, v2):
        v1 = np.array(v1)
        v2 = np.array(v2)
        return np.linalg.norm(v2 - v1)


    def get_feet_width(self):
        return self.length(self.joints["LAnkle"], self.joints["RAnkle"])


    def get_shoulder_plane(self):
        shoulder = self.joints["LShoulder"] - self.joints["RShoulder"]

        return self.angle(shoulder, self.normal)


    def get_hip_height(self):
        return self.length(self.joints["LHip"], self.joints["LAnkle"])


    def get_left_elbow_angle(self):
        up = self.joints["LElbow"] - self.joints["LShoulder"]
        down = self.joints["LElbow"]  - self.joints["LWrist"]

        return self.angle(up, down)


    def get_right_elbow_angle(self):
        up = self.joints["RElbow"] - self.joints["RShoulder"]
        down = self.joints["RElbow"] - self.joints["RWrist"]

        return self.angle(up, down)


    def get_left_ground_angle(self,):
        leg = self.joints["LKnee"] - self.joints["LAnkle"]
        return self.angle(leg, self.normal)


    def get_right_ground_angle(self,):
        leg = self.joints["RKnee"] - self.joints["RAnkle"]
        return self.angle(leg, self.normal)


    def get_left_knee_angle(self):
        up = self.joints["LKnee"] - self.joints["LHip"]
        down = self.joints["LKnee"] - self.joints["LAnkle"]

        return self.angle(up, down)


    def get_right_knee_angle(self):
        up = self.joints["RKnee"] - self.joints["RHip"]
        down = self.joints["RKnee"] - self.joints["RAnkle"]

        return self.angle(up, down)

def angle_demo():
    joints = json.load(open('examples/golf/Suzann Pettersen - 2013 0122 104829 Prosilica GE680C Face on right/pose-results.json'))
    joints = [x['keypoints'] for x in joints]

    for jts in joints:
        ang = Angles(jts)
        print(ang.joints)
        print(ang.values)

        plt.xlim([0,512])
        plt.ylim([480,0])
        showAnns(jts)
        plt.draw()
        plt.pause(0.1)
        plt.clf()

if __name__ == '__main__':
    angle_demo()