init model

README.md

@@ -0,0 +1,93 @@
+---
+license: apache-2.0
+datasets:
+- detection-datasets/coco
+language:
+- en
+metrics:
+- accuracy
+tags:
+- RyzenAI
+- pose estimation
+---
+
+# MoveNet
+
+MoveNet is an ultra fast and accurate model that detects 17 keypoints of a body. It released in [movenet.pytorch](https://github.com/fire717/movenet.pytorch/blob/master/README.md?plain=1)
+
+
+We develop a modified version that could be supported by [AMD Ryzen AI](https://ryzenai.docs.amd.com/).
+
+
+
+## How to use
+
+### Installation
+
+   Follow [Ryzen AI Installation](https://ryzenai.docs.amd.com/en/latest/inst.html) to prepare the environment for Ryzen AI.
+   Run the following script to install pre-requisites for this model.
+   ```bash
+   pip install -r requirements.txt 
+   ```
+
+
+### Data Preparation (optional: for accuracy evaluation)
+
+1.Download COCO dataset2017 from https://cocodataset.org/. (You need train2017.zip, val2017.zip and annotations.)Unzip to `./data/` like this:
+
+```
+├── data
+    ├── annotations (person_keypoints_train2017.json, person_keypoints_val2017.json, ...)
+    ├── train2017   (xx.jpg, xx.jpg,...)
+    └── val2017     (xx.jpg, xx.jpg,...)
+
+```
+
+
+2.Make data to our data format.
+ - Modify the path in line 282~287 in make_coco_data_17keypoints.py if needed
+ - run the code to pre-process the dataset
+```
+python make_coco_data_17keypoints.py
+```
+```
+Our data format: JSON file
+Keypoints order:['nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear', 
+    'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow', 'left_wrist', 
+    'right_wrist', 'left_hip', 'right_hip', 'left_knee', 'right_knee', 'left_ankle', 
+    'right_ankle']
+
+One item:
+[{"img_name": "0.jpg",
+  "keypoints": [x0,y0,z0,x1,y1,z1,...],
+  #z: 0 for no label, 1 for labeled but invisible, 2 for labeled and visible
+  "center": [x,y],
+  "bbox":[x0,y0,x1,y1],
+  "other_centers": [[x0,y0],[x1,y1],...],
+  "other_keypoints": [[[x0,y0],[x1,y1],...],[[x0,y0],[x1,y1],...],...], #lenth = num_keypoints
+ },
+ ...
+]
+```
+
+
+
+
+### Test & Evaluation
+
+ - Modify the DATASET_PATH in eval_onnx.py if needed
+ - Test accuracy of the quantized model
+  ```python
+  python eval_onnx.py --ipu --provider_config Path\To\vaip_config.json
+  ```
+
+### Performance
+
+|Metric |Accuracy on IPU|
+| :----:  | :----: |
+|accuracy | 79.745%|
+
+
+## Citation
+1.[model card](https://storage.googleapis.com/movenet/MoveNet.SinglePose%20Model%20Card.pdf)
+2.[movenet.pytorch](https://github.com/fire717/movenet.pytorch/blob/master/README.md?plain=1)

center_weight_origin.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88a437d41816cc6526e64f0014fb977d4d3f216bbec7daaea59492a1f3f1494a
+size 9296

eval_onnx.py

@@ -0,0 +1,533 @@
+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))

make_coco_data_17keypoints.py

@@ -0,0 +1,277 @@
+"""
+@Fire
+https://github.com/fire717
+"""
+import os
+import json
+import pickle
+import cv2
+import numpy as np
+
+
+
+
+"""
+segmentation格式取决于这个实例是一个单个的对象（即iscrowd=0，将使用polygons格式）
+还是一组对象（即iscrowd=1，将使用RLE格式
+
+iscrowd=1时（将标注一组对象，比如一群人）
+
+
+标注说明：x,y,v,x,y,v,...
+其中v：#0没有标注;1有标注不可见（被遮挡）;2有标注可见
+
+关键点顺序：'keypoints': ['nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear', 
+    'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow', 'left_wrist', 
+    'right_wrist', 'left_hip', 'right_hip', 'left_knee', 'right_knee', 'left_ankle', 
+    'right_ankle']
+
+"""
+
+
+def main(img_dir, labels_path, output_name, output_img_dir):
+
+    if not os.path.exists(output_img_dir):
+        os.makedirs(output_img_dir)
+
+
+    with open(labels_path, 'r') as f:
+        data = json.load(f)
+
+    #print("total: ", len(data)) 5
+    #print(data.keys())#['info', 'licenses', 'images', 'annotations', 'categories']
+    #print(len(data['annotations']), len(data['images']))#88153 40504
+    #print(data['categories'])
+    """
+    [{'supercategory': 'person', 'name': 'person', 
+    'skeleton': [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], [7, 13], 
+    [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3], [1, 2], [1, 3], [2, 4], [3, 5], 
+    4, 6], [5, 7]], 
+    'keypoints': ['nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear', 
+    'left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow', 'left_wrist', 
+    'right_wrist', 'left_hip', 'right_hip', 'left_knee', 'right_knee', 'left_ankle', 
+    'right_ankle'], 'id': 1}]
+    """
+    #print(data['images'][:3])#有filename和id
+
+    img_id_to_name = {}
+    img_name_to_id = {}
+    for item in data['images']:
+        idx = item['id']
+        name = item['file_name']
+        img_id_to_name[idx] = name
+        img_name_to_id[name] = idx
+    print(len(img_id_to_name))
+    
+
+    anno_by_imgname = {}
+    for annotation in data['annotations']:
+        name = img_id_to_name[annotation['image_id']]
+        if name in anno_by_imgname:
+            anno_by_imgname[name] += [annotation]
+        else:
+            anno_by_imgname[name] = [annotation]
+    print(len(anno_by_imgname))
+
+
+
+    new_label = []
+    for k,v in anno_by_imgname.items():
+        #filter out more than 3 people
+        if len(v)>3:
+            continue
+
+        # print(k)
+        # print(v)
+
+        img = cv2.imread(os.path.join(img_dir, k))
+        if img is None:
+            print(os.path.join(img_dir, k))
+            continue
+        h,w = img.shape[:2]
+        for idx,item in enumerate(v):
+            if item['iscrowd'] != 0:
+                continue
+
+            bbox = [int(x) for x in item['bbox']]#x,y,w,h
+            # cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[0]+bbox[2], bbox[1]+bbox[3]), (0,255,0), 2)
+
+            keypoints = item['keypoints']
+
+            # for i in range(len(keypoints)//3):
+            #     x = keypoints[i*3]
+            #     y = keypoints[i*3+1]
+            #     z = keypoints[i*3+2]#0没有标注;1有标注不可见（被遮挡）;2有标注可见
+            #     # print(x,y,z)
+            #     if z==1:
+            #         color = (255,0,0)
+            #     elif z==2:
+            #         color = (0,0,255)
+            #     else:
+            #         continue
+                # cv2.circle(img, (x, y), 4, color, 3)
+
+
+            # merge bbox and keypoints to get max bbox 
+            keypoints = np.array(keypoints).reshape((17,3))
+
+            keypoints_v = keypoints[keypoints[:,2]>0]
+            if len(keypoints_v)<8:#filter out keypoints not enough
+                continue
+            min_key_x = np.min(keypoints_v[:,0])
+            max_key_x = np.max(keypoints_v[:,0])
+            min_key_y = np.min(keypoints_v[:,1])
+            max_key_y = np.max(keypoints_v[:,1])
+
+            x0 = min(bbox[0], min_key_x)
+            x1 = max(bbox[0]+bbox[2], max_key_x)
+            y0 = min(bbox[1], min_key_y)
+            y1 = max(bbox[1]+bbox[3], max_key_y)
+            # cv2.rectangle(img, (x0, y0), (x1, y1), (0,255,255), 2)
+
+            # expand to square then expand
+            cx = (x0+x1)/2
+            cy = (y0+y1)/2
+            
+            half_size = ((x1-x0)+(y1-y0))/2 * EXPAND_RATIO
+            new_x0 = int(cx - half_size)
+            new_x1 = int(cx + half_size)
+            new_y0 = int(cy - half_size)
+            new_y1 = int(cy + half_size)
+
+            #pad where exceed edge
+            pad_top = 0
+            pad_left = 0
+            pad_right = 0
+            pad_bottom = 0
+            if new_x0 < 0:
+                pad_left = -new_x0+1
+            if new_y0 < 0:
+                pad_top = -new_y0+1
+            if new_x1 > w:
+                pad_right = new_x1-w+1
+            if new_y1 > h:
+                pad_bottom = new_y1-h+1
+
+            pad_img = np.zeros((h+pad_top+pad_bottom, w+pad_left+pad_right, 3))
+            pad_img[pad_top:pad_top+h,pad_left:pad_left+w] = img
+            new_x0 += pad_left
+            new_y0 += pad_top
+            new_x1 += pad_left
+            new_y1 += pad_top
+            # cv2.rectangle(pad_img, (new_x0, new_y0), (new_x1, new_y1), (0,255,0), 2)
+
+            # final save data
+            save_name = k[:-4]+"_"+str(idx)+".jpg"
+            new_w = new_x1-new_x0
+            new_h = new_y1-new_y0
+            save_img = pad_img[new_y0:new_y1,new_x0:new_x1]
+            save_bbox = [(bbox[0]+pad_left-new_x0)/new_w,
+                         (bbox[1]+pad_top-new_y0)/new_h,
+                         (bbox[0]+bbox[2]+pad_left-new_x0)/new_w,
+                         (bbox[1]+bbox[3]+pad_top-new_y0)/new_h
+                        ]
+            save_center = [(save_bbox[0]+save_bbox[2])/2,(save_bbox[1]+save_bbox[3])/2]
+
+            save_keypoints = []
+            for kid in range(len(keypoints)):
+                save_keypoints.extend([(int(keypoints[kid][0])+pad_left-new_x0)/new_w,
+                                       (int(keypoints[kid][1])+pad_top-new_y0)/new_h,
+                                       int(keypoints[kid][2])
+                                      ])
+            other_centers = []
+            other_keypoints = [[] for _ in range(17)]
+            for idx2,item2 in enumerate(v):
+                if item2['iscrowd'] != 0 or idx2==idx:
+                    continue
+                bbox2 = [int(x) for x in item2['bbox']]#x,y,w,h
+
+                save_bbox2 = [(bbox2[0]+pad_left-new_x0)/new_w,
+                             (bbox2[1]+pad_top-new_y0)/new_h,
+                             (bbox2[0]+bbox2[2]+pad_left-new_x0)/new_w,
+                             (bbox2[1]+bbox2[3]+pad_top-new_y0)/new_h
+                            ]
+                save_center2 = [(save_bbox2[0]+save_bbox2[2])/2,
+                                (save_bbox2[1]+save_bbox2[3])/2]
+                if save_center2[0]>0 and save_center2[0]<1 and save_center2[1]>0 and save_center2[1]<1:
+                    other_centers.append(save_center2)
+
+                keypoints2 = item2['keypoints']
+                keypoints2 = np.array(keypoints2).reshape((17,3))
+                for kid2 in range(17):
+                    if keypoints2[kid2][2]==0:
+                        continue
+                    kx = (keypoints2[kid2][0]+pad_left-new_x0)/new_w
+                    ky = (keypoints2[kid2][1]+pad_top-new_y0)/new_h
+                    if kx>0 and kx<1 and ky>0 and ky<1:
+                        other_keypoints[kid2].append([kx,ky])
+
+            save_item = {
+                         "img_name":save_name,
+                         "keypoints":save_keypoints,
+                         "center":save_center,
+                         "bbox":save_bbox,
+                         "other_centers":other_centers,
+                         "other_keypoints":other_keypoints,
+                        }
+            # for k,v in save_item.items():
+            #     print(type(v[0]))
+            # b
+            new_label.append(save_item)
+
+
+
+            ###visul for exam, comment when use
+            if SHOW_POINTS_ON_IMG:
+                cv2.circle(save_img, (int(save_center[0]*new_w), int(save_center[1]*new_h)), 4, (0,255,0), 3)
+                for show_kid in range(len(save_keypoints)//3):
+                    if save_keypoints[show_kid*3+2]==1:
+                        color = (255,0,0)
+                    elif save_keypoints[show_kid*3+2]==2:
+                        color = (0,0,255)
+                    else:
+                        continue
+                    cv2.circle(save_img, (int(save_keypoints[show_kid*3]*new_w), 
+                                int(save_keypoints[show_kid*3+1]*new_h)), 3, color, 2)
+                cv2.rectangle(save_img, (int(save_bbox[0]*new_w), int(save_bbox[1]*new_h)), 
+                        (int(save_bbox[2]*new_w), int(save_bbox[3]*new_h)), (0,255,0), 2)
+                for show_c in other_centers:
+                    cv2.circle(save_img, (int(show_c[0]*new_w), int(show_c[1]*new_h)), 4, (0,255,255), 3)
+                for show_ks in other_keypoints:
+                    for show_k in show_ks:
+                        cv2.circle(save_img, (int(show_k[0]*new_w), int(show_k[1]*new_h)), 3, (255,255,0), 2)
+
+
+            cv2.imwrite(os.path.join(output_img_dir, save_name), save_img)
+            
+            # print(save_item, save_img.shape)
+
+            # b
+        # cv2.imwrite(os.path.join("show", k), pad_img)
+
+    with open(output_name,'w') as f:  
+        json.dump(new_label, f, ensure_ascii=False)     
+    print('Total write ', len(new_label))
+
+
+if __name__ == '__main__':
+
+    #### PARAM ####
+
+    SHOW_POINTS_ON_IMG = False
+    #whether to show points on img for debug
+
+    EXPAND_RATIO = 1. 
+    #person body bbox expand range to image edge
+
+
+    output_img_dir = "./data/croped/imgs"
+    img_dir = "./data/val2017"
+    labels_path = "./data/annotations/person_keypoints_val2017.json"
+    output_name = "./data/croped/val2017.json"
+    main(img_dir, labels_path, output_name, output_img_dir)
+    img_dir = "./data/train2017"
+    labels_path = "./data/annotations/person_keypoints_train2017.json"
+    output_name = "./data/croped/train2017.json"
+    main(img_dir, labels_path, output_name, output_img_dir)
+    

movenet_int8.onnx

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

requirements.txt

@@ -0,0 +1,4 @@
+torch
+numpy
+opencv-python
+# onnxruntime