SignIDD_CodeFiles/plot_videos.py

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