def predict(image, params, model, model_params, heat_layers, paf_layers):
# print (image.shape)
heatmap_avg = np.zeros((image.shape[0], image.shape[1], heat_layers))
paf_avg = np.zeros((image.shape[0], image.shape[1], paf_layers))
multiplier = [x * model_params['boxsize'] / image.shape[0] for x in params['scale_search']]
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv2.resize(image, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest, model_params['stride'],
model_params['padValue'])
input_img = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 0, 1, 2)) # required shape (1, width, height, channels)
output_blobs = model.predict(input_img)
# extract outputs, resize, and remove padding
heatmap = np.squeeze(output_blobs[1]) # output 1 is heatmaps
heatmap = cv2.resize(heatmap, (0, 0), fx=model_params['stride'], fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
paf = np.squeeze(output_blobs[0]) # output 0 is PAFs
paf = cv2.resize(paf, (0, 0), fx=model_params['stride'], fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv2.resize(paf, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
#
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
return heatmap_avg, paf_avg
def handle_one_run2(oriImg, model_, model):
# for visualize
imageToTest = Variable(T.transpose(
T.transpose(T.unsqueeze(torch.from_numpy(oriImg).float(), 0), 2, 3), 1,
2),
volatile=True).cuda()
#print oriImg.shape
scale = model_['boxsize'] / float(oriImg.shape[0])
#print scale
imageToTest = cv2.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest,
model_['stride'],
model_['padValue'])
imageToTest_padded = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
feed = Variable(T.from_numpy(imageToTest_padded)).cuda()
out1, out2, out3, out4, out5, out5_1, out5_2, output1, output2, out6 = model(
feed)
return out1, out2, out3, out4, out5, out5_1, out5_2, output1, output2, out6
def get_mid_outputs(model_output):
picsize = [64, 64]
picsize_bf = [512, 512]
param, model_params = config_reader()
multiplier = [
x * model_params['boxsize'] / picsize_bf[0]
for x in param['scale_search']
]
heatmap_avg = np.zeros((picsize_bf[0], picsize_bf[1], 19))
paf_avg = np.zeros((picsize_bf[0], picsize_bf[1], 38))
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv2.resize(model_output, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model_params['stride'], model_params['padValue'])
input_img = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 0, 1, 2)) # required shape (1, width, height, channels)
#print("Input shape: " + str(input_img.shape))
output_blobs = model_output
#print("Output shape (heatmap): " + str(output_blobs[1].shape))
# extract outputs, resize, and remove padding
heatmap = np.squeeze(output_blobs[1]) # output 1 is heatmaps
heatmap = cv2.resize(heatmap, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv2.INTER_CUBIC)
paf = np.squeeze(output_blobs[0]) # output 0 is PAFs
paf = cv2.resize(paf, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv2.resize(paf, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
#pcm_cmp = pcm_cmp.reshape(picsize[0],picsize[1],19).astype("float32")
#paf_cmp = paf_cmp.reshape(picsize[0],picsize[1],38).astype("float32")
return heatmap_avg, paf_avg
def get_skeleton_multiple(self, images, scales):
oriImg = images[0]
multiplier = [
x * self.model['boxsize'] / oriImg.shape[0] for x in scales
]
heatmap_avg = [np.zeros(
(oriImg.shape[0], oriImg.shape[1], 19))] * len(images)
paf_avg = [np.zeros(
(oriImg.shape[0], oriImg.shape[1], 38))] * len(images)
for m in range(len(multiplier)):
scale = multiplier[m]
in_im = []
for i in range(len(images)):
oriImg = images[i]
imageToTest = cv.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, self.model['stride'], self.model['padValue'])
im = np.transpose(np.float32(imageToTest_padded[:, :, :]),
(2, 0, 1)) / 256 - 0.5
in_im.append(im)
self.net.blobs['data'].reshape(*(len(images), 3,
imageToTest_padded.shape[0],
imageToTest_padded.shape[1]))
self.net.blobs['data'].data[...] = in_im
output_blobs = self.net.forward()
for i in range(len(images)):
oriImg = images[i]
heatmap = np.transpose(
np.squeeze(self.net.blobs[output_blobs.keys()[1]].data)[
i, :, :, :], (1, 2, 0)) # output 1 is heatmaps
heatmap = cv.resize(heatmap, (0, 0),
fx=self.model['stride'],
fy=self.model['stride'],
interpolation=cv.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] -
pad[3], :]
heatmap = cv.resize(heatmap,
(oriImg.shape[1], oriImg.shape[0]),
interpolation=cv.INTER_CUBIC)
paf = np.transpose(
np.squeeze(self.net.blobs[output_blobs.keys()[0]].data)[
i, :, :, :], (1, 2, 0)) # output 0 is PAFs
paf = cv.resize(paf, (0, 0),
fx=self.model['stride'],
fy=self.model['stride'],
interpolation=cv.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv.resize(paf, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv.INTER_CUBIC)
heatmap_avg[i] = heatmap_avg[i] + heatmap / len(multiplier)
paf_avg[i] = paf_avg[i] + paf / len(multiplier)
return self.process_each_image(heatmap_avg, paf_avg, images)
def cnn_forward_pass(oriImg, params, model_params):
# on which scale to search, here single scale (=1) performs well enough (and is always faster)
params['scale_search'] = [1]
multiplier = [
x * model_params['boxsize'] / oriImg.shape[0]
for x in params['scale_search']
]
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
# calls to the CNN to create a Part Confidence Map (heatmap_avg) and Part Affinity Fields (paf_avg)
# by averaging the outputs on several scales
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv2.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model_params['stride'], model_params['padValue'])
input_img = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 0, 1, 2)) # required shape (1, width, height, channels)
output_blobs = model.predict(input_img)
# extract outputs, resize, and remove padding
heatmap = np.squeeze(output_blobs[1]) # output 1 is heatmaps
heatmap = cv2.resize(heatmap, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv2.INTER_CUBIC)
paf = np.squeeze(output_blobs[0]) # output 0 is PAFs
paf = cv2.resize(paf, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv2.resize(paf, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
return heatmap_avg, paf_avg
def transformation():
"""Do an inference on a single batch of data. In this sample server, we take data as CSV, convert
it to a pandas data frame for internal use and then convert the predictions back to CSV (which really
just means one prediction per line, since there's a single column.
"""
encoded = flask.request.data
nparr = np.frombuffer(encoded, np.uint8)
oriImg = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
param, model_params = config_reader()
scale = [x * model_params['boxsize'] / oriImg.shape[0] for x in param['scale_search']][3]
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
#paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
imageToTest = cv2.resize(oriImg, (0,0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest, model_params['stride'], model_params['padValue'])
input_img = np.transpose(np.float32(imageToTest_padded[:,:,:,np.newaxis]), (3,0,1,2)) # required shape (1, width, height, channels)
print("Input shape: " + str(input_img.shape))
output_blobs = ScoringService.predict(input_img)
heatmap = np.squeeze(output_blobs[1]) # output 1 is heatmaps
heatmap = cv2.resize(heatmap, (0,0), fx=model_params['stride'], fy=model_params['stride'], interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0]-pad[2], :imageToTest_padded.shape[1]-pad[3], :]
heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap
all_peaks = []
peak_counter = 0
for part in range(19-1):
map_ori = heatmap_avg[:,:,part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:,:] = map[:-1,:]
map_right = np.zeros(map.shape)
map_right[:-1,:] = map[1:,:]
map_up = np.zeros(map.shape)
map_up[:,1:] = map[:,:-1]
map_down = np.zeros(map.shape)
map_down[:,:-1] = map[:,1:]
peaks_binary = np.logical_and.reduce((map>=map_left, map>=map_right, map>=map_up, map>=map_down, map > param['thre1']))
peaks = list(zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0])) # note reverse
peaks_with_score = [x + (map_ori[x[1],x[0]],) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [peaks_with_score[i] + (id[i],) for i in range(len(id))]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
predictions = all_peaks
# Convert from numpy back to CSV
out = StringIO.StringIO()
pd.DataFrame({'results':predictions}).to_csv(out, header=False, index=False)
result = out.getvalue()
return flask.Response(response=result, status=200, mimetype='text/csv')
def preprocess(oriImg, model_params):
scale = float(model_params['boxsize']) / float(oriImg.shape[0])
imageToTest = cv2.resize(oriImg, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest, model_params['stride'],
model_params['padValue'])
input_img = np.transpose(np.float32(imageToTest_padded[:,:,:,np.newaxis]), (3,0,1,2)) # required shape (1, width, height, channels)
return input_img, pad
def Get_Net_res(oriImg, multiplier):
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
for m in range(len(multiplier)):
scale = multiplier[m]
print scale
imageToTest = cv.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model['stride'], model['padValue'])
print imageToTest_padded.shape
net.blobs['data'].reshape(*(1, 3, imageToTest_padded.shape[0],
imageToTest_padded.shape[1]))
# net.forward() # dry run
net.blobs['data'].data[...] = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
start_time = time.time()
output_blobs = net.forward()
print('At scale %d, The CNN took %.2f ms.' %
(m, 1000 * (time.time() - start_time)))
# extract outputs, resize, and remove padding
heatmap = np.transpose(
np.squeeze(net.blobs[output_blobs.keys()[1]].data),
(1, 2, 0)) # output 1 is heatmaps
heatmap = cv.resize(heatmap, (0, 0),
fx=model['stride'],
fy=model['stride'],
interpolation=cv.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv.INTER_CUBIC)
paf = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[0]].data),
(1, 2, 0)) # output 0 is PAFs
paf = cv.resize(paf, (0, 0),
fx=model['stride'],
fy=model['stride'],
interpolation=cv.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv.resize(paf, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
return heatmap_avg, paf_avg
def getSL(self, img, multipliers):
## Average confidence map and PAF
heatmap_avg = np.zeros(
(img.shape[0], img.shape[1], self.num_keypoints))
paf_avg = np.zeros((img.shape[0], img.shape[1], 2 * self.num_limbs))
for m in range(len(multipliers)):
scale = multipliers[m]
imageToTest = cv.resize(img, (0, 0),
fx=scale,
fy=scale,
interpolation=cv.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, self.model['stride'], self.model['padValue'])
self.net.blobs['data'].reshape(*(1, 3, imageToTest_padded.shape[0],
imageToTest_padded.shape[1]))
self.net.blobs['data'].data[...] = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
start_time = time.time()
output_blobs = self.net.forward()
# extract outputs, resize, and remove padding
heatmap = np.transpose(
np.squeeze(self.net.blobs[output_blobs.keys()[1]].data),
(1, 2, 0)) # output 1 is heatmaps
heatmap = cv.resize(heatmap, (0, 0),
fx=self.model['stride'],
fy=self.model['stride'],
interpolation=cv.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv.resize(heatmap, (img.shape[1], img.shape[0]),
interpolation=cv.INTER_CUBIC)
paf = np.transpose(
np.squeeze(self.net.blobs[output_blobs.keys()[0]].data),
(1, 2, 0)) # output 0 is PAFs
paf = cv.resize(paf, (0, 0),
fx=self.model['stride'],
fy=self.model['stride'],
interpolation=cv.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv.resize(paf, (img.shape[1], img.shape[0]),
interpolation=cv.INTER_CUBIC)
## Add to aggregate
heatmap_avg = heatmap_avg + heatmap / len(multipliers)
paf_avg = paf_avg + paf / len(multipliers)
return heatmap_avg, paf_avg
def _heatmap(self, img):
############################### Setup the heatmap vars. ########################
param, model_params = config_reader()
multiplier = [
x * model_params['boxsize'] / img.shape[0]
for x in param['scale_search']
]
############################### Get the heatmap ################################
heatmap_avg = np.zeros((img.shape[0], img.shape[1], 19))
paf_avg = np.zeros((img.shape[0], img.shape[1], 38))
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv2.resize(img, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model_params['stride'], model_params['padValue'])
input_img = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 0, 1, 2)) # required shape (1, width, height, channels)
output_blobs = self.model.predict(input_img)
# extract outputs, resize, and remove padding
heatmap = np.squeeze(output_blobs[1]) # output 1 is heatmaps
heatmap = cv2.resize(heatmap, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]),
interpolation=cv2.INTER_CUBIC)
paf = np.squeeze(output_blobs[0]) # output 0 is PAFs
paf = cv2.resize(paf, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv2.resize(paf, (img.shape[1], img.shape[0]),
interpolation=cv2.INTER_CUBIC)
# visualization
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
return (heatmap_avg, paf_avg)
def __call__(self, oriImg):
scale_search = [0.5, 1.0, 1.5, 2.0]
# scale_search = [0.5]
boxsize = 368
stride = 8
padValue = 128
thre = 0.05
multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 22))
# paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv2.resize(oriImg, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
im = np.ascontiguousarray(im)
data = torch.from_numpy(im).float()
if self.cuda_true:
data = data.cuda()
# data = data.permute([2, 0, 1]).unsqueeze(0).float()
with torch.no_grad():
if self.cuda_true:
output = self.model(data).cpu().numpy()
else:
output = self.model(data).numpy()
# extract outputs, resize, and remove padding
heatmap = np.transpose(np.squeeze(output), (1, 2, 0)) # output 1 is heatmaps
heatmap = cv2.resize(heatmap, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC)
heatmap_avg += heatmap / len(multiplier)
all_peaks = []
for part in range(21):
map_ori = heatmap_avg[:, :, part]
one_heatmap = gaussian_filter(map_ori, sigma=3)
binary = np.ascontiguousarray(one_heatmap > thre, dtype=np.uint8)
if np.sum(binary) == 0:
all_peaks.append([0, 0])
continue
label_img, label_numbers = label(binary, return_num=True, connectivity=binary.ndim)
max_index = np.argmax([np.sum(map_ori[label_img == i]) for i in range(1, label_numbers + 1)]) + 1
label_img[label_img != max_index] = 0
map_ori[label_img == 0] = 0
y, x = util.npmax(map_ori)
all_peaks.append([x, y])
return np.array(all_peaks)
def bodyDetect(oriImg, boxsize, detConf, person_net):
"""A body detector to detect body center
Args:
oriImg: The image data on which to detect body center.
boxsize: Scale height of image to boxsize.
detConf: Above which we think it will be a person, a hyper-parameter.
person_net: Caffe model object of the body detection net.
Returns:
imageToTest: Scaled image.
scale: imageToTest = oriImg * scale.
x, y: coordinates of body center
"""
##Scale height to match the boxsize
scale = boxsize / (oriImg.shape[0] * 1.0)
imageToTest = cv.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv.INTER_CUBIC)
##Pad image to fit detection input, that can be divided by 8
imageToTest_padded = util.padRightDownCorner(imageToTest)
person_net.blobs['image'].reshape(*(1, 3, imageToTest_padded.shape[0],
imageToTest_padded.shape[1]))
##Squeeze data range to [-0.5, 0.5]
person_net.blobs['image'].data[...] = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
output_blobs = person_net.forward()
##Remove single-dimensional entry
person_map = np.squeeze(person_net.blobs[output_blobs.keys()[0]].data)
##Location with high confidence with be returned, fx=8, because only 3 max pools in deploy file
person_map_resized = cv.resize(person_map, (0, 0),
fx=8,
fy=8,
interpolation=cv.INTER_CUBIC)
data_max = scipy.ndimage.filters.maximum_filter(person_map_resized, size=3)
maxima = (person_map_resized == data_max)
diff = (data_max > detConf)
maxima[diff == 0] = 0
x = np.nonzero(maxima)[1]
y = np.nonzero(maxima)[0]
return imageToTest, scale, x, y
def cal_heatmap(oriImg, net, param, model):
#test_image = 'test.jpg'
#test_image = '../sample_image/upper.jpg'
#test_image = '../sample_image/upper2.jpg'
#oriImg = cv.imread(test_image) # B,G,R order
#heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
#multiplier = [x * model['boxsize'] / oriImg.shape[0] for x in param['scale_search']]
scale = 1.0 * model['boxsize'] / oriImg.shape[0]
#print oriImg.shape
#print scale
imageToTest = cv.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest,
model['stride'],
model['padValue'])
net.blobs['data'].reshape(*(1, 3, imageToTest_padded.shape[0],
imageToTest_padded.shape[1]))
#net.forward() # dry run
net.blobs['data'].data[...] = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
#start_time = time.time()
output_blobs = net.forward()
#print('At scale %d, The CNN took %.2f ms.' % (m, 1000 * (time.time() - start_time)))
# extract outputs, resize, and remove padding
heatmap = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[1]].data),
(1, 2, 0)) # output 1 is heatmaps
heatmap = cv.resize(heatmap, (0, 0),
fx=model['stride'],
fy=model['stride'],
interpolation=cv.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv.INTER_CUBIC)
return heatmap
def run_at_multiple_scales(self, oriImg):
multiplier = [x * self.model['boxsize'] / oriImg.shape[0] for x in self.param['scale_search']]
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv.resize(oriImg, (0, 0), fx=scale, fy=scale, interpolation=cv.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest, self.model['stride'], self.model['padValue'])
print imageToTest_padded.shape
image = np.float32(imageToTest_padded[np.newaxis, :, :, :]) / 256 - 0.5
start_time = time.time()
output_blobs = self.sess.run(self.net.get_output(), feed_dict={self.input_node: image})
print('At scale %d, The CNN took %.2f ms.' % (m, 1000 * (time.time() - start_time)))
paf, heatmap = output_blobs
# extract outputs, resize, and remove padding
heatmap = heatmap[0, :, :, :] # output 1 is heatmaps
heatmap = cv.resize(heatmap, (0, 0), fx=self.model['stride'], fy=self.model['stride'], interpolation=cv.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv.INTER_CUBIC)
paf = paf[0, :, :, :]
paf = cv.resize(paf, (0, 0), fx=self.model['stride'], fy=self.model['stride'], interpolation=cv.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv.resize(paf, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv.INTER_CUBIC)
print "PAF ", paf.shape, np.min(paf), np.max(paf), np.mean(paf)
print "Heat ", heatmap.shape, np.min(heatmap), np.max(heatmap), np.mean(heatmap)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
return heatmap_avg, paf_avg
def collecting_images(input_image, params, model_params, buffs, buffs_for):
oriImg = input_image.copy() # B,G,R order
#oriImg = cv2.imread(input_image) # B,G,R order
multiplier = [
x * model_params['boxsize'] / oriImg.shape[0]
for x in params['scale_search']
]
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
#...........................................................................
#speeding up prediction
buffs[0].append(oriImg)
buffs[1].append(heatmap_avg)
buffs[2].append(paf_avg)
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv2.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model_params['stride'], model_params['padValue'])
input_img = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 0, 1, 2)) # required shape (1, width, height, channels)
buffs_for[0][m].append(imageToTest)
buffs_for[1][m].append(imageToTest_padded)
buffs_for[2][m].append(pad)
buffs_for[3][m].append(input_img)
return buffs, buffs_for
def handle_one(oriImg):
# for visualize
print(oriImg.shape)
canvas = np.copy(oriImg)
imageToTest = Variable(T.transpose(T.transpose(T.unsqueeze(torch.from_numpy(oriImg).float(),0),2,3),1,2),volatile=True)
print(oriImg.shape)
scale = model_['boxsize'] / float(oriImg.shape[0])
print(scale)
h = int(oriImg.shape[0]*scale)
w = int(oriImg.shape[1]*scale)
pad_h = 0 if (h%model_['stride']==0) else model_['stride'] - (h % model_['stride'])
pad_w = 0 if (w%model_['stride']==0) else model_['stride'] - (w % model_['stride'])
new_h = h+pad_h
new_w = w+pad_w
imageToTest = cv2.resize(oriImg, (0,0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest, model_['stride'], model_['padValue'])
imageToTest_padded = np.transpose(np.float32(imageToTest_padded[:,:,:,np.newaxis]), (3,2,0,1))/256 - 0.5
feed = Variable(T.from_numpy(imageToTest_padded))
time=datetime.datetime.now()
output1,output2 = model(feed.cuda())
print("time",(datetime.datetime.now()-time).microseconds)
heatmap = nn.UpsamplingBilinear2d((oriImg.shape[0], oriImg.shape[1]))(output2)
paf = nn.UpsamplingBilinear2d((oriImg.shape[0], oriImg.shape[1]))(output1)
print(heatmap.size())
print(paf.size())
print(type(heatmap))
heatmap_avg = T.transpose(T.transpose(heatmap[0],0,1),1,2).data.cpu().numpy()
paf_avg = T.transpose(T.transpose(paf[0],0,1),1,2).data.cpu().numpy()
all_peaks = []
peak_counter = 0
#maps =
for part in range(18):
map_ori = heatmap_avg[:,:,part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:,:] = map[:-1,:]
map_right = np.zeros(map.shape)
map_right[:-1,:] = map[1:,:]
map_up = np.zeros(map.shape)
map_up[:,1:] = map[:,:-1]
map_down = np.zeros(map.shape)
map_down[:,:-1] = map[:,1:]
peaks_binary = np.logical_and.reduce((map>=map_left, map>=map_right, map>=map_up, map>=map_down, map > param_['thre1']))
# peaks_binary = T.eq(
# peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])
peaks = list(zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0])) # note reverse
peaks_with_score = [x + (map_ori[x[1],x[0]],) for x in peaks]
id = list(range(peak_counter, peak_counter + len(peaks)))
peaks_with_score_and_id = [peaks_with_score[i] + (id[i],) for i in range(len(id))]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
connection_all = []
special_k = []
mid_num = 10
for k in range(len(mapIdx)):
score_mid = paf_avg[:,:,[x-19 for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0]-1]
candB = all_peaks[limbSeq[k][1]-1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if(nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0]*vec[0] + vec[1]*vec[1])
vec = np.divide(vec, norm)
startend = list(zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
np.linspace(candA[i][1], candB[j][1], num=mid_num)))
vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
for I in range(len(startend))])
vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
for I in range(len(startend))])
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
score_with_dist_prior = sum(score_midpts)/len(score_midpts) + min(0.5*oriImg.shape[0]/norm-1, 0)
criterion1 = len(np.nonzero(score_midpts > param_['thre2'])[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([i, j, score_with_dist_prior, score_with_dist_prior+candA[i][2]+candB[j][2]])
connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
connection = np.zeros((0,5))
for c in range(len(connection_candidate)):
i,j,s = connection_candidate[c][0:3]
if(i not in connection[:,3] and j not in connection[:,4]):
connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
if(len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
candidate = np.array([item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:,0]
partBs = connection_all[k][:,1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): #= 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): #1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if(subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
print("found = 2")
membership = ((subset[j1]>=0).astype(int) + (subset[j2]>=0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: #merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(candidate[connection_all[k][i,:2].astype(int), 2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
# delete some rows of subset which has few parts occur
deleteIdx = [];
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2]/subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
# canvas = cv2.imread(test_image) # B,G,R order
for i in range(18):
for j in range(len(all_peaks[i])):
cv2.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)
stickwidth = 4
for i in range(17):
for n in range(len(subset)):
index = subset[n][np.array(limbSeq[i])-1]
if -1 in index:
continue
cur_canvas = canvas.copy()
Y = candidate[index.astype(int), 0]
X = candidate[index.astype(int), 1]
mX = np.mean(X)
mY = np.mean(Y)
length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
polygon = cv2.ellipse2Poly((int(mY),int(mX)), (int(length/2), stickwidth), int(angle), 0, 360, 1)
cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
return canvas
def process(input_image, params, model_params):
oriImg = cv2.imread(input_image) # B,G,R order
multiplier = [
x * model_params['boxsize'] / oriImg.shape[0]
for x in params['scale_search']
]
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv2.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model_params['stride'], model_params['padValue'])
input_img = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 0, 1, 2)) # required shape (1, width, height, channels)
output_blobs = model.predict(input_img)
# extract outputs, resize, and remove padding
heatmap = np.squeeze(output_blobs[1]) # output 1 is heatmaps
heatmap = cv2.resize(heatmap, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv2.INTER_CUBIC)
paf = np.squeeze(output_blobs[0]) # output 0 is PAFs
paf = cv2.resize(paf, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv2.resize(paf, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
all_peaks = []
peak_counter = 0
for part in range(18):
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up, map >= map_down,
map > params['thre1']))
peaks = list(
zip(np.nonzero(peaks_binary)[1],
np.nonzero(peaks_binary)[0])) # note reverse
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [
peaks_with_score[i] + (id[i], ) for i in range(len(id))
]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
connection_all = []
special_k = []
mid_num = 10
for k in range(len(mapIdx)):
score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0] - 1]
candB = all_peaks[limbSeq[k][1] - 1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if (nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
# failure case when 2 body parts overlaps
if norm == 0:
continue
vec = np.divide(vec, norm)
startend = list(zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
np.linspace(candA[i][1], candB[j][1], num=mid_num)))
vec_x = np.array(
[score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
for I in range(len(startend))])
vec_y = np.array(
[score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
for I in range(len(startend))])
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(
vec_y, vec[1])
score_with_dist_prior = sum(
score_midpts) / len(score_midpts) + min(
0.5 * oriImg.shape[0] / norm - 1, 0)
criterion1 = len(
np.nonzero(score_midpts > params['thre2'])
[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([
i, j, score_with_dist_prior,
score_with_dist_prior + candA[i][2] + candB[j][2]
])
connection_candidate = sorted(connection_candidate,
key=lambda x: x[2],
reverse=True)
connection = np.zeros((0, 5))
for c in range(len(connection_candidate)):
i, j, s = connection_candidate[c][0:3]
if (i not in connection[:, 3] and j not in connection[:, 4]):
connection = np.vstack(
[connection, [candA[i][3], candB[j][3], s, i, j]])
if (len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
candidate = np.array([item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:, 0]
partBs = connection_all[k][:, 1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): # = 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): # 1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][
indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if (subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[partBs[i].astype(int),
2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
membership = ((subset[j1] >= 0).astype(int) +
(subset[j2] >= 0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: # merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[
partBs[i].astype(int), 2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(candidate[connection_all[k][i, :2].astype(int), 2]) + \
connection_all[k][i][2]
subset = np.vstack([subset, row])
# delete some rows of subset which has few parts occur
deleteIdx = []
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
canvas = cv2.imread(input_image) # B,G,R order
for i in range(18):
for j in range(len(all_peaks[i])):
cv2.circle(canvas,
all_peaks[i][j][0:2],
4,
colors[i],
thickness=-1)
stickwidth = 4
for i in range(17):
for n in range(len(subset)):
index = subset[n][np.array(limbSeq[i]) - 1]
if -1 in index:
continue
cur_canvas = canvas.copy()
Y = candidate[index.astype(int), 0]
X = candidate[index.astype(int), 1]
mX = np.mean(X)
mY = np.mean(Y)
length = ((X[0] - X[1])**2 + (Y[0] - Y[1])**2)**0.5
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
polygon = cv2.ellipse2Poly(
(int(mY), int(mX)), (int(length / 2), stickwidth), int(angle),
0, 360, 1)
cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
return canvas
def main(args):
global positions
parser = argparse.ArgumentParser(
description="Extract human pose from image(s) or video")
parser.add_argument("-i", "--inputFile")
parser.add_argument("-o", "--outputFile", default="output")
parser.add_argument("-n", "--name")
parser.add_argument("-t", "--duration",
help="duration in hh:mm:ss", default='00:00:05')
parser.add_argument("-s", "--startTime", help="start time in hh:mm:ss")
parser.add_argument("-r", "--frameRate",
help="frames per second", type=int, default=2)
parser.add_argument("-f", "--inputFolder", default=None)
args = parser.parse_args()
framesDir = 'frames'
# Get list of images
if args.inputFolder is not None: # Load from folder
images = os.listdir(os.path.join('.', args.inputFolder))
images = [os.path.join(args.inputFolder, img) for img in images]
if len(images) < 1:
print "Error: no images found in", args.inputFolder
sys.exit()
if args.name is not None:
images = [x for x in images if args.name in x]
oriImg = cv.imread(images[0]) # B,G,R order
elif args.inputFile is not None:
if any(ext in args.inputFile for ext in ['.mov', '.mp4', '.avi']):
images = movieToFrames(
args.inputFile, (args.frameRate, args.startTime, args.duration))
oriImg = cv.imread(images[0])
else:
oriImg = cv.imread(args.inputFile)
images = [args.inputFile]
currentFrame = 0
# Setup model
param, model = config_reader()
multiplier = [x * model['boxsize'] / oriImg.shape[0]
for x in param['scale_search']]
# if param['use_gpu']:
# caffe.set_mode_gpu()
# caffe.set_device(param['GPUdeviceNumber']) # set to your device!
# else:
caffe.set_mode_cpu()
net = caffe.Net(model['deployFile'], model['caffemodel'], caffe.TEST)
# Begin
positions = {}
for image in images:
positions[image] = []
oriImg = cv.imread(image)
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv.resize(oriImg, (0, 0), fx=scale,
fy=scale, interpolation=cv.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model['stride'], model['padValue'])
print imageToTest_padded.shape
net.blobs['data'].reshape(
*(1, 3, imageToTest_padded.shape[0], imageToTest_padded.shape[1]))
# net.forward() # dry run
net.blobs['data'].data[...] = np.transpose(np.float32(
imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
start_time = time.time()
output_blobs = net.forward()
print('At scale %d, The CNN took %.2f ms.' %
(m, 1000 * (time.time() - start_time)))
# extract outputs, resize, and remove padding
heatmap = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[
1]].data), (1, 2, 0)) # output 1 is heatmaps
heatmap = cv.resize(heatmap, (0, 0), fx=model['stride'], fy=model[
'stride'], interpolation=cv.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv.resize(heatmap, (oriImg.shape[1], oriImg.shape[
0]), interpolation=cv.INTER_CUBIC)
paf = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[
0]].data), (1, 2, 0)) # output 0 is PAFs
paf = cv.resize(paf, (0, 0), fx=model['stride'], fy=model[
'stride'], interpolation=cv.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] - pad[2],
:imageToTest_padded.shape[1] - pad[3], :]
paf = cv.resize(paf, (oriImg.shape[1], oriImg.shape[
0]), interpolation=cv.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
import scipy
print heatmap_avg.shape
# plt.imshow(heatmap_avg[:,:,2])
from scipy.ndimage.filters import gaussian_filter
all_peaks = []
peak_counter = 0
for part in range(19 - 1):
x_list = []
y_list = []
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up, map >= map_down, map > param['thre1']))
peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(
peaks_binary)[0]) # note reverse
peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [peaks_with_score[
i] + (id[i],) for i in range(len(id))]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
# find connection in the specified sequence, center 29 is in the position
# 15
limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10],
[10, 11], [2, 12], [12, 13], [
13, 14], [2, 1], [1, 15], [15, 17],
[1, 16], [16, 18], [3, 17], [6, 18]]
# the middle joints heatmap correpondence
mapIdx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44], [19, 20], [21, 22],
[23, 24], [25, 26], [27, 28], [29, 30], [
47, 48], [49, 50], [53, 54], [51, 52],
[55, 56], [37, 38], [45, 46]]
connection_all = []
special_k = []
mid_num = 10
for k in range(len(mapIdx)):
score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0] - 1]
candB = all_peaks[limbSeq[k][1] - 1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if(nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
vec = np.divide(vec, norm)
startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num),
np.linspace(candA[i][1], candB[j][1], num=mid_num))
vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0]
for I in range(len(startend))])
vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1]
for I in range(len(startend))])
score_midpts = np.multiply(
vec_x, vec[0]) + np.multiply(vec_y, vec[1])
try:
score_with_dist_prior = sum(
score_midpts) / len(score_midpts) + min(0.5 * oriImg.shape[0] / norm - 1, 0)
except ZeroDivisionError:
print "Zero Division Error Encountered"
score_with_dist_prior = 0
criterion1 = len(np.nonzero(score_midpts > param['thre2'])[
0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append(
[i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])
connection_candidate = sorted(
connection_candidate, key=lambda x: x[2], reverse=True)
connection = np.zeros((0, 5))
for c in range(len(connection_candidate)):
i, j, s = connection_candidate[c][0:3]
if(i not in connection[:, 3] and j not in connection[:, 4]):
connection = np.vstack(
[connection, [candA[i][3], candB[j][3], s, i, j]])
if(len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall
# configuration
subset = -1 * np.ones((0, 20))
candidate = np.array(
[item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:, 0]
partBs = connection_all[k][:, 1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): # = 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): # 1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if(subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[
j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
print "found = 2"
membership = ((subset[j1] >= 0).astype(
int) + (subset[j2] >= 0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: # merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[
j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(candidate[connection_all[k][i,
:2].astype(int), 2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
# delete some rows of subset which has few parts occur
deleteIdx = []
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
# visualize
colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],
[0, 255, 85], [0, 255, 170], [0, 255, 255], [
0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
canvas = cv.imread(image)
# visualize 2
stickwidth = 4
for i in range(17):
for n in range(len(subset)):
index = subset[n][np.array(limbSeq[i]) - 1]
if -1 in index:
continue
cur_canvas = canvas.copy()
Y = candidate[index.astype(int), 0]
X = candidate[index.astype(int), 1]
mX = np.mean(X)
mY = np.mean(Y)
length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
polygon = cv.ellipse2Poly((int(mY), int(mX)), (int(
length / 2), stickwidth), int(angle), 0, 360, 1)
cv.fillConvexPoly(cur_canvas, polygon, colors[i])
positions[image].append((
(int(mY), int(mX)),
int(length / 2),
int(angle), i))
canvas = cv.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
plt.imshow(canvas[:, :, [2, 1, 0]])
fig = matplotlib.pyplot.gcf()
fig.set_size_inches(12, 12)
outputPath = '{}{}{}.png'.format(args.inputFolder if args.inputFolder else '', args.outputFile,
'_' + str(currentFrame))
cv.imwrite(outputPath, canvas)
print "saved frame to", outputPath
currentFrame += 1
def processing(
image_path,
param,
parametre_search,
model,
net,
frame_number=0,
gpu=0
): #gpu=0 for cpu mode by default, if you want to enable gpu, put gpu=1
oriImg = cv.imread(image_path) # Reading the original image
multiplier = [
x * model['boxsize'] / oriImg.shape[0] for x in parametre_search
] #Computes the real scale by which the image will be multiplied, in the original implementation many
#scales were given in order to compare between them, in our case paramaetre_search is a list of a single element
#Initialise the heatmap and paf matrix
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
print("multiplier:", multiplier)
for m in range(len(multiplier)):
#Rescaling the image according to the multiplier
scale = multiplier[m]
imageToTest = cv.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model['stride'], model['padValue'])
print("scale:", scale)
net.blobs['data'].reshape(*(1, 3, imageToTest_padded.shape[0],
imageToTest_padded.shape[1]))
#net.forward() # dry run, loading the processed image (reshaped) as data in the correct format of input
net.blobs['data'].data[...] = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
start_time = time.time()
#computing in the network:
output_blobs = net.forward()
print('At scale %d, The CNN took %.2f ms.' %
(m, 1000 * (time.time() - start_time)))
# extract outputs, resize, and remove padding
print(list(output_blobs.keys()))
heatmap = np.transpose(
np.squeeze(net.blobs[list(output_blobs.keys())[1]].data),
(1, 2, 0)) # output 1 is heatmaps
heatmap = cv.resize(heatmap, (0, 0),
fx=model['stride'],
fy=model['stride'],
interpolation=cv.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv.INTER_CUBIC)
print("before squeeze",
net.blobs[list(output_blobs.keys())[0]].data.shape)
paf = np.transpose(
np.squeeze(net.blobs[list(output_blobs.keys())[0]].data),
(1, 2, 0)) # output 0 is PAFs
paf = cv.resize(paf, (0, 0),
fx=model['stride'],
fy=model['stride'],
interpolation=cv.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv.resize(paf, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv.INTER_CUBIC)
if m == 0:
heatmap_avg = heatmap / len(multiplier)
paf_avg = paf / len(multiplier)
else:
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
print("Scaling done")
import scipy
from scipy.ndimage.filters import gaussian_filter
all_peaks = []
peak_counter = 0
for part in range(19 - 1):
x_list = []
y_list = []
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up, map >= map_down,
map > param['thre1']))
peaks = zip(np.nonzero(peaks_binary)[1],
np.nonzero(peaks_binary)[0]) # note reverse
list_peaks = list(peaks)
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in list_peaks]
id = range(peak_counter, peak_counter + len(list_peaks))
peaks_with_score_and_id = [
peaks_with_score[i] + (id[i], ) for i in range(len(id))
]
all_peaks.append(peaks_with_score_and_id)
peak_counter = peak_counter + len(list_peaks)
# find connection in the specified sequence, center 29 is in the position 15
limbSeq = [[2,3], [2,6], [3,4], [4,5], [6,7], [7,8], [2,9], [9,10], \
[10,11], [2,12], [12,13], [13,14], [2,1], [1,15], [15,17], \
[1,16], [16,18], [3,17], [6,18]]
# the middle joints heatmap correpondence
mapIdx = [[31,32], [39,40], [33,34], [35,36], [41,42], [43,44], [19,20], [21,22], \
[23,24], [25,26], [27,28], [29,30], [47,48], [49,50], [53,54], [51,52], \
[55,56], [37,38], [45,46]]
connection_all = []
special_k = []
mid_num = 10
print(paf_avg.shape)
for k in range(len(mapIdx)):
score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0] - 1]
candB = all_peaks[limbSeq[k][1] - 1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if (nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
#if norm==0 : norm=1
vec = np.divide(vec, norm)
startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
np.linspace(candA[i][1], candB[j][1], num=mid_num))
list_startend = list(startend)
vec_x = np.array([score_mid[int(round(list_startend[I][1])), int(round(list_startend[I][0])), 0] \
for I in range(len(list_startend))])
vec_y = np.array([score_mid[int(round(list_startend[I][1])), int(round(list_startend[I][0])), 1] \
for I in range(len(list_startend))])
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(
vec_y, vec[1])
score_with_dist_prior = sum(
score_midpts) / len(score_midpts) + min(
0.5 * oriImg.shape[0] / norm - 1, 0)
criterion1 = len(
np.nonzero(score_midpts > param['thre2'])
[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([
i, j, score_with_dist_prior,
score_with_dist_prior + candA[i][2] + candB[j][2]
])
connection_candidate = sorted(connection_candidate,
key=lambda x: x[2],
reverse=True)
connection = np.zeros((0, 5))
for c in range(len(connection_candidate)):
i, j, s = connection_candidate[c][0:3]
if (i not in connection[:, 3] and j not in connection[:, 4]):
connection = np.vstack(
[connection, [candA[i][3], candB[j][3], s, i, j]])
if (len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
print("done")
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
candidate = np.array([item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:, 0]
partBs = connection_all[k][:, 1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): #= 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): #1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][
indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if (subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[partBs[i].astype(int),
2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
print("found = 2")
membership = ((subset[j1] >= 0).astype(int) +
(subset[j2] >= 0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: #merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[
partBs[i].astype(int), 2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(
candidate[connection_all[k][i, :2].astype(int),
2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
# delete some rows of subset which has few parts occur
# delete some rows of subset which has few parts occur
deleteIdx = []
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
# visualize
colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
[0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
cmap = matplotlib.cm.get_cmap('hsv')
canvas = cv.imread(image_path) # B,G,R order
for i in range(18):
rgba = np.array(cmap(1 - i / 18. - 1. / 36))
rgba[0:3] *= 255
for j in range(len(all_peaks[i])):
cv.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)
to_plot = cv.addWeighted(oriImg, 0.3, canvas, 0.7, 0)
# visualize 2
stickwidth = 4
for i in range(17):
for n in range(len(subset)):
index = subset[n][np.array(limbSeq[i]) - 1]
if -1 in index:
continue
cur_canvas = canvas.copy()
Y = candidate[index.astype(int), 0]
X = candidate[index.astype(int), 1]
mX = np.mean(X)
mY = np.mean(Y)
length = ((X[0] - X[1])**2 + (Y[0] - Y[1])**2)**0.5
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
polygon = cv.ellipse2Poly(
(int(mY), int(mX)), (int(length / 2), stickwidth), int(angle),
0, 360, 1)
cv.fillConvexPoly(cur_canvas, polygon, colors[i])
canvas = cv.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
plt.imshow(canvas[:, :, [2, 1, 0]])
fig = matplotlib.pyplot.gcf()
fig.set_size_inches(12, 12)
fig.savefig('../video_frames_processed/frame%d.jpg' % frame_number)
return (fig)
def predict(path, show=False):
test_image = path
oriImg = cv2.imread(test_image) # B,G,R order
# plt.imshow(oriImg[:,:,[2,1,0]])
# Load configuration
param, model_params = config_reader()
multiplier = [
x * model_params['boxsize'] / oriImg.shape[0]
for x in param['scale_search']
]
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv2.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model_params['stride'], model_params['padValue'])
input_img = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 0, 1, 2)) / 256 - 0.5
# required shape (1, width, height, channels)
# print("Input shape: " + str(input_img.shape))
output_blobs = model.predict(input_img)
# print("Output shape (heatmap): " + str(output_blobs[1].shape))
# extract outputs, resize, and remove padding
heatmap = np.squeeze(output_blobs[1]) # output 1 is heatmaps
heatmap = cv2.resize(heatmap, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv2.INTER_CUBIC)
paf = np.squeeze(output_blobs[0]) # output 0 is PAFs
paf = cv2.resize(paf, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv2.resize(paf, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
# peak
from scipy.ndimage.filters import gaussian_filter
all_peaks = []
peak_counter = 0
for part in range(19 - 1):
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up, map >= map_down,
map > param['thre1']))
peaks = list(
zip(np.nonzero(peaks_binary)[1],
np.nonzero(peaks_binary)[0])) # note reverse
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [
peaks_with_score[i] + (id[i], ) for i in range(len(id))
]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
# find connection in the specified sequence, center 29 is in the position 15
limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10],
[10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15],
[15, 17], [1, 16], [16, 18], [3, 17], [6, 18]]
# the middle joints heatmap correpondence
mapIdx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44],
[19, 20], [21, 22], [23, 24], [25, 26], [27, 28], [29, 30],
[47, 48], [49, 50], [53, 54], [51, 52], [55, 56], [37, 38],
[45, 46]]
# connection
connection_all = []
special_k = []
mid_num = 10
for k in range(len(mapIdx)):
score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0] - 1]
candB = all_peaks[limbSeq[k][1] - 1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if (nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
##################################
# 当A和B是同一个点的时候，vec=0，跳过
if (vec == np.array([0, 0])).all():
continue
##################################
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
vec = np.divide(vec, norm)
startend = list(
zip(np.linspace(candA[i][0], candB[j][0], num=mid_num),
np.linspace(candA[i][1], candB[j][1],
num=mid_num)))
vec_x = np.array([
score_mid[int(round(startend[I][1])),
int(round(startend[I][0])), 0]
for I in range(len(startend))
])
vec_y = np.array([
score_mid[int(round(startend[I][1])),
int(round(startend[I][0])), 1]
for I in range(len(startend))
])
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(
vec_y, vec[1])
score_with_dist_prior = sum(
score_midpts) / len(score_midpts) + min(
0.5 * oriImg.shape[0] / norm - 1, 0)
criterion1 = len(
np.nonzero(score_midpts > param['thre2'])
[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([
i, j, score_with_dist_prior,
score_with_dist_prior + candA[i][2] + candB[j][2]
])
connection_candidate = sorted(connection_candidate,
key=lambda x: x[2],
reverse=True)
connection = np.zeros((0, 5))
for c in range(len(connection_candidate)):
i, j, s = connection_candidate[c][0:3]
if (i not in connection[:, 3] and j not in connection[:, 4]):
connection = np.vstack(
[connection, [candA[i][3], candB[j][3], s, i, j]])
if (len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# subset
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
candidate = np.array([item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:, 0]
partBs = connection_all[k][:, 1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): #= 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): #1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][
indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if (subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[partBs[i].astype(int),
2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
# print ("found = 2")
membership = ((subset[j1] >= 0).astype(int) +
(subset[j2] >= 0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: #merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[
partBs[i].astype(int), 2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(
candidate[connection_all[k][i, :2].astype(int),
2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
# delete some rows of subset which has few parts occur
deleteIdx = []
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
# 创建human的dataframe
subset = pd.DataFrame(subset)
candidate = pd.DataFrame(candidate)
human = subset.T.copy()
human.drop([18, 19], axis=0, inplace=True)
for i in human:
human[i] = human[i].apply(lambda n: None if int(n) == -1 else list(
candidate.loc[int(n), [0, 1]]))
if show:
# visualize
colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0],
[170, 255, 0], [85, 255, 0], [0, 255, 0], [0, 255, 85],
[0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255],
[0, 0, 255], [85, 0, 255], [170, 0, 255], [255, 0, 255],
[255, 0, 170], [255, 0, 85]]
cmap = matplotlib.cm.get_cmap('hsv')
canvas = cv2.imread(test_image) # B,G,R order
for i in range(18):
rgba = np.array(cmap(1 - i / 18. - 1. / 36))
rgba[0:3] *= 255
for j in range(len(all_peaks[i])):
cv2.circle(canvas,
all_peaks[i][j][0:2],
4,
colors[i],
thickness=-1)
to_plot = cv2.addWeighted(oriImg, 0.3, canvas, 0.7, 0)
plt.imshow(to_plot[:, :, [2, 1, 0]])
fig = matplotlib.pyplot.gcf()
fig.set_size_inches(12, 12)
return human
def process(input_image, params, model_params):
''' Start of finding the Key points of full body using Open Pose.'''
oriImg = cv2.imread(input_image) # B,G,R order
multiplier = [
x * model_params['boxsize'] / oriImg.shape[0]
for x in params['scale_search']
]
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
for m in range(1):
scale = multiplier[m]
imageToTest = cv2.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model_params['stride'], model_params['padValue'])
input_img = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 0, 1, 2)) # required shape (1, width, height, channels)
output_blobs = model.predict(input_img)
heatmap = np.squeeze(output_blobs[1]) # output 1 is heatmaps
heatmap = cv2.resize(heatmap, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv2.INTER_CUBIC)
paf = np.squeeze(output_blobs[0]) # output 0 is PAFs
paf = cv2.resize(paf, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv2.resize(paf, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
all_peaks = [] #To store all the key points which a re detected.
peak_counter = 0
prinfTick(1) #prints time required till now.
for part in range(18):
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up, map >= map_down,
map > params['thre1']))
peaks = list(
zip(np.nonzero(peaks_binary)[1],
np.nonzero(peaks_binary)[0])) # note reverse
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [
peaks_with_score[i] + (id[i], ) for i in range(len(id))
]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
connection_all = []
special_k = []
mid_num = 10
prinfTick(2) #prints time required till now.
print()
position = checkPosition(all_peaks) #check position of spine.
checkKneeling(all_peaks) #check whether kneeling oernot
checkHandFold(all_peaks) #check whether hands are folding or not.
canvas1 = draw(input_image, all_peaks) #show the image.
return canvas1, position
def extract_feat(model, oriImg, param):
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
# first figure shows padded images
multiplier = [
x * model['boxsize'] / oriImg.shape[0] for x in param['scale_search']
]
# We use 1080 (8G memory), so we set 3 scales here.
for m in range(3): #range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model['stride'], model['padValue'])
#print imageToTest_padded.shape
net.blobs['data'].reshape(*(1, 3, imageToTest_padded.shape[0],
imageToTest_padded.shape[1]))
#net.forward() # dry run
net.blobs['data'].data[...] = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
#start_time = time.time()
output_blobs = net.forward()
#print('At scale %d, The CNN took %.2f ms.' % (m, 1000 * (time.time() - start_time)))
# extract outputs, resize, and remove padding
heatmap = np.transpose(
np.squeeze(net.blobs[output_blobs.keys()[1]].data),
(1, 2, 0)) # output 1 is heatmaps
heatmap = cv.resize(heatmap, (0, 0),
fx=model['stride'],
fy=model['stride'],
interpolation=cv.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
all_peaks = []
peak_counter = 0
for part in range(19 - 1):
x_list = []
y_list = []
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up, map >= map_down,
map > param['thre1']))
peaks = zip(np.nonzero(peaks_binary)[1],
np.nonzero(peaks_binary)[0]) # note reverse
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [
peaks_with_score[i] + (id[i], ) for i in range(len(id))
]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
return all_peaks, heatmap_avg
def get(self, test_image, idx):
oriImg = cv.imread(test_image) # B,G,R order
param, model = config_reader()
multiplier = [
x * model['boxsize'] / oriImg.shape[0]
for x in param['scale_search']
]
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 15))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 26))
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model['stride'], model['padValue'])
#print imageToTest_padded.shape
self.net.blobs['data'].reshape(*(1, 3, imageToTest_padded.shape[0],
imageToTest_padded.shape[1]))
#net.forward() # dry run
self.net.blobs['data'].data[...] = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
start_time = time.time()
output_blobs = self.net.forward()
#print('At scale %d, The CNN took %.2f ms.' % (m, 1000 * (time.time() - start_time)))
# extract outputs, resize, and remove padding
heatmap = np.transpose(
np.squeeze(self.net.blobs[output_blobs.keys()[1]].data),
(1, 2, 0)) # output 1 is heatmaps
heatmap = cv.resize(heatmap, (0, 0),
fx=model['stride'],
fy=model['stride'],
interpolation=cv.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv.INTER_CUBIC)
paf = np.transpose(
np.squeeze(self.net.blobs[output_blobs.keys()[0]].data),
(1, 2, 0)) # output 0 is PAFs
paf = cv.resize(paf, (0, 0),
fx=model['stride'],
fy=model['stride'],
interpolation=cv.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv.resize(paf, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
from scipy.ndimage.filters import gaussian_filter
all_peaks = []
peak_counter = 0
for part in range(15 - 1):
x_list = []
y_list = []
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up,
map >= map_down, map > param['thre1']))
peaks = zip(
np.nonzero(peaks_binary)[1],
np.nonzero(peaks_binary)[0]) # note reverse
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [
peaks_with_score[i] + (id[i], ) for i in range(len(id))
]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
# find connection in the specified sequence, center 29 is in the position 15
limbSeq = [[13, 14], [14, 4], [4, 5], [5, 6], [14, 1], [1, 2], [2, 3],
[14, 10], [10, 11], [11, 12], [14, 7], [7, 8], [8, 9]]
# the middle joints heatmap correpondence
mapIdx = [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13],
[14, 15], [16, 17], [18, 19], [20, 21], [22, 23], [24, 25]]
connection_all = []
special_k = []
mid_num = 10
for k in range(len(mapIdx)):
score_mid = paf_avg[:, :, [x for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0] - 1]
candB = all_peaks[limbSeq[k][1] - 1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if (nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
if norm == 0:
continue
vec = np.divide(vec, norm)
startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
np.linspace(candA[i][1], candB[j][1], num=mid_num))
vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
for I in range(len(startend))])
vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
for I in range(len(startend))])
score_midpts = np.multiply(
vec_x, vec[0]) + np.multiply(vec_y, vec[1])
score_with_dist_prior = sum(score_midpts) / len(
score_midpts) + min(
0.5 * oriImg.shape[0] / norm - 1, 0)
criterion1 = len(
np.nonzero(score_midpts > param['thre2'])
[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([
i, j, score_with_dist_prior,
score_with_dist_prior + candA[i][2] +
candB[j][2]
])
connection_candidate = sorted(connection_candidate,
key=lambda x: x[2],
reverse=True)
connection = np.zeros((0, 5))
for c in range(len(connection_candidate)):
i, j, s = connection_candidate[c][0:3]
if (i not in connection[:, 3]
and j not in connection[:, 4]):
connection = np.vstack(
[connection, [candA[i][3], candB[j][3], s, i, j]])
if (len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
candidate = np.array(
[item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:, 0]
partBs = connection_all[k][:, 1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): #= 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): #1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][
indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if (subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[
partBs[i].astype(int),
2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
#print "found = 2"
membership = ((subset[j1] >= 0).astype(int) +
(subset[j2] >= 0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: #merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[
partBs[i].astype(int),
2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(
candidate[connection_all[k][i, :2].astype(int),
2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
# delete some rows of subset which has few parts occur
deleteIdx = []
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
'''
# visualize
colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
[0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
cmap = matplotlib.cm.get_cmap('hsv')
canvas = cv.imread(test_image) # B,G,R order
for i in range(14):
rgba = np.array(cmap(1 - i/18. - 1./36))
rgba[0:3] *= 255
for j in range(len(all_peaks[i])):
cv.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)
to_plot = cv.addWeighted(oriImg, 0.3, canvas, 0.7, 0)
plt.imshow(to_plot[:,:,[2,1,0]])
fig = matplotlib.pyplot.gcf()
fig.set_size_inches(12, 12)
# visualize 2
stickwidth = 4
for i in range(13):
for n in range(len(subset)):
index = subset[n][np.array(limbSeq[i])-1]
if -1 in index:
continue
cur_canvas = canvas.copy()
Y = candidate[index.astype(int), 0]
X = candidate[index.astype(int), 1]
mX = np.mean(X)
mY = np.mean(Y)
length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
polygon = cv.ellipse2Poly((int(mY),int(mX)), (int(length/2), stickwidth), int(angle), 0, 360, 1)
cv.fillConvexPoly(cur_canvas, polygon, colors[i % 13])
canvas = cv.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
plt.imshow(canvas[:,:,[2,1,0]])
fig = matplotlib.pyplot.gcf()
fig.set_size_inches(12, 12)
plt.savefig(str(idx))
'''
ret = {
'test_image': test_image,
'subset': subset,
'candidate': candidate
}
return ret
def apply_model(oriImg, model, multiplier):
stride = 8
height, width, _ = oriImg.shape
normed_img = np.array(oriImg, dtype=np.float32)
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 25),
dtype=np.float32)
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv2.resize(normed_img, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imgToTest_padded, pad = util.padRightDownCorner(
imageToTest, stride, 128)
input_img = np.transpose(
np.float32(imgToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 255 - 0.5 # required shape (1, c, h, w)
input_var = torch.autograd.Variable(torch.from_numpy(input_img).cuda())
# get the features
heat1, heat2, heat3, heat4, heat5, heat6 = model(input_var)
# get the heatmap
heatmap = heat6.data.cpu().numpy()
heatmap = np.transpose(np.squeeze(heatmap), (1, 2, 0)) # (h, w, c)
heatmap = cv2.resize(heatmap, (0, 0),
fx=stride,
fy=stride,
interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imgToTest_padded.shape[0] -
pad[2], :imgToTest_padded.shape[1] - pad[3], :]
heatmap = cv2.resize(heatmap, (width, height),
interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
all_peaks = [] # all of the possible points by classes.
peak_counter = 0
thre1 = 0.1
for part in range(25 - 1):
x_list = []
y_list = []
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up, map >= map_down,
map > thre1))
peaks = zip(np.nonzero(peaks_binary)[1],
np.nonzero(peaks_binary)[0]) # note reverse
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [
peaks_with_score[i] + (id[i], ) for i in range(len(id))
]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
# sort by score
for i in range(24):
all_peaks[i] = sorted(all_peaks[i],
key=lambda ele: ele[2],
reverse=True)
canvas = oriImg.copy()
# draw points
for i in range(24):
for j in range(len(all_peaks[i])):
if j is 0:
cv2.circle(canvas,
all_peaks[i][j][0:2],
4, [0, 0, 255],
thickness=-1)
else:
cv2.circle(canvas,
all_peaks[i][j][0:2],
4, [255, 0, 0],
thickness=-1)
keypoints = -1 * np.ones((24, 3))
for i in range(24):
if len(all_peaks[i]) == 0:
continue
else:
keypoints[i, 0], keypoints[i, 1], keypoints[
i, 2] = all_peaks[i][0][0], all_peaks[i][0][1], 1
return keypoints, canvas
def do_work(test_image,save_name):
oriImg = cv2.imread(test_image) # B,G,R order
multiplier = [x * model['boxsize'] / oriImg.shape[0] for x in param['scale_search']]
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
print 'len of multipler'
print len(multiplier)
for m in range(len(multiplier)):
scale = multiplier[m]
print scale
imageToTest = cv2.resize(oriImg, (0,0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest, model['stride'], model['padValue'])
print imageToTest_padded.shape
net.blobs['data'].reshape(*(1, 3, imageToTest_padded.shape[0], imageToTest_padded.shape[1]))
#net.forward() # dry run
net.blobs['data'].data[...] = np.transpose(np.float32(imageToTest_padded[:,:,:,np.newaxis]), (3,2,0,1))/256 - 0.5;
start_time = time.time()
output_blobs = net.forward()
print('At scale %d, The CNN took %.2f ms.' % (m, 1000 * (time.time() - start_time)))
start_time = time.time()
# extract outputs, resize, and remove padding
heatmap = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[1]].data), (1,2,0)) # output 1 is heatmaps
heatmap = cv2.resize(heatmap, (0,0), fx=model['stride'], fy=model['stride'], interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0]-pad[2], :imageToTest_padded.shape[1]-pad[3], :]
heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC)
paf = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[0]].data), (1,2,0)) # output 0 is PAFs
paf = cv2.resize(paf, (0,0), fx=model['stride'], fy=model['stride'], interpolation=cv2.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0]-pad[2], :imageToTest_padded.shape[1]-pad[3], :]
paf = cv2.resize(paf, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
print('2 At scale %d, The CNN took %.2f ms.' % (m, 1000 * (time.time() - start_time)))
start_time = time.time()
all_peaks = []
peak_counter = 0
for part in range(19-1):
x_list = []
y_list = []
map_ori = heatmap_avg[:,:,part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:,:] = map[:-1,:]
map_right = np.zeros(map.shape)
map_right[:-1,:] = map[1:,:]
map_up = np.zeros(map.shape)
map_up[:,1:] = map[:,:-1]
map_down = np.zeros(map.shape)
map_down[:,:-1] = map[:,1:]
peaks_binary = np.logical_and.reduce((map>=map_left, map>=map_right, map>=map_up, map>=map_down, map > param['thre1']))
peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]) # note reverse
peaks_with_score = [x + (map_ori[x[1],x[0]],) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [peaks_with_score[i] + (id[i],) for i in range(len(id))]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
connection_all = []
special_k = []
mid_num = 10
for k in range(len(mapIdx)):
score_mid = paf_avg[:,:,[x-19 for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0]-1]
candB = all_peaks[limbSeq[k][1]-1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if(nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0]*vec[0] + vec[1]*vec[1])
vec = np.divide(vec, norm)
startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
np.linspace(candA[i][1], candB[j][1], num=mid_num))
vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
for I in range(len(startend))])
vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
for I in range(len(startend))])
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
score_with_dist_prior = sum(score_midpts)/len(score_midpts) + min(0.5*oriImg.shape[0]/norm-1, 0)
criterion1 = len(np.nonzero(score_midpts > param['thre2'])[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([i, j, score_with_dist_prior, score_with_dist_prior+candA[i][2]+candB[j][2]])
connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
connection = np.zeros((0,5))
for c in range(len(connection_candidate)):
i,j,s = connection_candidate[c][0:3]
if(i not in connection[:,3] and j not in connection[:,4]):
connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
if(len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
candidate = np.array([item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:,0]
partBs = connection_all[k][:,1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): #= 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): #1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if(subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
print "found = 2"
membership = ((subset[j1]>=0).astype(int) + (subset[j2]>=0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: #merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(candidate[connection_all[k][i,:2].astype(int), 2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
# delete some rows of subset which has few parts occur
deleteIdx = [];
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2]/subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
cmap = matplotlib.cm.get_cmap('hsv')
canvas = cv2.imread(test_image) # B,G,R order
for i in range(18):
rgba = np.array(cmap(1 - i/18. - 1./36))
rgba[0:3] *= 255
for j in range(len(all_peaks[i])):
cv2.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)
# visualize 2
stickwidth = 4
for i in range(17):
for n in range(len(subset)):
index = subset[n][np.array(limbSeq[i])-1]
if -1 in index:
continue
cur_canvas = canvas.copy()
Y = candidate[index.astype(int), 0]
X = candidate[index.astype(int), 1]
mX = np.mean(X)
mY = np.mean(Y)
length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
polygon = cv2.ellipse2Poly((int(mY),int(mX)), (int(length/2), stickwidth), int(angle), 0, 360, 1)
cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
cv2.imwrite(save_name,canvas)
print('3 At scale %d, The CNN took %.2f ms.' % (m, 1000 * (time.time() - start_time)))
param_, model_ = config_reader()
test_image = '../sample_image/ski.jpg'
#test_image = 'a.jpg'
oriImg = cv2.imread(test_image) # B,G,R order
multiplier = [x * model_['boxsize'] / oriImg.shape[0] for x in param_['scale_search']]
#heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
#paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv2.resize(oriImg, (0,0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest, model_['stride'], model_['padValue'])
feed = np.transpose(np.float32(imageToTest_padded[:,:,:,np.newaxis]), (3,2,0,1))/256 - 0.5
output1,output2 = model(Variable(torch.from_numpy(feed).float(), volatile=True))
output1 = output1.cpu().data.numpy()
output2 = output2.cpu().data.numpy()
net.blobs['data'].reshape(*(1, 3, imageToTest_padded.shape[0], imageToTest_padded.shape[1]))
net.blobs['data'].data[...] = feed
output_blobs = net.forward()
output1_ = net.blobs[output_blobs.keys()[0]].data
output2_ = net.blobs[output_blobs.keys()[1]].data
def __call__(self, oriImg):
scale_search = [0.5, 1.0, 1.5, 2.0]
boxsize = 368
stride = 8
padValue = 128
thre1 = 0.1
thre2 = 0.05
multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 16))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 28))
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv2.resize(oriImg, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
im = np.ascontiguousarray(im)
data = torch.from_numpy(im).float()
if torch.cuda.is_available():
data = data.cuda()
# data = data.permute([2, 0, 1]).unsqueeze(0).float()
with torch.no_grad():
Mconv7_stage6_L1, Mconv7_stage6_L2 = self.model(data)
Mconv7_stage6_L1 = Mconv7_stage6_L1.cpu().numpy()
Mconv7_stage6_L2 = Mconv7_stage6_L2.cpu().numpy()
# extract outputs, resize, and remove padding
heatmap = np.transpose(np.squeeze(Mconv7_stage6_L2), (1, 2, 0)) # output 1 is heatmaps
heatmap = cv2.resize(heatmap, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC)
# paf = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[0]].data), (1, 2, 0)) # output 0 is PAFs
paf = np.transpose(np.squeeze(Mconv7_stage6_L1), (1, 2, 0)) # output 0 is PAFs
paf = cv2.resize(paf, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv2.resize(paf, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC)
heatmap_avg += heatmap_avg + heatmap / len(multiplier)
paf_avg += + paf / len(multiplier)
all_peaks = []
peak_counter = 0
num_parts = heatmap_avg.shape[-1]
for part in range(num_parts - 1):
map_ori = heatmap_avg[:, :, part]
one_heatmap = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(one_heatmap.shape)
map_left[1:, :] = one_heatmap[:-1, :]
map_right = np.zeros(one_heatmap.shape)
map_right[:-1, :] = one_heatmap[1:, :]
map_up = np.zeros(one_heatmap.shape)
map_up[:, 1:] = one_heatmap[:, :-1]
map_down = np.zeros(one_heatmap.shape)
map_down[:, :-1] = one_heatmap[:, 1:]
peaks_binary = np.logical_and.reduce(
(one_heatmap >= map_left, one_heatmap >= map_right, one_heatmap >= map_up, one_heatmap >= map_down, one_heatmap > thre1))
peaks = list(zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0])) # note reverse
peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]
peak_id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [peaks_with_score[i] + (peak_id[i],) for i in range(len(peak_id))]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
# find connection in the specified sequence
limbSeq = [[1, 2], [2, 3], [3, 4], [4, 5], [2, 6], [6, 7], [7, 8], [2, 15],
[15, 12], [12, 13], [13, 14], [15, 9], [9, 10], [10, 11]]
# the middle joints heatmap correpondence
mapIdx = [[16, 17], [18, 19], [20, 21], [22, 23], [24, 25], [26, 27], [28, 29], [30, 31],
[38, 39], [40, 41], [42, 43], [32, 33], [34, 35], [36, 37]]
connection_all = []
special_k = []
mid_num = 10
for k in range(len(mapIdx)):
score_mid = paf_avg[:, :, [x - 16 for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0] - 1]
candB = all_peaks[limbSeq[k][1] - 1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if (nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
vec = np.divide(vec, norm)
startend = list(zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
np.linspace(candA[i][1], candB[j][1], num=mid_num)))
vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
for I in range(len(startend))])
vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
for I in range(len(startend))])
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
score_with_dist_prior = sum(score_midpts) / len(score_midpts) + min(
0.5 * oriImg.shape[0] / norm - 1, 0)
criterion1 = len(np.nonzero(score_midpts > thre2)[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append(
[i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])
connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
connection = np.zeros((0, 5))
for c in range(len(connection_candidate)):
i, j, s = connection_candidate[c][0:3]
if (i not in connection[:, 3] and j not in connection[:, 4]):
connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
if (len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
candidate = np.array([item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:, 0]
partBs = connection_all[k][:, 1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): # = 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): # 1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if subset[j][indexB] != partBs[i]:
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
membership = ((subset[j1] >= 0).astype(int) + (subset[j2] >= 0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: # merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(candidate[connection_all[k][i, :2].astype(int), 2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
# delete some rows of subset which has few parts occur
deleteIdx = []
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
# subset: n*20 array, 0-17 is the index in candidate, 18 is the total score, 19 is the total parts
# candidate: x, y, score, id
return candidate, subset
def infer(image_path):
weights_path = "model/keras/model.h5"
input_shape = (None, None, 3)
img_input = Input(shape=input_shape)
stages = 6
np_branch1 = 38
np_branch2 = 19
img_normalized = Lambda(lambda x: x / 256 - 0.5)(img_input) # [-0.5, 0.5]
# VGG
stage0_out = vgg_block(img_normalized)
# stage 1
stage1_branch1_out = stage1_block(stage0_out, np_branch1, 1)
stage1_branch2_out = stage1_block(stage0_out, np_branch2, 2)
x = Concatenate()([stage1_branch1_out, stage1_branch2_out, stage0_out])
# stage t >= 2
for sn in range(2, stages + 1):
stageT_branch1_out = stageT_block(x, np_branch1, sn, 1)
stageT_branch2_out = stageT_block(x, np_branch2, sn, 2)
if (sn < stages):
x = Concatenate()(
[stageT_branch1_out, stageT_branch2_out, stage0_out])
model = Model(img_input, [stageT_branch1_out, stageT_branch2_out])
model.load_weights(weights_path)
#%matplotlib inline
import cv2
import matplotlib
import pylab as plt
import numpy as np
import util
test_image = image_path
#test_image = 'sample_images/input.jpg'
oriImg = cv2.imread(test_image)
#resize the image - haritha
#oriImg = cv2.resize(oriImg, (0,0), fx=0.1, fy=0.1)
# B,G,R order
plt.imshow(oriImg[:, :, [2, 1, 0]])
param, model_params = config_reader()
multiplier = [
x * model_params['boxsize'] / oriImg.shape[0]
for x in param['scale_search']
]
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
# first figure shows padded images
f, axarr = plt.subplots(1, len(multiplier))
f.set_size_inches((20, 5))
# second figure shows heatmaps
f2, axarr2 = plt.subplots(1, len(multiplier))
f2.set_size_inches((20, 5))
# third figure shows PAFs
f3, axarr3 = plt.subplots(2, len(multiplier))
f3.set_size_inches((20, 10))
# updated the range to 1 to increase performance - haritha
for m in range(1):
scale = multiplier[m]
imageToTest = cv2.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model_params['stride'], model_params['padValue'])
axarr[m].imshow(imageToTest_padded[:, :, [2, 1, 0]])
axarr[m].set_title('Input image: scale %d' % m)
input_img = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 0, 1, 2)) # required shape (1, width, height, channels)
print("Input shape: " + str(input_img.shape))
#calculate output
output_blobs = model.predict(input_img)
print("Output shape (heatmap): " + str(output_blobs[1].shape))
# extract outputs, resize, and remove padding
heatmap = np.squeeze(output_blobs[1]) # output 1 is heatmaps
heatmap = cv2.resize(heatmap, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv2.INTER_CUBIC)
heatmap_avg = heatmap
paf = np.squeeze(output_blobs[0]) # output 0 is PAFs
paf = cv2.resize(paf, (0, 0),
fx=model_params['stride'],
fy=model_params['stride'],
interpolation=cv2.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv2.resize(paf, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv2.INTER_CUBIC)
##visualization
axarr2[m].imshow(oriImg[:, :, [2, 1, 0]])
ax2 = axarr2[m].imshow(heatmap[:, :, 3], alpha=.5) # right elbow
axarr2[m].set_title('Heatmaps (Relb): scale %d' % m)
axarr3.flat[m].imshow(oriImg[:, :, [2, 1, 0]])
ax3x = axarr3.flat[m].imshow(paf[:, :, 16], alpha=.5) # right elbow
axarr3.flat[m].set_title('PAFs (x comp. of Rwri to Relb): scale %d' %
m)
axarr3.flat[len(multiplier) + m].imshow(oriImg[:, :, [2, 1, 0]])
ax3y = axarr3.flat[len(multiplier) + m].imshow(paf[:, :, 17],
alpha=.5) # right wrist
axarr3.flat[len(multiplier) + m].set_title(
'PAFs (y comp. of Relb to Rwri): scale %d' % m)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
f2.subplots_adjust(right=0.93)
cbar_ax = f2.add_axes([0.95, 0.15, 0.01, 0.7])
_ = f2.colorbar(ax2, cax=cbar_ax)
f3.subplots_adjust(right=0.93)
cbar_axx = f3.add_axes([0.95, 0.57, 0.01, 0.3])
_ = f3.colorbar(ax3x, cax=cbar_axx)
cbar_axy = f3.add_axes([0.95, 0.15, 0.01, 0.3])
_ = f3.colorbar(ax3y, cax=cbar_axy)
plt.imshow(oriImg[:, :, [2, 1, 0]])
plt.imshow(heatmap_avg[:, :, 3], alpha=.5)
plt.imshow(heatmap_avg[:, :, 4], alpha=.5)
fig = matplotlib.pyplot.gcf()
cax = matplotlib.pyplot.gca()
fig.set_size_inches(20, 20)
fig.subplots_adjust(right=0.93)
cbar_ax = fig.add_axes([0.95, 0.15, 0.01, 0.7])
_ = fig.colorbar(ax2, cax=cbar_ax)
from numpy import ma
U = paf_avg[:, :, 16] * -1
V = paf_avg[:, :, 17]
X, Y = np.meshgrid(np.arange(U.shape[1]), np.arange(U.shape[0]))
M = np.zeros(U.shape, dtype='bool')
M[U**2 + V**2 < 0.5 * 0.5] = True
U = ma.masked_array(U, mask=M)
V = ma.masked_array(V, mask=M)
# 1
plt.figure()
plt.imshow(oriImg[:, :, [2, 1, 0]], alpha=.5)
s = 5
Q = plt.quiver(X[::s, ::s],
Y[::s, ::s],
U[::s, ::s],
V[::s, ::s],
scale=50,
headaxislength=4,
alpha=.5,
width=0.001,
color='r')
fig = matplotlib.pyplot.gcf()
fig.set_size_inches(20, 20)
from scipy.ndimage.filters import gaussian_filter
all_peaks = []
peak_counter = 0
for part in range(19 - 1):
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up, map >= map_down,
map > param['thre1']))
peaks = list(
zip(np.nonzero(peaks_binary)[1],
np.nonzero(peaks_binary)[0])) # note reverse
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [
peaks_with_score[i] + (id[i], ) for i in range(len(id))
]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
# find connection in the specified sequence, center 29 is in the position 15
limbSeq = [[2,3], [2,6], [3,4], [4,5], [6,7], [7,8], [2,9], [9,10], \
[10,11], [2,12], [12,13], [13,14], [2,1], [1,15], [15,17], \
[1,16], [16,18], [3,17], [6,18]]
# the middle joints heatmap correpondence
mapIdx = [[31,32], [39,40], [33,34], [35,36], [41,42], [43,44], [19,20], [21,22], \
[23,24], [25,26], [27,28], [29,30], [47,48], [49,50], [53,54], [51,52], \
[55,56], [37,38], [45,46]]
connection_all = []
special_k = []
mid_num = 10
for k in range(len(mapIdx)):
score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0] - 1]
candB = all_peaks[limbSeq[k][1] - 1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if (nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
# failure case when 2 body parts overlaps
if norm == 0:
continue
vec = np.divide(vec, norm)
startend = list(zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
np.linspace(candA[i][1], candB[j][1], num=mid_num)))
vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
for I in range(len(startend))])
vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
for I in range(len(startend))])
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(
vec_y, vec[1])
score_with_dist_prior = sum(
score_midpts) / len(score_midpts) + min(
0.5 * oriImg.shape[0] / norm - 1, 0)
criterion1 = len(
np.nonzero(score_midpts > param['thre2'])
[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([
i, j, score_with_dist_prior,
score_with_dist_prior + candA[i][2] + candB[j][2]
])
connection_candidate = sorted(connection_candidate,
key=lambda x: x[2],
reverse=True)
connection = np.zeros((0, 5))
for c in range(len(connection_candidate)):
i, j, s = connection_candidate[c][0:3]
if (i not in connection[:, 3] and j not in connection[:, 4]):
connection = np.vstack(
[connection, [candA[i][3], candB[j][3], s, i, j]])
if (len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
candidate = np.array([item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:, 0]
partBs = connection_all[k][:, 1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): #= 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): #1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][
indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if (subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[partBs[i].astype(int),
2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
print("found = 2")
membership = ((subset[j1] >= 0).astype(int) +
(subset[j2] >= 0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: #merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[
partBs[i].astype(int), 2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(
candidate[connection_all[k][i, :2].astype(int),
2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
# delete some rows of subset which has few parts occur
deleteIdx = []
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
# visualize
colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
[0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
#what does this line do?
#cmap = matplotlib.cm.get_cmap('hsv')
canvas = cv2.resize(cv2.imread(test_image), (0, 0), fx=0.1,
fy=0.1) # B,G,R order
# for i in range(1):
# rgba = np.array(cmap(1 - i/18. - 1./36))
# rgba[0:3] *= 255
#show wrist (the 4th index is the right wrist)
if (len(all_peaks[4]) > 0):
cv2.circle(canvas, all_peaks[4][0][0:2], 4, colors[1], thickness=-1)
#print(all_peaks)
return (all_peaks[3][1])
def handle_one(oriImg):
# for visualize
canvas = np.copy(oriImg)
imageToTest = Variable(T.transpose(T.transpose(T.unsqueeze(torch.from_numpy(oriImg).float(),0),2,3),1,2),volatile=True).cuda()
print oriImg.shape
scale = model_['boxsize'] / float(oriImg.shape[0])
print scale
h = int(oriImg.shape[0]*scale)
w = int(oriImg.shape[1]*scale)
pad_h = 0 if (h%model_['stride']==0) else model_['stride'] - (h % model_['stride'])
pad_w = 0 if (w%model_['stride']==0) else model_['stride'] - (w % model_['stride'])
new_h = h+pad_h
new_w = w+pad_w
imageToTest = cv2.resize(oriImg, (0,0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest, model_['stride'], model_['padValue'])
imageToTest_padded = np.transpose(np.float32(imageToTest_padded[:,:,:,np.newaxis]), (3,2,0,1))/256 - 0.5
feed = Variable(T.from_numpy(imageToTest_padded)).cuda()
output1,output2 = model(feed)
heatmap = nn.UpsamplingBilinear2d((oriImg.shape[0], oriImg.shape[1])).cuda()(output2)
paf = nn.UpsamplingBilinear2d((oriImg.shape[0], oriImg.shape[1])).cuda()(output1)
print heatmap.size()
print paf.size()
print type(heatmap)
heatmap_avg = T.transpose(T.transpose(heatmap[0],0,1),1,2).data.cpu().numpy()
paf_avg = T.transpose(T.transpose(paf[0],0,1),1,2).data.cpu().numpy()
all_peaks = []
peak_counter = 0
#maps =
for part in range(18):
map_ori = heatmap_avg[:,:,part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:,:] = map[:-1,:]
map_right = np.zeros(map.shape)
map_right[:-1,:] = map[1:,:]
map_up = np.zeros(map.shape)
map_up[:,1:] = map[:,:-1]
map_down = np.zeros(map.shape)
map_down[:,:-1] = map[:,1:]
peaks_binary = np.logical_and.reduce((map>=map_left, map>=map_right, map>=map_up, map>=map_down, map > param_['thre1']))
# peaks_binary = T.eq(
# peaks = zip(T.nonzero(peaks_binary)[0],T.nonzero(peaks_binary)[0])
peaks = zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]) # note reverse
peaks_with_score = [x + (map_ori[x[1],x[0]],) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [peaks_with_score[i] + (id[i],) for i in range(len(id))]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
connection_all = []
special_k = []
mid_num = 10
for k in range(len(mapIdx)):
score_mid = paf_avg[:,:,[x-19 for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0]-1]
candB = all_peaks[limbSeq[k][1]-1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if(nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0]*vec[0] + vec[1]*vec[1])
vec = np.divide(vec, norm)
startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
np.linspace(candA[i][1], candB[j][1], num=mid_num))
vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
for I in range(len(startend))])
vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
for I in range(len(startend))])
score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
score_with_dist_prior = sum(score_midpts)/len(score_midpts) + min(0.5*oriImg.shape[0]/norm-1, 0)
criterion1 = len(np.nonzero(score_midpts > param_['thre2'])[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([i, j, score_with_dist_prior, score_with_dist_prior+candA[i][2]+candB[j][2]])
connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
connection = np.zeros((0,5))
for c in range(len(connection_candidate)):
i,j,s = connection_candidate[c][0:3]
if(i not in connection[:,3] and j not in connection[:,4]):
connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
if(len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
candidate = np.array([item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:,0]
partBs = connection_all[k][:,1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): #= 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): #1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if(subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
print "found = 2"
membership = ((subset[j1]>=0).astype(int) + (subset[j2]>=0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: #merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(candidate[connection_all[k][i,:2].astype(int), 2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
# delete some rows of subset which has few parts occur
deleteIdx = [];
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2]/subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
# canvas = cv2.imread(test_image) # B,G,R order
for i in range(18):
for j in range(len(all_peaks[i])):
cv2.circle(canvas, all_peaks[i][j][0:2], 4, colors[i], thickness=-1)
stickwidth = 4
for i in range(17):
for n in range(len(subset)):
index = subset[n][np.array(limbSeq[i])-1]
if -1 in index:
continue
cur_canvas = canvas.copy()
Y = candidate[index.astype(int), 0]
X = candidate[index.astype(int), 1]
mX = np.mean(X)
mY = np.mean(Y)
length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
polygon = cv2.ellipse2Poly((int(mY),int(mX)), (int(length/2), stickwidth), int(angle), 0, 360, 1)
cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
return canvas
for m in range(len(multiplier)):
scale = multiplier[m]
h = int(oriImg.shape[0] * scale)
w = int(oriImg.shape[1] * scale)
pad_h = 0 if (h % model_['stride']
== 0) else model_['stride'] - (h % model_['stride'])
pad_w = 0 if (w % model_['stride']
== 0) else model_['stride'] - (w % model_['stride'])
new_h = h + pad_h
new_w = w + pad_w
imageToTest = cv2.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, model_['stride'], model_['padValue'])
imageToTest_padded = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
feed = Variable(T.from_numpy(imageToTest_padded)).cuda()
output1, output2 = model(feed)
heatmap = nn.UpsamplingBilinear2d(
(oriImg.shape[0], oriImg.shape[1])).cuda()(output2)
paf = nn.UpsamplingBilinear2d(
(oriImg.shape[0], oriImg.shape[1])).cuda()(output1)
heatmap_avg[m] = heatmap[0].data
paf_avg[m] = paf[0].data
print("oriImg shape: ", oriImg.shape)
multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
print("multiplier: ", multiplier)
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
print("heatmap_avg: ", heatmap_avg.shape, "paf_avg: ", paf_avg.shape)
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv.INTER_CUBIC)
print("imageToTest: ", imageToTest.shape)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride,
padValue)
print("imageToTest_padded: ", imageToTest_padded.shape, "pad: ", pad)
im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
im = np.ascontiguousarray(im)
start_time = time.time()
output = pred_net.eval({pred_net.arguments[0]: [im]})
print("Mconv7_stage6_L1: ", output[Mconv7_stage6_L1].shape,
"Mconv7_stage6_L2: ", output[Mconv7_stage6_L2].shape)
# print output[Mconv7_stage6_L2]
print('At scale %.2f, The CNN took %.2f ms.' %
(scale_search[m], 1000 * (time.time() - start_time)))
# extract outputs, resize, and remove padding
# heatmap = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[1]].data), (1, 2, 0)) # output 1 is heatmaps
heatmap = np.transpose(np.squeeze(output[Mconv7_stage6_L2]),
def inference_and_compute_loss(self, image, annotation):
"""
Perform forward pass on image and compute loss
Inputs
------
image: np.ndarray [H x W]
annotation: dict
Dictionary of MPII annotation containing keys such as 'joint_self', etc.
"""
# Convert joint information to ground truth
people = util.get_openpose_annotation(annotation)
gt_heatmaps = util.get_gaussian_maps(people, image.shape)[:, :, :14]
gt_pafs = util.get_pafs(people, image.shape)
# Get masks
image_name = annotation['img_paths']
if not os.path.exists(os.path.join('masks', 'mask_' + image_name)):
print('Mask does not exist for {}'.format(image_name))
return None
mask = cv2.imread(os.path.join('masks', 'mask_' + image_name))
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
if mask.max() > 1:
mask = np.divide(mask, 255)
scale_search = [0.5, 1.0, 1.5, 2.0]
boxsize = 368
stride = 8
padValue = 128
thre1 = 0.1
thre2 = 0.05
multiplier = [x * boxsize / image.shape[0] for x in scale_search]
heatmap_avg = np.zeros((image.shape[0], image.shape[1], 16))
paf_avg = np.zeros((image.shape[0], image.shape[1], 28))
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv2.resize(image, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
im = np.ascontiguousarray(im)
data = torch.from_numpy(im).float()
if torch.cuda.is_available():
data = data.cuda()
# data = data.permute([2, 0, 1]).unsqueeze(0).float()
with torch.no_grad():
Mconv7_stage6_L1, Mconv7_stage6_L2 = self.model(data)
Mconv7_stage6_L1 = Mconv7_stage6_L1.cpu().numpy()
Mconv7_stage6_L2 = Mconv7_stage6_L2.cpu().numpy()
# extract outputs, resize, and remove padding
heatmap = np.transpose(np.squeeze(Mconv7_stage6_L2), (1, 2, 0)) # output 1 is heatmaps
heatmap = cv2.resize(heatmap, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
# paf = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[0]].data), (1, 2, 0)) # output 0 is PAFs
paf = np.transpose(np.squeeze(Mconv7_stage6_L1), (1, 2, 0)) # output 0 is PAFs
paf = cv2.resize(paf, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv2.resize(paf, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_CUBIC)
heatmap_avg += heatmap_avg + heatmap / len(multiplier)
paf_avg += + paf / len(multiplier)
# Resize outputs to the ground truth size
output_heatmaps = cv2.resize(heatmap_avg[:, :, :14], tuple(reversed(gt_heatmaps.shape[:2])))
output_pafs = cv2.resize(paf_avg, tuple(reversed(gt_heatmaps.shape[:2])))
# Apply mask and compute loss
mask = cv2.resize(mask, (tuple(reversed(gt_heatmaps.shape[:2]))))
mask = np.expand_dims(mask, axis=2)
heatmap_sse = np.multiply(mask, (gt_heatmaps - output_heatmaps) ** 2)
pafs_sse = np.multiply(mask, (gt_pafs - output_pafs) ** 2)
loss = np.sum(heatmap_sse) + np.sum(pafs_sse)
return loss
def KeypointDetection(self, oriImg_path, render_image_path):
caffe.set_mode_gpu()
caffe.set_device(self.param['GPUdeviceNumber'])
oriImg = cv.imread(oriImg_path)
print "image size:", oriImg.shape
multiplier = [
x * self.model['boxsize'] / oriImg.shape[0]
for x in self.param['scale_search']
]
# get heatmap and paf map
heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
for m in range(len(multiplier)):
scale = multiplier[m]
imageToTest = cv.resize(oriImg, (0, 0),
fx=scale,
fy=scale,
interpolation=cv.INTER_CUBIC)
imageToTest_padded, pad = util.padRightDownCorner(
imageToTest, self.model['stride'], self.model['padValue'])
print imageToTest_padded.shape
self.net.blobs['data'].reshape(*(1, 3, imageToTest_padded.shape[0],
imageToTest_padded.shape[1]))
# self.net.forward() # dry run
self.net.blobs['data'].data[...] = np.transpose(
np.float32(imageToTest_padded[:, :, :, np.newaxis]),
(3, 2, 0, 1)) / 256 - 0.5
start_time = time.time()
output_blobs = self.net.forward()
print('At scale %d, The CNN took %.2f ms.' %
(m, 1000 * (time.time() - start_time)))
# extract outputs, resize, and remove padding
heatmap = np.transpose(
np.squeeze(self.net.blobs[output_blobs.keys()[1]].data),
(1, 2, 0)) # output 1 is heatmaps
heatmap = cv.resize(heatmap, (0, 0),
fx=self.model['stride'],
fy=self.model['stride'],
interpolation=cv.INTER_CUBIC)
heatmap = heatmap[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
heatmap = cv.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv.INTER_CUBIC)
paf = np.transpose(
np.squeeze(self.net.blobs[output_blobs.keys()[0]].data),
(1, 2, 0)) # output 0 is PAFs
paf = cv.resize(paf, (0, 0),
fx=self.model['stride'],
fy=self.model['stride'],
interpolation=cv.INTER_CUBIC)
paf = paf[:imageToTest_padded.shape[0] -
pad[2], :imageToTest_padded.shape[1] - pad[3], :]
paf = cv.resize(paf, (oriImg.shape[1], oriImg.shape[0]),
interpolation=cv.INTER_CUBIC)
heatmap_avg = heatmap_avg + heatmap / len(multiplier)
paf_avg = paf_avg + paf / len(multiplier)
# --------------------------------------------
# get all peaks
print heatmap_avg.shape
# plt.imshow(heatmap_avg[:,:,2])
all_peaks = []
peak_counter = 0
for part in range(19 - 1):
x_list = []
y_list = []
map_ori = heatmap_avg[:, :, part]
map = gaussian_filter(map_ori, sigma=3)
map_left = np.zeros(map.shape)
map_left[1:, :] = map[:-1, :]
map_right = np.zeros(map.shape)
map_right[:-1, :] = map[1:, :]
map_up = np.zeros(map.shape)
map_up[:, 1:] = map[:, :-1]
map_down = np.zeros(map.shape)
map_down[:, :-1] = map[:, 1:]
peaks_binary = np.logical_and.reduce(
(map >= map_left, map >= map_right, map >= map_up,
map >= map_down, map > self.param['thre1']))
peaks = zip(
np.nonzero(peaks_binary)[1],
np.nonzero(peaks_binary)[0]) # note reverse
peaks_with_score = [x + (map_ori[x[1], x[0]], ) for x in peaks]
id = range(peak_counter, peak_counter + len(peaks))
peaks_with_score_and_id = [
peaks_with_score[i] + (id[i], ) for i in range(len(id))
]
all_peaks.append(peaks_with_score_and_id)
peak_counter += len(peaks)
# --------------------------------------------
# find connection in the specified sequence, center 29 is in the position 15
limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
[10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
[1, 16], [16, 18], [3, 17], [6, 18]]
# the middle joints heatmap correpondence
mapIdx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44], [19, 20], [21, 22], \
[23, 24], [25, 26], [27, 28], [29, 30], [47, 48], [49, 50], [53, 54], [51, 52], \
[55, 56], [37, 38], [45, 46]]
connection_all = []
special_k = []
mid_num = 10
for k in range(len(mapIdx)):
score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
candA = all_peaks[limbSeq[k][0] - 1]
candB = all_peaks[limbSeq[k][1] - 1]
nA = len(candA)
nB = len(candB)
indexA, indexB = limbSeq[k]
if (nA != 0 and nB != 0):
connection_candidate = []
for i in range(nA):
for j in range(nB):
vec = np.subtract(candB[j][:2], candA[i][:2])
norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
vec = np.divide(vec, norm)
startend = zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
np.linspace(candA[i][1], candB[j][1], num=mid_num))
vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
for I in range(len(startend))])
vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
for I in range(len(startend))])
score_midpts = np.multiply(
vec_x, vec[0]) + np.multiply(vec_y, vec[1])
score_with_dist_prior = sum(score_midpts) / len(
score_midpts) + min(
0.5 * oriImg.shape[0] / norm - 1, 0)
criterion1 = len(
np.nonzero(score_midpts > self.param['thre2'])
[0]) > 0.8 * len(score_midpts)
criterion2 = score_with_dist_prior > 0
if criterion1 and criterion2:
connection_candidate.append([
i, j, score_with_dist_prior,
score_with_dist_prior + candA[i][2] +
candB[j][2]
])
connection_candidate = sorted(connection_candidate,
key=lambda x: x[2],
reverse=True)
connection = np.zeros((0, 5))
for c in range(len(connection_candidate)):
i, j, s = connection_candidate[c][0:3]
if (i not in connection[:, 3]
and j not in connection[:, 4]):
connection = np.vstack(
[connection, [candA[i][3], candB[j][3], s, i, j]])
if (len(connection) >= min(nA, nB)):
break
connection_all.append(connection)
else:
special_k.append(k)
connection_all.append([])
# -----------------------------------------------
# get all person keypoint detection
# last number in each row is the total parts number of that person
# the second last number in each row is the score of the overall configuration
subset = -1 * np.ones((0, 20))
candidate = np.array(
[item for sublist in all_peaks for item in sublist])
for k in range(len(mapIdx)):
if k not in special_k:
partAs = connection_all[k][:, 0]
partBs = connection_all[k][:, 1]
indexA, indexB = np.array(limbSeq[k]) - 1
for i in range(len(connection_all[k])): # = 1:size(temp,1)
found = 0
subset_idx = [-1, -1]
for j in range(len(subset)): # 1:size(subset,1):
if subset[j][indexA] == partAs[i] or subset[j][
indexB] == partBs[i]:
subset_idx[found] = j
found += 1
if found == 1:
j = subset_idx[0]
if (subset[j][indexB] != partBs[i]):
subset[j][indexB] = partBs[i]
subset[j][-1] += 1
subset[j][-2] += candidate[
partBs[i].astype(int),
2] + connection_all[k][i][2]
elif found == 2: # if found 2 and disjoint, merge them
j1, j2 = subset_idx
print "found = 2"
membership = ((subset[j1] >= 0).astype(int) +
(subset[j2] >= 0).astype(int))[:-2]
if len(np.nonzero(membership == 2)[0]) == 0: # merge
subset[j1][:-2] += (subset[j2][:-2] + 1)
subset[j1][-2:] += subset[j2][-2:]
subset[j1][-2] += connection_all[k][i][2]
subset = np.delete(subset, j2, 0)
else: # as like found == 1
subset[j1][indexB] = partBs[i]
subset[j1][-1] += 1
subset[j1][-2] += candidate[
partBs[i].astype(int),
2] + connection_all[k][i][2]
# if find no partA in the subset, create a new subset
elif not found and k < 17:
row = -1 * np.ones(20)
row[indexA] = partAs[i]
row[indexB] = partBs[i]
row[-1] = 2
row[-2] = sum(
candidate[connection_all[k][i, :2].astype(int),
2]) + connection_all[k][i][2]
subset = np.vstack([subset, row])
# -----------------------------------------------------
# delete bad keypoint detection
deleteIdx = []
for i in range(len(subset)):
if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
deleteIdx.append(i)
subset = np.delete(subset, deleteIdx, axis=0)
# -----------------------------------------------------
# visualize to render limbs of each person
colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
[0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
[170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
stickwidth = 4
canvas = oriImg.copy()
for i in range(17):
for n in range(len(subset)): # each subset
index = subset[n][np.array(limbSeq[i]) - 1]
if -1 in index:
continue
cur_canvas = canvas.copy()
Y = candidate[index.astype(int), 0]
X = candidate[index.astype(int), 1]
mX = np.mean(X)
mY = np.mean(Y)
length = ((X[0] - X[1])**2 + (Y[0] - Y[1])**2)**0.5
angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
polygon = cv.ellipse2Poly((int(mY), int(mX)),
(int(length / 2), stickwidth),
int(angle), 0, 360, 1)
cv.fillConvexPoly(cur_canvas, polygon, colors[i])
canvas = cv.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
# plt.imshow(canvas[:, :, [2, 1, 0]])
# fig = matplotlib.pyplot.gcf()
# fig.set_size_inches(12, 12)
cv.imwrite(render_image_path, canvas)
all_keypoints = []
print subset
for i in range(0, len(subset)):
keypoints = []
for j in range(0, 18):
index = subset[i][j].astype(int)
if index == -1:
keypoints.append([0, 0, -1])
continue
X = candidate[index, 0]
Y = candidate[index, 1]
keypoints.append([X, Y, 1])
print keypoints
all_keypoints.append(keypoints)
print len(all_keypoints)
# -----------------------------------------------------
# return center person keypoints
centerIndex = self.getCenterKeypointsIndex(oriImg, all_keypoints)
if centerIndex != -1:
return all_keypoints[centerIndex]
else:
return []
