Upload 2 files

eval_onnx.py

@@ -1,3 +1,4 @@
+
 import onnxruntime as rt
 import numpy as np
 import json
@@ -14,7 +15,7 @@ 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
+DATASET_PATH = '/group/dphi_algo_scratch_02/ziheng/datasets/coco/croped'  # 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
 
@@ -50,6 +51,7 @@ def getAccRight(dist, th = 5/IMG_SIZE):
     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):
@@ -63,6 +65,7 @@ def myAcc(output, target):
     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]
@@ -91,11 +94,15 @@ def maxPoint(heatmap, center=True):
     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*_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
@@ -103,47 +110,47 @@ def maxPoint(heatmap, center=True):
     # bv
     return x,y
 
+# Function for decoding MoveNet output data
 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):
+            #nchw!!!!!!!!!!!!!!!!!
+
             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)
+            #### 根据center得到关键点回归位置，然后加权heatmap
             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
@@ -152,10 +159,12 @@ def movenetDecode(data, kps_mask=None,mode='output', num_joints = 17,
             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
@@ -163,57 +172,69 @@ def movenetDecode(data, kps_mask=None,mode='output', num_joints = 17,
             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])
+            # b
+                
         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):
+            #nchw!!!!!!!!!!!!!!!!!
             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
+
+            # print(reg_x, reg_y)
             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
+            # print(offset_x,offset_y)
             res_x = (reg_x+offset_x)/(img_size//4)
             res_y = (reg_y+offset_y)/(img_size//4)
+
+            #不存在的点设为-1 后续不参与acc计算
             res_x[kps_mask[:,n]==0] = -1
             res_y[kps_mask[:,n]==0] = -1
             res.extend([res_x, res_y])
+            # b
+                
         res = np.concatenate(res,axis=1) #bs*14
+
     return res
 
+# Function to convert labeled keypoints to heatmaps for keypoints
 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]
+    # print(keypoints_range)
     min_x = np.min(keypoints_range[:,0])
     min_y = np.min(keypoints_range[:,1])
     max_x = np.max(keypoints_range[:,0])
@@ -226,40 +247,36 @@ def label2heatmap(keypoints, other_keypoints, img_size):
         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) 
+
+        x = int(keypoints[i]*img_size//4) #取值应该是0-47
         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
-    '''
+# Function to generate a heatmap for a specific keypoint
+def generate_heatmap(x, y, other_keypoints, size, sigma):
     #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
@@ -270,6 +287,8 @@ def generate_heatmap(x, y, other_keypoints, size, sigma):
             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
@@ -277,60 +296,55 @@ def generate_heatmap(x, y, other_keypoints, size, sigma):
         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
 
+# Function to convert labeled keypoints to a center 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
 
+# Function to convert labeled keypoints to regression maps
 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)
+    # print(keypoints)
     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
@@ -345,55 +359,57 @@ def label2reg(keypoints, cx, cy, img_size):
                     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
 
+# Function to convert labeled keypoints to offset maps
 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
 
+# Custom Dataset class for handling data loading and preprocessing
 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,
@@ -405,50 +421,70 @@ class TensorDataset(Dataset):
         """
         # [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):
+            ##0没有标注;1有标注不可见（被遮挡）;2有标注可见
             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,
@@ -458,33 +494,27 @@ def getDataLoader(mode, input_data):
                                         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",
@@ -514,20 +544,30 @@ if __name__ == '__main__':
     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})
+        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))
+    print('[Info] acc: {:.3f}% \n'.format(100. * acc))

movenet_int8.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9e3243a492e8886d29d9dcdd2831da3c69eb18bc76185f2db565cfe8ddcd58d8
-size 7681846
+oid sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855
+size 0