NTU 60 and NTU 120¶

import glob
import tqdm
import numpy
import matplotlib.pyplot as plt
import torch
import pathlib
import numpy as np
import os
import pickle
import joblib


# ---------------------------------------------------------
# 0. Adapted from st-gcn repo
# ---------------------------------------------------------
def read_skeleton(file):
    with open(file, 'r') as f:
        skeleton_sequence = {}
        skeleton_sequence['numFrame'] = int(f.readline())
        skeleton_sequence['frameInfo'] = []
        for t in range(skeleton_sequence['numFrame']):
            frame_info = {}
            frame_info['numBody'] = int(f.readline())
            frame_info['bodyInfo'] = []
            for m in range(frame_info['numBody']):
                body_info = {}
                body_info_key = [
                    'bodyID', 'clipedEdges', 'handLeftConfidence',
                    'handLeftState', 'handRightConfidence', 'handRightState',
                    'isResticted', 'leanX', 'leanY', 'trackingState'
                ]
                body_info = {
                    k: float(v)
                    for k, v in zip(body_info_key, f.readline().split())
                }
                body_info['numJoint'] = int(f.readline())
                body_info['jointInfo'] = []
                for v in range(body_info['numJoint']):
                    joint_info_key = [
                        'x', 'y', 'z', 'depthX', 'depthY', 'colorX', 'colorY',
                        'orientationW', 'orientationX', 'orientationY',
                        'orientationZ', 'trackingState'
                    ]
                    joint_info = {
                        k: float(v)
                        for k, v in zip(joint_info_key, f.readline().split())
                    }
                    body_info['jointInfo'].append(joint_info)
                frame_info['bodyInfo'].append(body_info)
            skeleton_sequence['frameInfo'].append(frame_info)
    return skeleton_sequence


def read_xyz(file, max_body=2, num_joint=25):
    seq_info = read_skeleton(file)
    data = np.zeros((3, seq_info['numFrame'], num_joint, max_body))
    for n, f in enumerate(seq_info['frameInfo']):
        for m, b in enumerate(f['bodyInfo']):
            for j, v in enumerate(b['jointInfo']):
                if m < max_body and j < num_joint:
                    data[:, n, j, m] = [v['x'], v['y'], v['z']]
                else:
                    pass
    data = np.moveaxis(data, [0, 1, 2, 3], [3, 0, 2, 1])
    # duration, body, joint, component
    return data


# ---------------------------------------------------------
# 1. Txt files to individual original numpy files
# ---------------------------------------------------------
def generate_np_files():
    assert pathlib.Path('./nturgb+d_skeletons/').exists()
    assert pathlib.Path('./ignored_sk.txt').exists()

    all_sk = glob.glob('nturgb+d_skeletons/*.skeleton')
    ignored_sk = ['nturgb+d_skeletons/' + r.rstrip('\n') + '.skeleton' for r in open('ignored_sk.txt', 'r').readlines()]
    all_sk = [ski for ski in all_sk if ski not in ignored_sk]

    pathlib.Path('./nturgb+d_skeletons_np_raw').mkdir(parents=True, exist_ok=True)

    # read first shape
    s0 = read_xyz(all_sk[0]).shape

    # save as numpy arrays for quicker access
    for idx, skf in enumerate(tqdm.tqdm(all_sk)):
        data = read_xyz(skf)
        # except duration, all the shapes dims should match
        assert data.shape[1:] == s0[1:]
        np.save(skf.replace('nturgb+d_skeletons/', 'nturgb+d_skeletons_np_raw/').replace('.skeleton', ''), data)


# ---------------------------------------------------------
# 2. Resize individual numpy files to the same length and save as one dataset
# ---------------------------------------------------------
def generate_resized_dataset_file(final_length=100):
    npy = sorted(glob.glob('nturgb+d_skeletons_np_raw/*'))

    pathlib.Path('./nturgb+d_skeletons_np_{}'.format(final_length)).mkdir(parents=True, exist_ok=True)

    all_tensors = []
    all_scpra = []
    all_meta = []

    for n in tqdm.tqdm(npy):
        d = numpy.load(n)
        L0 = d.shape[0]
        nn = n.split('/')[-1].split('.npy')[0]
        S, C, P, R, A, filename = int(nn[1:4]), int(nn[5:8]), int(nn[9:12]), int(nn[13:16]), int(nn[17:20]), nn
        dr = torch.nn.functional.interpolate(input=torch.from_numpy(d).reshape(-1, 2 * 25 * 3).transpose(0, 1).unsqueeze(0),
                                             size=final_length).squeeze(0).transpose(1, 0).reshape(-1, 2, 25, 3).numpy()
        numpy.save(n.replace('nturgb+d_skeletons_np_raw', 'nturgb+d_skeletons_np_{}'.format(final_length)), dr)
        all_tensors += [dr]
        all_scpra += [numpy.array([S, C, P, R, A])]
        all_meta += [(filename, 'original_len_{}'.format(L0), 'resized_len_{}'.format(final_length))]

    all_scpra = numpy.array(all_scpra)
    all_tensors = numpy.array(all_tensors)

    data = {
        'x': all_tensors,
        'scpra': all_scpra,
        'meta': all_meta
    }
    return data


