111 lines
3.6 KiB
Python
111 lines
3.6 KiB
Python
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 |