generate_NTU60_occ

# THis file should be used under [CrosSCLR](https://github.com/LinguoLi/CrosSCLR).
# This file is used to generate NTU-60 data from occluded skeleton files, 
# including data, data_nan, missing_joints and missing_joints_distribution.
# After generation, similar preprocessing is used based on CrosSCLR.
import os
import sys
import pickle
import argparse
import numpy as np
from numpy.lib.format import open_memmap

from utils.ntu_read_skeleton import read_xyz_all

training_subjects = [
    1, 2, 4, 5, 8, 9, 13, 14, 15, 16, 17, 18, 19, 25, 27, 28, 31, 34, 35, 38
]
training_cameras = [2, 3]
max_body = 2
num_joint = 25
max_frame = 300
toolbar_width = 30

def print_toolbar(rate, annotation=''):
    # setup toolbar
    sys.stdout.write("{}[".format(annotation))
    for i in range(toolbar_width):
        if i * 1.0 / toolbar_width > rate:
            sys.stdout.write(' ')
        else:
            sys.stdout.write('-')
        sys.stdout.flush()
    sys.stdout.write(']\r')


def end_toolbar():
    sys.stdout.write("\n")


def gendata(data_path,
            out_path,
            ignored_sample_path=None,
            benchmark='xview',
            part='eval'):
    if ignored_sample_path != None:
        with open(ignored_sample_path, 'r') as f:
            ignored_samples = [
                line.strip() + '.skeleton' for line in f.readlines()
            ]
    else:
        ignored_samples = []
    sample_name = []
    sample_label = []
    missing_joints = []
    missing_joints_distribution = []
    dir_list = os.listdir(data_path)
    dir_list.sort()
    for filename in dir_list:
        if filename in ignored_samples:
            continue
        action_class = int(
            filename[filename.find('A') + 1:filename.find('A') + 4])
        subject_id = int(
            filename[filename.find('P') + 1:filename.find('P') + 4])
        camera_id = int(
            filename[filename.find('C') + 1:filename.find('C') + 4])

        if benchmark == 'xview':
            istraining = (camera_id in training_cameras)
        elif benchmark == 'xsub':
            istraining = (subject_id in training_subjects)
        else:
            raise ValueError()

        if part == 'train':
            issample = istraining
        elif part == 'val':
            issample = not (istraining)
        else:
            raise ValueError()

        if issample:
            sample_name.append(filename)
            sample_label.append(action_class - 1)

    with open('{}/{}_label.pkl'.format(out_path, part), 'wb') as f:
        pickle.dump((sample_name, list(sample_label)), f)

    f_data = open_memmap(
        '{}/{}_data.npy'.format(out_path, part),
        dtype='float32',
        mode='w+',
        shape=(len(sample_label), 3, max_frame, num_joint, max_body))

    f_data_nan = open_memmap(
        '{}/{}_data_nan.npy'.format(out_path, part),
        dtype='float32',
        mode='w+',
        shape=(len(sample_label), 3, max_frame, num_joint, max_body))
    
    for i, s in enumerate(sample_name):
        print_toolbar(i * 1.0 / len(sample_label),
                      '({:>5}/{:<5}) Processing {:>5}-{:<5} data: '.format(
                          i + 1, len(sample_name), benchmark, part))
        data,data_nan,missing,missing_distribution = read_xyz_all(
            os.path.join(data_path, s), max_body=max_body, num_joint=num_joint)
        f_data[i, :, 0:data.shape[1], :, :] = data
        f_data_nan[i, :, 0:data.shape[1], :, :] = data_nan
        missing_joints.append(missing)
    missing_joints_distribution = np.array(missing_distribution)
    with open('{}/{}_missing_joints_distribution.pkl'.format(out_path, part), 'wb') as f:
        pickle.dump(missing_joints_distribution, f)
    with open('{}/{}_missing_joints.pkl'.format(out_path, part), 'wb') as f:
        pickle.dump(missing_joints, f)
    end_toolbar()

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='NTU-RGB-D Data Converter.')
    parser.add_argument(
        '--data_path', default='/cvhci/temp/ychen2/data_occ/NTU60_occ_skeleton')
    parser.add_argument(
        '--ignored_sample_path',
        default='resource/NTU-RGB-D/NTU_RGBD60_samples_with_missing_skeletons.txt')
    parser.add_argument('--out_folder', default='/cvhci/temp/ychen2/data_occ/NTU-RGB-D-60-occ')

    benchmark = ['xsub', 'xview']
    part = ['train', 'val']
    arg = parser.parse_args()

    for b in benchmark:
        for p in part:
            out_path = os.path.join(arg.out_folder, b)
            if not os.path.exists(out_path):
                os.makedirs(out_path)
            gendata(
                arg.data_path,
                out_path,
                arg.ignored_sample_path,
                benchmark=b,
                part=p)