Spaces:
Running
Running
| import os | |
| import cv2 | |
| import math | |
| import torch | |
| import numpy as np | |
| from dtw import dtw | |
| from constants import PAD_TOKEN | |
| from helpers import getSkeletalModelStructure | |
| # This is the format of the 3D data, outputted from the Inverse Kinematics model | |
| # def getSkeletalModelStructure(): | |
| # # Definition of skeleton model structure: | |
| # # The structure is an n-tuple of: | |
| # # | |
| # # (index of a start point, index of an end point, index of a bone) | |
| # # | |
| # # E.g., this simple skeletal model | |
| # # | |
| # # (0) | |
| # # | | |
| # # | | |
| # # 0 | |
| # # | | |
| # # | | |
| # # (2)--1--(1)--1--(3) | |
| # # | | | |
| # # | | | |
| # # 2 2 | |
| # # | | | |
| # # | | | |
| # # (4) (5) | |
| # # | |
| # # has this structure: | |
| # # | |
| # # ( | |
| # # (0, 1, 0), | |
| # # (1, 2, 1), | |
| # # (1, 3, 1), | |
| # # (2, 4, 2), | |
| # # (3, 5, 2), | |
| # # ) | |
| # # | |
| # # Warning 1: The structure has to be a tree. | |
| # # Warning 2: The order isn't random. The order is from a root to lists. | |
| # # | |
| # return ( | |
| # # head | |
| # (0, 1, 0), | |
| # # left shoulder | |
| # (1, 2, 1), | |
| # # left arm | |
| # (2, 3, 2), | |
| # # (3, 4, 3), | |
| # # Changed to avoid wrist, go straight to hands | |
| # (3, 29, 3), | |
| # # right shoulder | |
| # (1, 5, 1), | |
| # # right arm | |
| # (5, 6, 2), | |
| # # (6, 7, 3), | |
| # # Changed to avoid wrist, go straight to hands | |
| # (6, 8, 3), | |
| # # left hand - wrist | |
| # # (7, 8, 4), | |
| # # left hand - palm | |
| # (8, 9, 5), | |
| # (8, 13, 9), | |
| # (8, 17, 13), | |
| # (8, 21, 17), | |
| # (8, 25, 21), | |
| # # left hand - 1st finger | |
| # (9, 10, 6), | |
| # (10, 11, 7), | |
| # (11, 12, 8), | |
| # # left hand - 2nd finger | |
| # (13, 14, 10), | |
| # (14, 15, 11), | |
| # (15, 16, 12), | |
| # # left hand - 3rd finger | |
| # (17, 18, 14), | |
| # (18, 19, 15), | |
| # (19, 20, 16), | |
| # # left hand - 4th finger | |
| # (21, 22, 18), | |
| # (22, 23, 19), | |
| # (23, 24, 20), | |
| # # left hand - 5th finger | |
| # (25, 26, 22), | |
| # (26, 27, 23), | |
| # (27, 28, 24), | |
| # # right hand - wrist | |
| # # (4, 29, 4), | |
| # # right hand - palm | |
| # (29, 30, 5), | |
| # (29, 34, 9), | |
| # (29, 38, 13), | |
| # (29, 42, 17), | |
| # (29, 46, 21), | |
| # # right hand - 1st finger | |
| # (30, 31, 6), | |
| # (31, 32, 7), | |
| # (32, 33, 8), | |
| # # right hand - 2nd finger | |
| # (34, 35, 10), | |
| # (35, 36, 11), | |
| # (36, 37, 12), | |
| # # right hand - 3rd finger | |
| # (38, 39, 14), | |
| # (39, 40, 15), | |
| # (40, 41, 16), | |
| # # right hand - 4th finger | |
| # (42, 43, 18), | |
| # (43, 44, 19), | |
| # (44, 45, 20), | |
| # # right hand - 5th finger | |
| # (46, 47, 22), | |
| # (47, 48, 23), | |
| # (48, 49, 24), | |
| # ) | |
| # Plot a video given a tensor of joints, a file path, video name and references/sequence ID | |
| def plot_video(joints, | |
| file_path, | |
| video_name, | |
| references=None, | |
| skip_frames=1, | |
| sequence_ID=None): | |
| # Create video template | |
| FPS = (25 // skip_frames) | |
| # ipdb.set_trace() | |
| video_file = file_path + "/{}.mp4".format(sequence_ID.split(".")[0]) | |
| video_path, video_name = os.path.split(video_file) | |
| if not os.path.exists(video_path): | |
| os.mkdir(video_path) | |
| fourcc = cv2.VideoWriter_fourcc(*'mp4v') | |
| if references is None: | |
| video = cv2.VideoWriter(video_file, fourcc, float(FPS), (650, 650), True) | |
| elif references is not None: | |
| video = cv2.VideoWriter(video_file, fourcc, float(FPS), (1300, 650), True) # Long | |
| num_frames = 0 | |
| for (j, frame_joints) in enumerate(joints): | |
| # Reached padding | |
| if PAD_TOKEN in frame_joints.astype('str').tolist(): | |
| continue | |
| # Initialise frame of white | |
| frame = np.ones((650, 650, 3), np.uint8) * 255 | |
| # Cut off the percent_tok, multiply by 3 to restore joint size | |
| # TODO - Remove the *3 if the joints weren't divided by 3 in data creation | |
| frame_joints = frame_joints[:-1] * 3 | |
| # Reduce the frame joints down to 2D for visualisation - Frame joints 2d shape is (48,2) | |
| frame_joints_2d = np.reshape(frame_joints, (50, 3))[:, :2] | |
| # Draw the frame given 2D joints | |
| draw_frame_2D(frame, frame_joints_2d) | |
| cv2.putText(frame, "Predicted Sign Pose", (180, 600), cv2.FONT_HERSHEY_SIMPLEX, 1, | |
| (0, 0, 0), 2) | |
| # If reference is provided, create and concatenate on the end | |
| if references is not None: | |
| # Extract the reference joints | |
| ref_joints = references[j] | |
| # Initialise frame of white | |
| ref_frame = np.ones((650, 650, 3), np.uint8) * 255 | |
| # Cut off the percent_tok and multiply each joint by 3 (as was reduced in training files) | |
| ref_joints = ref_joints[:-1] * 3 | |
| # Reduce the frame joints down to 2D- Frame joints 2d shape is (48,2) | |
| ref_joints_2d = np.reshape(ref_joints, (50, 3))[:, :2] | |
| # Draw these joints on the frame | |
| draw_frame_2D(ref_frame, ref_joints_2d) | |
| cv2.putText(ref_frame, "Ground Truth Pose", (190, 600), cv2.FONT_HERSHEY_SIMPLEX, 1, | |
| (0, 0, 0), 2) | |
| frame = np.concatenate((frame, ref_frame), axis=1) | |
| sequence_ID_write = "Sequence ID: " + sequence_ID.split("/")[-1] | |
| cv2.putText(frame, sequence_ID_write, (150, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, | |
| (0, 0, 0), 2) | |
| # Write the video frame | |
| video.write(frame) | |
| num_frames += 1 | |
| # Release the video | |
| video.release() | |
| # Draw a line between two points, if they are positive points | |
| def draw_line(im, joint1, joint2, c=(0, 0, 255),t=1, width=3): | |
| thresh = -100 | |
| if joint1[0] > thresh and joint1[1] > thresh and joint2[0] > thresh and joint2[1] > thresh: | |
| center = (int((joint1[0] + joint2[0]) / 2), int((joint1[1] + joint2[1]) / 2)) | |
| length = int(math.sqrt(((joint1[0] - joint2[0]) ** 2) + ((joint1[1] - joint2[1]) ** 2))/2) | |
| angle = math.degrees(math.atan2((joint1[0] - joint2[0]),(joint1[1] - joint2[1]))) | |
| cv2.ellipse(im, center, (width,length), -angle,0.0,360.0, c, -1) | |
| # Draw the frame given 2D joints that are in the Inverse Kinematics format | |
| def draw_frame_2D(frame, joints): | |
| # Line to be between the stacked | |
| draw_line(frame, [1, 650], [1, 1], c=(0,0,0), t=1, width=1) | |
| # Give an offset to center the skeleton around | |
| offset = [350, 250] | |
| # Get the skeleton structure details of each bone, and size | |
| skeleton = getSkeletalModelStructure() | |
| skeleton = np.array(skeleton) | |
| number = skeleton.shape[0] | |
| # Increase the size and position of the joints | |
| joints = joints * 10 * 12 * 2 | |
| joints = joints + np.ones((50, 2)) * offset | |
| # Loop through each of the bone structures, and plot the bone | |
| for j in range(number): | |
| c = get_bone_colour(skeleton,j) | |
| draw_line(frame, [joints[skeleton[j, 0]][0], joints[skeleton[j, 0]][1]], | |
| [joints[skeleton[j, 1]][0], joints[skeleton[j, 1]][1]], c=c, t=1, width=1) | |
| # get bone colour given index | |
| def get_bone_colour(skeleton,j): | |
| return (0, 0, 0) | |
| # Apply DTW to the produced sequence, so it can be visually compared to the reference sequence | |
| def alter_DTW_timing(pred_seq,ref_seq): | |
| # Define a cost function | |
| euclidean_norm = lambda x, y: np.sum(np.abs(x - y)) | |
| # Cut the reference down to the max count value | |
| _ , ref_max_idx = torch.max(ref_seq[:, -1], 0) | |
| if ref_max_idx == 0: ref_max_idx += 1 | |
| # Cut down frames by counter | |
| ref_seq = ref_seq[:ref_max_idx,:].cpu().numpy() | |
| # Cut the hypothesis down to the max count value | |
| _, hyp_max_idx = torch.max(pred_seq[:, -1], 0) | |
| if hyp_max_idx == 0: hyp_max_idx += 1 | |
| # Cut down frames by counter | |
| pred_seq = pred_seq[:hyp_max_idx,:].cpu().numpy() | |
| #pred_seq = pred_seq[:ref_max_idx, :].cpu().numpy() | |
| # Run DTW on the reference and predicted sequence | |
| d, cost_matrix, acc_cost_matrix, path = dtw(ref_seq[:,:-1], pred_seq[:,:-1], dist=euclidean_norm) | |
| # Normalise the dtw cost by sequence length | |
| d = d / acc_cost_matrix.shape[0] | |
| # Initialise new sequence | |
| new_pred_seq = np.zeros_like(ref_seq) | |
| # j tracks the position in the reference sequence | |
| j = 0 | |
| skips = 0 | |
| squeeze_frames = [] | |
| for (i, pred_num) in enumerate(path[0]): | |
| if i == len(path[0]) - 1: | |
| break | |
| if path[1][i] == path[1][i + 1]: | |
| skips += 1 | |
| # If a double coming up | |
| if path[0][i] == path[0][i + 1]: | |
| squeeze_frames.append(pred_seq[i - skips]) | |
| j += 1 | |
| # Just finished a double | |
| elif path[0][i] == path[0][i - 1]: | |
| new_pred_seq[pred_num] = avg_frames(squeeze_frames) | |
| squeeze_frames = [] | |
| else: | |
| new_pred_seq[pred_num] = pred_seq[i - skips] | |
| return new_pred_seq, ref_seq, d | |
| # Find the average of the given frames | |
| def avg_frames(frames): | |
| frames_sum = np.zeros_like(frames[0]) | |
| for frame in frames: | |
| frames_sum += frame | |
| avg_frame = frames_sum / len(frames) | |
| return avg_frame |