|
|
"""
|
|
|
# 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 contextlib
|
|
|
import copy
|
|
|
import os
|
|
|
|
|
|
import faster_coco_eval.core.mask as mask_util
|
|
|
import numpy as np
|
|
|
import torch
|
|
|
from faster_coco_eval import COCO, COCOeval_faster
|
|
|
|
|
|
from ...core import register
|
|
|
from ...misc import dist_utils
|
|
|
|
|
|
__all__ = [
|
|
|
"CocoEvaluator",
|
|
|
]
|
|
|
|
|
|
|
|
|
@register()
|
|
|
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 = coco_gt
|
|
|
self.iou_types = iou_types
|
|
|
|
|
|
self.coco_eval = {}
|
|
|
for iou_type in iou_types:
|
|
|
self.coco_eval[iou_type] = COCOeval_faster(
|
|
|
coco_gt, iouType=iou_type, print_function=print, separate_eval=True
|
|
|
)
|
|
|
|
|
|
self.img_ids = []
|
|
|
self.eval_imgs = {k: [] for k in iou_types}
|
|
|
|
|
|
def cleanup(self):
|
|
|
self.coco_eval = {}
|
|
|
for iou_type in self.iou_types:
|
|
|
self.coco_eval[iou_type] = COCOeval_faster(
|
|
|
self.coco_gt, iouType=iou_type, print_function=print, separate_eval=True
|
|
|
)
|
|
|
self.img_ids = []
|
|
|
self.eval_imgs = {k: [] for k in self.iou_types}
|
|
|
|
|
|
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)
|
|
|
coco_eval = self.coco_eval[iou_type]
|
|
|
|
|
|
|
|
|
with open(os.devnull, "w") as devnull:
|
|
|
with contextlib.redirect_stdout(devnull):
|
|
|
coco_dt = self.coco_gt.loadRes(results) if results else COCO()
|
|
|
coco_eval.cocoDt = coco_dt
|
|
|
coco_eval.params.imgIds = list(img_ids)
|
|
|
coco_eval.evaluate()
|
|
|
|
|
|
self.eval_imgs[iou_type].append(
|
|
|
np.array(coco_eval._evalImgs_cpp).reshape(
|
|
|
len(coco_eval.params.catIds),
|
|
|
len(coco_eval.params.areaRng),
|
|
|
len(coco_eval.params.imgIds),
|
|
|
)
|
|
|
)
|
|
|
|
|
|
def synchronize_between_processes(self):
|
|
|
for iou_type in self.iou_types:
|
|
|
img_ids, eval_imgs = merge(self.img_ids, self.eval_imgs[iou_type])
|
|
|
|
|
|
coco_eval = self.coco_eval[iou_type]
|
|
|
coco_eval.params.imgIds = img_ids
|
|
|
coco_eval._paramsEval = copy.deepcopy(coco_eval.params)
|
|
|
coco_eval._evalImgs_cpp = eval_imgs
|
|
|
|
|
|
def accumulate(self):
|
|
|
for coco_eval in self.coco_eval.values():
|
|
|
coco_eval.accumulate()
|
|
|
|
|
|
def summarize(self):
|
|
|
for iou_type, coco_eval in self.coco_eval.items():
|
|
|
print("IoU metric: {}".format(iou_type))
|
|
|
coco_eval.summarize()
|
|
|
|
|
|
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
|
|
|
|
|
|
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 = dist_utils.all_gather(img_ids)
|
|
|
all_eval_imgs = dist_utils.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.extend(p)
|
|
|
|
|
|
merged_img_ids = np.array(merged_img_ids)
|
|
|
merged_eval_imgs = np.concatenate(merged_eval_imgs, axis=2).ravel()
|
|
|
|
|
|
|
|
|
|
|
|
merged_img_ids, idx = np.unique(merged_img_ids, return_index=True)
|
|
|
|
|
|
return merged_img_ids.tolist(), merged_eval_imgs.tolist()
|
|
|
|
