evaluation/context_eval.py

import os
import sys
import json
import argparse
import pandas as pd
from tqdm import tqdm 
from pycocotools_ovod.semantic_matching import is_semantic_match, gt_cat_match_path

gt_dataset = None
preds = None
cat_id_to_name = None
WHOLE_DATASET_PATH = './gts/det/semantics/union3_test.json'


def get_obj_name(cat):
    return cat.split('-')[0]


def is_category_interactable(cat):
    if isinstance(cat, str):
        return not cat.endswith('-n')
    elif isinstance(cat, dict):
        return not cat['name'].endswith('-n')
    else:
        raise ValueError("Invalid input type")


def iou(bbox1, bbox2):
    x1, y1, w1, h1 = bbox1
    x2, y2, w2, h2 = bbox2
    union = w1 * h1 + w2 * h2
    inter = max(0, min(x1 + w1, x2 + w2) - max(x1, x2)) * \
        max(0, min(y1 + h1, y2 + h2) - max(y1, y2))
    return inter / (union - inter)


def best_match_gt(pred, anns):
    if len(anns) == 0:
        return 0, None
    best_iou = -1
    best_match = None
    for ann in anns:
        iou_score = iou(pred['bbox'], ann['bbox'])
        if iou_score > best_iou:
            best_iou = iou_score
            best_match = ann
    return best_iou, best_match


def match_cats(gt_cats, preds, eval_dimension):
    if os.path.exists(gt_cat_match_path):
        os.remove(gt_cat_match_path)
    dt_cats = set()
    for pred in preds:
        dt_cats.add(pred['category_id'])
    dt_cats = list(dt_cats)
    dt_cats.sort()
    gt_cats.sort()
    # dt_cats = gt_cats
    gt_cat_match = {gt_cat: [] for gt_cat in gt_cats}
    print('matching dt cats to gt cats...', file=sys.stderr)
    for gt_cat in tqdm(gt_cats):
        for dt_cat in dt_cats:
            if is_semantic_match(gt_cat, dt_cat, eval_dimension=eval_dimension):
                gt_cat_match[gt_cat].append(dt_cat)
        # print(gt_cat, gt_cat_match[gt_cat])
    with open(gt_cat_match_path, 'w') as f:
        json.dump(gt_cat_match, f)


def eval_category(imgs_anns, imgs_preds, iou_threshold):
    global cat_id_to_name
    tp = 0 # predicion matching interactable annotation
    fp = 0 # prediction matching non-interactable annotation
    tn = 0 # not matched non-interactable annotation
    fn = 0 # not matched interactable annotation
    bg = 0 # background
    for img_id in imgs_anns:
        anns_match_flags = {ann['id']: False for ann in imgs_anns[img_id]}
        if img_id in imgs_preds:
            for pred in imgs_preds[img_id]:
                iou_score, best_match_ann = best_match_gt(pred, imgs_anns[img_id])
                if iou_score < iou_threshold:
                    bg += 1
                    continue
                if is_category_interactable(cat_id_to_name[best_match_ann['category_id']]):
                    tp += 1
                else:
                    fp += 1
                anns_match_flags[best_match_ann['id']] = True
        fn += sum([not anns_match_flags[ann['id']] for ann in imgs_anns[img_id]
                  if is_category_interactable(cat_id_to_name[ann['category_id']])])
        tn += sum([not anns_match_flags[ann['id']] for ann in imgs_anns[img_id]
                    if not is_category_interactable(cat_id_to_name[ann['category_id']])])
    return tp, fp, tn, fn, bg


