Visualization evaluation result(s)

[duongminhhieu16]

I hope these lines work for someone to visualize image (result from inference/evaluation)

import json

from pycocotools.coco import COCO
import os
from PIL import Image
from matplotlib import pyplot as plt
import numpy as np

def getAnnId(imgId, results):
    """
    Get annotation for specific id
    :param  imgId   (int)    : id of specific image
            results (list)   : list of result dictionaries
    :return anns    (list)   : list of annotations for imgId
    """
    
    values = [result["image_id"] for result in results]
    ann_indices = np.where(np.array(values) == imgId)[0]

    anns = [results[ann_idx] for ann_idx in ann_indices]
    return anns

def visualize_result(db_path=None, annotation_path=None, result_path=None, img_id=None, save_dir=None):
    """
    Visualize result of specific image id
    :param  db_path         (str)    : database path
            annotation_path (list)   : list of result dictionaries
            result_path     (str)    : evaluation result path
            img_id          (int)    : image id
            save_dir       (str)     : save image to directory
    :return anns    (list)   : list of annotations for imgId
    """
    # coco_annotation_file_path = annotation_path
    # db = db_path
    
    # check if path exists
    if db_path == None or os.path.exists(db_path) is False:
        print("Database path doesn't exist. Database path: ", db_path)
        return None
    if annotation_path == None or os.path.exists(annotation_path) is False:
        print("Annotation path doesn't exist. Annotation path: ", annotation_path)
        return None
    if result_path == None or os.path.exists(result_path) is False:
        print("Result path doesn't exist. Result path: ", result_path)
        return None
    
    coco_annotation = COCO(annotation_file=annotation_path)

    # read evaluation result
    with open(result_path, "r") as f:
        # each image has image_id. category_id, keypoints, score, bbox
        results = json.load(f)
    
        # img_id = 87038 #results[12]["image_id"]
        # print(img_id)
        # load image by id
        img_info = coco_annotation.loadImgs(img_id)
        img_path = os.path.join(db_path, img_info[0]['file_name'])
        img = Image.open(img_path)
        plt.imshow(np.asarray(img))
        
        # get all annotations in coco dataset
        # ann_ids = coco_annotation.getAnnIds(imgIds=img_id, iscrowd=None)
        # anns = coco_annotation.loadAnns(ann_ids)

        # get annotations for specific id
        anns = getAnnId(imgId=img_id, results=results)
        
        # draw annotations
        coco_annotation.showAnns(anns, draw_bbox=True)
        
        # save result
        if os.path.exists(save_dir) is False:
            os.makedirs(save_dir)
        path = os.path.join(save_dir, img_info[0]['file_name'])
        plt.savefig(path)
        print("Save image to: ", path)
        
        # show result
        plt.show()