Plot Ev2Hands-R data

plot_data.py

# Adapted from `demo.py`

import sys; sys.path.append('..') # add settings.
import os
import torch
import cv2
import time
import pyrender
import numpy as np
import trimesh

import arg_parser

from torch.utils.data import DataLoader
from dataset import Ev2HandRDataset
from model import TEHNetWrapper
from settings import OUTPUT_HEIGHT, OUTPUT_WIDTH, MAIN_CAMERA, REAL_TEST_DATA_PATH

import plotly.graph_objects as go
import cv2 as cv
import tqdm
from pathlib import Path

# Ref: https://github.com/3d-hand-shape/hand-graph-cnn/blob/2f59da3e7d9d78b8de0549b9600df6f44ca90fa0/hand_shape_pose/util/vis.py#L62
color_hand_joints = [[1.0, 0.0, 0.0],
                     [0.0, 0.4, 0.0], [0.0, 0.6, 0.0], [0.0, 0.8, 0.0], [0.0, 1.0, 0.0],  # thumb
                     [0.0, 0.0, 0.6], [0.0, 0.0, 1.0], [0.2, 0.2, 1.0], [0.4, 0.4, 1.0],  # index
                     [0.0, 0.4, 0.4], [0.0, 0.6, 0.6], [0.0, 0.8, 0.8], [0.0, 1.0, 1.0],  # middle
                     [0.4, 0.4, 0.0], [0.6, 0.6, 0.0], [0.8, 0.8, 0.0], [1.0, 1.0, 0.0],  # ring
                     [0.4, 0.0, 0.4], [0.6, 0.0, 0.6], [0.8, 0.0, 0.8], [1.0, 0.0, 1.0]]  # little
def draw_2d_skeleton(image, pose_uv):
    """
    :param image: H x W x 3
    :param pose_uv: 21 x 2
    wrist,
    thumb_mcp, thumb_pip, thumb_dip, thumb_tip
    index_mcp, index_pip, index_dip, index_tip,
    middle_mcp, middle_pip, middle_dip, middle_tip,
    ring_mcp, ring_pip, ring_dip, ring_tip,
    little_mcp, little_pip, little_dip, little_tip
    :return:
    """
    assert pose_uv.shape[0] == 21
    skeleton_overlay = np.copy(image)

    marker_sz = 6-4
    line_wd = 3-2
    root_ind = 0

    for joint_ind in range(pose_uv.shape[0]):
        joint = pose_uv[joint_ind, 0].astype('int32'), pose_uv[joint_ind, 1].astype('int32')
        cv.circle(
            skeleton_overlay, joint,
            radius=marker_sz, color=color_hand_joints[joint_ind] * np.array(255), thickness=-1,
            lineType=cv.CV_AA if cv.__version__.startswith('2') else cv.LINE_AA)

        if joint_ind == 0:
            continue
        elif joint_ind % 4 == 1:
            root_joint = pose_uv[root_ind, 0].astype('int32'), pose_uv[root_ind, 1].astype('int32')
            cv.line(
                skeleton_overlay, root_joint, joint,
                color=color_hand_joints[joint_ind] * np.array(255), thickness=int(line_wd),
                lineType=cv.CV_AA if cv.__version__.startswith('2') else cv.LINE_AA)
        else:
            joint_2 = pose_uv[joint_ind - 1, 0].astype('int32'), pose_uv[joint_ind - 1, 1].astype('int32')
            cv.line(
                skeleton_overlay, joint_2, joint,
                color=color_hand_joints[joint_ind] * np.array(255), thickness=int(line_wd),
                lineType=cv.CV_AA if cv.__version__.startswith('2') else cv.LINE_AA)

    return skeleton_overlay


def main():    
    arg_parser.demo()
    os.makedirs('outputs', exist_ok=True)

    batch_size = int(os.environ['BATCH_SIZE'])
    video_fps = 25

    data = REAL_TEST_DATA_PATH
    subject_idx = 1
    validation_dataset = Ev2HandRDataset(f"{data}/subject_{subject_idx}_event.pickle", demo=True, augment=False)
    validation_loader = DataLoader(validation_dataset, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)

    out_dir = Path(f"outputs/data_plots/subject_{subject_idx}")
    out_dir.mkdir(parents=True, exist_ok=True)
    video = cv2.VideoWriter(f'{out_dir}/video.mp4', cv2.VideoWriter_fourcc(*'avc1'), video_fps, (3 * OUTPUT_WIDTH, OUTPUT_HEIGHT))
    out_img_idx = 0
    n_batches = 10
    pbar = tqdm.tqdm(total=n_batches)
    for bdx, batch in enumerate(validation_loader):
        for idx in range(len(batch)):
            event_frame = batch['event_frame'][idx].cpu().numpy().astype(dtype=np.uint8)
            
            frame_index = batch['frame_index'][idx]
            input_rgb = cv.imread(f"./datasets/RGB_Data/subject_{subject_idx}_rgb/{frame_index:07d}.png")
            
            skeleton_overlay = event_frame.copy()
            skeleton_overlay = draw_2d_skeleton(skeleton_overlay, batch['left']['j2d'][idx].cpu().numpy().astype(int))
            skeleton_overlay = draw_2d_skeleton(skeleton_overlay, batch['right']['j2d'][idx].cpu().numpy().astype(int))

            input_rgb_resized = cv.resize(input_rgb, (event_frame.shape[1], event_frame.shape[0]))  # input_rgb.shape = (1080, 1920, 3); input_rgb_resized.shape = (260, 346, 3)
            img_stack = np.hstack([input_rgb_resized, event_frame, skeleton_overlay])             
            
            cv.imwrite(f"{out_dir}/{out_img_idx:07d}.png", img_stack)
            out_img_idx += 1
            video.write(img_stack)

        if bdx >= n_batches:
            break

        pbar.update(1)
    pbar.close()

    
    video.release()


if __name__ == '__main__':
    main()