init

utils/functions_shield.py

@@ -0,0 +1,133 @@
+import cv2
+import numpy as np
+import os
+
+# Read shield with alpha channel
+def readAlpha(path):
+    src = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+    src = cv2.cvtColor(src, cv2.COLOR_BGRA2RGBA)
+    # b, g, r, alpha = cv2.split(src)
+    # dst1 = cv2.merge((b, g, r))
+    # alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2BGR)
+    # dst1_64 = dst1.astype(np.float64)
+    # alpha_64 = alpha.astype(np.float64)
+    # alpha_64 *= 1./255.
+    # dst2_64 = np.multiply(dst1_64, alpha_64)
+    # dst2 = dst2_64.astype(np.uint8)
+    return src
+
+# Read shield with alpha channel and points
+def readShield(path, step_board=0):
+    pict = readAlpha(path)
+    h_pict = pict.shape[0]
+    w_pict = pict.shape[1]
+    pts_pict = np.array(
+        [[step_board, step_board], [w_pict - step_board, step_board], [w_pict - step_board, h_pict - step_board],
+         [step_board, h_pict - step_board]])
+    return pict, pts_pict
+
+def ReadShields(shield_album):
+    picts_paths = []
+    f = open(shield_album)
+    for g in f:
+        if g[-1] == '\n':
+            picts_paths.append(g[:-1])
+        else:
+            picts_paths.append(g)
+    f.close()
+    # print(picts_paths)
+    picts = []
+    pts_picts = []
+    for shield_path in picts_paths:
+        pict, pts_pict = readShield(shield_path)
+        picts.append(pict)
+        pts_picts.append(pts_pict)
+    return picts, pts_picts
+
+def ReadShieldsAni(shield_album):
+    picts_paths = []
+    # filelist = os.listdir(shield_album)
+    # for g in filelist:
+    #     picts_paths.append(shield_album + "/" + g)
+    # # print(picts_paths)
+    picts = []
+    pts_picts = []
+    for i in range(0,55):
+        str_i = str(i).zfill(4)
+        pict, pts_pict = readShield(shield_album + "/frame_"+str_i+".png")
+        picts.append(pict)
+        pts_picts.append(pts_pict)
+    return picts, pts_picts
+
+# Overaling transformed shield on frame
+def overlaying(pict, img, pts_pict, pts_frame, part=1.0):
+    # Homography
+    # h, status = cv2.findHomography(pts_pict, pts_frame)
+
+    
+    
+#     x1 = int(pts_frame[0][1])
+#     if x1 < 0:
+#         x1 = 0
+#     x2 = int(pts_frame[2][1])
+#     y1 = int(pts_frame[0][0])
+#     if y1 < 0:
+#         y1 = 0
+#     y2 = int(pts_frame[2][0])
+
+    (y1, x1), (y2, x2) = pts_frame.min(axis=0), pts_frame.max(axis=0)
+    if (x2 - x1) > 50 and (y2 - y1) > 50:
+        try:
+            area = img[x1:x2, y1:y2]
+            # rows, cols, channels = area.shape
+
+            b, g, r, alpha = cv2.split(pict)
+            pict_rgb = cv2.merge((r, g, b))
+
+            h, status = cv2.findHomography(pts_pict, pts_frame)
+            imgWarp = cv2.warpPerspective(pict_rgb, h, (img.shape[1], img.shape[0]))
+            alphaWarp = cv2.warpPerspective(alpha, h, (img.shape[1], img.shape[0]))
+            pict_area = imgWarp[x1:x2, y1:y2]
+            alpha_area = alphaWarp[x1:x2, y1:y2]
+
+            # pict = cv2.resize(pict, (cols, rows))
+
+            znam = 1./255.
+            alpha_area = cv2.cvtColor(alpha_area, cv2.COLOR_GRAY2BGR)
+            pict_32 = pict_area.astype(np.float32)
+            area_32 = area.astype(np.float32)
+            alpha_32 = alpha_area.astype(np.float32)
+            alpha_32 *= znam
+            if part != 1.:
+                alpha_32 *= part
+            alpha_nega_32 = 1. - alpha_32
+            dst1_32 = np.multiply(pict_32, alpha_32)
+            dst2_32 = np.multiply(area_32, alpha_nega_32)
+            dst_32 = dst1_32 + dst2_32
+            img[x1:x2, y1:y2] = dst_32.astype(np.uint8)
+
+            # Warp transform shield
+            # imgWarp = cv2.warpPerspective(pict_rgb, h, (img.shape[1], img.shape[0]))
+            # alphaWarp = cv2.warpPerspective(alpha, h, (img.shape[1], img.shape[0]))
+
+            # Placing sheild on frame
+            # imgAug = cv2.addWeighted(img, 1.0, imgWarp, 0.5, 0.0)
+        except:
+            pass
+
+    return img
+
+def applistShields(album_paths):
+    picts_paths = []
+    picts_names = []
+    f = open(album_paths)
+    for g in f:
+        if g[-1] == '\n':
+            picts_paths.append(g[:-1])
+        else:
+            picts_paths.append(g)
+    f.close()
+    for path in picts_paths:
+        name = path.split('.')[0].split("/")[1]
+        picts_names.append(name)
+    return picts_paths, picts_names

utils/karussel_24kps.csv

@@ -0,0 +1,25 @@
+,x,y,z
+1,-11.728358247427778,-219.4412122023117,0.0
+2,-15.92420132696894,20.938476833288533,0.0
+3,-24.15779771912961,46.96169512870693,0.0
+4,-86.61374177228646,-142.95474716218305,0.0
+5,-134.74977744065188,-261.8997851845144,0.0
+6,-188.2631533958743,-315.8749516397929,0.0
+7,-191.57404333357664,-166.7104853355808,0.0
+8,-210.8939138893729,-113.29941343370568,0.0
+9,-231.00716853514677,-54.9516886489786,0.0
+10,-216.46512728143847,3.293732082401192,0.0
+11,-187.73016312616568,99.26925935365503,0.0
+12,-248.253432074131,163.27655515540138,0.0
+13,-72.81535484187529,200.9923483599607,0.0
+14,70.59797139345626,250.17179384400777,0.0
+15,126.43140175736619,171.23113224430733,0.0
+16,165.9457279162327,135.85291821065778,0.0
+17,194.32083871787557,90.37926554238746,0.0
+18,216.53917157107745,-7.535726594794127,0.0
+19,239.72723685724318,-79.97407022773838,0.0
+20,197.07800219673948,-148.61110675229477,0.0
+21,200.24438188899256,-248.98029544904819,0.0
+22,133.21163830661362,-298.2409418903916,0.0
+23,69.21865584018204,-197.5190565076657,0.0
+24,55.01178651022396,-112.90129912872368,0.0

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