coco_evaluate.py

import contextlib
import copy
import os
from typing import Dict, List, Union

import numpy as np
import torch

from .coco import COCO
from .cocoeval import COCOeval
from .utils import (
    _TYPING_BOX,
    _TYPING_PREDICTIONS,
    convert_to_xywh,
    create_common_coco_eval,
)

_SUPPORTED_TYPES = ["bbox"]


class COCOEvaluator(object):
    """
    Class to perform evaluation for the COCO dataset.
    """

    def __init__(self, coco_gt: COCO, iou_types: List[str] = ["bbox"]):
        """
        Initializes COCOEvaluator with the ground truth COCO dataset and IoU types.

        Args:
            coco_gt: The ground truth COCO dataset.
            iou_types: Intersection over Union (IoU) types for evaluation (Supported: "bbox").
        """
        self.coco_gt = copy.deepcopy(coco_gt)

        self.coco_eval = {}
        for iou_type in iou_types:
            assert iou_type in _SUPPORTED_TYPES, ValueError(
                f"IoU type not supported {iou_type}"
            )
            self.coco_eval[iou_type] = COCOeval(self.coco_gt, iouType=iou_type)

        self.iou_types = iou_types
        self.img_ids = []
        self.eval_imgs = {k: [] for k in iou_types}

    def update(self, predictions: _TYPING_PREDICTIONS) -> None:
        """
        Update the evaluator with new predictions.

        Args:
            predictions: The predictions to update.
        """
        img_ids = list(np.unique(list(predictions.keys())))
        self.img_ids.extend(img_ids)

        for iou_type in self.iou_types:
            results = self.prepare(predictions, iou_type)

            # suppress pycocotools prints
            with open(os.devnull, "w") as devnull:
                with contextlib.redirect_stdout(devnull):
                    coco_dt = COCO.loadRes(self.coco_gt, results) if results else COCO()
            coco_eval = self.coco_eval[iou_type]

            coco_eval.cocoDt = coco_dt
            coco_eval.params.imgIds = list(img_ids)
            eval_imgs = coco_eval.evaluate()
            self.eval_imgs[iou_type].append(eval_imgs)

    def synchronize_between_processes(self) -> None:
        """
        Synchronizes evaluation images between processes.
        """
        for iou_type in self.iou_types:
            self.eval_imgs[iou_type] = np.concatenate(self.eval_imgs[iou_type], 2)
            create_common_coco_eval(
                self.coco_eval[iou_type], self.img_ids, self.eval_imgs[iou_type]
            )

    def accumulate(self) -> None:
        """
        Accumulates the evaluation results.
        """
        for coco_eval in self.coco_eval.values():
            coco_eval.accumulate()

    def summarize(self) -> None:
        """
        Prints the IoU metric and summarizes the evaluation results.
        """
        for iou_type, coco_eval in self.coco_eval.items():
            print("IoU metric: {}".format(iou_type))
            coco_eval.summarize()

    def prepare(
        self, predictions: _TYPING_PREDICTIONS, iou_type: str
    ) -> List[Dict[str, Union[int, _TYPING_BOX, float]]]:
        """
        Prepares the predictions for COCO detection.

        Args:
            predictions: The predictions to prepare.
            iou_type: The Intersection over Union (IoU) type for evaluation.

        Returns:
            A dictionary with the prepared predictions.
        """
        if iou_type == "bbox":
            return self.prepare_for_coco_detection(predictions)
        else:
            raise ValueError(f"IoU type not supported {iou_type}")

    def _post_process_stats(
        self, stats, coco_eval_object, iou_type="bbox"
    ) -> Dict[str, float]:
        """
        Prepares the predictions for COCO detection.

        Args:
            predictions: The predictions to prepare.
            iou_type: The Intersection over Union (IoU) type for evaluation.

        Returns:
            A dictionary with the prepared predictions.
        """
        if iou_type not in _SUPPORTED_TYPES:
            raise ValueError(f"iou_type '{iou_type}' not supported")

        current_max_dets = coco_eval_object.params.maxDets

        index_to_title = {
            "bbox": {
                0: f"AP-IoU=0.50:0.95-area=all-maxDets={current_max_dets[2]}",
                1: f"AP-IoU=0.50-area=all-maxDets={current_max_dets[2]}",
                2: f"AP-IoU=0.75-area=all-maxDets={current_max_dets[2]}",
                3: f"AP-IoU=0.50:0.95-area=small-maxDets={current_max_dets[2]}",
                4: f"AP-IoU=0.50:0.95-area=medium-maxDets={current_max_dets[2]}",
                5: f"AP-IoU=0.50:0.95-area=large-maxDets={current_max_dets[2]}",
                6: f"AR-IoU=0.50:0.95-area=all-maxDets={current_max_dets[0]}",
                7: f"AR-IoU=0.50:0.95-area=all-maxDets={current_max_dets[1]}",
                8: f"AR-IoU=0.50:0.95-area=all-maxDets={current_max_dets[2]}",
                9: f"AR-IoU=0.50:0.95-area=small-maxDets={current_max_dets[2]}",
                10: f"AR-IoU=0.50:0.95-area=medium-maxDets={current_max_dets[2]}",
                11: f"AR-IoU=0.50:0.95-area=large-maxDets={current_max_dets[2]}",
            },
            "keypoints": {
                0: "AP-IoU=0.50:0.95-area=all-maxDets=20",
                1: "AP-IoU=0.50-area=all-maxDets=20",
                2: "AP-IoU=0.75-area=all-maxDets=20",
                3: "AP-IoU=0.50:0.95-area=medium-maxDets=20",
                4: "AP-IoU=0.50:0.95-area=large-maxDets=20",
                5: "AR-IoU=0.50:0.95-area=all-maxDets=20",
                6: "AR-IoU=0.50-area=all-maxDets=20",
                7: "AR-IoU=0.75-area=all-maxDets=20",
                8: "AR-IoU=0.50:0.95-area=medium-maxDets=20",
                9: "AR-IoU=0.50:0.95-area=large-maxDets=20",
            },
        }

        output_dict: Dict[str, float] = {}
        for index, stat in enumerate(stats):
            output_dict[index_to_title[iou_type][index]] = stat

        return output_dict

    def get_results(self) -> Dict[str, Dict[str, float]]:
        """
        Gets the results of the COCO evaluation.

        Returns:
            A dictionary with the results of the COCO evaluation.
        """
        output_dict = {}

        for iou_type, coco_eval in self.coco_eval.items():
            if iou_type == "segm":
                iou_type = "bbox"
            output_dict[f"iou_{iou_type}"] = self._post_process_stats(
                coco_eval.stats, coco_eval, iou_type
            )
        return output_dict

    def prepare_for_coco_detection(
        self, predictions: _TYPING_PREDICTIONS
    ) -> List[Dict[str, Union[int, _TYPING_BOX, float]]]:
        """
        Prepares the predictions for COCO detection.

        Args:
            predictions: The predictions to prepare.

        Returns:
            A list of dictionaries with the prepared predictions.
        """
        coco_results = []
        for original_id, prediction in predictions.items():
            if len(prediction) == 0:
                continue

            boxes = prediction["boxes"]
            if len(boxes) == 0:
                continue

            if not isinstance(boxes, torch.Tensor):
                boxes = torch.as_tensor(boxes)
                boxes = boxes.tolist()

            scores = prediction["scores"]
            if not isinstance(scores, list):
                scores = scores.tolist()

            labels = prediction["labels"]
            if not isinstance(labels, list):
                labels = prediction["labels"].tolist()

            coco_results.extend(
                [
                    {
                        "image_id": original_id,
                        "category_id": labels[k],
                        "bbox": box,
                        "score": scores[k],
                    }
                    for k, box in enumerate(boxes)
                ]
            )
        return coco_results