def __call__(self, img_input):
t = time.time()
img_input = np.transpose(
img_input, axes=[1, 0, 2]).copy() if self.T else img_input
img_batch = self._gen_input_batch(img_input, self._box_size,
self.scales)
# inference
hm, xm, ym, zm = self.sess.run(
[self.heatmap, self.x_heatmap, self.y_heatmap, self.z_heatmap],
{self.input_crops: img_batch})
# average scale outputs
hm_size = self._box_size // self._hm_factor
hm_avg = np.zeros((hm_size, hm_size, self._joints_num))
xm_avg = np.zeros((hm_size, hm_size, self._joints_num))
ym_avg = np.zeros((hm_size, hm_size, self._joints_num))
zm_avg = np.zeros((hm_size, hm_size, self._joints_num))
for i in range(len(self.scales)):
rescale = 1.0 / self.scales[i]
scaled_hm = utils.img_scale(hm[i, :, :, :], rescale)
scaled_x_hm = utils.img_scale(xm[i, :, :, :], rescale)
scaled_y_hm = utils.img_scale(ym[i, :, :, :], rescale)
scaled_z_hm = utils.img_scale(zm[i, :, :, :], rescale)
mid = [scaled_hm.shape[0] // 2, scaled_hm.shape[1] // 2]
hm_avg += scaled_hm[mid[0] - hm_size // 2:mid[0] + hm_size // 2,
mid[1] - hm_size // 2:mid[1] + hm_size // 2, :]
xm_avg += scaled_x_hm[mid[0] - hm_size // 2:mid[0] + hm_size // 2,
mid[1] - hm_size // 2:mid[1] +
hm_size // 2, :]
ym_avg += scaled_y_hm[mid[0] - hm_size // 2:mid[0] + hm_size // 2,
mid[1] - hm_size // 2:mid[1] +
hm_size // 2, :]
zm_avg += scaled_z_hm[mid[0] - hm_size // 2:mid[0] + hm_size // 2,
mid[1] - hm_size // 2:mid[1] +
hm_size // 2, :]
hm_avg /= len(self.scales)
xm_avg /= len(self.scales)
ym_avg /= len(self.scales)
zm_avg /= len(self.scales)
joints_2d = utils.extract_2d_joints_from_heatmap(
hm_avg, self._box_size, self._hm_factor)
joints_3d = utils.extract_3d_joints_from_heatmap(
joints_2d, xm_avg, ym_avg, zm_avg, self._box_size, self._hm_factor)
joints_2d, joints_3d = self._joint_filter(joints_2d, joints_3d)
# if self.T:
# joints_2d = joints_2d[:, ::-1]
# joints_3d = joints_3d[:, [1, 0, 2]]
print('FPS: {:>2.2f}'.format(1 / (time.time() - t)))
if self.plot:
# 2d plotting
frame_square = utils.img_scale_squareify(img_input, self._box_size)
frame_square = utils.draw_limbs_2d(frame_square, joints_2d,
self._joint_parents)
cv2.imshow('2D Prediction', frame_square)
# 3d plotting
self.imshow_3d(self.ax_3d, joints_3d, self._joint_parents)
return joints_2d, joints_3d
def run(self):
# initial BB by HOG detection
self.BB_init()
success, frame = self.cameraCapture.read()
frame = frame.T if self.T else frame
while success and cv2.waitKey(1) == -1:
found, w = self.hog.detectMultiScale(frame)
if len(found) > 0:
self._draw_BB_rect(frame, found[np.argmax(w)])
scale = 400 / self.H
frame = cv2.resize(frame, (0, 0), fx=scale, fy=scale)
cv2.imshow(self._box_init_window_name, frame)
if self._clicked:
self._clicked = False
cv2.destroyWindow(self._box_init_window_name)
break
success, frame = self.cameraCapture.read()
frame = frame.T if self.T else frame
x, y, w, h = self.rect
plt.show()
# main loop
success, frame = self.cameraCapture.read()
frame = frame.T if self.T else frame
while success and cv2.waitKey(1) == -1:
t = time.time()
# crop bounding box from the raw frame
frame_cropped = frame[y:y + h, x:x + w, :]
# crop --> one sqrare input img for CNN
self.frame_square = utils.img_scale_squareify(
frame_cropped, self.box_size)
# one sqrare input img --> a batch of sqrare input imgs
self._create_input_batch()
# sqrare input img batch --CNN net--> 2d and 3d skeleton joint coordinates
self._run_net()
# filter to smooth the joint coordinate results
self._joint_coord_filter()
## plot ##
# 2d plotting
self.frame_square = utils.draw_limbs_2d(self.frame_square,
self.joints_2d,
self.joint_parents)
cv2.imshow('2D Prediction', self.frame_square)
# 3d plotting
self._imshow_3d()
print('FPS: {:>2.2f}'.format(1 / (time.time() - t)))
success, frame = self.cameraCapture.read()
frame = frame.T if self.T else frame
self._exit()
def _gen_input_batch(img_input, box_size, scales):
# any input image --> sqrared input image acceptable for the model
img_square = utils.img_scale_squareify(img_input, box_size)
# generate multi-scale input batch
input_batch = []
for scale in scales:
img = utils.img_scale_padding(img_square,
scale) if scale < 1 else img_square
input_batch.append(img)
# input image range: [0, 255) --> [-0.4, 0.6)
input_batch = np.asarray(input_batch, dtype=np.float32) / 255 - 0.4
return input_batch
1, 15, 1, 2, 3, 1, 5, 6, 14, 8, 9, 14, 11, 12, 14, 14, 1, 4, 7, 10, 13
]
with tf.Session() as sess:
saver = tf.train.import_meta_graph('./models/tf_model/vnect_tf.meta')
saver.restore(sess, tf.train.latest_checkpoint('./models/tf_model/'))
graph = tf.get_default_graph()
input_batch = graph.get_tensor_by_name('Placeholder:0')
heatmap = graph.get_tensor_by_name('split_2:0')
x_heatmap = graph.get_tensor_by_name('split_2:1')
y_heatmap = graph.get_tensor_by_name('split_2:2')
z_heatmap = graph.get_tensor_by_name('split_2:3')
img = cv2.imread('./pic/test_pic.jpg')
img_square = utils.img_scale_squareify(img, box_size)
img_square = img_square[np.newaxis, ...]
hm, xm, ym, zm = sess.run([heatmap, x_heatmap, y_heatmap, z_heatmap],
{input_batch: img_square / 255 - 0.4})
joints_2d = utils.extract_2d_joints_from_heatmap(hm[0, ...], box_size,
hm_factor)
for i in range(21):
if i == 0:
himg = hm[0, :, :, i]
ximg = xm[0, :, :, i]
yimg = ym[0, :, :, i]
zimg = zm[0, :, :, i]
else: