Visualize coco detection output with fp tp missing boxes

visualize_coco.py

""" only part codes """

# read groundtruth and predictions
dataDir= Path('datasets/coco')
valImageRoot = dataDir / 'val2017'
groudtruth_path = dataDir / 'annotations/instances_val2017.json'
prediction_path = 'output/inference/coco_instances_results.json'

cocoGt=COCO(groudtruth_path)
cocoDt = cocoGt.loadRes(prediction_path)
imgIds=sorted(cocoGt.getImgIds())

# some function definition

def draw_boxes(image, boxes, labels, scores, color):
    image_draw = copy.deepcopy(image)
    draw = ImageDraw.Draw(image_draw)
    color = ImageColor.getrgb(color)
    font_path = '/usr/share/fonts/truetype/ubuntu-font-family/UbuntuMono-R.ttf'
    for i in range(len(boxes)):
        box = boxes[i]
        draw.rectangle([box[0], box[1], box[2], box[3]], outline=color, width=2)
        text = ''
        font_obj = ImageFont.truetype(font_path, size=14)
        if labels is not None:
            text += labels[i]
        if scores is not None:
            text += f':{int(round(scores[i]*100))}%'
        if text:
            draw.text(
                (box[0], box[1]), text=text,
                fill=color, font=font_obj
            )
    return image_draw

# Visual a single image
imgId = random.choice(imgIds)
# extract predictions and groundtruth
dt = get_image_data(imgId, cocoDt)
if len(dt.boxes) == 0:
    dt_boxes = np.array([0,0,0,0])
    dt_classes = np.array([0])
    dt_scores = np.array([0])
else:
    dt_boxes = xywh2xyxy(np.array(dt.boxes))
    dt_classes = np.array(dt.classes)
    dt_scores = np.array(dt.scores)
dt_labels = cocoDt.loadCats(ids=dt_classes)
dt_labels = np.array([x['name'] for x in dt_labels])
dt_mask = dt_scores > 0.5
dt = edict({
    'boxes': dt_boxes[dt_mask],
    'labels': dt_labels[dt_mask],
    'scores': dt_scores[dt_mask],
})

gt = get_image_data(imgId, cocoGt)
gt_boxes = xywh2xyxy(np.array(gt.boxes))
gt_classes = np.array(gt.classes)
gt_scores = np.array(gt.scores)
gt_labels = cocoDt.loadCats(ids=gt_classes)
gt_labels = np.array([x['name'] for x in gt_labels])
gt = edict({
    'boxes': gt_boxes,
    'labels': gt_labels,
    'scores': None
})

# calculate colors of each boxes:
ious = pairwise_iou(torch.tensor(dt.boxes), torch.tensor(gt.boxes))
ious = ious.numpy()
# TP
maxv = np.amax(ious, axis=1)
ix, iy = np.nonzero(ious > 0.5)
tp_lst = []
fp_lst = []
for i in range(len(dt.boxes)):
    matched = iy[ix==i]
    matched_labels = gt.labels[matched].tolist()
    if dt.labels[i] in matched_labels:
        tp_lst.append(i)
    else:
        fp_lst.append(i)
tp_boxes = edict({
    'boxes': dt.boxes[tp_lst],
    'labels': dt.labels[tp_lst],
    'scores': dt.scores[tp_lst]
})
fp_boxes = edict({
    'boxes': dt.boxes[fp_lst],
    'scores': dt.scores[fp_lst],
    'labels': dt.labels[fp_lst]
})
# missing
ms_lst = []
for i in range(len(gt.boxes)):
    matched = ix[iy==i]
    matched_labels = dt.labels[matched].tolist()
    if gt.labels[i] not in matched_labels:
        ms_lst.append(i)
ms_boxes = edict({
    'boxes': gt.boxes[ms_lst],
    'labels': gt.labels[ms_lst],
    'scores': None
})

image = cocoGt.loadImgs([imgId])[0]

image_path = valImageRoot / image['file_name']
image = Image.open(image_path)
# TP
image_boxes = draw_boxes(image, tp_boxes.boxes, tp_boxes.labels, tp_boxes.scores, 'green')
image_boxes = draw_boxes(image_boxes, fp_boxes.boxes, fp_boxes.labels, fp_boxes.scores, 'red')
image_boxes = draw_boxes(image_boxes, ms_boxes.boxes, ms_boxes.labels, None, 'yellow')

image_boxes