change folders

Karussell/utils/utils.py

@@ -0,0 +1,111 @@
+import cv2 as cv
+import pandas as pd
+import math
+import numpy as np
+from tqdm.notebook import tqdm
+
+
+def encode(ann):
+    keypoints_2d = pd.DataFrame(np.array(ann['keypoints']).reshape(-1, 3),
+                                columns=['x','y','conf'], index=range(1, 25))
+    bbox = np.array(ann['bbox'])
+    score = ann['score']
+    return keypoints_2d, bbox, score
+
+
+def fit(imageSize, keypoints_2d, keypoints_3d, focus=1):
+
+    objectPoints = keypoints_3d.loc[keypoints_2d.index].values
+    imagePoints = keypoints_2d[['x', 'y']].values.astype('float')
+
+    n = len(imagePoints)
+    fx = fy = focus*np.hypot(*imageSize)
+    cx = imageSize[0]/2
+    cy = imageSize[1]/2
+    distCoeffs = np.zeros(4, np.float32)
+    if n < 6:
+        raise ValueError('Number of keypoints must be > 5')
+
+    cameraMatrix = np.float32([[fx,0, cx],
+                               [0, fy,cy],
+                               [0, 0, 1]])
+
+    _, rvecs, tvecs = cv.solvePnP(objectPoints, imagePoints, cameraMatrix, distCoeffs, flags=cv.SOLVEPNP_ITERATIVE )
+    return rvecs, tvecs, cameraMatrix, distCoeffs
+
+
+def rvec2euler(rvec):
+    rvec_matrix = cv.Rodrigues(rvec)[0]
+    proj_matrix = np.hstack((rvec_matrix, np.zeros_like(rvec)))
+    euler_angles = cv.decomposeProjectionMatrix(proj_matrix)[6]
+    return euler_angles.flatten()
+
+# Checks if a matrix is a valid rotation matrix.
+def isRotationMatrix(R) :
+    Rt = np.transpose(R)
+    shouldBeIdentity = np.dot(Rt, R)
+    I = np.identity(3, dtype = R.dtype)
+    n = np.linalg.norm(I - shouldBeIdentity)
+    return n < 1e-6
+
+# Calculates rotation matrix to euler angles
+# The result is the same as MATLAB except the order
+# of the euler angles ( x and z are swapped ).
+def rotationMatrixToEulerAngles(R) :
+
+    assert(isRotationMatrix(R))
+
+    sy = math.sqrt(R[0,0] * R[0,0] +  R[1,0] * R[1,0])
+
+    singular = sy < 1e-6
+
+    if  not singular :
+        x = math.atan2(R[2,1] , R[2,2])
+        y = math.atan2(-R[2,0], sy)
+        z = math.atan2(R[1,0], R[0,0])
+    else :
+        x = math.atan2(-R[1,2], R[1,1])
+        y = math.atan2(-R[2,0], sy)
+        z = 0
+
+    return np.array([x, y, z])
+
+def rvec2euler_Fridman(rvec):
+    R = cv.Rodrigues(rvec)[0]
+    euler_angles = np.rad2deg(rotationMatrixToEulerAngles(R))
+    return euler_angles
+
+
+def pose_estimation(img_json, keypoints_3d, ImageSize=(1920, 1080), treshhold=0.01, focus=1, fridman=False):
+    poses = []
+    for i, ann in enumerate(img_json):
+        keypoints_2d, bbox, score = encode(ann)
+
+        conf_keypoints_2d = keypoints_2d[keypoints_2d['conf']>treshhold]
+        if len(conf_keypoints_2d)>5:
+            rvec, tvec, camMatrx, dist = fit(ImageSize, conf_keypoints_2d, keypoints_3d, focus=focus)
+            if not fridman:
+                euler_angles = rvec2euler(rvec)
+            else:
+                euler_angles = rvec2euler_Fridman(rvec) #alternative
+
+            pose = {'Euler angles':[-euler_angles[2], -euler_angles[0]+90, -euler_angles[1]],
+                    'xyz coords':[tvec[2][0],tvec[0][0],-tvec[1][0]]}
+            poses.append(pose)
+    return poses
+
+
+def image_pose_estimation(img_json, img_name, *args, **kvargs):
+    poses = pose_estimation(img_json, *args, **kvargs)
+    UE4_json = {'data':{'filename':img_name, 'poses':poses}}
+    return UE4_json
+
+
+def video_pose_estimation(video_json, video_name, *args, **kvargs):
+    video_poses = []
+    for frame in tqdm(video_json):
+        num = frame['frame']
+        frame_poses = pose_estimation(frame['predictions'], *args, **kvargs)
+        video_poses.append({'frame':num, 'poses':frame_poses})
+    UE4_json_video = {'filename':video_name, 'data':video_poses}
+    return UE4_json_video