import onnxruntime as rt import numpy as np import json import torch import cv2 import os from torch.utils.data.dataset import Dataset import random import math import argparse # Constants and paths defining model, image, and dataset specifics MODEL_DIR = './movenet_int8.onnx' # Path to the MoveNet model IMG_SIZE = 192 # Image size used for processing FEATURE_MAP_SIZE = 48 # Feature map size used in the model CENTER_WEIGHT_ORIGIN_PATH = './center_weight_origin.npy' # Path to center weight origin file DATASET_PATH = 'your_dataset_path' # Base path for the dataset EVAL_LABLE_PATH = os.path.join(DATASET_PATH, "val2017.json") # Path to validation labels JSON file EVAL_IMG_PATH = os.path.join(DATASET_PATH, 'imgs') # Path to validation images def getDist(pre, labels): """ Calculate the Euclidean distance between predicted and labeled keypoints. Args: pre: Predicted keypoints [batchsize, 14] labels: Labeled keypoints [batchsize, 14] Returns: dist: Distance between keypoints [batchsize, 7] """ pre = pre.reshape([-1, 17, 2]) labels = labels.reshape([-1, 17, 2]) res = np.power(pre[:,:,0]-labels[:,:,0],2)+np.power(pre[:,:,1]-labels[:,:,1],2) return res def getAccRight(dist, th = 5/IMG_SIZE): """ Compute accuracy for each keypoint based on a threshold. Args: dist: Distance between keypoints [batchsize, 7] th: Threshold for accuracy computation Returns: res: Accuracy per keypoint [7,] representing the count of correct predictions """ res = np.zeros(dist.shape[1], dtype=np.int64) for i in range(dist.shape[1]): res[i] = sum(dist[:,i]47] = 47 reg_x[reg_x<0] = 0 reg_y[reg_y>47] = 47 reg_y[reg_y<0] = 0 score = heatmaps[dim0,dim1+n,reg_y,reg_x] offset_x = offsets[dim0,dim1+n*2,reg_y,reg_x]#*img_size//4 offset_y = offsets[dim0,dim1+n*2+1,reg_y,reg_x]#*img_size//4 res_x = (reg_x+offset_x)/(img_size//4) res_y = (reg_y+offset_y)/(img_size//4) res_x[score47] = 47 reg_x[reg_x<0] = 0 reg_y[reg_y>47] = 47 reg_y[reg_y<0] = 0 offset_x = offsets[dim0,dim1+n*2,reg_y,reg_x]#*img_size//4 offset_y = offsets[dim0,dim1+n*2+1,reg_y,reg_x]#*img_size//4 res_x = (reg_x+offset_x)/(img_size//4) res_y = (reg_y+offset_y)/(img_size//4) res_x[kps_mask[:,n]==0] = -1 res_y[kps_mask[:,n]==0] = -1 res.extend([res_x, res_y]) res = np.concatenate(res,axis=1) #bs*14 return res def label2heatmap(keypoints, other_keypoints, img_size): ''' Convert labeled keypoints to heatmaps for keypoints. Args: keypoints: Target person's keypoints other_keypoints: Other people's keypoints img_size: Size of the image Returns: heatmaps: Heatmaps for keypoints sigma: Value used for heatmap generation ''' #keypoints: target person #other_keypoints: other people's keypoints need to be add to the heatmap heatmaps = [] keypoints_range = np.reshape(keypoints,(-1,3)) keypoints_range = keypoints_range[keypoints_range[:,2]>0] min_x = np.min(keypoints_range[:,0]) min_y = np.min(keypoints_range[:,1]) max_x = np.max(keypoints_range[:,0]) max_y = np.max(keypoints_range[:,1]) area = (max_y-min_y)*(max_x-min_x) sigma = 3 if area < 0.16: sigma = 3 elif area < 0.3: sigma = 5 else: sigma = 7 for i in range(0,len(keypoints),3): if keypoints[i+2]==0: heatmaps.append(np.zeros((img_size//4, img_size//4))) continue x = int(keypoints[i]*img_size//4) y = int(keypoints[i+1]*img_size//4) if x==img_size//4:x=(img_size//4-1) if y==img_size//4:y=(img_size//4-1) if x>img_size//4 or x<0:x=-1 if y>img_size//4 or y<0:y=-1 heatmap = generate_heatmap(x, y, other_keypoints[i//3], (img_size//4, img_size//4),sigma) heatmaps.append(heatmap) heatmaps = np.array(heatmaps, dtype=np.float32) return heatmaps,sigma def generate_heatmap(x, y, other_keypoints, size, sigma): ''' Generate a heatmap for a specific keypoint. Args: x, y: Absolute position of the keypoint other_keypoints: Position of other keypoints size: Size of the heatmap sigma: Value used for heatmap generation Returns: heatmap: Generated heatmap for the keypoint ''' #x,y abs postion #other_keypoints positive position sigma+=6 heatmap = np.zeros(size) if x<0 or y<0 or x>=size[0] or y>=size[1]: return heatmap tops = [[x,y]] if len(other_keypoints)>0: #add other people's keypoints for i in range(len(other_keypoints)): x = int(other_keypoints[i][0]*size[0]) y = int(other_keypoints[i][1]*size[1]) if x==size[0]:x=(size[0]-1) if y==size[1]:y=(size[1]-1) if x>size[0] or x<0 or y>size[1] or y<0: continue tops.append([x,y]) for top in tops: #heatmap[top[1]][top[0]] = 1 x,y = top x0 = max(0,x-sigma//2) x1 = min(size[0],x+sigma//2) y0 = max(0,y-sigma//2) y1 = min(size[1],y+sigma//2) for map_y in range(y0, y1): for map_x in range(x0, x1): d2 = ((map_x - x) ** 2 + (map_y - y) ** 2)**0.5 if d2<=sigma//2: heatmap[map_y, map_x] += math.exp(-d2/(sigma//2)*3) if heatmap[map_y, map_x] > 1: heatmap[map_y, map_x] = 1 # heatmap[heatmap<0.1] = 0 return heatmap def label2center(cx, cy, other_centers, img_size, sigma): ''' Convert labeled keypoints to a center heatmap. Args: cx, cy: Center coordinates other_centers: Other people's centers img_size: Size of the image sigma: Value used for heatmap generation Returns: heatmaps: Heatmap representing the center ''' heatmaps = [] heatmap = generate_heatmap(cx, cy, other_centers, (img_size//4, img_size//4),sigma+2) heatmaps.append(heatmap) heatmaps = np.array(heatmaps, dtype=np.float32) return heatmaps def label2reg(keypoints, cx, cy, img_size): ''' Convert labeled keypoints to regression maps. Args: keypoints: Labeled keypoints cx, cy: Center coordinates img_size: Size of the image Returns: heatmaps: Regression maps for keypoints ''' heatmaps = np.zeros((len(keypoints)//3*2, img_size//4, img_size//4), dtype=np.float32) for i in range(len(keypoints)//3): if keypoints[i*3+2]==0: continue x = keypoints[i*3]*img_size//4 y = keypoints[i*3+1]*img_size//4 if x==img_size//4:x=(img_size//4-1) if y==img_size//4:y=(img_size//4-1) if x>img_size//4 or x<0 or y>img_size//4 or y<0: continue reg_x = x-cx reg_y = y-cy for j in range(cy-2,cy+3): if j<0 or j>img_size//4-1: continue for k in range(cx-2,cx+3): if k<0 or k>img_size//4-1: continue if cximg_size//4 or small_x<0 or small_y>img_size//4 or small_y<0: continue heatmaps[i*2][small_y][small_x] = offset_x#/(img_size//4) heatmaps[i*2+1][small_y][small_x] = offset_y#/(img_size//4) return heatmaps class TensorDataset(Dataset): ''' Custom Dataset class for handling data loading and preprocessing ''' def __init__(self, data_labels, img_dir, img_size, data_aug=None): self.data_labels = data_labels self.img_dir = img_dir self.data_aug = data_aug self.img_size = img_size self.interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4] def __getitem__(self, index): item = self.data_labels[index] """ item = { "img_name":save_name, "keypoints":save_keypoints, "center":save_center, "other_centers":other_centers, "other_keypoints":other_keypoints, } """ # [name,h,w,keypoints...] img_path = os.path.join(self.img_dir, item["img_name"]) img = cv2.imread(img_path, cv2.IMREAD_COLOR) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = cv2.resize(img, (self.img_size, self.img_size), interpolation=random.choice(self.interp_methods)) #### Data Augmentation if self.data_aug is not None: img, item = self.data_aug(img, item) img = img.astype(np.float32) img = np.transpose(img,axes=[2,0,1]) keypoints = item["keypoints"] center = item['center'] other_centers = item["other_centers"] other_keypoints = item["other_keypoints"] kps_mask = np.ones(len(keypoints)//3) for i in range(len(keypoints)//3): if keypoints[i*3+2]==0: kps_mask[i] = 0 heatmaps,sigma = label2heatmap(keypoints, other_keypoints, self.img_size) #(17, 48, 48) cx = min(max(0,int(center[0]*self.img_size//4)),self.img_size//4-1) cy = min(max(0,int(center[1]*self.img_size//4)),self.img_size//4-1) centers = label2center(cx, cy, other_centers, self.img_size, sigma) #(1, 48, 48) regs = label2reg(keypoints, cx, cy, self.img_size) #(14, 48, 48) offsets = label2offset(keypoints, cx, cy, regs, self.img_size)#(14, 48, 48) labels = np.concatenate([heatmaps,centers,regs,offsets],axis=0) img = img / 127.5 - 1.0 return img, labels, kps_mask, img_path def __len__(self): return len(self.data_labels) # Function to get data loader based on mode (e.g., evaluation) def getDataLoader(mode, input_data): ''' Function to get data loader based on mode (e.g., evaluation). Args: mode: Mode of data loader (e.g., 'eval') input_data: Input data Returns: data_loader: DataLoader for specified mode ''' if mode=="eval": val_loader = torch.utils.data.DataLoader( TensorDataset(input_data[0], EVAL_IMG_PATH, IMG_SIZE, ), batch_size=1, shuffle=False, num_workers=0, pin_memory=False) return val_loader # Class for managing data and obtaining evaluation data loader class Data(): ''' Class for managing data and obtaining evaluation data loader. ''' def __init__(self): pass def getEvalDataloader(self): with open(EVAL_LABLE_PATH, 'r') as f: data_label_list = json.loads(f.readlines()[0]) print("[INFO] Total images: ", len(data_label_list)) input_data = [data_label_list] data_loader = getDataLoader("eval", input_data) return data_loader # Configs for onnx inference session def make_parser(): ''' Create parser for MoveNet ONNX runtime inference. Returns: parser: Argument parser for MoveNet inference ''' parser = argparse.ArgumentParser("movenet onnxruntime inference") parser.add_argument( "--ipu", action="store_true", help="Use IPU for inference.", ) parser.add_argument( "--provider_config", type=str, default="vaip_config.json", help="Path of the config file for seting provider_options.", ) return parser.parse_args() if __name__ == '__main__': args = make_parser() if args.ipu: providers = ["VitisAIExecutionProvider"] provider_options = [{"config_file": args.provider_config}] else: providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] provider_options = None # Get evaluation data loader using the Data class data = Data() data_loader = data.getEvalDataloader() # Load MoveNet model using ONNX runtime model = rt.InferenceSession(MODEL_DIR, providers=providers, provider_options=provider_options) correct = 0 total = 0 # Loop through the data loader for evaluation for batch_idx, (imgs, labels, kps_mask, img_names) in enumerate(data_loader): if batch_idx%100 == 0: print('Finish ',batch_idx) imgs = imgs.detach().cpu().numpy() imgs = imgs.transpose((0,2,3,1)) output = model.run(['1548_transpose','1607_transpose','1665_transpose','1723_transpose'],{'blob.1':imgs}) output[0] = output[0].transpose((0,3,1,2)) output[1] = output[1].transpose((0,3,1,2)) output[2] = output[2].transpose((0,3,1,2)) output[3] = output[3].transpose((0,3,1,2)) pre = movenetDecode(output, kps_mask,mode='output',img_size=IMG_SIZE) gt = movenetDecode(labels, kps_mask,mode='label',img_size=IMG_SIZE) #n acc = myAcc(pre, gt) correct += sum(acc) total += len(acc) # Compute and print accuracy based on evaluated data acc = correct/total print('[Info] acc: {:.3f}% \n'.format(100. * acc))