File size: 3,170 Bytes
6d35aff | 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 |
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)
|