movenet / eval_onnx.py
zhengrongzhang's picture
init model
21794d5
raw
history blame
19.3 kB
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]<th)
return res
def myAcc(output, target):
'''
Compute accuracy across keypoints.
Args:
output: Predicted keypoints
target: Labeled keypoints
Returns:
cate_acc: Categorical accuracy [7,] representing the count of correct predictions per keypoint
'''
# [h, ls, rs, lb, rb, lr, rr]
# output[:,6:10] = output[:,6:10]+output[:,2:6]
# output[:,10:14] = output[:,10:14]+output[:,6:10]
# Calculate distance between predicted and labeled keypoints
dist = getDist(output, target)
# Calculate accuracy for each keypoint
cate_acc = getAccRight(dist)
return cate_acc
# Predefined numpy arrays and weights for calculations
_range_weight_x = np.array([[x for x in range(FEATURE_MAP_SIZE)] for _ in range(FEATURE_MAP_SIZE)])
_range_weight_y = _range_weight_x.T
_center_weight = np.load(CENTER_WEIGHT_ORIGIN_PATH).reshape(FEATURE_MAP_SIZE,FEATURE_MAP_SIZE)
def maxPoint(heatmap, center=True):
"""
Find the coordinates of maximum values in a heatmap.
Args:
heatmap: Input heatmap data
center: Flag to indicate whether to consider center-weighted points
Returns:
x, y: Coordinates of maximum values in the heatmap
"""
if len(heatmap.shape) == 3:
batch_size,h,w = heatmap.shape
c = 1
elif len(heatmap.shape) == 4:
# n,c,h,w
batch_size,c,h,w = heatmap.shape
if center:
heatmap = heatmap*_center_weight
heatmap = heatmap.reshape((batch_size,c, -1)) #64,c, cfg['feature_map_size']xcfg['feature_map_size']
max_id = np.argmax(heatmap,2)#64,c, 1
y = max_id//w
x = max_id%w
# bv
return x,y
def movenetDecode(data, kps_mask=None,mode='output', num_joints = 17,
img_size=192, hm_th=0.1):
'''
Decode MoveNet output data to predicted keypoints.
Args:
data: MoveNet output data
kps_mask: Keypoints mask
mode: Mode of decoding ('output' or 'label')
num_joints: Number of joints/keypoints
img_size: Image size
hm_th: Threshold for heatmap processing
Returns:
res: Decoded keypoints
'''
##data [64, 7, 48, 48] [64, 1, 48, 48] [64, 14, 48, 48] [64, 14, 48, 48]
#kps_mask [n, 7]
if mode == 'output':
batch_size = data[0].shape[0]
heatmaps = data[0]
heatmaps[heatmaps < hm_th] = 0
centers = data[1]
regs = data[2]
offsets = data[3]
cx,cy = maxPoint(centers)
dim0 = np.arange(batch_size,dtype=np.int32).reshape(batch_size,1)
dim1 = np.zeros((batch_size,1),dtype=np.int32)
res = []
for n in range(num_joints):
reg_x_origin = (regs[dim0,dim1+n*2,cy,cx]+0.5).astype(np.int32)
reg_y_origin = (regs[dim0,dim1+n*2+1,cy,cx]+0.5).astype(np.int32)
reg_x = reg_x_origin+cx
reg_y = reg_y_origin+cy
### for post process
reg_x = np.reshape(reg_x, (reg_x.shape[0],1,1))
reg_y = np.reshape(reg_y, (reg_y.shape[0],1,1))
reg_x = reg_x.repeat(FEATURE_MAP_SIZE,1).repeat(FEATURE_MAP_SIZE,2)
reg_y = reg_y.repeat(FEATURE_MAP_SIZE,1).repeat(FEATURE_MAP_SIZE,2)
range_weight_x = np.reshape(_range_weight_x,(1,FEATURE_MAP_SIZE,FEATURE_MAP_SIZE)).repeat(reg_x.shape[0],0)
range_weight_y = np.reshape(_range_weight_y,(1,FEATURE_MAP_SIZE,FEATURE_MAP_SIZE)).repeat(reg_x.shape[0],0)
tmp_reg_x = (range_weight_x-reg_x)**2
tmp_reg_y = (range_weight_y-reg_y)**2
tmp_reg = (tmp_reg_x+tmp_reg_y)**0.5+1.8#origin 1.8
tmp_reg = heatmaps[:,n,...]/tmp_reg
tmp_reg = tmp_reg[:,np.newaxis,:,:]
reg_x,reg_y = maxPoint(tmp_reg, center=False)
reg_x[reg_x>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[score<hm_th] = -1
res_y[score<hm_th] = -1
res.extend([res_x, res_y])
res = np.concatenate(res,axis=1) #bs*14
elif mode == 'label':
kps_mask = kps_mask.detach().cpu().numpy()
data = data.detach().cpu().numpy()
batch_size = data.shape[0]
heatmaps = data[:,:17,:,:]
centers = data[:,17:18,:,:]
regs = data[:,18:52,:,:]
offsets = data[:,52:,:,:]
cx,cy = maxPoint(centers)
dim0 = np.arange(batch_size,dtype=np.int32).reshape(batch_size,1)
dim1 = np.zeros((batch_size,1),dtype=np.int32)
res = []
for n in range(num_joints):
reg_x_origin = (regs[dim0,dim1+n*2,cy,cx]+0.5).astype(np.int32)
reg_y_origin = (regs[dim0,dim1+n*2+1,cy,cx]+0.5).astype(np.int32)
reg_x = reg_x_origin+cx
reg_y = reg_y_origin+cy
reg_x[reg_x>47] = 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 cx<img_size//4/2-1:
heatmaps[i*2][j][k] = reg_x-(cx-k)#/(img_size//4)
else:
heatmaps[i*2][j][k] = reg_x+(cx-k)#/(img_size//4)
if cy<img_size//4/2-1:
heatmaps[i*2+1][j][k] = reg_y-(cy-j)#/(img_size//4)
else:
heatmaps[i*2+1][j][k] = reg_y+(cy-j)
return heatmaps
def label2offset(keypoints, cx, cy, regs, img_size):
'''
Convert labeled keypoints to offset maps.
Args:
keypoints: Labeled keypoints
cx, cy: Center coordinates
regs: Regression maps
img_size: Size of the image
Returns:
heatmaps: Offset 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
large_x = int(keypoints[i*3]*img_size)
large_y = int(keypoints[i*3+1]*img_size)
small_x = int(regs[i*2,cy,cx]+cx)
small_y = int(regs[i*2+1,cy,cx]+cy)
offset_x = large_x/4-small_x
offset_y = large_y/4-small_y
if small_x==img_size//4:small_x=(img_size//4-1)
if small_y==img_size//4:small_y=(img_size//4-1)
if small_x>img_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()
output = model.run(['1548','1607','1665','1723'],{'blob.1':imgs})
pre = movenetDecode(output, kps_mask,mode='output',img_size=IMG_SIZE)
gt = movenetDecode(labels, kps_mask,mode='label',img_size=IMG_SIZE)
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))