text: # Copyright (c) Microsoft Corporation.
  # Licensed under the MIT License.
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import os.path as osp
  import numpy as np
  import pickle
  
  from dataset.joints_dataset import JointsDataset
  
  
  class H36MDataset(JointsDataset):
  
      def __init__(self, cfg, image_set, is_train, transform=None):
          super().__init__(cfg, image_set, is_train, transform)
          self.actual_joints = {
              0: 'root',
              1: 'rhip',
              2: 'rkne',
              3: 'rank',
              4: 'lhip',
              5: 'lkne',
              6: 'lank',
              7: 'belly',
              8: 'neck',
              9: 'nose',
              10: 'head',
              11: 'lsho',
              12: 'lelb',
              13: 'lwri',
              14: 'rsho',
              15: 'relb',
              16: 'rwri'
          }
  
          anno_file = osp.join(self.root, 'h36m', 'annot',
                               'h36m_{}.pkl'.format(image_set))
          self.db = self.load_db(anno_file)
  
          self.u2a_mapping = super().get_mapping()
          super().do_mapping()
  
          self.grouping = self.get_group(self.db)
          self.group_size = len(self.grouping)
  
      def index_to_action_names(self):
          return {
              2: 'Direction',
              3: 'Discuss',
              4: 'Eating',
              5: 'Greet',
              6: 'Phone',
              7: 'Photo',
              8: 'Pose',
              9: 'Purchase',
              10: 'Sitting',
              11: 'SittingDown',
              12: 'Smoke',
              13: 'Wait',
              14: 'WalkDog',
              15: 'Walk',
              16: 'WalkTwo'
          }
  
      def load_db(self, dataset_file):
          with open(dataset_file, 'rb') as f:
              dataset = pickle.load(f)
              return dataset
  
      def get_group(self, db):
          grouping = {}
          nitems = len(db)
          for i in range(nitems):
              keystr = self.get_key_str(db[i])
              camera_id = db[i]['camera_id']
              if keystr not in grouping:
                  grouping[keystr] = [-1, -1, -1, -1]
              grouping[keystr][camera_id] = i
  
          filtered_grouping = []
          for _, v in grouping.items():
              if np.all(np.array(v) != -1):
                  filtered_grouping.append(v)
  
          if not self.is_train:
              filtered_grouping = filtered_grouping[::64]
  
          return filtered_grouping
  
      def __getitem__(self, idx):
          item = np.random.choice(self.grouping[idx])
          return super().__getitem__(item)
  
      def __len__(self):
          return self.group_size
  
      def get_key_str(self, datum):
          return 's_{:02}_act_{:02}_subact_{:02}_imgid_{:06}'.format(
              datum['subject'], datum['action'], datum['subaction'],
              datum['image_id'])
  