Attached below is my code for generating the confidence maps for the joint locations, it is in the same vein as the maps generated for the paper Convolutional Pose Machines and OpenPose. I am wondering if anyone has input on how to speed up the process, as when I try to train my own variant of OpenPose or Convolutional Pose machines, my generator times out for batch sizes larger than one.... Just as a note so everyone knows the structure of the keypoint data, as I am training on MS COCO: for each image there can be N people and each person has 18 keypoints associated with them, and each keypoint comes with a list [y,x,o] where the x,y are self explanatory and o is an inclusion/occlusion indicator.
import numpy as np import cv2 import os def GeneratorMaskFiller(ImgPath,keys,string_ids,bs,mode="train",aug=None): os.chdir(ImgPath) while True: loop_count=0 imgs= masks= for i in range(loop_count * bs,bs+loop_count * bs): img=cv2.imread(string_ids[i]+'.jpg') resized_img = cv2.resize(img,(224,224)) imgs.append(resized_img) length = len(keys[i]) num_people=int(length/18) sigma=7 zeros=np.zeros((img.shape,img.shape,18,num_people)) for part in range(18): for x in range(img.shape): for y in range(img.shape): for people in range(num_people): zeros[x][y][part][people]=np.exp((-(pow(x-keys[loop_count][part+people*18],2)+pow(y-keys[loop_count][part+people*18],2)))/sigma) mask=np.zeros((img.shape,img.shape,18)) for x in range(img.shape): for y in range(img.shape): for part in range(18): mask[x][y][part]=max(zeros[x,y,part,:]) mask=cv2.resize(mask,(224,224)) masks.append(mask) imgs = np.asarray(imgs) masks=np.asarray(masks) loop_count+=1 if loop_count == len(keys): loop_count=0 yield (imgs,masks)
Edits**: For clarification the the argument keys is the a list of lists, where each element of the list corresponds to the keypoints for each image, so say there are three individuals in one image then there are 54 elements in that list where each element is a list of three elements. Out of the 18 masks each mask corresponds to a specific joint and will have disks centered at the true joint location for dependent on the number of people in the image and if the joints are visible. I am curious if there is a way for me to remove the multiple for loops I currently have in my code to generate these ground truth confident maps.