Rk6S0sm1

1import logging
2import tempfile
3import os
4import torch
5from collections import OrderedDict
6from tqdm import tqdm
7
8from maskrcnn_benchmark.modeling.roi_heads.mask_head.inference import Masker
9from maskrcnn_benchmark.structures.bounding_box import BoxList
10from maskrcnn_benchmark.structures.boxlist_ops import boxlist_iou
11
12
13def do_coco_evaluation(
14    dataset,
15    predictions,
16    box_only,
17    output_folder,
18    iou_types,
19    expected_results,
20    expected_results_sigma_tol,
21):
22    logger = logging.getLogger("maskrcnn_benchmark.inference")
23
24    if box_only:
25        logger.info("Evaluating bbox proposals")
26        areas = {"all": "", "small": "s", "medium": "m", "large": "l"}
27        res = COCOResults("box_proposal")
28        for limit in [100, 1000]:
29            for area, suffix in areas.items():
30                stats = evaluate_box_proposals(
31                    predictions, dataset, area=area, limit=limit
32                )
33                key = "AR{}@{:d}".format(suffix, limit)
34                res.results["box_proposal"][key] = stats["ar"].item()
35        logger.info(res)
36        check_expected_results(res, expected_results, expected_results_sigma_tol)
37        if output_folder:
38            torch.save(res, os.path.join(output_folder, "box_proposals.pth"))
39        return
40    logger.info("Preparing results for COCO format")
41    coco_results = {}
42    if "bbox" in iou_types:
43        logger.info("Preparing bbox results")
44        coco_results["bbox"] = prepare_for_coco_detection(predictions, dataset)
45    if "segm" in iou_types:
46        logger.info("Preparing segm results")
47        coco_results["segm"] = prepare_for_coco_segmentation(predictions, dataset)
48    if 'keypoints' in iou_types:
49        logger.info('Preparing keypoints results')
50        coco_results['keypoints'] = prepare_for_coco_keypoint(predictions, dataset)
51
52    results = COCOResults(*iou_types)
53    logger.info("Evaluating predictions")
54    for iou_type in iou_types:
55        with tempfile.NamedTemporaryFile() as f:
56            file_path = f.name
57            if output_folder:
58                file_path = os.path.join(output_folder, iou_type + ".json")
59            res = evaluate_predictions_on_coco(
60                dataset.coco, coco_results[iou_type], file_path, iou_type
61            )
62            results.update(res)
63    logger.info(results)
64    check_expected_results(results, expected_results, expected_results_sigma_tol)
65    if output_folder:
66        torch.save(results, os.path.join(output_folder, "coco_results.pth"))
67    return results, coco_results
68
69
70def prepare_for_coco_detection(predictions, dataset):
71    # assert isinstance(dataset, COCODataset)
72    coco_results = []
73    for image_id, prediction in enumerate(predictions):
74        original_id = dataset.id_to_img_map[image_id]
75        if len(prediction) == 0:
76            continue
77
78        img_info = dataset.get_img_info(image_id)
79        image_width = img_info["width"]
80        image_height = img_info["height"]
81        prediction = prediction.resize((image_width, image_height))
82        prediction = prediction.convert("xywh")
83
84        boxes = prediction.bbox.tolist()
85        scores = prediction.get_field("scores").tolist()
86        labels = prediction.get_field("labels").tolist()
87
88        mapped_labels = [dataset.contiguous_category_id_to_json_id[i] for i in labels]
89
90        coco_results.extend(
91            [
92                {
93                    "image_id": original_id,
94                    "category_id": mapped_labels[k],
95                    "bbox": box,
96                    "score": scores[k],
97                }
98                for k, box in enumerate(boxes)
99            ]
100        )
101    return coco_results
102
103
104def prepare_for_coco_segmentation(predictions, dataset):
105    import pycocotools.mask as mask_util
106    import numpy as np
107
108    masker = Masker(threshold=0.5, padding=1)
109    # assert isinstance(dataset, COCODataset)
110    coco_results = []
111    for image_id, prediction in tqdm(enumerate(predictions)):
112        original_id = dataset.id_to_img_map[image_id]
113        if len(prediction) == 0:
114            continue
115
116        img_info = dataset.get_img_info(image_id)
117        image_width = img_info["width"]
118        image_height = img_info["height"]
119        prediction = prediction.resize((image_width, image_height))
120        masks = prediction.get_field("mask")
121        # t = time.time()
122        # Masker is necessary only if masks haven't been already resized.
123        if list(masks.shape[-2:]) != [image_height, image_width]:
124            masks = masker(masks.expand(1, -1, -1, -1, -1), prediction)
125            masks = masks[0]
126        # logger.info('Time mask: {}'.format(time.time() - t))
127        # prediction = prediction.convert('xywh')
128
129        # boxes = prediction.bbox.tolist()
130        scores = prediction.get_field("scores").tolist()
131        labels = prediction.get_field("labels").tolist()
132
133        # rles = prediction.get_field('mask')
134
135        rles = [
136            mask_util.encode(np.array(mask[0, :, :, np.newaxis], order="F"))[0]
137            for mask in masks
138        ]
139        for rle in rles:
140            rle["counts"] = rle["counts"].decode("utf-8")
141
142        mapped_labels = [dataset.contiguous_category_id_to_json_id[i] for i in labels]
143
144        coco_results.extend(
145            [
146                {
147                    "image_id": original_id,
148                    "category_id": mapped_labels[k],
149                    "segmentation": rle,
150                    "score": scores[k],
151                }
152                for k, rle in enumerate(rles)
153            ]
154        )
155    return coco_results
156
157
158def prepare_for_coco_keypoint(predictions, dataset):
159    # assert isinstance(dataset, COCODataset)
160    coco_results = []
161    for image_id, prediction in enumerate(predictions):
162        original_id = dataset.id_to_img_map[image_id]
163        if len(prediction.bbox) == 0:
164            continue
165
166        # TODO replace with get_img_info?
167        image_width = dataset.coco.imgs[original_id]['width']
168        image_height = dataset.coco.imgs[original_id]['height']
169        prediction = prediction.resize((image_width, image_height))
170        prediction = prediction.convert('xywh')
171
172        boxes = prediction.bbox.tolist()
173        scores = prediction.get_field('scores').tolist()
174        labels = prediction.get_field('labels').tolist()
175        keypoints = prediction.get_field('keypoints')
176        keypoints = keypoints.resize((image_width, image_height))
177        keypoints = keypoints.keypoints.view(keypoints.keypoints.shape[0], -1).tolist()
178
179        mapped_labels = [dataset.contiguous_category_id_to_json_id[i] for i in labels]
180
181        coco_results.extend([{
182            'image_id': original_id,
183            'category_id': mapped_labels[k],
184            'keypoints': keypoint,
185            'score': scores[k]} for k, keypoint in enumerate(keypoints)])
186    return coco_results
187
188# inspired from Detectron
189def evaluate_box_proposals(
190    predictions, dataset, thresholds=None, area="all", limit=None
191):
192    """Evaluate detection proposal recall metrics. This function is a much
193    faster alternative to the official COCO API recall evaluation code. However,
194    it produces slightly different results.
195    """
196    # Record max overlap value for each gt box
197    # Return vector of overlap values
198    areas = {
199        "all": 0,
200        "small": 1,
201        "medium": 2,
202        "large": 3,
203        "96-128": 4,
204        "128-256": 5,
205        "256-512": 6,
206        "512-inf": 7,
207    }
208    area_ranges = [
209        [0 ** 2, 1e5 ** 2],  # all
210        [0 ** 2, 32 ** 2],  # small
211        [32 ** 2, 96 ** 2],  # medium
212        [96 ** 2, 1e5 ** 2],  # large
213        [96 ** 2, 128 ** 2],  # 96-128
214        [128 ** 2, 256 ** 2],  # 128-256
215        [256 ** 2, 512 ** 2],  # 256-512
216        [512 ** 2, 1e5 ** 2],
217    ]  # 512-inf
218    assert area in areas, "Unknown area range: {}".format(area)
219    area_range = area_ranges[areas[area]]
220    gt_overlaps = []
221    num_pos = 0
222
223    for image_id, prediction in enumerate(predictions):
224        original_id = dataset.id_to_img_map[image_id]
225
226        img_info = dataset.get_img_info(image_id)
227        image_width = img_info["width"]
228        image_height = img_info["height"]
229        prediction = prediction.resize((image_width, image_height))
230
231        # sort predictions in descending order
232        # TODO maybe remove this and make it explicit in the documentation
233        inds = prediction.get_field("objectness").sort(descending=True)[1]
234        prediction = prediction[inds]
235
236        ann_ids = dataset.coco.getAnnIds(imgIds=original_id)
237        anno = dataset.coco.loadAnns(ann_ids)
238        gt_boxes = [obj["bbox"] for obj in anno if obj["iscrowd"] == 0]
239        gt_boxes = torch.as_tensor(gt_boxes).reshape(-1, 4)  # guard against no boxes
240        gt_boxes = BoxList(gt_boxes, (image_width, image_height), mode="xywh").convert(
241            "xyxy"
242        )
243        gt_areas = torch.as_tensor([obj["area"] for obj in anno if obj["iscrowd"] == 0])
244
245        if len(gt_boxes) == 0:
246            continue
247
248        valid_gt_inds = (gt_areas >= area_range[0]) & (gt_areas <= area_range[1])
249        gt_boxes = gt_boxes[valid_gt_inds]
250
251        num_pos += len(gt_boxes)
252
253        if len(gt_boxes) == 0:
254            continue
255
256        if len(prediction) == 0:
257            continue
258
259        if limit is not None and len(prediction) > limit:
260            prediction = prediction[:limit]
261
262        overlaps = boxlist_iou(prediction, gt_boxes)
263
264        _gt_overlaps = torch.zeros(len(gt_boxes))
265        for j in range(min(len(prediction), len(gt_boxes))):
266            # find which proposal box maximally covers each gt box
267            # and get the iou amount of coverage for each gt box
268            max_overlaps, argmax_overlaps = overlaps.max(dim=0)
269
270            # find which gt box is 'best' covered (i.e. 'best' = most iou)
271            gt_ovr, gt_ind = max_overlaps.max(dim=0)
272            assert gt_ovr >= 0
273            # find the proposal box that covers the best covered gt box
274            box_ind = argmax_overlaps[gt_ind]
275            # record the iou coverage of this gt box
276            _gt_overlaps[j] = overlaps[box_ind, gt_ind]
277            assert _gt_overlaps[j] == gt_ovr
278            # mark the proposal box and the gt box as used
279            overlaps[box_ind, :] = -1
280            overlaps[:, gt_ind] = -1
281
282        # append recorded iou coverage level
283        gt_overlaps.append(_gt_overlaps)
284    gt_overlaps = torch.cat(gt_overlaps, dim=0)
285    gt_overlaps, _ = torch.sort(gt_overlaps)
286
287    if thresholds is None:
288        step = 0.05
289        thresholds = torch.arange(0.5, 0.95 + 1e-5, step, dtype=torch.float32)
290    recalls = torch.zeros_like(thresholds)
291    # compute recall for each iou threshold
292    for i, t in enumerate(thresholds):
293        recalls[i] = (gt_overlaps >= t).float().sum() / float(num_pos)
294    # ar = 2 * np.trapz(recalls, thresholds)
295    ar = recalls.mean()
296    return {
297        "ar": ar,
298        "recalls": recalls,
299        "thresholds": thresholds,
300        "gt_overlaps": gt_overlaps,
301        "num_pos": num_pos,
302    }
303
304
305def evaluate_predictions_on_coco(
306    coco_gt, coco_results, json_result_file, iou_type="bbox"
307):
308    import json
309
310    with open(json_result_file, "w") as f:
311        json.dump(coco_results, f)
312
313    from pycocotools.coco import COCO
314    from pycocotools.cocoeval import COCOeval
315
316    coco_dt = coco_gt.loadRes(str(json_result_file)) if coco_results else COCO()
317
318    # coco_dt = coco_gt.loadRes(coco_results)
319    coco_eval = COCOeval(coco_gt, coco_dt, iou_type)
320    coco_eval.evaluate()
321    coco_eval.accumulate()
322    coco_eval.summarize()
323    return coco_eval
324
325
326class COCOResults(object):
327    METRICS = {
328        "bbox": ["AP", "AP50", "AP75", "APs", "APm", "APl"],
329        "segm": ["AP", "AP50", "AP75", "APs", "APm", "APl"],
330        "box_proposal": [
331            "AR@100",
332            "ARs@100",
333            "ARm@100",
334            "ARl@100",
335            "AR@1000",
336            "ARs@1000",
337            "ARm@1000",
338            "ARl@1000",
339        ],
340        "keypoints": ["AP", "AP50", "AP75", "APm", "APl"],
341    }
342
343    def __init__(self, *iou_types):
344        allowed_types = ("box_proposal", "bbox", "segm", "keypoints")
345        assert all(iou_type in allowed_types for iou_type in iou_types)
346        results = OrderedDict()
347        for iou_type in iou_types:
348            results[iou_type] = OrderedDict(
349                [(metric, -1) for metric in COCOResults.METRICS[iou_type]]
350            )
351        self.results = results
352
353    def update(self, coco_eval):
354        if coco_eval is None:
355            return
356        from pycocotools.cocoeval import COCOeval
357
358        assert isinstance(coco_eval, COCOeval)
359        s = coco_eval.stats
360        iou_type = coco_eval.params.iouType
361        res = self.results[iou_type]
362        metrics = COCOResults.METRICS[iou_type]
363        for idx, metric in enumerate(metrics):
364            res[metric] = s[idx]
365
366    def __repr__(self):
367        results = '\n'
368        for task, metrics in self.results.items():
369            results += 'Task: {}\n'.format(task)
370            metric_names = metrics.keys()
371            metric_vals = ['{:.4f}'.format(v) for v in metrics.values()]
372            results += (', '.join(metric_names) + '\n')
373            results += (', '.join(metric_vals) + '\n')
374        return results
375
376
377def check_expected_results(results, expected_results, sigma_tol):
378    if not expected_results:
379        return
380
381    logger = logging.getLogger("maskrcnn_benchmark.inference")
382    for task, metric, (mean, std) in expected_results:
383        actual_val = results.results[task][metric]
384        lo = mean - sigma_tol * std
385        hi = mean + sigma_tol * std
386        ok = (lo < actual_val) and (actual_val < hi)
387        msg = (
388            "{} > {} sanity check (actual vs. expected): "
389            "{:.3f} vs. mean={:.4f}, std={:.4}, range=({:.4f}, {:.4f})"
390        ).format(task, metric, actual_val, mean, std, lo, hi)
391        if not ok:
392            msg = "FAIL: " + msg
393            logger.error(msg)
394        else:
395            msg = "PASS: " + msg
396            logger.info(msg)