def read_tfrecord_test():
demo = TFRecordDataset()
iterator, next_element = demo.generateDataset(
tfrecord_path="./tfrecord_dataset/train.tfrecords", batch_size=64)
sess = tf.Session()
for i in range(1000):
sess.run(iterator.initializer)
while True:
try:
images, binned_labels, cont_labels = sess.run(next_element)
print(binned_labels[0], cont_labels[0])
resultimg = images[0]
resultimg = cv2.cvtColor(resultimg, cv2.COLOR_RGB2BGR)
resultimg = cv2.resize(resultimg, (400, 400))
#utils.plot_pose_cube(nimg, cont_labels[0], cont_labels[1], cont_labels[2], tdx=None, tdy=None, size=150.)
resultimg = resultimg * 128.0 + 127.5
resultimg = resultimg.astype(np.uint8)
utils.plot_pose_cube(resultimg,
cont_labels[0][0],
cont_labels[0][1],
cont_labels[0][2],
tdx=200,
tdy=200,
size=150.)
#utils.draw_axis(resultimg, cont_labels[0][0], cont_labels[0][1], cont_labels[0][2], tdx=200, tdy=200, size = 100)
cv2.imshow('test', resultimg)
cv2.waitKey(0)
except tf.errors.OutOfRangeError:
print("End of dataset")
break
def test(self, test_dataset, show_result=False):
yaw_error = .0
pitch_error = .0
roll_error = .0
landmark_error = .0
total_time = .0
total_samples = 0
test_dataset.set_normalization(False)
for images, [batch_yaw, batch_pitch, batch_roll, batch_landmark] in test_dataset:
start_time = time.time()
batch_yaw_pred, batch_pitch_pred, batch_roll_pred, batch_landmark_pred = self.predict_batch(images, normalize=True)
total_time += time.time() - start_time
total_samples += np.array(images).shape[0]
batch_yaw = batch_yaw[:, 0]
batch_pitch = batch_pitch[:, 0]
batch_roll = batch_roll[:, 0]
# Mean absolute error
yaw_error += np.sum(np.abs(batch_yaw - batch_yaw_pred))
pitch_error += np.sum(np.abs(batch_pitch - batch_pitch_pred))
roll_error += np.sum(np.abs(batch_roll - batch_roll_pred))
landmark_error += np.sum(np.abs(batch_landmark - batch_landmark_pred))
# Show result
if show_result:
for i in range(images.shape[0]):
image = images[i]
yaw = batch_yaw_pred[i]
pitch = batch_pitch_pred[i]
roll = batch_roll_pred[i]
landmark = batch_landmark_pred[i]
image = utils.draw_landmark(image, landmark)
image = utils.plot_pose_cube(image, yaw, pitch, roll, tdx=image.shape[1] // 2, tdy=image.shape[0] // 2, size=80)
cv2.imshow("Test result", image)
cv2.waitKey(0)
avg_time = total_time / total_samples
avg_fps = 1.0 / avg_time
print("### MAE: ")
print("- Yaw MAE: {}".format(yaw_error / total_samples))
print("- Pitch MAE: {}".format(pitch_error / total_samples))
print("- Roll MAE: {}".format(roll_error / total_samples))
print("- Head pose MAE: {}".format((yaw_error + pitch_error + roll_error) / total_samples / 3))
print("- Landmark MAE: {}".format(landmark_error / total_samples / 10))
print("- Avg. FPS: {}".format(avg_fps))
roll_predicted = F.softmax(roll)
# Get continuous predictions in degrees.
yaw_predicted = torch.sum(
yaw_predicted.data * idx_tensor) * 3 - 99
pitch_predicted = torch.sum(
pitch_predicted.data * idx_tensor) * 3 - 99
roll_predicted = torch.sum(
roll_predicted.data * idx_tensor) * 3 - 99
# Print new frame with cube and axis
txt_out.write(
str(frame_num) + ' %f %f %f\n' %
(yaw_predicted, pitch_predicted, roll_predicted))
utils.plot_pose_cube(frame,
yaw_predicted,
pitch_predicted,
roll_predicted, (x_min + x_max) / 2,
(y_min + y_max) / 2,
size=bbox_width)
utils.draw_axis(frame,
yaw_predicted,
pitch_predicted,
roll_predicted,
tdx=(x_min + x_max) / 2,
tdy=(y_min + y_max) / 2,
size=bbox_height / 2)
# Plot expanded bounding box
# cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0,255,0), 1)
#cv2.waitKey(0)
cv2.putText(frame, str(yaw_predicted), (00, 250),
cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2,
def inference():
net = caffe.Net(args.caffe_prototxt_path, args.caffe_model_path,
caffe.TEST)
input_size = [int(v.strip()) for v in args.input_size.split(",")]
witdh = input_size[0]
height = input_size[1]
# priors = define_img_size(input_size)
net.blobs['input'].reshape(1, 3, height, witdh)
result_path = args.results_path
imgs_path = args.imgs_path
if not os.path.exists(result_path):
os.makedirs(result_path)
listdir = os.listdir(imgs_path)
for file_path in listdir:
img_path = os.path.join(imgs_path, file_path)
img_ori = cv2.imread(img_path)
img_det = cv2.cvtColor(img_ori, cv2.COLOR_BGR2RGB)
result = det.detect_face(img_det)
img_h, img_w = img_ori.shape[:2]
for i in range(len(result)):
box = result[i]['box']
cls = result[i]['cls']
pts = result[i]['pts']
x1, y1, x2, y2 = box
cv2.rectangle(img_ori, (x1, y1), (x2, y2), (255, 255, 25))
w = x2 - x1 + 1
h = y2 - y1 + 1
size_w = int(max([w, h]) * 0.9)
size_h = int(max([w, h]) * 0.9)
cx = x1 + w // 2
cy = y1 + h // 2
x1 = cx - size_w // 2
x2 = x1 + size_w
y1 = cy - int(size_h * 0.4)
y2 = y1 + size_h
left = 0
top = 0
bottom = 0
right = 0
if x1 < 0:
left = -x1
if y1 < 0:
top = -y1
if x1 >= img_w:
right = x2 - img_w
if y1 >= img_h:
bottom = y2 - img_h
x1 = max(0, x1)
y1 = max(0, y1)
x2 = min(img_w, x2)
y2 = min(img_h, y2)
cropped = img_ori[y1:y2, x1:x2]
print(top, bottom, left, right)
cropped = cv2.copyMakeBorder(cropped, top, bottom, left, right,
cv2.BORDER_CONSTANT, 0)
tmp_batch = np.zeros([1, 3, height, witdh], dtype=np.float32)
cropped = cv2.resize(cropped, (witdh, height))
image = (cropped - image_mean) / image_std
tmp_batch[0, :, :, :] = image.transpose(2, 0, 1)
net.blobs['input'].data[...] = tmp_batch
time_time = time.time()
res = net.forward()
poses = res['pose'][0]
landms = res['landms'][0]
poses = poses * 180 / np.pi
landms = landms
print("inference time: {} s".format(
round(time.time() - time_time, 4)))
for i in range(98):
lx, ly = (int(landms[i * 2] * size_w + x1 - left),
int(landms[i * 2 + 1] * size_h + y1 - bottom))
cv2.circle(img_ori, (lx, ly), 1, (0, 255, 255), 2)
plot_pose_cube(img_ori,
poses[0],
poses[1],
poses[2],
tdx=pts['nose'][0],
tdy=pts['nose'][1],
size=(x2 - x1) // 2)
cv2.imwrite(os.path.join(result_path, file_path), img_ori)
# print("result_pic is written to {}".format(os.path.join(result_path, file_path)))
# cv2.imshow("show", img_ori)
# cv2.waitKey(-1)
cv2.destroyAllWindows()
pitch_predicted = utils.softmax_temperature(pitch.data, 1)
roll_predicted = utils.softmax_temperature(roll.data, 1)
yaw_predicted = torch.sum(yaw_predicted * idx_tensor, 1).cpu() * 3 - 99
pitch_predicted = torch.sum(pitch_predicted * idx_tensor, 1).cpu() * 3 - 99
roll_predicted = torch.sum(roll_predicted * idx_tensor, 1).cpu() * 3 - 99
# Mean absolute error
yaw_error += torch.sum(torch.abs(yaw_predicted - label_yaw))
pitch_error += torch.sum(torch.abs(pitch_predicted - label_pitch))
roll_error += torch.sum(torch.abs(roll_predicted - label_roll))
# Save first image in batch with pose cube or axis.
if args.save_viz:
name = name[0]
if args.dataset == 'BIWI':
cv2_img = cv2.imread(os.path.join(args.data_dir, name + '_rgb.png'))
else:
cv2_img = cv2.imread(os.path.join(args.data_dir, name + '.jpg'))
if args.batch_size == 1:
error_string = 'y %.2f, p %.2f, r %.2f' % (torch.sum(torch.abs(yaw_predicted - label_yaw)), torch.sum(torch.abs(pitch_predicted - label_pitch)), torch.sum(torch.abs(roll_predicted - label_roll)))
cv2.putText(cv2_img, error_string, (30, cv2_img.shape[0]- 30), fontFace=1, fontScale=1, color=(0,0,255), thickness=2)
utils.plot_pose_cube(cv2_img, yaw_predicted[0], pitch_predicted[0], roll_predicted[0], size=100)
utils.draw_axis(cv2_img, yaw_predicted[0], pitch_predicted[0], roll_predicted[0], tdx = 200, tdy= 200, size=100)
cv2.imwrite(os.path.join('output/images', name + '.jpg'), cv2_img)
print('Test error in degrees of the model on the ' + str(total) + ' test images. Yaw: %.4f, Pitch: %.4f, Roll: %.4f'% (yaw_error / total,
pitch_error / total, roll_error / total))
end = timeit.timeit()
print "The inference ran for % seconds." % (end - start)
