text: import os
  from collections import OrderedDict
  
  import json_tricks as json
  import numpy as np
  
  from mmpose.core.evaluation.top_down_eval import (keypoint_auc, keypoint_epe,
                                                    keypoint_pck_accuracy)
  from ...registry import DATASETS
  from .animal_base_dataset import AnimalBaseDataset
  
  
  @DATASETS.register_module()
  class AnimalFlyDataset(AnimalBaseDataset):
      """AnimalFlyDataset for animal pose estimation.
  
      `Fast animal pose estimation using deep neural networks'
      Nature methods'2019. More details can be found in the `paper
      <https://www.biorxiv.org/content/
        biorxiv/early/2018/05/25/331181.full.pdf>`__ .
  
      The dataset loads raw features and apply specified transforms
      to return a dict containing the image tensors and other information.
  
      Vinegar Fly keypoint indexes::
  
          0: "head",
          1: "eyeL",
          2: "eyeR",
          3: "neck",
          4: "thorax",
          5: "abdomen",
          6: "forelegR1",
          7: "forelegR2",
          8: "forelegR3",
          9: "forelegR4",
          10: "midlegR1",
          11: "midlegR2",
          12: "midlegR3",
          13: "midlegR4",
          14: "hindlegR1",
          15: "hindlegR2",
          16: "hindlegR3",
          17: "hindlegR4",
          18: "forelegL1",
          19: "forelegL2",
          20: "forelegL3",
          21: "forelegL4",
          22: "midlegL1",
          23: "midlegL2",
          24: "midlegL3",
          25: "midlegL4",
          26: "hindlegL1",
          27: "hindlegL2",
          28: "hindlegL3",
          29: "hindlegL4",
          30: "wingL",
          31: "wingR"
  
      Args:
          ann_file (str): Path to the annotation file.
          img_prefix (str): Path to a directory where images are held.
              Default: None.
          data_cfg (dict): config
          pipeline (list[dict | callable]): A sequence of data transforms.
          test_mode (bool): Store True when building test or
              validation dataset. Default: False.
      """
  
      def __init__(self,
                   ann_file,
                   img_prefix,
                   data_cfg,
                   pipeline,
                   test_mode=False):
  
          super().__init__(
              ann_file, img_prefix, data_cfg, pipeline, test_mode=test_mode)
  
          self.ann_info['use_different_joint_weights'] = False
          assert self.ann_info['num_joints'] == 32
          self.ann_info['joint_weights'] = \
              np.ones((self.ann_info['num_joints'], 1), dtype=np.float32)
  
          self.ann_info['flip_pairs'] = [[1, 2], [6, 18], [7, 19], [8, 20],
                                         [9, 21], [10, 22], [11, 23], [12, 24],
                                         [13, 25], [14, 26], [15, 27], [16, 28],
                                         [17, 29], [30, 31]]
  
          self.dataset_name = 'fly'
          self.db = self._get_db()
  
          print(f'=> num_images: {self.num_images}')
          print(f'=> load {len(self.db)} samples')
  
      def _get_db(self):
          """Load dataset."""
          gt_db = []
          bbox_id = 0
          num_joints = self.ann_info['num_joints']
          for img_id in self.img_ids:
  
              ann_ids = self.coco.getAnnIds(imgIds=img_id, iscrowd=False)
              objs = self.coco.loadAnns(ann_ids)
  
              for obj in objs:
                  if max(obj['keypoints']) == 0:
                      continue
                  joints_3d = np.zeros((num_joints, 3), dtype=np.float32)
                  joints_3d_visible = np.zeros((num_joints, 3), dtype=np.float32)
  
                  keypoints = np.array(obj['keypoints']).reshape(-1, 3)
                  joints_3d[:, :2] = keypoints[:, :2]
                  joints_3d_visible[:, :2] = np.minimum(1, keypoints[:, 2:3])
  
                  # the ori image is 192x192
                  center, scale = self._xywh2cs(0, 0, 192, 192, 0.8)
  
                  image_file = os.path.join(self.img_prefix,
                                            self.id2name[img_id])
  
                  gt_db.append({
                      'image_file': image_file,
                      'center': center,
                      'scale': scale,
                      'rotation': 0,
                      'joints_3d': joints_3d,
                      'joints_3d_visible': joints_3d_visible,
                      'dataset': self.dataset_name,
                      'bbox': obj['bbox'],
                      'bbox_score': 1,
                      'bbox_id': bbox_id
                  })
                  bbox_id = bbox_id + 1
          gt_db = sorted(gt_db, key=lambda x: x['bbox_id'])
  
          return gt_db
  
      def _report_metric(self, res_file, metrics, pck_thr=0.2, auc_nor=30):
          """Keypoint evaluation.
  
          Args:
              res_file (str): Json file stored prediction results.
              metrics (str | list[str]): Metric to be performed.
                  Options: 'PCK', 'PCKh', 'AUC', 'EPE'.
              pck_thr (float): PCK threshold, default as 0.2.
              pckh_thr (float): PCKh threshold, default as 0.7.
              auc_nor (float): AUC normalization factor, default as 30 pixel.
  
          Returns:
              List: Evaluation results for evaluation metric.
          """
          info_str = []
  
          with open(res_file, 'r') as fin:
              preds = json.load(fin)
          assert len(preds) == len(self.db)
  
          outputs = []
          gts = []
          masks = []
          threshold_bbox = []
  
          for pred, item in zip(preds, self.db):
              outputs.append(np.array(pred['keypoints'])[:, :-1])
              gts.append(np.array(item['joints_3d'])[:, :-1])
              masks.append((np.array(item['joints_3d_visible'])[:, 0]) > 0)
              if 'PCK' in metrics:
                  bbox = np.array(item['bbox'])
                  bbox_thr = np.max(bbox[2:])
                  threshold_bbox.append(np.array([bbox_thr, bbox_thr]))
  
          outputs = np.array(outputs)
          gts = np.array(gts)
          masks = np.array(masks)
          threshold_bbox = np.array(threshold_bbox)
  
          if 'PCK' in metrics:
              _, pck, _ = keypoint_pck_accuracy(outputs, gts, masks, pck_thr,
                                                threshold_bbox)
              info_str.append(('PCK', pck))
  
          if 'AUC' in metrics:
              info_str.append(('AUC', keypoint_auc(outputs, gts, masks,
                                                   auc_nor)))
  
          if 'EPE' in metrics:
              info_str.append(('EPE', keypoint_epe(outputs, gts, masks)))
  
          return info_str
  
      def evaluate(self, outputs, res_folder, metric='PCK', **kwargs):
          """Evaluate Fly keypoint results. The pose prediction results will be
          saved in `${res_folder}/result_keypoints.json`.
  
          Note:
              batch_size: N
              num_keypoints: K
              heatmap height: H
              heatmap width: W
  
          Args:
              outputs (list(preds, boxes, image_path, output_heatmap))
                  :preds (np.ndarray[N,K,3]): The first two dimensions are
                      coordinates, score is the third dimension of the array.
                  :boxes (np.ndarray[N,6]): [center[0], center[1], scale[0]
                      , scale[1],area, score]
                  :image_paths (list[str]): For example, ['Test/source/0.jpg']
                  :output_heatmap (np.ndarray[N, K, H, W]): model outpus.
  
              res_folder (str): Path of directory to save the results.
              metric (str | list[str]): Metric to be performed.
                  Options: 'PCK', 'AUC', 'EPE'.
  
          Returns:
              dict: Evaluation results for evaluation metric.
          """
          metrics = metric if isinstance(metric, list) else [metric]
          allowed_metrics = ['PCK', 'AUC', 'EPE']
          for metric in metrics:
              if metric not in allowed_metrics:
                  raise KeyError(f'metric {metric} is not supported')
  
          res_file = os.path.join(res_folder, 'result_keypoints.json')
  
          kpts = []
          for output in outputs:
              preds = output['preds']
              boxes = output['boxes']
              image_paths = output['image_paths']
              bbox_ids = output['bbox_ids']
  
              batch_size = len(image_paths)
              for i in range(batch_size):
                  image_id = self.name2id[image_paths[i][len(self.img_prefix):]]
  
                  kpts.append({
                      'keypoints': preds[i].tolist(),
                      'center': boxes[i][0:2].tolist(),
                      'scale': boxes[i][2:4].tolist(),
                      'area': float(boxes[i][4]),
                      'score': float(boxes[i][5]),
                      'image_id': image_id,
                      'bbox_id': bbox_ids[i]
                  })
          kpts = self._sort_and_unique_bboxes(kpts)
  
          self._write_keypoint_results(kpts, res_file)
          info_str = self._report_metric(res_file, metrics)
          name_value = OrderedDict(info_str)
  
          return name_value
  