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
#################################################################

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298