text: # Copyright (c) 2018-present, Facebook, Inc.
  # All rights reserved.
  #
  # This source code is licensed under the license found in the
  # LICENSE file in the root directory of this source tree.
  #
  
  import argparse
  import os
  import re
  import sys
  from glob import glob
  
  import numpy as np
  from data_utils import suggest_metadata, suggest_pose_importer
  
  sys.path.append('../')
  from itertools import groupby
  
  subjects = ['Train/S1', 'Train/S2', 'Train/S3', 'Validate/S1', 'Validate/S2', 'Validate/S3']
  
  cam_map = {
      'C1': 0,
      'C2': 1,
      'C3': 2,
  }
  
  # Frame numbers for train/test split
  # format: [start_frame, end_frame[ (inclusive, exclusive)
  index = {
      'Train/S1': {
          'Walking 1': (590, 1203),
          'Jog 1': (367, 740),
          'ThrowCatch 1': (473, 945),
          'Gestures 1': (395, 801),
          'Box 1': (385, 789),
      },
      'Train/S2': {
          'Walking 1': (438, 876),
          'Jog 1': (398, 795),
          'ThrowCatch 1': (550, 1128),
          'Gestures 1': (500, 901),
          'Box 1': (382, 734),
      },
      'Train/S3': {
          'Walking 1': (448, 939),
          'Jog 1': (401, 842),
          'ThrowCatch 1': (493, 1027),
          'Gestures 1': (533, 1102),
          'Box 1': (512, 1021),
      },
      'Validate/S1': {
          'Walking 1': (5, 590),
          'Jog 1': (5, 367),
          'ThrowCatch 1': (5, 473),
          'Gestures 1': (5, 395),
          'Box 1': (5, 385),
      },
      'Validate/S2': {
          'Walking 1': (5, 438),
          'Jog 1': (5, 398),
          'ThrowCatch 1': (5, 550),
          'Gestures 1': (5, 500),
          'Box 1': (5, 382),
      },
      'Validate/S3': {
          'Walking 1': (5, 448),
          'Jog 1': (5, 401),
          'ThrowCatch 1': (5, 493),
          'Gestures 1': (5, 533),
          'Box 1': (5, 512),
      },
  }
  
  # Frames to skip for each video (synchronization)
  sync_data = {
      'S1': {
          'Walking 1': (82, 81, 82),
          'Jog 1': (51, 51, 50),
          'ThrowCatch 1': (61, 61, 60),
          'Gestures 1': (45, 45, 44),
          'Box 1': (57, 57, 56),
      },
      'S2': {
          'Walking 1': (115, 115, 114),
          'Jog 1': (100, 100, 99),
          'ThrowCatch 1': (127, 127, 127),
          'Gestures 1': (122, 122, 121),
          'Box 1': (119, 119, 117),
      },
      'S3': {
          'Walking 1': (80, 80, 80),
          'Jog 1': (65, 65, 65),
          'ThrowCatch 1': (79, 79, 79),
          'Gestures 1': (83, 83, 82),
          'Box 1': (1, 1, 1),
      },
      'S4': {}
  }
  
  if __name__ == '__main__':
      if os.path.basename(os.getcwd()) != 'data':
          print('This script must be launched from the "data" directory')
          exit(0)
  
      parser = argparse.ArgumentParser(description='HumanEva dataset converter')
  
      parser.add_argument('-p', '--path', default='', type=str, metavar='PATH', help='path to the processed HumanEva dataset')
      parser.add_argument('--convert-3d', action='store_true', help='convert 3D mocap data')
      parser.add_argument('--convert-2d', default='', type=str, metavar='PATH', help='convert user-supplied 2D detections')
      parser.add_argument('-o', '--output', default='', type=str, metavar='PATH', help='output suffix for 2D detections (e.g. detectron_pt_coco)')
  
      args = parser.parse_args()
  
      if not args.convert_2d and not args.convert_3d:
          print('Please specify one conversion mode')
          exit(0)
  
      if args.path:
          print('Parsing HumanEva dataset from', args.path)
          output = {}
          output_2d = {}
          frame_mapping = {}
  
          from scipy.io import loadmat
  
          num_joints = None
  
          for subject in subjects:
              output[subject] = {}
              output_2d[subject] = {}
              split, subject_name = subject.split('/')
              if subject_name not in frame_mapping:
                  frame_mapping[subject_name] = {}
  
              file_list = glob(args.path + '/' + subject + '/*.mat')
              for f in file_list:
                  action = os.path.splitext(os.path.basename(f))[0]
  
                  # Use consistent naming convention
                  canonical_name = action.replace('_', ' ')
  
                  hf = loadmat(f)
                  positions = hf['poses_3d']
                  positions_2d = hf['poses_2d'].transpose(1, 0, 2, 3)  # Ground-truth 2D poses
                  assert positions.shape[0] == positions_2d.shape[0] and positions.shape[1] == positions_2d.shape[2]
                  assert num_joints is None or num_joints == positions.shape[1], "Joint number inconsistency among files"
                  num_joints = positions.shape[1]
  
                  # Sanity check for the sequence length
                  assert positions.shape[0] == index[subject][canonical_name][1] - index[subject][canonical_name][0]
  
                  # Split corrupted motion capture streams into contiguous chunks
                  # e.g. 012XX567X9 is split into "012", "567", and "9".
                  all_chunks = [list(v) for k, v in groupby(positions, lambda x: np.isfinite(x).all())]
                  all_chunks_2d = [list(v) for k, v in groupby(positions_2d, lambda x: np.isfinite(x).all())]
                  assert len(all_chunks) == len(all_chunks_2d)
                  current_index = index[subject][canonical_name][0]
                  chunk_indices = []
                  for i, chunk in enumerate(all_chunks):
                      next_index = current_index + len(chunk)
                      name = canonical_name + ' chunk' + str(i)
                      if np.isfinite(chunk).all():
                          output[subject][name] = np.array(chunk, dtype='float32') / 1000
                          output_2d[subject][name] = list(np.array(all_chunks_2d[i], dtype='float32').transpose(1, 0, 2, 3))
                      chunk_indices.append((current_index, next_index, np.isfinite(chunk).all(), split, name))
                      current_index = next_index
                  assert current_index == index[subject][canonical_name][1]
                  if canonical_name not in frame_mapping[subject_name]:
                      frame_mapping[subject_name][canonical_name] = []
                  frame_mapping[subject_name][canonical_name] += chunk_indices
  
          metadata = suggest_metadata('humaneva' + str(num_joints))
          output_filename = 'data_3d_' + metadata['layout_name']
          output_prefix_2d = 'data_2d_' + metadata['layout_name'] + '_'
  
          if args.convert_3d:
              print('Saving...')
              np.savez_compressed(output_filename, positions_3d=output)
              np.savez_compressed(output_prefix_2d + 'gt', positions_2d=output_2d, metadata=metadata)
              print('Done.')
  
      else:
          print('Please specify the dataset source')
          exit(0)
  
      if args.convert_2d:
          if not args.output:
              print('Please specify an output suffix (e.g. detectron_pt_coco)')
              exit(0)
  
          import_func = suggest_pose_importer(args.output)
          metadata = suggest_metadata(args.output)
  
          print('Parsing 2D detections from', args.convert_2d)
  
          output = {}
          file_list = glob(args.convert_2d + '/S*/*.avi.npz')
          for f in file_list:
              path, fname = os.path.split(f)
              subject = os.path.basename(path)
              assert subject.startswith('S'), subject + ' does not look like a subject directory'
  
              m = re.search('(.*) \\((.*)\\)', fname.replace('_', ' '))
              action = m.group(1)
              camera = m.group(2)
              camera_idx = cam_map[camera]
  
              keypoints = import_func(f)
              assert keypoints.shape[1] == metadata['num_joints']
  
              if action in sync_data[subject]:
                  sync_offset = sync_data[subject][action][camera_idx] - 1
              else:
                  sync_offset = 0
  
              if subject in frame_mapping and action in frame_mapping[subject]:
                  chunks = frame_mapping[subject][action]
                  for (start_idx, end_idx, labeled, split, name) in chunks:
                      canonical_subject = split + '/' + subject
                      if not labeled:
                          canonical_subject = 'Unlabeled/' + canonical_subject
                      if canonical_subject not in output:
                          output[canonical_subject] = {}
                      kps = keypoints[start_idx + sync_offset:end_idx + sync_offset]
                      assert len(kps) == end_idx - start_idx, "Got len {}, expected {}".format(len(kps), end_idx - start_idx)
  
                      if name not in output[canonical_subject]:
                          output[canonical_subject][name] = [None, None, None]
  
                      output[canonical_subject][name][camera_idx] = kps.astype('float32')
              else:
                  canonical_subject = 'Unlabeled/' + subject
                  if canonical_subject not in output:
                      output[canonical_subject] = {}
                  if action not in output[canonical_subject]:
                      output[canonical_subject][action] = [None, None, None]
                  output[canonical_subject][action][camera_idx] = keypoints.astype('float32')
  
          print('Saving...')
          np.savez_compressed(output_prefix_2d + args.output, positions_2d=output, metadata=metadata)
          print('Done.')
  