cocoeval

import numpy as np
from pycocotools.cocoeval import COCOeval


class CustomCOCOeval(COCOeval):
    def __init__(self, cocoGt=None, cocoDt=None, iouType='keypoints', kpt_oks_sigmas=None):
        super(CustomCOCOeval, self).__init__(cocoGt, cocoDt, iouType)
        if kpt_oks_sigmas is not None:
            self.params.kpt_oks_sigmas = kpt_oks_sigmas

    def computeIoU(self, imgId, catId):
        p = self.params
        if p.useCats:
            gt = self._gts[imgId, catId]
            dt = self._dts[imgId, catId]
        else:
            gt = [g for cId in p.catIds for g in self._gts[imgId, cId]]
            dt = [d for cId in p.catIds for d in self._dts[imgId, cId]]
        if len(gt) == 0 and len(dt) == 0:
            return []
        ious = np.zeros((len(dt), len(gt)))

        for j, gt_obj in enumerate(gt):
            gx, gy, gwidth, gheight = gt_obj['bbox']
            gx2, gy2 = gx + gwidth, gy + gheight
            gt_area = gwidth * gheight
            for i, dt_obj in enumerate(dt):
                dx, dy, dwidth, dheight = dt_obj['bbox']
                dx2, dy2 = dx + dwidth, dy + dheight
                dt_area = dwidth * dheight

                inter_w = min(dx2, gx2) - max(dx, gx)
                inter_h = min(dy2, gy2) - max(dy, gy)
                if inter_w <= 0 or inter_h <= 0:
                    continue

                inter_area = inter_w * inter_h
                union_area = dt_area + gt_area - inter_area
                ious[i, j] = inter_area / union_area
        return ious

    def computeOks(self, imgId, catId):
        p = self.params
        gts = self._gts[imgId, catId]
        dts = self._dts[imgId, catId]
        inds = np.argsort([-d['score'] for d in dts], kind='mergesort')
        dts = [dts[i] for i in inds]
        if len(dts) > p.maxDets[-1]:
            dts = dts[0:p.maxDets[-1]]
        if len(gts) == 0 or len(dts) == 0:
            return []
        ious = np.zeros((len(dts), len(gts)))
        sigmas = p.kpt_oks_sigmas
        vars = (sigmas * 2) ** 2
        k = len(sigmas)
        for j, gt in enumerate(gts):
            g = np.array(gt['keypoints'])
            xg = g[0::3];
            yg = g[1::3];
            vg = g[2::3]
            k1 = np.count_nonzero(vg > 0)
            bb = gt['bbox']
            x0 = bb[0] - bb[2];
            x1 = bb[0] + bb[2] * 2
            y0 = bb[1] - bb[3];
            y1 = bb[1] + bb[3] * 2
            for i, dt in enumerate(dts):
                d = np.array(dt['keypoints'])
                xd = d[0::3];
                yd = d[1::3]
                if k1 > 0:
                    dx = xd - xg
                    dy = yd - yg
                else:
                    z = np.zeros((k))
                    dx = np.max((z, x0 - xd), axis=0) + np.max((z, xd - x1), axis=0)
                    dy = np.max((z, y0 - yd), axis=0) + np.max((z, yd - y1), axis=0)
                e = (dx ** 2 + dy ** 2) / vars / (gt['area'] + np.spacing(1)) / 2
                if k1 > 0:
                    e = e[vg > 0]
                ious[i, j] = np.sum(np.exp(-e)) / e.shape[0]
        return ious


def evaluate_model(coco_gt, coco_dt, kpt_oks_sigmas=None):
    coco_eval = CustomCOCOeval(coco_gt, coco_dt, 'keypoints', kpt_oks_sigmas=kpt_oks_sigmas)
    coco_eval.evaluate()
    coco_eval.accumulate()
    coco_eval.summarize()
    return coco_eval.stats