def main(args):
    global gt_dataset, preds, cat_id_to_name

    with open(args.gt, 'r') as f:
        gt_dataset = json.load(f)
    with open(args.pred, 'r') as f:
        preds = json.load(f)
        
    if args.num_cat:
        with open(WHOLE_DATASET_PATH, 'r') as f:
            whole_dataset = json.load(f)
        cat_id_to_name = {cat['id']: cat['name'] for cat in whole_dataset['categories']}
        for pred in preds:
            pred['category_id'] = cat_id_to_name[pred['category_id']]
    

    cat_id_to_name = {cat['id']: cat['name']
                      for cat in gt_dataset['categories']}

    obj_names = [cat['name'] for cat in gt_dataset['categories']
                 if is_category_interactable(cat)]
    
    match_cats(obj_names, preds, args.dimension)

    # find annotations by object name and image id
    objs_imgs_anns = {obj_name: {} for obj_name in obj_names}
    for ann in gt_dataset['annotations']:
        obj_name = get_obj_name(cat_id_to_name[ann['category_id']])
        if ann['image_id'] not in objs_imgs_anns[obj_name]:
            objs_imgs_anns[obj_name][ann['image_id']] = []
        objs_imgs_anns[obj_name][ann['image_id']].append(ann)
        
    for obj_name in obj_names:
        if objs_imgs_anns[obj_name] == {}:
            print(f'No annotation for {obj_name}')
            obj_names.remove(obj_name)

    # find predictions by object name and image id
    objs_imgs_preds = {obj_name: {} for obj_name in obj_names}
    for obj_name in tqdm(obj_names, total=len(obj_names)):
        for pred in preds:
            if is_semantic_match(obj_name, pred['category_id'], eval_dimension=args.dimension):
                if pred['image_id'] not in objs_imgs_preds[obj_name]:
                    objs_imgs_preds[obj_name][pred['image_id']] = []
                objs_imgs_preds[obj_name][pred['image_id']].append(pred)

    result_list = []
    P_avg, R_avg, f1_avg, bgr_avg = 0, 0, 0, 0
    tp_avg, fp_avg, tn_avg, fn_avg, bg_avg = 0, 0, 0, 0, 0

    for obj_name in obj_names:
        tp, fp, tn, fn, bg = eval_category(objs_imgs_anns[obj_name], objs_imgs_preds[obj_name], args.iou)
        tp_avg += tp
        fp_avg += fp
        tn_avg += tn
        fn_avg += fn
        bg_avg += bg
        precision = tp / (tp + fp) if tp + fp > 0 else 0
        P_avg += precision
        recall = tp / (tp + fn) if tp + fn > 0 else 0
        R_avg += recall
        f1 = 2 * precision * recall / (precision + recall) if precision + recall > 0 else 0
        f1_avg += f1
        bg_rate = bg / (tp + fp + bg) if tp + fp + bg > 0 else 0
        bgr_avg += bg_rate
        result_list.append([obj_name, precision, recall, f1, bg_rate, tp, fp, tn, fn, bg])
        
    P_avg /= len(obj_names)
    R_avg /= len(obj_names)
    f1_avg /= len(obj_names)
    bgr_avg /= len(obj_names)
    tp_avg /= len(obj_names)
    fp_avg /= len(obj_names)
    tn_avg /= len(obj_names)
    fn_avg /= len(obj_names)
    bg_avg /= len(obj_names)
    
    result_list.append(['average', P_avg, R_avg, f1_avg, bgr_avg, tp_avg, fp_avg, tn_avg, fn_avg, bg_avg])
        
    df = pd.DataFrame(result_list, columns=['object', 'precision', 'recall', 'f1', 'bg_rate', 'tp', 'fp', 'tn', 'fn', 'bg'])        
    df.to_csv(args.output, index=False)
    
    if os.path.exists(gt_cat_match_path):
        os.remove(gt_cat_match_path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-g', '--gt', type=str, required=True)
    parser.add_argument('-p', '--pred', type=str, required=True)
    parser.add_argument('-o', '--output', type=str)
    parser.add_argument('-d', '--dimension', type=str, default='s')
    parser.add_argument('-n', '--num_cat', action='store_true')
    parser.add_argument('-i', '--iou', type=float)
    args = parser.parse_args()
    main(args)