Source code for gluoncv.utils.metrics.coco_detection

"""MS COCO Detection Evaluate Metrics."""
from __future__ import absolute_import

import sys
try:
    from StringIO import StringIO
except ImportError:
    from io import StringIO
import os
from os import path as osp
import warnings
import numpy as np
import mxnet as mx
try:
    from mxnet.metric import EvalMetric
except ImportError:
    from mxnet.gluon.metric import EvalMetric


[docs]class COCODetectionMetric(EvalMetric): """Detection metric for COCO bbox task. Parameters ---------- dataset : instance of gluoncv.data.COCODetection The validation dataset. save_prefix : str Prefix for the saved JSON results. use_time : bool Append unique datetime string to created JSON file name if ``True``. cleanup : bool Remove created JSON file if ``True``. score_thresh : float Detection results with confident scores smaller than ``score_thresh`` will be discarded before saving to results. data_shape : tuple of int, default is None If `data_shape` is provided as (height, width), we will rescale bounding boxes when saving the predictions. This is helpful when SSD/YOLO box predictions cannot be rescaled conveniently. Note that the data_shape must be fixed for all validation images. post_affine : a callable function with input signature (orig_w, orig_h, out_w, out_h) If not None, the bounding boxes will be affine transformed rather than simply scaled. """ def __init__(self, dataset, save_prefix, use_time=True, cleanup=False, score_thresh=0.05, data_shape=None, post_affine=None): super(COCODetectionMetric, self).__init__('COCOMeanAP') self.dataset = dataset self._img_ids = sorted(dataset.coco.getImgIds()) self._current_id = 0 self._cleanup = cleanup self._results = [] self._score_thresh = score_thresh if isinstance(data_shape, (tuple, list)): assert len(data_shape) == 2, "Data shape must be (height, width)" elif not data_shape: data_shape = None else: raise ValueError("data_shape must be None or tuple of int as (height, width)") self._data_shape = data_shape if post_affine is not None: assert self._data_shape is not None, "Using post affine transform requires data_shape" self._post_affine = post_affine else: self._post_affine = None if use_time: import datetime t = datetime.datetime.now().strftime('_%Y_%m_%d_%H_%M_%S') else: t = '' self._filename = osp.abspath(osp.expanduser(save_prefix) + t + '.json') try: f = open(self._filename, 'w') except IOError as e: raise RuntimeError("Unable to open json file to dump. What(): {}".format(str(e))) else: f.close() def __del__(self): if self._cleanup: try: os.remove(self._filename) except IOError as err: warnings.warn(str(err))
[docs] def reset(self): self._current_id = 0 self._results = []
def _update(self): """Use coco to get real scores. """ if not self._current_id == len(self._img_ids): warnings.warn( 'Recorded {} out of {} validation images, incomplete results'.format( self._current_id, len(self._img_ids))) if not self._results: # in case of empty results, push a dummy result self._results.append({'image_id': self._img_ids[0], 'category_id': 0, 'bbox': [0, 0, 0, 0], 'score': 0}) import json try: with open(self._filename, 'w') as f: json.dump(self._results, f) except IOError as e: raise RuntimeError("Unable to dump json file, ignored. What(): {}".format(str(e))) pred = self.dataset.coco.loadRes(self._filename) gt = self.dataset.coco # lazy import pycocotools from ...data.mscoco.utils import try_import_pycocotools try_import_pycocotools() from pycocotools.cocoeval import COCOeval coco_eval = COCOeval(gt, pred, 'bbox') coco_eval.evaluate() coco_eval.accumulate() self._coco_eval = coco_eval return coco_eval
[docs] def get(self): """Get evaluation metrics. """ # Metric printing adapted from detectron/json_dataset_evaluator. def _get_thr_ind(coco_eval, thr): ind = np.where((coco_eval.params.iouThrs > thr - 1e-5) & (coco_eval.params.iouThrs < thr + 1e-5))[0][0] iou_thr = coco_eval.params.iouThrs[ind] assert np.isclose(iou_thr, thr) return ind # call real update try: coco_eval = self._update() except IndexError: # invalid model may result in empty JSON results, skip it return ['mAP',], ['0.0',] IoU_lo_thresh = 0.5 IoU_hi_thresh = 0.95 ind_lo = _get_thr_ind(coco_eval, IoU_lo_thresh) ind_hi = _get_thr_ind(coco_eval, IoU_hi_thresh) # precision has dims (iou, recall, cls, area range, max dets) # area range index 0: all area ranges # max dets index 2: 100 per image precision = coco_eval.eval['precision'][ind_lo:(ind_hi + 1), :, :, 0, 2] ap_default = np.mean(precision[precision > -1]) names, values = [], [] names.append('~~~~ Summary metrics ~~~~\n') # catch coco print string, don't want directly print here _stdout = sys.stdout sys.stdout = StringIO() coco_eval.summarize() coco_summary = sys.stdout.getvalue() sys.stdout = _stdout values.append(str(coco_summary).strip()) for cls_ind, cls_name in enumerate(self.dataset.classes): precision = coco_eval.eval['precision'][ ind_lo:(ind_hi + 1), :, cls_ind, 0, 2] ap = np.mean(precision[precision > -1]) names.append(cls_name) values.append('{:.1f}'.format(100 * ap)) # put mean AP at last, for comparing perf names.append('~~~~ MeanAP @ IoU=[{:.2f},{:.2f}] ~~~~\n'.format( IoU_lo_thresh, IoU_hi_thresh)) values.append('{:.1f}'.format(100 * ap_default)) return names, values
# pylint: disable=arguments-differ, unused-argument
[docs] def update(self, pred_bboxes, pred_labels, pred_scores, *args, **kwargs): """Update internal buffer with latest predictions. Note that the statistics are not available until you call self.get() to return the metrics. Parameters ---------- pred_bboxes : mxnet.NDArray or numpy.ndarray Prediction bounding boxes with shape `B, N, 4`. Where B is the size of mini-batch, N is the number of bboxes. pred_labels : mxnet.NDArray or numpy.ndarray Prediction bounding boxes labels with shape `B, N`. pred_scores : mxnet.NDArray or numpy.ndarray Prediction bounding boxes scores with shape `B, N`. """ from ...data.transforms.bbox import affine_transform def as_numpy(a): """Convert a (list of) mx.NDArray into numpy.ndarray""" if isinstance(a, (list, tuple)): out = [x.asnumpy() if isinstance(x, mx.nd.NDArray) else x for x in a] return np.concatenate(out, axis=0) elif isinstance(a, mx.nd.NDArray): a = a.asnumpy() return a for pred_bbox, pred_label, pred_score in zip( *[as_numpy(x) for x in [pred_bboxes, pred_labels, pred_scores]]): valid_pred = np.where(pred_label.flat >= 0)[0] pred_bbox = pred_bbox[valid_pred, :].astype(np.float) pred_label = pred_label.flat[valid_pred].astype(int) pred_score = pred_score.flat[valid_pred].astype(np.float) imgid = self._img_ids[self._current_id] self._current_id += 1 affine_mat = None if self._data_shape is not None: entry = self.dataset.coco.loadImgs(imgid)[0] orig_height = entry['height'] orig_width = entry['width'] height_scale = float(orig_height) / self._data_shape[0] width_scale = float(orig_width) / self._data_shape[1] if self._post_affine is not None: affine_mat = self._post_affine(orig_width, orig_height, self._data_shape[1], self._data_shape[0]) else: height_scale, width_scale = (1., 1.) # for each bbox detection in each image for bbox, label, score in zip(pred_bbox, pred_label, pred_score): if label not in self.dataset.contiguous_id_to_json: # ignore non-exist class continue if score < self._score_thresh: continue category_id = self.dataset.contiguous_id_to_json[label] # rescale bboxes/affine transform bboxes if affine_mat is not None: bbox[0:2] = affine_transform(bbox[0:2], affine_mat) bbox[2:4] = affine_transform(bbox[2:4], affine_mat) else: bbox[[0, 2]] *= width_scale bbox[[1, 3]] *= height_scale # convert [xmin, ymin, xmax, ymax] to [xmin, ymin, w, h] bbox[2:4] -= (bbox[:2] - 1) self._results.append({'image_id': imgid, 'category_id': category_id, 'bbox': bbox[:4].tolist(), 'score': score})