# AUTOGENERATED! DO NOT EDIT! File to edit: ../metrics.ipynb.

# %% auto 0
__all__ = ['PCK', 'get_visible_part_mean_absolute_reprojection_error', 'get_mAA', 'fraction_of_gt_corrs_consisent_with_F']

# %% ../metrics.ipynb 3
import numpy as np

def PCK(pts: np.array, gt_pts: np.array, ths = np.arange(20)):
    '''Function to calculate  Probability of Correct Keypoint (PCK) given the error threshold'''
    assert pts.shape[0] == gt_pts.shape[0]
    assert pts.shape[1] == 2
    diff  = np.linalg.norm(pts - gt_pts, ord=2, axis=1)
    out = []
    for th in ths:
        correct = diff <= th
        out.append(correct.astype(np.float32).mean())
    return np.array(out)

# %% ../metrics.ipynb 4
import numpy as np
import cv2

def get_visible_part_mean_absolute_reprojection_error(img1_shape, img2_shape, H_gt, H):
    '''We reproject the image 1 mask to image2 and back to get the visible part mask.
    Then we average the reprojection absolute error over that area'''
    h,w = img1_shape[:2]
    mask1 = np.ones((h,w))
    mask1in2 = cv2.warpPerspective(mask1, H_gt, img2_shape[::-1])
    mask1inback = cv2.warpPerspective(mask1in2, np.linalg.inv(H_gt), img1_shape[:2][::-1]) > 0
    xi = np.arange(w)
    yi = np.arange(h)
    xg, yg = np.meshgrid(xi,yi)
    coords = np.concatenate([xg.reshape(*xg.shape,1), yg.reshape(*yg.shape,1)], axis=-1)
    shape_orig = coords.shape
    xy_rep_gt = cv2.perspectiveTransform(coords.reshape(-1, 1,2).astype(np.float32), H_gt.astype(np.float32)).squeeze(1)
    xy_rep_estimated = cv2.perspectiveTransform(coords.reshape(-1, 1,2).astype(np.float32),
                                                H.astype(np.float32)).squeeze(1)
    #error = np.abs(xy_rep_gt-xy_rep_estimated).sum(axis=1).reshape(xg.shape) * mask1inback
    error = np.sqrt(((xy_rep_gt-xy_rep_estimated)**2).sum(axis=1)).reshape(xg.shape) * mask1inback
    mean_error = error.sum() / mask1inback.sum()
    return mean_error


def get_mAA(mAE: np.array, ths=np.logspace(np.log2(1.0), np.log2(20), 10, base=2.0)):
    out = []
    for th in ths:
        correct = mAE <= th
        out.append(correct.astype(np.float32).mean())
    return out

# %% ../metrics.ipynb 8
import kornia.geometry as KG
import torch
def fraction_of_gt_corrs_consisent_with_F(Fm: np.array, gt_corrs: np.array, ths=np.arange(20)):
    assert Fm.shape == (3, 3)
    assert gt_corrs.shape[1] == 4
    pts_torch = torch.from_numpy(gt_corrs)[None]
    F_torch = torch.from_numpy(Fm)[None]
    errors = KG.symmetrical_epipolar_distance(pts_torch[:,:,:2],
                                              pts_torch[:,:,2:],
                                              F_torch, squared=False).detach().cpu().numpy().reshape(-1, 1)
    out = []
    for th in ths:
        correct = errors <= th
        out.append(correct.astype(np.float32).mean())
    return out
