|
|
| import json |
| import argparse |
| import pandas as pd |
| from pycocotools.coco import COCO |
| import pycocotools.mask as maskUtils |
|
|
| gt_cat_match_path = 'tmp_gt_cat_match.json' |
| tmp_ann_path = 'tmp_ann.json' |
|
|
|
|
| def do_evaluate(args): |
| coco_gt = COCO(args.gt,) |
| coco_dt = coco_gt.loadRes(args.dt,) |
| img_list = coco_gt.getImgIds() |
| gt_masks = {} |
| dt_masks = {} |
| for ann in coco_gt.dataset['annotations']: |
| seg = ann['segmentation'] |
| RLEs = maskUtils.frPyObjects(seg, 540, 960) |
| RLE = maskUtils.merge(RLEs) |
| if ann['image_id'] not in gt_masks: |
| gt_masks[ann['image_id']] = RLE |
| else: |
| gt_masks[ann['image_id']] = maskUtils.merge([gt_masks[ann['image_id']], RLE]) |
| for ann in coco_dt.dataset['annotations']: |
| seg = ann['segmentation'] |
| if type(seg['counts']) != str: |
| RLEs = maskUtils.frPyObjects(seg, 540, 960) |
| RLE = maskUtils.merge(RLEs) |
| else: |
| RLE = seg |
| if ann['image_id'] not in dt_masks: |
| dt_masks[ann['image_id']] = RLE |
| else: |
| dt_masks[ann['image_id']] = maskUtils.merge([dt_masks[ann['image_id']], RLE]) |
| precisions = {} |
| recalls = {} |
| for img_id in img_list: |
| if img_id not in gt_masks or img_id not in dt_masks: |
| precisions[img_id] = 0 |
| recalls[img_id] = 0 |
| continue |
| tp_mask = maskUtils.merge([gt_masks[img_id], dt_masks[img_id]], intersect=True) |
| gt_area = maskUtils.area(gt_masks[img_id]) |
| dt_area = maskUtils.area(dt_masks[img_id]) |
| tp_area = maskUtils.area(tp_mask) |
| fp_area = dt_area - tp_area |
| fn_area = gt_area - tp_area |
| precisions[img_id] = tp_area / (tp_area + fp_area + 1e-8) |
| recalls[img_id] = tp_area / (tp_area + fn_area + 1e-8) |
|
|
| return precisions, recalls |
|
|
| def id2name(args): |
| with open(args.gt, 'r') as f: |
| gt = json.load(f) |
| with open(args.dt, 'r') as f: |
| dt = json.load(f) |
| cat_id_name = {cat['id']: cat['name'] for cat in gt['categories']} |
| for res in dt: |
| res['category_id'] = cat_id_name[res['category_id']] |
| with open(tmp_ann_path, 'w') as f: |
| json.dump(dt, f) |
| args.dt = tmp_ann_path |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description="Evaluate Metrics from the predictions and Ground Truths") |
| parser.add_argument('-gt', '--gt', type=str, help='path to ground truth json', required=True) |
| parser.add_argument('-dt', '--dt', type=str, help='path to detection json', required=True) |
| parser.add_argument('-l', '--log', type=str, default="evaluation.log") |
| parser.add_argument('-n', '--name_id', action="store_true", help="Change category id to the corresponding name") |
| args = parser.parse_args() |
|
|
| if args.name_id: |
| id2name(args) |
| precisions, recalls = do_evaluate(args) |
| precision = sum(precisions.values()) / len(precisions) |
| recall = sum(recalls.values()) / len(recalls) |
| f1 = 2 * precision * recall / (precision + recall + 1e-8) |
| result = {'precision': precision, 'recall': recall, 'f1': f1} |
| pd.DataFrame(result, index=[0]).to_csv(args.log, index=False) |
|
|
