File size: 35,512 Bytes
a249588
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
# Copyright (c) OpenMMLab. All rights reserved.
import datetime
import os.path as osp
import tempfile
from collections import OrderedDict, defaultdict
from typing import Dict, Optional, Sequence
import traceback

import numpy as np
from mmengine.evaluator import BaseMetric
from mmengine.fileio import dump, get_local_path, load
from mmengine.logging import MessageHub, MMLogger, print_log
from xtcocotools.coco import COCO
from xtcocotools.cocoeval import COCOeval

from mmpose.registry import METRICS
from mmpose.structures.bbox import bbox_xyxy2xywh
from mmpose.structures.keypoint import find_min_padding_exact, fix_bbox_aspect_ratio
from ..functional import (oks_nms, soft_oks_nms, transform_ann, transform_pred,
                          transform_sigmas)

import cv2
import os

import matplotlib.pyplot as plt
from matplotlib import rc

from xtcocotools.mask import _mask as maskUtils


@METRICS.register_module()
class CocoMetric(BaseMetric):
    """COCO pose estimation task evaluation metric.

    Evaluate AR, AP, and mAP for keypoint detection tasks. Support COCO
    dataset and other datasets in COCO format. Please refer to
    `COCO keypoint evaluation <https://cocodataset.org/#keypoints-eval>`__
    for more details.

    Args:
        ann_file (str, optional): Path to the coco format annotation file.
            If not specified, ground truth annotations from the dataset will
            be converted to coco format. Defaults to None
        use_area (bool): Whether to use ``'area'`` message in the annotations.
            If the ground truth annotations (e.g. CrowdPose, AIC) do not have
            the field ``'area'``, please set ``use_area=False``.
            Defaults to ``True``
        iou_type (str): The same parameter as `iouType` in
            :class:`xtcocotools.COCOeval`, which can be ``'keypoints'``, or
            ``'keypoints_crowd'`` (used in CrowdPose dataset).
            Defaults to ``'keypoints'``
        score_mode (str): The mode to score the prediction results which
            should be one of the following options:

                - ``'bbox'``: Take the score of bbox as the score of the
                    prediction results.
                - ``'bbox_keypoint'``: Use keypoint score to rescore the
                    prediction results.
                - ``'bbox_rle'``: Use rle_score to rescore the
                    prediction results.

            Defaults to ``'bbox_keypoint'`
        keypoint_score_thr (float): The threshold of keypoint score. The
            keypoints with score lower than it will not be included to
            rescore the prediction results. Valid only when ``score_mode`` is
            ``bbox_keypoint``. Defaults to ``0.2``
        nms_mode (str): The mode to perform Non-Maximum Suppression (NMS),
            which should be one of the following options:

                - ``'oks_nms'``: Use Object Keypoint Similarity (OKS) to
                    perform NMS.
                - ``'soft_oks_nms'``: Use Object Keypoint Similarity (OKS)
                    to perform soft NMS.
                - ``'none'``: Do not perform NMS. Typically for bottomup mode
                    output.

            Defaults to ``'oks_nms'`
        nms_thr (float): The Object Keypoint Similarity (OKS) threshold
            used in NMS when ``nms_mode`` is ``'oks_nms'`` or
            ``'soft_oks_nms'``. Will retain the prediction results with OKS
            lower than ``nms_thr``. Defaults to ``0.9``
        format_only (bool): Whether only format the output results without
            doing quantitative evaluation. This is designed for the need of
            test submission when the ground truth annotations are absent. If
            set to ``True``, ``outfile_prefix`` should specify the path to
            store the output results. Defaults to ``False``
        pred_converter (dict, optional): Config dictionary for the prediction
            converter. The dictionary has the same parameters as
            'KeypointConverter'. Defaults to None.
        gt_converter (dict, optional): Config dictionary for the ground truth
            converter. The dictionary has the same parameters as
            'KeypointConverter'. Defaults to None.
        outfile_prefix (str | None): The prefix of json files. It includes
            the file path and the prefix of filename, e.g., ``'a/b/prefix'``.
            If not specified, a temp file will be created. Defaults to ``None``
        collect_device (str): Device name used for collecting results from
            different ranks during distributed training. Must be ``'cpu'`` or
            ``'gpu'``. Defaults to ``'cpu'``
        prefix (str, optional): The prefix that will be added in the metric
            names to disambiguate homonymous metrics of different evaluators.
            If prefix is not provided in the argument, ``self.default_prefix``
            will be used instead. Defaults to ``None``
    """
    default_prefix: Optional[str] = 'coco'

    def __init__(self,
                 ann_file: Optional[str] = None,
                 use_area: bool = True,
                 iou_type: str = 'keypoints',
                 score_mode: str = 'bbox_keypoint',
                 score_thresh_type: str = 'score',
                 keypoint_score_thr: float = 0.2,
                 nms_mode: str = 'oks_nms',
                 nms_thr: float = 0.9,
                 format_only: bool = False,
                 pred_converter: Dict = None,
                 gt_converter: Dict = None,
                 outfile_prefix: Optional[str] = None,
                 collect_device: str = 'cpu',
                 prefix: Optional[str] = None,
                 extended: list = [False],
                 match_by_bbox: list = [False],
                 ignore_border_points: list = [False],
                 ignore_stats: list = [],
                 padding: float = 1.25) -> None:
        super().__init__(collect_device=collect_device, prefix=prefix)
        self.ann_file = ann_file
        # initialize coco helper with the annotation json file
        # if ann_file is not specified, initialize with the converted dataset
        if ann_file is not None:
            with get_local_path(ann_file) as local_path:
                self.coco = COCO(local_path)
        else:
            self.coco = None

        self.use_area = use_area
        self.iou_type = iou_type

        allowed_score_modes = ['bbox', 'bbox_keypoint', 'bbox_rle', 'keypoint']
        if score_mode not in allowed_score_modes:
            raise ValueError(
                "`score_mode` should be one of 'bbox', 'bbox_keypoint', "
                f"'bbox_rle', but got {score_mode}")
        self.score_mode = score_mode
        self.keypoint_score_thr = keypoint_score_thr
        if score_thresh_type not in ['score', 'prob']:
            raise ValueError(
                "'score_thresh_type' should be one of 'score' or 'prob'"
            )
        self.score_thresh_type = score_thresh_type

        allowed_nms_modes = ['oks_nms', 'soft_oks_nms', 'none']
        if nms_mode not in allowed_nms_modes:
            raise ValueError(
                "`nms_mode` should be one of 'oks_nms', 'soft_oks_nms', "
                f"'none', but got {nms_mode}")
        self.nms_mode = nms_mode
        self.nms_thr = nms_thr

        if format_only:
            assert outfile_prefix is not None, '`outfile_prefix` can not be '\
                'None when `format_only` is True, otherwise the result file '\
                'will be saved to a temp directory which will be cleaned up '\
                'in the end.'
        elif ann_file is not None:
            # do evaluation only if the ground truth annotations exist
            assert 'annotations' in load(ann_file), \
                'Ground truth annotations are required for evaluation '\
                'when `format_only` is False.'

        self.format_only = format_only
        self.outfile_prefix = outfile_prefix
        self.pred_converter = pred_converter
        self.gt_converter = gt_converter

        len_params = max(len(extended), len(match_by_bbox))
        if len(extended) == 1 and len_params > 1:
            extended = extended * len_params
        if len(match_by_bbox) == 1 and len_params > 1:
            match_by_bbox = match_by_bbox * len_params
        assert len(extended) == len(match_by_bbox), \
            'The length of `extended` and `match_by_bbox` should be the same.'
        assert len(extended) >= 1, \
            'The length of `extended` and `match_by_bbox` should be at least 1.'
        self.extended = extended
        self.match_by_bbox = match_by_bbox
        self.ignore_border_points = ignore_border_points

        self.ignore_stats = ignore_stats
        self.prob_thr = -1
        self.has_probability = True
        self.padding = padding

        self._compute_min_padding_in_coco()

    @property
    def dataset_meta(self) -> Optional[dict]:
        """Optional[dict]: Meta info of the dataset."""
        return self._dataset_meta

    @dataset_meta.setter
    def dataset_meta(self, dataset_meta: dict) -> None:
        """Set the dataset meta info to the metric."""
        if self.gt_converter is not None:
            dataset_meta['sigmas'] = transform_sigmas(
                dataset_meta['sigmas'], self.gt_converter['num_keypoints'],
                self.gt_converter['mapping'])
            dataset_meta['num_keypoints'] = len(dataset_meta['sigmas'])
        self._dataset_meta = dataset_meta

        if self.coco is None:
            message = MessageHub.get_current_instance()
            ann_file = message.get_info(
                f"{dataset_meta['dataset_name']}_ann_file", None)
            if ann_file is not None:
                with get_local_path(ann_file) as local_path:
                    self.coco = COCO(local_path)
                print_log(
                    f'CocoMetric for dataset '
                    f"{dataset_meta['dataset_name']} has successfully "
                    f'loaded the annotation file from {ann_file}', 'current')

    def _compute_min_padding_in_coco(self):
        """Compute the minimum padding in COCO format."""
        if self.coco is None:
            return
        
        for _, ann in self.coco.anns.items():
            if 'pad_to_contain' in ann.keys():
                continue

            kpts = np.array(ann['keypoints']).reshape(-1, 3)
            bbox = np.array(ann['bbox']).flatten()
            min_padding = find_min_padding_exact(bbox, kpts)
            ann['pad_to_contain'] = min_padding

        return

    def process(self, data_batch: Sequence[dict],
                data_samples: Sequence[dict]) -> None:
        """Process one batch of data samples and predictions. The processed
        results should be stored in ``self.results``, which will be used to
        compute the metrics when all batches have been processed.

        Args:
            data_batch (Sequence[dict]): A batch of data
                from the dataloader.
            data_samples (Sequence[dict]): A batch of outputs from
                the model, each of which has the following keys:

                - 'id': The id of the sample
                - 'img_id': The image_id of the sample
                - 'pred_instances': The prediction results of instance(s)
        """
        self.results_len = len(self.results)
        for data_sample in data_samples:
            if 'pred_instances' not in data_sample:
                raise ValueError(
                    '`pred_instances` are required to process the '
                    f'predictions results in {self.__class__.__name__}. ')

            # keypoints.shape: [N, K, 2],
            # N: number of instances, K: number of keypoints
            # for topdown-style output, N is usually 1, while for
            # bottomup-style output, N is the number of instances in the image
            keypoints = data_sample['pred_instances']['keypoints']
            N, K, _ = keypoints.shape
            # [N, K], the scores for all keypoints of all instances
            keypoint_scores = data_sample['pred_instances']['keypoint_scores']
            assert keypoint_scores.shape == keypoints.shape[:2]
            
            if 'keypoints_visible' in data_sample['pred_instances']:
                keypoints_visible = data_sample['pred_instances']['keypoints_visible']
            else:
                keypoints_visible = keypoint_scores.copy()
            
            if 'keypoints_probs' in data_sample['pred_instances']:
                keypoints_probs = data_sample['pred_instances']['keypoints_probs']
                # keypoints_probs = keypoint_scores.copy()
            else:
                self.has_probability = False
                keypoints_probs = keypoint_scores.copy()

            if 'keypoints_oks' in data_sample['pred_instances']:
                keypoints_oks = data_sample['pred_instances']['keypoints_oks']
            else:
                keypoints_oks = keypoint_scores.copy()

            if 'keypoints_error' in data_sample['pred_instances']:
                keypoints_error = data_sample['pred_instances']['keypoints_error']
            else:
                keypoints_error = keypoint_scores.copy()

            if K == 21:
                # Translate 21 keypoints to 17 keypoints by ignoring the last 4 keypoints
                keypoints = keypoints[:, :17, :]
                keypoint_scores = keypoint_scores[:, :17]
                keypoints_visible = keypoints_visible[:, :17]
                keypoints_probs = keypoints_probs[:, :17]
                keypoints_oks = keypoints_oks[:, :17]
                keypoints_error = keypoints_error[:, :17]

            elif K != 17:
                raise ValueError('The number of keypoints should be 17 or 21, '
                                    f'but got {K}.')

            assert keypoints.shape[1] == 17, f'Number of keypoints should be 17 but got {keypoints.shape}'
            assert keypoint_scores.shape[1] == 17, f'Number of keypoint scores should be 17 but got {keypoint_scores.shape}'
            assert keypoints_visible.shape[1] == 17, f'Number of visible keypoints should be 17 but got {keypoints_visible.shape}'
            assert keypoints_probs.shape[1] == 17, f'Number of keypoint probs should be 17 but got {keypoints_probs.shape}'
            assert keypoints_oks.shape[1] == 17, f'Number of keypoint oks should be 17 but got {keypoints_oks.shape}'
            assert keypoints_error.shape[1] == 17, f'Number of keypoint error should be 17 but got {keypoints_error.shape}'
            assert heatmaps.shape[1] == 17, f'Number of heatmaps should be 17 but got {heatmaps.shape}'

            # parse prediction results
            pred = dict()
            pred['id'] = data_sample['id']
            pred['img_id'] = data_sample['img_id']

            pred['keypoints'] = keypoints
            pred['keypoint_scores'] = keypoint_scores
            pred['keypoints_visible'] = keypoints_visible
            pred['keypoint_probs'] = keypoints_probs
            pred['keypoint_oks'] = keypoints_oks
            pred['keypoint_error'] = keypoints_error
            pred['category_id'] = data_sample.get('category_id', 1)
            if 'bboxes' in data_sample['pred_instances']:
                pred['bbox'] = bbox_xyxy2xywh(
                    data_sample['pred_instances']['bboxes'])

            if 'bbox_scores' in data_sample['pred_instances']:
                # some one-stage models will predict bboxes and scores
                # together with keypoints
                bbox_scores = data_sample['pred_instances']['bbox_scores']
            elif ('bbox_scores' not in data_sample['gt_instances']
                  or len(data_sample['gt_instances']['bbox_scores']) !=
                  len(keypoints)):
                # bottom-up models might output different number of
                # instances from annotation
                bbox_scores = np.ones(len(keypoints))
            else:
                # top-down models use detected bboxes, the scores of which
                # are contained in the gt_instances
                bbox_scores = data_sample['gt_instances']['bbox_scores']
            pred['bbox_scores'] = bbox_scores

            # get area information
            if 'bbox_scales' in data_sample['gt_instances']:
                pred['areas'] = np.prod(
                    data_sample['gt_instances']['bbox_scales'], axis=1)

            # parse gt
            gt = dict()
            if self.coco is None:
                gt['width'] = data_sample['ori_shape'][1]
                gt['height'] = data_sample['ori_shape'][0]
                gt['img_id'] = data_sample['img_id']
                if self.iou_type == 'keypoints_crowd':
                    assert 'crowd_index' in data_sample, \
                        '`crowd_index` is required when `self.iou_type` is ' \
                        '`keypoints_crowd`'
                    gt['crowd_index'] = data_sample['crowd_index']
                assert 'raw_ann_info' in data_sample, \
                    'The row ground truth annotations are required for ' \
                    'evaluation when `ann_file` is not provided'
                anns = data_sample['raw_ann_info']
                gt['raw_ann_info'] = anns if isinstance(anns, list) else [anns]

            # add converted result to the results list
            self.results.append((pred, gt))
        processed_len = len(self.results) - self.results_len
        if processed_len != len(data_samples):
            print(f'Warning: {processed_len} samples are processed, ')
            print(f'but {len(data_samples)} samples are provided.')
        
    def gt_to_coco_json(self, gt_dicts: Sequence[dict],
                        outfile_prefix: str) -> str:
        """Convert ground truth to coco format json file.

        Args:
            gt_dicts (Sequence[dict]): Ground truth of the dataset. Each dict
                contains the ground truth information about the data sample.
                Required keys of the each `gt_dict` in `gt_dicts`:
                    - `img_id`: image id of the data sample
                    - `width`: original image width
                    - `height`: original image height
                    - `raw_ann_info`: the raw annotation information
                Optional keys:
                    - `crowd_index`: measure the crowding level of an image,
                        defined in CrowdPose dataset
                It is worth mentioning that, in order to compute `CocoMetric`,
                there are some required keys in the `raw_ann_info`:
                    - `id`: the id to distinguish different annotations
                    - `image_id`: the image id of this annotation
                    - `category_id`: the category of the instance.
                    - `bbox`: the object bounding box
                    - `keypoints`: the keypoints cooridinates along with their
                        visibilities. Note that it need to be aligned
                        with the official COCO format, e.g., a list with length
                        N * 3, in which N is the number of keypoints. And each
                        triplet represent the [x, y, visible] of the keypoint.
                    - `iscrowd`: indicating whether the annotation is a crowd.
                        It is useful when matching the detection results to
                        the ground truth.
                There are some optional keys as well:
                    - `area`: it is necessary when `self.use_area` is `True`
                    - `num_keypoints`: it is necessary when `self.iou_type`
                        is set as `keypoints_crowd`.
            outfile_prefix (str): The filename prefix of the json files. If the
                prefix is "somepath/xxx", the json file will be named
                "somepath/xxx.gt.json".
        Returns:
            str: The filename of the json file.
        """
        image_infos = []
        annotations = []
        img_ids = []
        ann_ids = []

        for gt_dict in gt_dicts:
            # filter duplicate image_info
            if gt_dict['img_id'] not in img_ids:
                image_info = dict(
                    id=gt_dict['img_id'],
                    width=gt_dict['width'],
                    height=gt_dict['height'],
                )
                if self.iou_type == 'keypoints_crowd':
                    image_info['crowdIndex'] = gt_dict['crowd_index']

                image_infos.append(image_info)
                img_ids.append(gt_dict['img_id'])

            # filter duplicate annotations
            for ann in gt_dict['raw_ann_info']:
                if ann is None:
                    # during evaluation on bottom-up datasets, some images
                    # do not have instance annotation
                    continue

                annotation = dict(
                    id=ann['id'],
                    image_id=ann['image_id'],
                    category_id=ann['category_id'],
                    bbox=ann['bbox'],
                    keypoints=ann['keypoints'],
                    iscrowd=ann['iscrowd'],
                )
                if self.use_area:
                    assert 'area' in ann, \
                        '`area` is required when `self.use_area` is `True`'
                    annotation['area'] = ann['area']

                if self.iou_type == 'keypoints_crowd':
                    assert 'num_keypoints' in ann, \
                        '`num_keypoints` is required when `self.iou_type` ' \
                        'is `keypoints_crowd`'
                    annotation['num_keypoints'] = ann['num_keypoints']

                annotations.append(annotation)
                ann_ids.append(ann['id'])

        info = dict(
            date_created=str(datetime.datetime.now()),
            description='Coco json file converted by mmpose CocoMetric.')
        coco_json = dict(
            info=info,
            images=image_infos,
            categories=self.dataset_meta['CLASSES'],
            licenses=None,
            annotations=annotations,
        )
        converted_json_path = f'{outfile_prefix}.gt.json'
        dump(coco_json, converted_json_path, sort_keys=True, indent=4)
        return converted_json_path

    def compute_metrics(self, results: list) -> Dict[str, float]:
        """Compute the metrics from processed results.

        Args:
            results (list): The processed results of each batch.

        Returns:
            Dict[str, float]: The computed metrics. The keys are the names of
            the metrics, and the values are corresponding results.
        """
        logger: MMLogger = MMLogger.get_current_instance()

        # split prediction and gt list
        preds, gts = zip(*results)

        tmp_dir = None
        if self.outfile_prefix is None:
            tmp_dir = tempfile.TemporaryDirectory()
            outfile_prefix = osp.join(tmp_dir.name, 'results')
        else:
            outfile_prefix = self.outfile_prefix

        if self.coco is None:
            # use converted gt json file to initialize coco helper
            logger.info('Converting ground truth to coco format...')
            coco_json_path = self.gt_to_coco_json(
                gt_dicts=gts, outfile_prefix=outfile_prefix)
            self.coco = COCO(coco_json_path)
        if self.gt_converter is not None:
            for id_, ann in self.coco.anns.items():
                self.coco.anns[id_] = transform_ann(
                    ann, self.gt_converter['num_keypoints'],
                    self.gt_converter['mapping'])

        kpts = defaultdict(list)

        # group the preds by img_id
        for pred in preds:
            img_id = pred['img_id']

            if self.pred_converter is not None:
                pred = transform_pred(pred,
                                      self.pred_converter['num_keypoints'],
                                      self.pred_converter['mapping'])

            for idx, keypoints in enumerate(pred['keypoints']):
                
                instance = {
                    'id': pred['id'],
                    'img_id': pred['img_id'],
                    'category_id': pred['category_id'],
                    'keypoints': keypoints,
                    'keypoint_scores': pred['keypoint_scores'][idx],
                    'bbox_score': pred['bbox_scores'][idx],
                    'keypoints_visible': pred['keypoints_visible'][idx],
                    'keypoint_probs': pred['keypoint_probs'][idx],
                    'keypoint_oks': pred['keypoint_oks'][idx],
                    'keypoint_error': pred['keypoint_error'][idx],
                }
                
                # breakpoint()
                if 'bbox' in pred:
                    instance['bbox'] = pred['bbox'][idx]
                    diagonal = np.sqrt(
                        instance['bbox'][2]**2 + instance['bbox'][3]**2)
                if 'areas' in pred:
                    instance['area'] = pred['areas'][idx]
                    diagonal = np.sqrt(instance['area'])
                else:
                    # use keypoint to calculate bbox and get area
                    area = (
                        np.max(keypoints[:, 0]) - np.min(keypoints[:, 0])) * (
                            np.max(keypoints[:, 1]) - np.min(keypoints[:, 1]))
                    instance['area'] = area
                    diagonal = np.sqrt(area)
                
                kpts[img_id].append(instance)

        # sort keypoint results according to id and remove duplicate ones
        kpts = self._sort_and_unique_bboxes(kpts, key='id')

        # score the prediction results according to `score_mode`
        # and perform NMS according to `nms_mode`
        valid_kpts = defaultdict(list)
        if self.pred_converter is not None:
            num_keypoints = self.pred_converter['num_keypoints']
        else:
            num_keypoints = self.dataset_meta['num_keypoints']
        for img_id, instances in kpts.items():
            for instance in instances:
                # concatenate the keypoint coordinates and scores
                instance['keypoints'] = np.concatenate([
                    instance['keypoints'], instance['keypoint_probs'][:, None]
                ],
                                                       axis=-1)
                if self.score_mode == 'bbox':
                    instance['score'] = instance['bbox_score']
                elif self.score_mode == 'keypoint':
                    instance['score'] = np.mean(instance['keypoint_scores'])
                else:
                    bbox_score = instance['bbox_score']
                    if self.score_mode == 'bbox_rle':
                        keypoint_scores = instance['keypoint_scores']
                        instance['score'] = float(bbox_score +
                                                  np.mean(keypoint_scores) +
                                                  np.max(keypoint_scores))

                    else:  # self.score_mode == 'bbox_keypoint':
                        mean_kpt_score = 0
                        valid_num = 0
                        for kpt_idx in range(num_keypoints):
                            kpt_score = instance['keypoint_scores'][kpt_idx]
                            kpt_prob = instance['keypoint_probs'][kpt_idx]
                            kpt_thresh = kpt_score if self.score_thresh_type == 'score' else kpt_prob
                            if kpt_thresh > self.keypoint_score_thr:
                                mean_kpt_score += kpt_score
                                valid_num += 1
                        if valid_num != 0:
                            mean_kpt_score /= valid_num
                        instance['score'] = bbox_score * mean_kpt_score
            # perform nms
            if self.nms_mode == 'none':
                valid_kpts[img_id] = instances
            else:
                nms = oks_nms if self.nms_mode == 'oks_nms' else soft_oks_nms
                keep = nms(
                    instances,
                    self.nms_thr,
                    sigmas=self.dataset_meta['sigmas'])
                valid_kpts[img_id] = [instances[_keep] for _keep in keep]

        # convert results to coco style and dump into a json file
        self.results2json(valid_kpts, outfile_prefix=outfile_prefix)

        # only format the results without doing quantitative evaluation
        if self.format_only:
            logger.info('results are saved in '
                        f'{osp.dirname(outfile_prefix)}')
            return {}

        eval_results = OrderedDict()
        
        # mAP evaluation results
        logger.info(f'Evaluating {self.__class__.__name__}...')
        self.prob_thr = 0.51

        # Localization evaluation results
        info_str = self._do_python_keypoint_eval(outfile_prefix)
        name_value = OrderedDict(info_str)
        eval_results.update(name_value)


        logger.info('Number of values per dataset: {}'.format(len(eval_results)))

        if tmp_dir is not None:
            tmp_dir.cleanup()
        return eval_results

    def results2json(self, keypoints: Dict[int, list],
                     outfile_prefix: str) -> str:
        """Dump the keypoint detection results to a COCO style json file.

        Args:
            keypoints (Dict[int, list]): Keypoint detection results
                of the dataset.
            outfile_prefix (str): The filename prefix of the json files. If the
                prefix is "somepath/xxx", the json files will be named
                "somepath/xxx.keypoints.json",

        Returns:
            str: The json file name of keypoint results.
        """
        # the results with category_id
        cat_results = []

        for _, img_kpts in keypoints.items():
            _keypoints = np.array(
                [img_kpt['keypoints'] for img_kpt in img_kpts])
            num_keypoints = self.dataset_meta['num_keypoints']
            # collect all the person keypoints in current image
            _keypoints = _keypoints.reshape(-1, num_keypoints * 3)

            result = []
            for img_kpt, keypoint in zip(img_kpts, _keypoints):
                res = {
                    'image_id': img_kpt['img_id'],
                    'category_id': img_kpt['category_id'],
                    'keypoints': keypoint.tolist(),
                    'score': float(img_kpt['score']),
                }
                if 'bbox' in img_kpt:
                    res['bbox'] = img_kpt['bbox'].tolist()
                if 'keypoint_probs' in img_kpt:
                    res['probs'] = img_kpt['keypoint_probs'].tolist()
                result.append(res)

            cat_results.extend(result)

        res_file = f'{outfile_prefix}.keypoints.json'
        dump(cat_results, res_file, sort_keys=True, indent=4)

    def _do_python_keypoint_eval(self, outfile_prefix: str) -> list:
        """Do keypoint evaluation using COCOAPI.

        Args:
            outfile_prefix (str): The filename prefix of the json files. If the
                prefix is "somepath/xxx", the json files will be named
                "somepath/xxx.keypoints.json",

        Returns:
            list: a list of tuples. Each tuple contains the evaluation stats
            name and corresponding stats value.
        """
        res_file = f'{outfile_prefix}.keypoints.json'
        coco_det = self.coco.loadRes(res_file)
        sigmas = self.dataset_meta['sigmas']

        info_str = []
        for extended_oks, match_by_bbox, ignore_border_points in zip(
            self.extended, self.match_by_bbox, self.ignore_border_points
        ):
            prefix = ""
            suffix = ""
            if match_by_bbox:
                prefix = "bbox_" + prefix
            if extended_oks:
                prefix = "Ex_" + prefix
            if ignore_border_points:
                suffix = suffix + "_NoBrd"

            conf_thr = self.prob_thr
            print("+"*80)
            print("COCO Eval params: Bbox {:5s}, ExOKS {:5s}".format(
                str(match_by_bbox), str(extended_oks)
            ), end="")
            if extended_oks:
                print(" with conf_thr: {:.2f} (has probability: {})".format(conf_thr, self.has_probability), end="")
            print()

            coco_eval = COCOeval(
                self.coco,
                coco_det,
                iouType=self.iou_type,
                sigmas=sigmas,
                use_area=self.use_area,
                extended_oks=extended_oks,
                match_by_bbox=match_by_bbox,
                confidence_thr=conf_thr,
                padding=self.padding,
                ignore_near_bbox=ignore_border_points
            )
            coco_eval.params.useSegm = None
            coco_eval.evaluate()
            coco_eval.accumulate()
            coco_eval.summarize()

            try:
                stats_names = coco_eval.stats_names
            except AttributeError:
                if self.iou_type == 'keypoints_crowd':
                    stats_names = [
                        'AP', 'AP .5', 'AP .75', 'AR', 'AR .5', 'AR .75',
                        'AP(E)', 'AP(M)', 'AP(H)'
                    ]
                else:
                    stats_names = [
                        'AP', 'AP .5', 'AP .75', 'AP (M)', 'AP (L)', 'AR',
                        'AR .5', 'AR .75', 'AR (M)', 'AR (L)'
                    ]
            i_str = list(zip(stats_names, coco_eval.stats))
            ignore_stats = self.ignore_stats
            # if match_by_bbox or extended_oks:
            #     ignore_stats.extend(['AP (M)', 'AP (L)', 'AR (M)', 'AR (L)', 'AR'])
            i_str = [(k, v) for k, v in i_str if k not in self.ignore_stats]
            i_str = [(f'{prefix}{k}', v) for k, v in i_str]
            i_str = [(f'{k}{suffix}', v) for k, v in i_str]

            info_str.extend(i_str)

        return info_str

    def _sort_and_unique_bboxes(self,
                                kpts: Dict[int, list],
                                key: str = 'id') -> Dict[int, list]:
        """Sort keypoint detection results in each image and remove the
        duplicate ones. Usually performed in multi-batch testing.

        Args:
            kpts (Dict[int, list]): keypoint prediction results. The keys are
                '`img_id`' and the values are list that may contain
                keypoints of multiple persons. Each element in the list is a
                dict containing the ``'key'`` field.
                See the argument ``key`` for details.
            key (str): The key name in each person prediction results. The
                corresponding value will be used for sorting the results.
                Default: ``'id'``.

        Returns:
            Dict[int, list]: The sorted keypoint detection results.
        """
        for img_id, persons in kpts.items():
            # deal with bottomup-style output
            if isinstance(kpts[img_id][0][key], Sequence):
                return kpts
            num = len(persons)
            kpts[img_id] = sorted(kpts[img_id], key=lambda x: x[key])
            for i in range(num - 1, 0, -1):
                if kpts[img_id][i][key] == kpts[img_id][i - 1][key]:
                    del kpts[img_id][i]

        return kpts