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)