"""
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
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)
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)
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)
def evaluate(self):
'''
Run per image evaluation on given images and store results (a list of dict) in self.evalImgs
:return: None
'''
p = self.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))
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()
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
]
evalImgs = np.asarray(evalImgs).reshape(len(catIds), len(p.areaRng), len(p.imgIds))
self._paramsEval = copy.deepcopy(self.params)
return p.imgIds, evalImgs
