keypoint_utils.py

from numba import njit, prange

@njit(fastmath=True, cache=True)
def score_mask_numba_fast(pred, expected, ground, pixels_on_lines):
    h, w = pred.shape
    pp = 0
    po = 0
    
    for y in prange(h):
        for x in range(w):
            p_val = pred[y, x]
            g_val = ground[y, x]
            e_val = expected[y, x]
            p = (p_val != 0) & (g_val != 0)
            e = e_val != 0
            pp += p
            po += p & e
    if pp == 0:
        return 0.0
    pr = pp - po
    total = pixels_on_lines + pp - po
    if total == 0 or pr * 10 > total * 9:
        return 0.0
    return po / (pixels_on_lines + 1e-8)