# ---------------------------------------------------------
# 3. I/O
# ---------------------------------------------------------
def write_data(data, filepath):
    """Save the dataset to a file. Note: data is a dict with keys 'x', ..."""
    with open(filepath, 'wb') as output_file:
        joblib.dump(data, output_file)


def load_ntu_data(filepath='./nturgbd__len100__data.pckl'):
    """
    Returns gesture sequences (X) and their associated labels (SCPRA) and meta information (meta).
    X.shape = (56578, 100, 2, 25, 3) = (batch, duration, body, joint, joint_component)
    SCPRA.shape = (56578, 5) = (batch, labels_s_c_p_r_and_a)
    len(Y) = 56578, with each line in Y being: (original_filename, original_len, resized_len=100)
    """
    file = open(filepath, 'rb')
    data = joblib.load(file)
    file.close()
    return data['x'], data['scpra'], data['meta']


# ---------------------------------------------------------
# 4. main
# ---------------------------------------------------------
if __name__ == '__main__':
    final_length = 100
    filename = 'nturgbd__len{}__data.pckl'.format(final_length)

    # 1. Generate individual numpy files
    #generate_np_files()

    # 2. Resize sequences and generate one dataset file
    #data = generate_resized_dataset_file(final_length=final_length)
    #write_data(data, filename)

    # 3. Check everything went as we wanted
    x, scpra, meta = load_ntu_data(filename)
    y = scpra[:, -1]
    print('NTU (xyz version) dataset shape:', x.shape, scpra.shape)
    print('Y labels values are between {} (included) and {} (included).'.format(int(y.min()), int(y.max())))

    sanity_check = True
    idx_sanity_check = 2222

    sanity_check_filename = meta[idx_sanity_check][0]
    x_res = x[idx_sanity_check]
    x_nonres = read_xyz('./nturgb+d_skeletons/' + sanity_check_filename + '.skeleton')
    print('Shape of same sequence before and after resize:', x_nonres.shape, '--->', x_res.shape)

    if sanity_check:
        # shapes: duration, body, joint, joint_component
        x_res_viz_alternated = x_res[:, 0, :, :].reshape(x_res.shape[0], -1)
        x_nonres_viz_alternated = x_nonres[:, 0, :, :].reshape(x_nonres.shape[0], -1)
        x_res_viz_block = numpy.hstack([x_res_viz_alternated[:, ::3], x_res_viz_alternated[:, 1::3], x_res_viz_alternated[:, 2::3]])
        x_nonres_viz_block = numpy.hstack([x_nonres_viz_alternated[:, ::3], x_nonres_viz_alternated[:, 1::3], x_nonres_viz_alternated[:, 2::3]])
        fig, (ax1, ax2) = plt.subplots(1, 2)
        fig.suptitle('Same sequence resized or not (viz: alternated)')
        ax1.imshow(x_nonres_viz_alternated)
        ax1.set_title('Original')
        ax2.imshow(x_res_viz_alternated)
        ax2.set_title('Resized')
        plt.show()
        fig, (ax1, ax2) = plt.subplots(1, 2)
        fig.suptitle('Same sequence resized or not (viz: block)')
        ax1.imshow(x_nonres_viz_block)
        ax1.set_title('Original')
        ax2.imshow(x_res_viz_block)
        ax2.set_title('Resized')
        plt.show()

# TODO

import joblib

def load_ntu_data(filepath='./nturgbd__len100__data.pckl'):
    """
    Returns gesture sequences (X) and their associated labels (SCPRA) and meta information (meta).
    X.shape = (56578, 100, 2, 25, 3) = (batch, duration, body, joint, joint_component)
    SCPRA.shape = (56578, 5) = (batch, labels_s_c_p_r_and_a)
    len(Y) = 56578, with each line in Y being: (original_filename, original_len, resized_len=100)
    """
    file = open(filepath, 'rb')
    data = joblib.load(file)
    file.close()
    return data['x'], data['scpra'], data['meta']


filename = './nturgbd__len100__data.pckl'
x, scpra, meta = load_ntu_data(filename)
y = scpra[:, -1]

