compute_rmse

import os
import json
import time
import math
import pandas as pd
import numpy as np
from scipy.spatial import cKDTree

def _find_closest_neighbors(xy_true, xy_pred, k=5):
    n_preds = xy_pred.shape[0]
    tree = cKDTree(xy_pred)
    dist, inds = tree.query(xy_true, k=k)
    idx = np.argsort(dist[:, 0])
    neighbors = np.full(len(xy_true), -1, dtype=int)
    picked = set()
    for i, ind in enumerate(inds[idx]):
        for j in ind:
            if j not in picked:
                picked.add(j)
                neighbors[idx[i]] = j
                break
        if len(picked) == n_preds:
            break
    return neighbors

pred_json ='/media/wagnchogn/data_2tb/pose_estimation/logocap-main/tools/output/logocap-hrnet-w48-mouse/eval_results-val2017.json'
gt_json = '/media/wagnchogn/data_2tb/pose_estimation/logocap-main/data/mouse/annotations/person_keypoints_val2017.json'
#df = pd.read_csv(r'')
with open(gt_json,"r") as gt:
    gts = json.load(gt)
gts_data = gts['annotations']

with open(pred_json,'r') as f:
    preds = json.load(f)

keypoints_num = 12
mouses = 3
pred_bpts = []
gt_bpts = []
imgs_num = int(len(gts_data)/3)

dist = np.full((imgs_num,keypoints_num*3),np.nan)
conf = np.full_like(dist,np.nan)
relativeimagenames = []
for bpt in range(keypoints_num):
    for i in range(imgs_num):
        pred_coords = []
        pred_scores = []
        true_coords = []
        true_vises = []
        indexs = []
        for j in range(i * mouses, (i + 1) * mouses):
            pred = preds[j]['keypoints'][bpt*3:(bpt+1)*3]
            gt = gts_data[j]['keypoints'][bpt*3:(bpt+1)*3]
            pred_x = pred[0]
            pred_y = pred[1]
            pred_score = pred[2]
            true_x = gt[0]
            true_y = gt[1]
            ture_vis = gt[2]
            if math.isnan(true_x):
                pred_x = 0
                pred_y = 0
                true_x = 0
                true_y = 0
            if math.isnan(true_y):
                pred_x = 0
                pred_y = 0
                true_x = 0
                true_y = 0
            pred_coord = [pred_x,pred_y]
            true_coord = [true_x,true_y]
            pred_coords.append(pred_coord)
            pred_scores.append(pred_score)
            true_coords.append(true_coord)
            true_vises.append(ture_vis)

        pred_coords = np.array(pred_coords)
        true_coords = np.array(true_coords)

        neithbors = _find_closest_neighbors(true_coords,pred_coords,k=3)

        found = neithbors != -1
        #print(bpt,i,indexs)

        #time.sleep(1)
        min_dist = np.linalg.norm(true_coords[found]-pred_coords[neithbors[found]],
                                   axis=1)
        indexs = [bpt*3,bpt*3+1,bpt*3+2]
        sl = i,np.array(indexs)
        dist[sl] = min_dist

        #a =  np.array(pred_scores)[neithbors[found]].squeeze()
        conf[sl] = np.array(pred_scores)[neithbors[found]].squeeze()
        #print(min_dist)
    relativeimagenames.append(bpt)
multibodyparts = ['snout', 'leftear', 'rightear', 'shoulder', 'spine1',
        'spine2', 'spine3', 'spine4', 'tailbase', 'tail1',
        'tail2', 'tailend']
mouse = ['mus1','mus2','mus3']
df = None
a = np.empty((imgs_num, 2))
a[:] = np.nan
'''
for prfxindex, prefix in enumerate(mouse):
    for c, bp in enumerate(multibodyparts):
        cols = pd.MultiIndex.from_product(
            [[prefix], [bp]],
            names=[ "individuals", "bodyparts"],
        )
        #index = pd.MultiIndex.from_tuples(relativeimagenames)
        frame = pd.DataFrame(a, columns=cols)
        df = pd.concat([df, frame], axis=1)
'''
dist_df = pd.DataFrame(dist)
conf_df = pd.DataFrame(conf)
#dist_df.columns = pd.MultiIndex.from_product([['a'],dist_df.columns])
#conf_df.columns = pd.MultiIndex.from_product([['a'],conf_df.columns])
#dist_df = pd.DataFrame(dist)
'''
df_joint = pd.concat([dist_df,conf_df],
                     keys=["rmse","conf"],
                     names = ["metrics"],
                     axis=1)
'''
#error = dist_df.xs("rmse",level="metrics",axis=1)
error_test = np.nanmean(dist_df)
print(error_test)