efficientdet/effdet/evaluator.py

import torch
import torch.distributed as dist
import abc
import json
import logging
import time
import numpy as np

from .distributed import synchronize, is_main_process, all_gather_container
from pycocotools.cocoeval import COCOeval

# FIXME experimenting with speedups for OpenImages eval, it's slow
#import pyximport; py_importer, pyx_importer = pyximport.install(pyimport=True)
import effdet.evaluation.detection_evaluator as tfm_eval
#pyximport.uninstall(py_importer, pyx_importer)

_logger = logging.getLogger(__name__)


__all__ = ['CocoEvaluator', 'PascalEvaluator', 'OpenImagesEvaluator', 'create_evaluator']


class Evaluator:

    def __init__(self, distributed=False, pred_yxyx=False):
        self.distributed = distributed
        self.distributed_device = None
        self.pred_yxyx = pred_yxyx
        self.img_indices = []
        self.predictions = []

    def add_predictions(self, detections, target):
        if self.distributed:
            if self.distributed_device is None:
                # cache for use later to broadcast end metric
                self.distributed_device = detections.device
            synchronize()
            detections = all_gather_container(detections)
            img_indices = all_gather_container(target['img_idx'])
            if not is_main_process():
                return
        else:
            img_indices = target['img_idx']

        detections = detections.cpu().numpy()
        img_indices = img_indices.cpu().numpy()
        for img_idx, img_dets in zip(img_indices, detections):
            self.img_indices.append(img_idx)
            self.predictions.append(img_dets)

    def _coco_predictions(self):
        # generate coco-style predictions
        coco_predictions = []
        coco_ids = []
        for img_idx, img_dets in zip(self.img_indices, self.predictions):
            img_id = self._dataset.img_ids[img_idx]
            coco_ids.append(img_id)
            if self.pred_yxyx:
                # to xyxy
                img_dets[:, 0:4] = img_dets[:, [1, 0, 3, 2]]
            # to xywh
            img_dets[:, 2] -= img_dets[:, 0]
            img_dets[:, 3] -= img_dets[:, 1]
            for det in img_dets:
                score = float(det[4])
                if score < .001:  # stop when below this threshold, scores in descending order
                    break
                coco_det = dict(
                    image_id=int(img_id),
                    bbox=det[0:4].tolist(),
                    score=score,
                    category_id=int(det[5]))
                coco_predictions.append(coco_det)
        return coco_predictions, coco_ids

    @abc.abstractmethod
    def evaluate(self):
        pass

    def save(self, result_file):
        # save results in coco style, override to save in a alternate form
        if not self.distributed or dist.get_rank() == 0:
            assert len(self.predictions)
            coco_predictions, coco_ids = self._coco_predictions()
            json.dump(coco_predictions, open(result_file, 'w'), indent=4)


class CocoEvaluator(Evaluator):

    def __init__(self, dataset, neptune=None, distributed=False, pred_yxyx=False):
        super().__init__(distributed=distributed, pred_yxyx=pred_yxyx)
        self._dataset = dataset.parser
        self.coco_api = dataset.parser.coco
        self.neptune = neptune

    def reset(self):
        self.img_indices = []
        self.predictions = []

    def evaluate(self):
        if not self.distributed or dist.get_rank() == 0:
            assert len(self.predictions)
            coco_predictions, coco_ids = self._coco_predictions()
            json.dump(coco_predictions, open('./temp.json', 'w'), indent=4)
            results = self.coco_api.loadRes('./temp.json')
            coco_eval = COCOeval(self.coco_api, results, 'bbox')
            coco_eval.params.imgIds = coco_ids  # score only ids we've used
            coco_eval.evaluate()
            coco_eval.accumulate()
            coco_eval.summarize()
            metric = coco_eval.stats[0]  # mAP 0.5-0.95
            if self.neptune:
                self.neptune.log_metric('valid/mAP/0.5-0.95IOU', metric)
                self.neptune.log_metric('valid/mAP/0.5IOU', coco_eval.stats[1])
            if self.distributed:
                dist.broadcast(torch.tensor(metric, device=self.distributed_device), 0)
        else:
            metric = torch.tensor(0, device=self.distributed_device)
            dist.broadcast(metric, 0)
            metric = metric.item()
        self.reset()
        return metric


class TfmEvaluator(Evaluator):
    """ Tensorflow Models Evaluator Wrapper """
    def __init__(
            self, dataset, neptune=None, distributed=False, pred_yxyx=False, 
            evaluator_cls=tfm_eval.ObjectDetectionEvaluator):
        super().__init__(distributed=distributed, pred_yxyx=pred_yxyx)
        self._evaluator = evaluator_cls(categories=dataset.parser.cat_dicts)
        self._eval_metric_name = self._evaluator._metric_names[0]
        self._dataset = dataset.parser
        self.neptune = neptune

    def reset(self):
        self._evaluator.clear()
        self.img_indices = []
        self.predictions = []

    def evaluate(self):
        if not self.distributed or dist.get_rank() == 0:
            for img_idx, img_dets in zip(self.img_indices, self.predictions):
                gt = self._dataset.get_ann_info(img_idx)
                self._evaluator.add_single_ground_truth_image_info(img_idx, gt)

                bbox = img_dets[:, 0:4] if self.pred_yxyx else img_dets[:, [1, 0, 3, 2]]
                det = dict(bbox=bbox, score=img_dets[:, 4], cls=img_dets[:, 5])
                self._evaluator.add_single_detected_image_info(img_idx, det)

            metrics = self._evaluator.evaluate()
            _logger.info('Metrics:')
            for k, v in metrics.items():
                _logger.info(f'{k}: {v}')
                if self.neptune:
                    key = 'valid/mAP/' + str(k).split('/')[-1]
                    self.neptune.log_metric(key, v)
                
            map_metric = metrics[self._eval_metric_name]
            if self.distributed:
                dist.broadcast(torch.tensor(map_metric, device=self.distributed_device), 0)
        else:
            map_metric = torch.tensor(0, device=self.distributed_device)
            wait = dist.broadcast(map_metric, 0, async_op=True)
            while not wait.is_completed():
                # wait without spinning the cpu @ 100%, no need for low latency here
                time.sleep(0.5)
            map_metric = map_metric.item()
        self.reset()
        return map_metric


class PascalEvaluator(TfmEvaluator):

    def __init__(self, dataset, neptune=None, distributed=False, pred_yxyx=False):
        super().__init__(
            dataset, neptune, distributed=distributed, pred_yxyx=pred_yxyx, evaluator_cls=tfm_eval.PascalDetectionEvaluator)


class OpenImagesEvaluator(TfmEvaluator):

    def __init__(self, dataset, distributed=False, pred_yxyx=False):
        super().__init__(
            dataset, distributed=distributed, pred_yxyx=pred_yxyx, evaluator_cls=tfm_eval.OpenImagesDetectionEvaluator)


def create_evaluator(name, dataset, neptune=None, distributed=False, pred_yxyx=False):
    # FIXME support OpenImages Challenge2019 metric w/ image level label consideration
    if 'coco' in name:
        return CocoEvaluator(dataset, neptune, distributed=distributed, pred_yxyx=pred_yxyx)
    elif 'openimages' in name:
        return OpenImagesEvaluator(dataset, distributed=distributed, pred_yxyx=pred_yxyx)
    else:
        return CocoEvaluator(dataset, neptune, distributed=distributed, pred_yxyx=pred_yxyx)
        #return PascalEvaluator(dataset, neptune, distributed=distributed, pred_yxyx=pred_yxyx)