Detailed explanation of coco_eval.py

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
COCO evaluator that works in distributed mode.

Mostly copy-paste from https://github.com/pytorch/vision/blob/edfd5a7/references/detection/coco_eval.py
The difference is that there is less copy-pasting from pycocotools
in the end of the file, as python3 can suppress prints with contextlib
"""
import os
import contextlib#用于创建和管理上下文管理器。
import copy
import numpy as np
import torch

from pycocotools.cocoeval import COCOeval #用于评估 COCO 数据集的预测结果。
from pycocotools.coco import COCO
import pycocotools.mask as mask_util#用于进行掩码相关的操作。

from util.misc import all_gather#用于在分布式环境中将每个进程中的数据收集到所有进程中。

'''用于评估 COCO 数据集的结果.这个类的目的是简化 COCO 数据集的评估过程。
它在初始化时创建了多个 COCOeval 对象，以支持不同的 iou_types，并提供了一些成员变量来存储评估结果。这样，我们就可以方便地对预测结果进行评估并获取相应的指标。
在类的构造函数 __init__ 中，接收两个参数 coco_gt 和 iou_types。
首先，对 iou_types 进行类型检查，确保它是一个列表或元组。
然后，使用 copy.deepcopy() 函数对 coco_gt 进行深拷贝，将其赋值给 self.coco_gt。这样可以避免对原始的 coco_gt 对象进行修改。
创建一个空字典 self.coco_eval，用于存储不同 iou_type 对应的 COCOeval 对象。
使用 for 循环遍历 iou_types 列表，为每个 iou_type 创建一个 COCOeval 对象，并将其添加到 self.coco_eval 字典中。
初始化 self.img_ids 为空列表，用于存储图像的标识符。
初始化 self.eval_imgs 为一个字典，键为 iou_types 中的每个值，值为一个空列表。它用于存储不同 iou_types 的评估结果。'''
class CocoEvaluator(object):
    def __init__(self, coco_gt, iou_types):
        assert isinstance(iou_types, (list, tuple))
        coco_gt = copy.deepcopy(coco_gt)
        self.coco_gt = coco_gt

        self.iou_types = iou_types
        self.coco_eval = {
    
    }
        for iou_type in iou_types:
            self.coco_eval[iou_type] = COCOeval(coco_gt, iouType=iou_type)

        self.img_ids = []
        self.eval_imgs = {
    
    k: [] for k in iou_types}
    '''首先，从 predictions 字典中提取所有唯一的图像 ID，并将其转换为列表。然后使用 extend 方法将这些图像 ID 添加到 self.img_ids 列表中。
    接下来，使用 for 循环遍历 self.iou_types 列表中的每个 iou_type：
    使用 self.prepare 方法准备预测结果，得到 results，其中 iou_type 参数指定了要使用的 IoU 类型。
    通过创建 COCO 对象或使用 COCO.loadRes 加载预测结果，创建一个 coco_dt 对象。如果没有预测结果，将创建一个空的 COCO 对象。
    获取该 iou_type 对应的 COCOeval 对象 coco_eval。
    将 coco_dt 赋值给 coco_eval.cocoDt，表示要使用这个预测结果进行评估。
    将 img_ids 转换为列表，并将其赋值给 coco_eval.params.imgIds，表示要对哪些图像进行评估。
    调用 evaluate 函数进行评估，获取评估的结果，其中返回值 img_ids 是评估过的图像 ID，eval_imgs 是包含各种评估指标的字典。
    将 eval_imgs 添加到 self.eval_imgs[iou_type] 列表中，保存评估结果。
    通过调用 update 方法，我们可以将预测结果添加到评估器中，并进行相应的评估。'''
    def update(self, predictions):
        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)
            img_ids, eval_imgs = evaluate(coco_eval)

            self.eval_imgs[iou_type].append(eval_imgs)
    '''它首先遍历 self.iou_types 列表中的每个 iou_type。
    
    然后，对于每个 iou_type，它执行以下操作：

    使用 np.concatenate 方法将 self.eval_imgs[iou_type] 列表中的所有元素沿第三个维度（即列维度）进行连接，得到一个新的数组。这样可以将多个进程的评估结果合并成一个大的数组。
    调用 create_common_coco_eval 函数，将合并后的评估结果、图像 ID 和相应的 COCOeval 对象作为参数传递给该函数。这个函数用于创建一个统一的 COCOeval 对象，其中包含合并后的评估结果和相关参数。
    通过调用 synchronize_between_processes 方法，可以将多个进程中的评估结果合并为一个统一的结果对象，并准备进行最终的评估。'''
    def synchronize_between_processes(self):
        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])
    '''它使用一个 for 循环遍历 self.coco_eval.values()，其中 self.coco_eval 是一个字典，键是 iou_type，值是相应的 COCOeval 对象。
    然后，对于每个 coco_eval，它调用 coco_eval.accumulate() 方法，对该 COCOeval 对象进行累积操作。
    这个方法会将每张图像的评估结果加入到总体的评估结果中，以便最终计算出整体的评估指标。    
    通过调用 accumulate 方法，可以对评估器中的评估结果进行累积，从而得到全局的评估结果。'''
    def accumulate(self):
        for coco_eval in self.coco_eval.values():
            coco_eval.accumulate()
    '''它使用一个 for 循环遍历 self.coco_eval.items()，其中 self.coco_eval 是一个字典，键是 iou_type，值是相应的 COCOeval 对象。
    然后，对于每个 iou_type 和相应的 coco_eval，它先打印输出当前正在处理的 IoU 指标：
    接着，调用 coco_eval.summarize() 方法，对该 COCOeval 对象进行总结操作。这个方法会计算并打印出各种评估指标，如准确率、召回率、平均精确度等。
    通过调用 summarize 方法，可以对评估器中的评估结果进行总结，并将结果打印输出，以便查看模型在不同 IoU 指标下的表现情况。'''
    def summarize(self):
        for iou_type, coco_eval in self.coco_eval.items():
            print("IoU metric: {}".format(iou_type))
            coco_eval.summarize()
    '''通过调用 prepare 方法，可以根据不同的 iou_type 准备相应的评估数据，以便后续进行评估操作。'''
    def prepare(self, predictions, iou_type):
        if iou_type == "bbox":
            return self.prepare_for_coco_detection(predictions)
        elif iou_type == "segm":
            return self.prepare_for_coco_segmentation(predictions)
        elif iou_type == "keypoints":
            return self.prepare_for_coco_keypoint(predictions)
        else:
            raise ValueError("Unknown iou type {}".format(iou_type))

    def prepare_for_coco_detection(self, predictions):
        coco_results = []
        for original_id, prediction in predictions.items():
            if len(prediction) == 0:
                continue

            boxes = prediction["boxes"]
            boxes = convert_to_xywh(boxes).tolist()
            scores = prediction["scores"].tolist()
            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
    '''首先，它创建一个空列表 coco_results，用于存储转换后的结果。
    
    然后，对于每个原始图像的预测结果，它首先检查预测结果是否为空，如果为空则跳过继续下一次循环。
    
    接着，从预测结果中提取出边界框 (boxes)、置信度分数 (scores) 和类别标签 (labels)，并将它们转换为 Python 列表格式。
    
    接下来，使用列表解析将每个边界框的相关信息（包括图像ID、类别ID、边界框坐标和置信度分数）添加到 coco_results 列表中。
    
    最后，返回转换后的 coco_results 列表。
    
    通过调用 prepare_for_coco_detection 方法，可以将目标检测模型的预测结果转换为符合 COCO 格式的评估数据，以便后续的评估操作。'''
    def prepare_for_coco_segmentation(self, predictions):
        coco_results = []
        for original_id, prediction in predictions.items():
            if len(prediction) == 0:
                continue

            scores = prediction["scores"]
            labels = prediction["labels"]
            masks = prediction["masks"]

            masks = masks > 0.5

            scores = prediction["scores"].tolist()
            labels = prediction["labels"].tolist()

            rles = [
                mask_util.encode(np.array(mask[0, :, :, np.newaxis], dtype=np.uint8, order="F"))[0]
                for mask in masks
            ]
            for rle in rles:
                rle["counts"] = rle["counts"].decode("utf-8")

            coco_results.extend(
                [
                    {
    
    
                        "image_id": original_id,
                        "category_id": labels[k],
                        "segmentation": rle,
                        "score": scores[k],
                    }
                    for k, rle in enumerate(rles)
                ]
            )
        return coco_results

    def prepare_for_coco_keypoint(self, predictions):
        coco_results = []
        for original_id, prediction in predictions.items():
            if len(prediction) == 0:
                continue

            boxes = prediction["boxes"]
            boxes = convert_to_xywh(boxes).tolist()
            scores = prediction["scores"].tolist()
            labels = prediction["labels"].tolist()
            keypoints = prediction["keypoints"]
            keypoints = keypoints.flatten(start_dim=1).tolist()

            coco_results.extend(
                [
                    {
    
    
                        "image_id": original_id,
                        "category_id": labels[k],
                        'keypoints': keypoint,
                        "score": scores[k],
                    }
                    for k, keypoint in enumerate(keypoints)
                ]
            )
        return coco_results


def convert_to_xywh(boxes):
    xmin, ymin, xmax, ymax = boxes.unbind(1)
    return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=1)


def merge(img_ids, eval_imgs):
    all_img_ids = all_gather(img_ids)
    all_eval_imgs = all_gather(eval_imgs)

    merged_img_ids = []
    for p in all_img_ids:
        merged_img_ids.extend(p)

    merged_eval_imgs = []
    for p in all_eval_imgs:
        merged_eval_imgs.append(p)

    merged_img_ids = np.array(merged_img_ids)
    merged_eval_imgs = np.concatenate(merged_eval_imgs, 2)

    # keep only unique (and in sorted order) images
    merged_img_ids, idx = np.unique(merged_img_ids, return_index=True)
    merged_eval_imgs = merged_eval_imgs[..., idx]

    return merged_img_ids, merged_eval_imgs


def create_common_coco_eval(coco_eval, img_ids, eval_imgs):
    img_ids, eval_imgs = merge(img_ids, eval_imgs)
    img_ids = list(img_ids)
    eval_imgs = list(eval_imgs.flatten())

    coco_eval.evalImgs = eval_imgs
    coco_eval.params.imgIds = img_ids
    coco_eval._paramsEval = copy.deepcopy(coco_eval.params)


#################################################################
# From pycocotools, just removed the prints and fixed
# a Python3 bug about unicode not defined
#################################################################


def evaluate(self):
    '''
    Run per image evaluation on given images and store results (a list of dict) in self.evalImgs
    :return: None
    '''
    # tic = time.time()
    # print('Running per image evaluation...')
    p = self.params
    # add backward compatibility if useSegm is specified in params
    if p.useSegm is not None:
        p.iouType = 'segm' if p.useSegm == 1 else 'bbox'
        print('useSegm (deprecated) is not None. Running {} evaluation'.format(p.iouType))
    # print('Evaluate annotation type *{}*'.format(p.iouType))
    p.imgIds = list(np.unique(p.imgIds))
    if p.useCats:
        p.catIds = list(np.unique(p.catIds))
    p.maxDets = sorted(p.maxDets)
    self.params = p

    self._prepare()
    # loop through images, area range, max detection number
    catIds = p.catIds if p.useCats else [-1]

    if p.iouType == 'segm' or p.iouType == 'bbox':
        computeIoU = self.computeIoU
    elif p.iouType == 'keypoints':
        computeIoU = self.computeOks
    self.ious = {
    
    
        (imgId, catId): computeIoU(imgId, catId)
        for imgId in p.imgIds
        for catId in catIds}
    coco_eval.accumulate()

    evaluateImg = self.evaluateImg
    maxDet = p.maxDets[-1]
    evalImgs = [
        evaluateImg(imgId, catId, areaRng, maxDet)
        for catId in catIds
        for areaRng in p.areaRng
        for imgId in p.imgIds
    ]
    # this is NOT in the pycocotools code, but could be done outside
    evalImgs = np.asarray(evalImgs).reshape(len(catIds), len(p.areaRng), len(p.imgIds))
    self._paramsEval = copy.deepcopy(self.params)
    # toc = time.time()
    # print('DONE (t={:0.2f}s).'.format(toc-tic))
    return p.imgIds, evalImgs

#################################################################
# end of straight copy from pycocotools, just removing the prints
#################################################################