mmpose.datasets.datasets.animal.animal_locust_dataset 源代码

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 ...builder import DATASETS
from .animal_base_dataset import AnimalBaseDataset


[文档]@DATASETS.register_module() class AnimalLocustDataset(AnimalBaseDataset): """AnimalLocustDataset for animal pose estimation. `DeepPoseKit, a software toolkit for fast and robust animal pose estimation using deep learning' Elife'2019. More details can be found in the `paper. The dataset loads raw features and apply specified transforms to return a dict containing the image tensors and other information. Desert Locust keypoint indexes:: 0: "head", 1: "neck", 2: "thorax", 3: "abdomen1", 4: "abdomen2", 5: "anttipL", 6: "antbaseL", 7: "eyeL", 8: "forelegL1", 9: "forelegL2", 10: "forelegL3", 11: "forelegL4", 12: "midlegL1", 13: "midlegL2", 14: "midlegL3", 15: "midlegL4", 16: "hindlegL1", 17: "hindlegL2", 18: "hindlegL3", 19: "hindlegL4", 20: "anttipR", 21: "antbaseR", 22: "eyeR", 23: "forelegR1", 24: "forelegR2", 25: "forelegR3", 26: "forelegR4", 27: "midlegR1", 28: "midlegR2", 29: "midlegR3", 30: "midlegR4", 31: "hindlegR1", 32: "hindlegR2", 33: "hindlegR3", 34: "hindlegR4" 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'] == 35 self.ann_info['joint_weights'] = \ np.ones((self.ann_info['num_joints'], 1), dtype=np.float32) self.ann_info['flip_pairs'] = [[5, 20], [6, 21], [7, 22], [8, 23], [9, 24], [10, 25], [11, 26], [12, 27], [13, 28], [14, 29], [15, 30], [16, 31], [17, 32], [18, 33], [19, 34]] self.dataset_name = 'locust' 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 160x160 center, scale = self._xywh2cs(0, 0, 160, 160, 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