# TODO

import itertools

# NTU 120
ntu_120_daily_actions = {
    1: 'Drink water',
    2: 'Eat meal',
    3: 'Brush teeth',
    4: 'Brush hair',
    5: 'Drop',
    6: 'Pick up',
    7: 'Throw',
    8: 'Sit down',
    9: 'Stand up',
    10: 'Clapping',
    11: 'Reading',
    12: 'Writing',
    13: 'Tear up paper',
    14: 'Put on jacket',
    15: 'Take off jacket',
    16: 'Put on a shoe',
    17: 'Take off a shoe',
    18: 'Put on glasses',
    19: 'Take off glasses',
    20: 'Put on a hat/cap',
    21: 'Take off a hat/cap',
    22: 'Cheer up',
    23: 'Hand waving',
    24: 'Kicking something',
    25: 'Reach into pocket',
    26: 'Hopping',
    27: 'Jump up',
    28: 'Phone call',
    29: 'Play with phone/tablet',
    30: 'Type on a keyboard',
    31: 'Point to something',
    32: 'Taking a selfie',
    33: 'Check time (from watch',
    34: 'Rub two hands',
    35: 'Nod head/bow',
    36: 'Shake head',
    37: 'Wipe face',
    38: 'Salute',
    39: 'Put palms together',
    40: 'Cross hands in front',
    61: 'Put on headphone',
    62: 'Take off headphone',
    63: 'Shoot at basket',
    64: 'Bounce ball',
    65: 'Tennis bat swing',
    66: 'Juggle table tennis ball',
    67: 'Hush',
    68: 'Flick hair',
    69: 'Thumb up',
    70: 'Thumb down',
    71: 'Make OK sign',
    72: 'Make victory sign',
    73: 'Staple book',
    74: 'Counting money',
    75: 'Cutting nails',
    76: 'Cutting paper',
    77: 'Snap fingers',
    78: 'Open bottle',
    79: 'Sniff/smell',
    80: 'Squat down',
    81: 'Toss a coin',
    82: 'Fold paper',
    83: 'Ball up paper',
    84: 'Play magic cube',
    85: 'Apply cream on face',
    86: 'Apply cream on hand',
    87: 'Put on bag',
    88: 'Take off bag',
    89: 'Put object into bag',
    90: 'Take object out of bag',
    91: 'Open a box',
    92: 'Move heavy objects',
    93: 'Shake fist',
    94: 'Throw up cap/hat',
    95: 'Capitulate',
    96: 'Cross arms',
    97: 'Arm circles',
    98: 'Arm swings',
    99: 'Run on the spot',
    100: 'Butt kicks',
    101: 'Cross toe touch',
    102: 'Side kick'
}
ntu_120_medical_actions = {
    41: 'Sneeze/cough',
    42: 'Staggering',
    43: 'Falling down',
    44: 'Headache',
    45: 'Chest pain',
    46: 'Back pain',
    47: 'Neck pain',
    48: 'Nausea/vomiting',
    49: 'Fan self',
    103: 'Yawn',
    104: 'Stretch oneself',
    105: 'Blow nose'
}
ntu_120_two_person_actions = {
    50: 'Punch/slap',
    51: 'Kicking',
    52: 'Pushing',
    53: 'Pat on back',
    54: 'Point finger',
    55: 'Hugging',
    56: 'Giving object',
    57: 'Touch pocket',
    58: 'Shaking hands',
    59: 'Walking towards',
    60: 'Walking apart',
    106: 'Hit with object',
    107: 'Wield knife',
    108: 'Knock over',
    109: 'Grab stuff',
    110: 'Shoot with gun',
    111: 'Step on foot',
    112: 'High-five',
    113: 'Cheers and drink',
    114: 'Carry object',
    115: 'Take a photo',
    116: 'Follow',
    117: 'Whisper',
    118: 'Exchange things',
    119: 'Support somebody',
    120: 'Rock-paper-scissors',
}
ntu_120_actions = dict(itertools.chain.from_iterable(d.items() for d in (ntu_120_daily_actions, ntu_120_medical_actions, ntu_120_two_person_actions)))

# NTU 60
ntu_60_daily_actions = {k:ntu_120_daily_actions[k] for k in ntu_120_daily_actions if k<=60}
ntu_60_medical_actions = {k:ntu_120_medical_actions[k] for k in ntu_120_medical_actions if k<=60}
ntu_60_two_person_actions = {k:ntu_120_two_person_actions[k] for k in ntu_120_two_person_actions if k<=60}
ntu_60_actions = {k:ntu_120_actions[k] for k in ntu_120_actions if k<=60}

# TODO