def __init__(self, input_size, output_size): self.graph = tf.Graph() with self.graph.as_default(): self.model = teacher_model.CPM_Model(input_size, output_size, 6, 21, img_type='RGB', is_training=False) sess_config = tf.ConfigProto() sess_config.gpu_options.per_process_gpu_memory_fraction = 0.7 sess_config.gpu_options.allow_growth = True sess_config.allow_soft_placement = True self.sess = tf.Session(config=sess_config, graph=self.graph) self._init_vars() self.saver = tf.train.Saver()
def main(argv): tracker = tracking_module.SelfTracker( [FLAGS.webcam_height, FLAGS.webcam_width], FLAGS.input_size) model = cpm_hand.CPM_Model( input_size=FLAGS.input_size, # 导入model1的模型结构 heatmap_size=FLAGS.heatmap_size, stages=FLAGS.cpm_stages, joints=FLAGS.num_of_joints, img_type=FLAGS.color_channel, is_training=False) print("成功导入1的模型") saver = tf.train.Saver() output_node = tf.get_default_graph().get_tensor_by_name( name=FLAGS.output_node_names) with tf.Session() as sess: saver.restore(sess, './models/weights/cpm_hand') # 恢复model1的参数 print("成功导入1的model参数") tf.train.import_meta_graph( './classify/modelSave/model.ckpt.meta') # 导入model3的模型结构 print("成功导入3的model结构") all_vars = tf.trainable_variables() # print(all_vars[61:]) saver_fenlei = tf.train.Saver( all_vars[62:]) # 这个saver1只导入第62个变量及以后变量,是model3的专属变量 saver_fenlei.restore( sess, tf.train.latest_checkpoint('./classify/modelSave/')) print("成功导入3的model参数") x = tf.get_default_graph().get_tensor_by_name("x:0") logits = tf.get_default_graph().get_tensor_by_name("logits_eval:0") # Create kalman filters if FLAGS.use_kalman: kalman_filter_array = [ cv2.KalmanFilter(4, 2) for _ in range(FLAGS.num_of_joints) ] for _, joint_kalman_filter in enumerate(kalman_filter_array): joint_kalman_filter.transitionMatrix = np.array( [[1, 0, 1, 0], [0, 1, 0, 1], [0, 0, 1, 0], [0, 0, 0, 1]], np.float32) joint_kalman_filter.measurementMatrix = np.array( [[1, 0, 0, 0], [0, 1, 0, 0]], np.float32) joint_kalman_filter.processNoiseCov = np.array( [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], np.float32) * FLAGS.kalman_noise else: kalman_filter_array = None while True: data = [] global receive_flag, do_again, again_time files = os.listdir(local_dir + '/storePic') for f in files: if f == 'flag_ok': receive_flag = 1 if receive_flag or (do_again and again_time < 2): os.system('rm ./storePic/flag_ok') receive_flag = 0 do_again = 0 for i in range(num_of_photo): full_img = cv2.imread("./storePic/test" + str(i) + ".jpg") cv2.waitKey(1) cv2.namedWindow("yuantu") cv2.imshow("yuantu", full_img) cv2.moveWindow("yuantu", 0, 0) test_img = tracker.tracking_by_joints( full_img, joint_detections=joint_detections) crop_full_scale = tracker.input_crop_ratio test_img_copy = test_img.copy() # White balance test_img_wb = utils.img_white_balance(test_img, 5) test_img_input = normalize_and_centralize_img(test_img_wb) t1 = time.time() stage_heatmap_np = sess.run( [output_node], feed_dict={model.input_images: test_img_input}) print('FPS: %.2f' % (1 / (time.time() - t1))) local_img = visualize_result(full_img, stage_heatmap_np, kalman_filter_array, tracker, crop_full_scale, test_img_copy) cv2.namedWindow("local_image") cv2.namedWindow("global_image") cv2.imshow('local_img', local_img.astype(np.uint8)) # 训练用图 cv2.imshow('global_img', full_img.astype(np.uint8)) # 单人大框 cv2.moveWindow("local_image", 1000, 0) cv2.moveWindow("global_image", 0, 2000) cv2.imwrite("./blackPic/black" + str(i) + ".jpg", local_img) local_img = io.imread("./blackPic/black" + str(i) + ".jpg") img = transform.resize(local_img, (w, h)) data.append(np.asarray(img)) print(len(data)) feed_dict = {x: data} classification_result = sess.run(logits, feed_dict) # 打印出预测矩阵 print(classification_result) # 打印出预测矩阵每一行最大值的索引 print(tf.argmax(classification_result, 1).eval()) # 根据索引通过字典对应花的分类 output = tf.argmax(classification_result, 1).eval() for k in range(len(output)): print("第" + str(k + 1) + "帧预测:" + dict[output[k]]) print("#########################################") output = output.tolist() num_of_none = find_most(output) if num_of_none > int((num_of_photo / 2.0)): again_time += 1 do_again = 1 else: again_time = 0 do_again = 0 cv2.destroyAllWindows() print('waiting...')
def main(argv): tf_device = '/gpu:0' with tf.device(tf_device): """Build graph """ if FLAGS.color_channel == 'RGB': input_data = tf.placeholder(dtype=tf.float32, shape=[None, FLAGS.input_size, FLAGS.input_size, 3], name='input_image') else: input_data = tf.placeholder(dtype=tf.float32, shape=[None, FLAGS.input_size, FLAGS.input_size, 1], name='input_image') center_map = tf.placeholder(dtype=tf.float32, shape=[None, FLAGS.input_size, FLAGS.input_size, 1], name='center_map') model = cpm_hand.CPM_Model(FLAGS.stages, FLAGS.joints + 1) model.build_model(input_data, center_map, 1) saver = tf.train.Saver() """Create session and restore weights """ sess = tf.Session() sess.run(tf.global_variables_initializer()) if FLAGS.model_path.endswith('pkl'): model.load_weights_from_file(FLAGS.model_path, sess, False) else: saver.restore(sess, FLAGS.model_path) test_center_map = cpm_utils.gaussian_img(FLAGS.input_size, FLAGS.input_size, FLAGS.input_size / 2, FLAGS.input_size / 2, FLAGS.cmap_radius) test_center_map = np.reshape(test_center_map, [1, FLAGS.input_size, FLAGS.input_size, 1]) # Check weights for variable in tf.trainable_variables(): with tf.variable_scope('', reuse=True): var = tf.get_variable(variable.name.split(':0')[0]) print(variable.name, np.mean(sess.run(var))) if not FLAGS.DEMO_TYPE.endswith(('png', 'jpg')): cam = cv2.VideoCapture(FLAGS.cam_num) # Create kalman filters if FLAGS.KALMAN_ON: kalman_filter_array = [cv2.KalmanFilter(4, 2) for _ in range(FLAGS.joints)] for _, joint_kalman_filter in enumerate(kalman_filter_array): joint_kalman_filter.transitionMatrix = np.array([[1, 0, 1, 0], [0, 1, 0, 1], [0, 0, 1, 0], [0, 0, 0, 1]], np.float32) joint_kalman_filter.measurementMatrix = np.array([[1, 0, 0, 0], [0, 1, 0, 0]], np.float32) joint_kalman_filter.processNoiseCov = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], np.float32) * FLAGS.kalman_noise else: kalman_filter_array = None with tf.device(tf_device): while True: t1 = time.time() if FLAGS.DEMO_TYPE.endswith(('png', 'jpg')): test_img = cpm_utils.read_image(FLAGS.DEMO_TYPE, [], FLAGS.input_size, 'IMAGE') else: test_img = cpm_utils.read_image([], cam, FLAGS.input_size, 'WEBCAM') test_img_resize = cv2.resize(test_img, (FLAGS.input_size, FLAGS.input_size)) print('img read time %f' % (time.time() - t1)) if FLAGS.color_channel == 'GRAY': test_img_resize = np.dot(test_img_resize[..., :3], [0.299, 0.587, 0.114]).reshape( (FLAGS.input_size, FLAGS.input_size, 1)) cv2.imshow('color', test_img.astype(np.uint8)) cv2.imshow('gray', test_img_resize.astype(np.uint8)) cv2.waitKey(1) test_img_input = test_img_resize / 256.0 - 0.5 test_img_input = np.expand_dims(test_img_input, axis=0) if FLAGS.DEMO_TYPE.endswith(('png', 'jpg')): # Inference t1 = time.time() predict_heatmap, stage_heatmap_np = sess.run([model.current_heatmap, model.stage_heatmap, ], feed_dict={'input_image:0': test_img_input, 'center_map:0': test_center_map}) # Show visualized image demo_img = visualize_result(test_img, FLAGS, stage_heatmap_np, kalman_filter_array) cv2.imshow('demo_img', demo_img.astype(np.uint8)) if cv2.waitKey(0) == ord('q'): break print('fps: %.2f' % (1 / (time.time() - t1))) elif FLAGS.DEMO_TYPE == 'MULTI': # Inference t1 = time.time() predict_heatmap, stage_heatmap_np = sess.run([model.current_heatmap, model.stage_heatmap, ], feed_dict={'input_image:0': test_img_input, 'center_map:0': test_center_map}) # Show visualized image demo_img = visualize_result(test_img, FLAGS, stage_heatmap_np, kalman_filter_array) cv2.imshow('demo_img', demo_img.astype(np.uint8)) if cv2.waitKey(1) == ord('q'): break print('fps: %.2f' % (1 / (time.time() - t1))) elif FLAGS.DEMO_TYPE == 'SINGLE': # Inference t1 = time.time() stage_heatmap_np = sess.run([model.stage_heatmap[5]], feed_dict={'input_image:0': test_img_input, 'center_map:0': test_center_map}) # Show visualized image demo_img = visualize_result(test_img, FLAGS, stage_heatmap_np, kalman_filter_array) cv2.imshow('current heatmap', (demo_img).astype(np.uint8)) if cv2.waitKey(1) == ord('q'): break print('fps: %.2f' % (1 / (time.time() - t1))) elif FLAGS.DEMO_TYPE == 'HM': # Inference t1 = time.time() stage_heatmap_np = sess.run([model.stage_heatmap[FLAGS.stages - 1]], feed_dict={'input_image:0': test_img_input, 'center_map:0': test_center_map}) print('fps: %.2f' % (1 / (time.time() - t1))) demo_stage_heatmap = stage_heatmap_np[len(stage_heatmap_np) - 1][0, :, :, 0:FLAGS.joints].reshape( (FLAGS.hmap_size, FLAGS.hmap_size, FLAGS.joints)) demo_stage_heatmap = cv2.resize(demo_stage_heatmap, (FLAGS.input_size, FLAGS.input_size)) vertical_imgs = [] tmp_img = None joint_coord_set = np.zeros((FLAGS.joints, 2)) for joint_num in range(FLAGS.joints): # Concat until 4 img if (joint_num % 4) == 0 and joint_num != 0: vertical_imgs.append(tmp_img) tmp_img = None demo_stage_heatmap[:, :, joint_num] *= (255 / np.max(demo_stage_heatmap[:, :, joint_num])) # Plot color joints if np.min(demo_stage_heatmap[:, :, joint_num]) > -50: joint_coord = np.unravel_index(np.argmax(demo_stage_heatmap[:, :, joint_num]), (FLAGS.input_size, FLAGS.input_size)) joint_coord_set[joint_num, :] = joint_coord color_code_num = (joint_num // 4) if joint_num in [0, 4, 8, 12, 16]: if PYTHON_VERSION == 3: joint_color = list( map(lambda x: x + 35 * (joint_num % 4), joint_color_code[color_code_num])) else: joint_color = map(lambda x: x + 35 * (joint_num % 4), joint_color_code[color_code_num]) cv2.circle(test_img, center=(joint_coord[1], joint_coord[0]), radius=3, color=joint_color, thickness=-1) else: if PYTHON_VERSION == 3: joint_color = list( map(lambda x: x + 35 * (joint_num % 4), joint_color_code[color_code_num])) else: joint_color = map(lambda x: x + 35 * (joint_num % 4), joint_color_code[color_code_num]) cv2.circle(test_img, center=(joint_coord[1], joint_coord[0]), radius=3, color=joint_color, thickness=-1) # Put text tmp = demo_stage_heatmap[:, :, joint_num].astype(np.uint8) tmp = cv2.putText(tmp, 'Min:' + str(np.min(demo_stage_heatmap[:, :, joint_num])), org=(5, 20), fontFace=cv2.FONT_HERSHEY_COMPLEX, fontScale=0.3, color=150) tmp = cv2.putText(tmp, 'Mean:' + str(np.mean(demo_stage_heatmap[:, :, joint_num])), org=(5, 30), fontFace=cv2.FONT_HERSHEY_COMPLEX, fontScale=0.3, color=150) tmp_img = np.concatenate((tmp_img, tmp), axis=0) \ if tmp_img is not None else tmp # Plot limbs for limb_num in range(len(limbs)): if np.min(demo_stage_heatmap[:, :, limbs[limb_num][0]]) > -2000 and np.min( demo_stage_heatmap[:, :, limbs[limb_num][1]]) > -2000: x1 = joint_coord_set[limbs[limb_num][0], 0] y1 = joint_coord_set[limbs[limb_num][0], 1] x2 = joint_coord_set[limbs[limb_num][1], 0] y2 = joint_coord_set[limbs[limb_num][1], 1] length = ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5 if length < 10000 and length > 5: deg = math.degrees(math.atan2(x1 - x2, y1 - y2)) polygon = cv2.ellipse2Poly((int((y1 + y2) / 2), int((x1 + x2) / 2)), (int(length / 2), 3), int(deg), 0, 360, 1) color_code_num = limb_num // 4 if PYTHON_VERSION == 3: limb_color = list( map(lambda x: x + 35 * (limb_num % 4), joint_color_code[color_code_num])) else: limb_color = map(lambda x: x + 35 * (limb_num % 4), joint_color_code[color_code_num]) cv2.fillConvexPoly(test_img, polygon, color=limb_color) if tmp_img is not None: tmp_img = np.lib.pad(tmp_img, ((0, vertical_imgs[0].shape[0] - tmp_img.shape[0]), (0, 0)), 'constant', constant_values=(0, 0)) vertical_imgs.append(tmp_img) # Concat horizontally output_img = None for col in range(len(vertical_imgs)): output_img = np.concatenate((output_img, vertical_imgs[col]), axis=1) if output_img is not None else \ vertical_imgs[col] output_img = output_img.astype(np.uint8) output_img = cv2.applyColorMap(output_img, cv2.COLORMAP_JET) test_img = cv2.resize(test_img, (300, 300), cv2.INTER_LANCZOS4) cv2.imshow('hm', output_img) cv2.moveWindow('hm', 2000, 200) cv2.imshow('rgb', test_img) cv2.moveWindow('rgb', 2000, 750) if cv2.waitKey(1) == ord('q'): break
def main(): with open(sys.argv[1]) as f: conf = json.load(f) images_folder = conf['images_folder'] heatmaps_output_folder = conf['heatmaps_output_folder'] masks_output_folder = conf['masks_output_folder'] images_output_folder = conf['images_output_folder'] model_path = conf['model_path'] #1. Read list of images from folder images_ids = [a for a in os.listdir(images_folder) if a.endswith(".jpg")] #2. tf_device = "/gpu:0" with tf.device(tf_device): input_data = tf.placeholder(dtype=tf.float32, shape=[None, input_size, input_size, 3], name='input_image') center_map = tf.placeholder(dtype=tf.float32, shape=[None, input_size, input_size, 1], name='center_map') model = cpm_hand.CPM_Model(stages, joints + 1) model.build_model(input_data, center_map, 1) saver = tf.train.Saver() sess = tf.Session() sess.run(tf.global_variables_initializer()) model.load_weights_from_file(model_path, sess, False) test_center_map = cpm_utils.gaussian_img(input_size, input_size, input_size / 2, input_size / 2, cmap_radius) test_center_map = np.reshape(test_center_map, [1, input_size, input_size, 1]) with tf.device(tf_device): for im_id in images_ids: print(im_id) im_path = os.path.join(images_folder, im_id) test_img = cpm_utils.read_image(im_path, [], input_size, 'IMAGE') test_img_resize = cv2.resize(test_img, (input_size, input_size)) test_img_input = test_img_resize / 256.0 - 0.5 test_img_input = np.expand_dims(test_img_input, axis=0) predict_heatmap, stage_heatmap_np = sess.run([ model.current_heatmap, model.stage_heatmap, ], feed_dict={ 'input_image:0': test_img_input, 'center_map:0': test_center_map }) last_heatmap = stage_heatmap_np[len(stage_heatmap_np) - 1][0, :, :, 0:joints].reshape( (hmap_size, hmap_size, joints)) heatmap_path = os.path.join(heatmaps_output_folder, im_id + ".npy") np.save(heatmap_path, last_heatmap) last_heatmap = cv2.resize(last_heatmap, (test_img.shape[1], test_img.shape[0])) demo_img, masked_img = visualize_result(test_img, last_heatmap) heatmap_path = os.path.join(heatmaps_output_folder, im_id + ".npy") mask_path = os.path.join(masks_output_folder, im_id) im_path = os.path.join(images_output_folder, im_id) imageio.imsave(mask_path, masked_img) imageio.imsave(im_path, demo_img)
def main(argv): """Build graph """ batch_x, batch_c, batch_y, batch_x_orig = tf_utils.read_batch_cpm(FLAGS.tfr_data_files, FLAGS.input_size, FLAGS.heatmap_size, FLAGS.num_of_joints, FLAGS.center_radius, FLAGS.batch_size) if FLAGS.color_channel == 'RGB': input_placeholder = tf.placeholder(dtype=tf.float32, shape=(FLAGS.batch_size, FLAGS.input_size, FLAGS.input_size, 3), name='input_placeholer') elif FLAGS.color_channel == 'GRAY': input_placeholder = tf.placeholder(dtype=tf.float32, shape=(FLAGS.batch_size, FLAGS.input_size, FLAGS.input_size, 1), name='input_placeholer') cmap_placeholder = tf.placeholder(dtype=tf.float32, shape=(FLAGS.batch_size, FLAGS.input_size, FLAGS.input_size, 1), name='cmap_placeholder') hmap_placeholder = tf.placeholder(dtype=tf.float32, shape=( FLAGS.batch_size, FLAGS.heatmap_size, FLAGS.heatmap_size, FLAGS.num_of_joints + 1), name='hmap_placeholder') model = cpm_hand.CPM_Model(FLAGS.stages, FLAGS.num_of_joints + 1) model.build_model(input_placeholder, cmap_placeholder, FLAGS.batch_size) model.build_loss(hmap_placeholder, FLAGS.lr, FLAGS.lr_decay_rate, FLAGS.lr_decay_step) print('=====Model Build=====\n') """Training """ with tf.Session() as sess: # Create dataset queue coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(coord=coord) ## Create summary tf_writer = tf.summary.FileWriter(FLAGS.log_dir, sess.graph, filename_suffix=FLAGS.log_file_name) ## Create model saver saver = tf.train.Saver(max_to_keep=None) # Init init = tf.global_variables_initializer() sess.run(init) # Restore weights if FLAGS.pretrained_model is not None: if FLAGS.pretrained_model.endswith('.pkl'): model.load_weights_from_file(FLAGS.pretrained_model, sess, finetune=True) # Check weights for variable in tf.trainable_variables(): with tf.variable_scope('', reuse=True): var = tf.get_variable(variable.name.split(':0')[0]) print(variable.name, np.mean(sess.run(var))) else: saver.restore(sess, FLAGS.pretrained_model) # check weights for variable in tf.trainable_variables(): with tf.variable_scope('', reuse=True): var = tf.get_variable(variable.name.split(':0')[0]) print(variable.name, np.mean(sess.run(var))) while True: # Read in batch data batch_x_np, batch_y_np, batch_c_np = sess.run([batch_x, batch_y, batch_c]) # Warp training images for img_num in range(batch_x_np.shape[0]): deg1 = (2 * np.random.rand() - 1) * 50 deg2 = (2 * np.random.rand() - 1) * 50 batch_x_np[img_num, ...] = cpm_utils.warpImage(batch_x_np[img_num, ...], 0, deg1, deg2, 1, 30) batch_y_np[img_num, ...] = cpm_utils.warpImage(batch_y_np[img_num, ...], 0, deg1, deg2, 1, 30) batch_y_np[img_num, :, :, FLAGS.num_of_joints] = np.ones(shape=(FLAGS.input_size, FLAGS.input_size)) - \ np.max( batch_y_np[img_num, :, :, 0:FLAGS.num_of_joints], axis=2) batch_c_np[img_num, ...] = cpm_utils.warpImage(batch_c_np[img_num, ...], 0, deg1, deg2, 1, 30).reshape( (FLAGS.input_size, FLAGS.input_size, 1)) # Convert image to grayscale if FLAGS.color_channel == 'GRAY': batch_x_gray_np = np.zeros((batch_x_np.shape[0], FLAGS.input_size, FLAGS.input_size, 1)) for img_num in range(batch_x_np.shape[0]): tmp = batch_x_np[img_num, ...] tmp += 0.5 tmp *= 255 tmp = np.dot(tmp[..., :3], [0.114, 0.587, 0.299]) tmp /= 255 tmp -= 0.5 batch_x_gray_np[img_num, ...] = tmp.reshape((FLAGS.input_size, FLAGS.input_size, 1)) batch_x_np = batch_x_gray_np # Recreate heatmaps gt_heatmap_np = cpm_utils.make_gaussian_batch(batch_y_np, FLAGS.heatmap_size, 3) # Update once stage_losses_np, total_loss_np, _, summary, current_lr, \ stage_heatmap_np, global_step = sess.run([model.stage_loss, model.total_loss, model.train_op, model.merged_summary, model.lr, model.stage_heatmap, model.global_step ], feed_dict={input_placeholder: batch_x_np, cmap_placeholder: batch_c_np, hmap_placeholder: gt_heatmap_np}) # Write logs tf_writer.add_summary(summary, global_step) # Draw intermediate results if global_step % 50 == 0: if FLAGS.color_channel == 'GRAY': demo_img = np.repeat(batch_x_np[0], 3, axis=2) demo_img += 0.5 elif FLAGS.color_channel == 'RGB': demo_img = batch_x_np[0] + 0.5 demo_stage_heatmaps = [] for stage in range(FLAGS.stages): demo_stage_heatmap = stage_heatmap_np[stage][0, :, :, 0:FLAGS.num_of_joints].reshape( (FLAGS.heatmap_size, FLAGS.heatmap_size, FLAGS.num_of_joints)) demo_stage_heatmap = cv2.resize(demo_stage_heatmap, (FLAGS.input_size, FLAGS.input_size)) demo_stage_heatmap = np.amax(demo_stage_heatmap, axis=2) demo_stage_heatmap = np.reshape(demo_stage_heatmap, (FLAGS.input_size, FLAGS.input_size, 1)) demo_stage_heatmap = np.repeat(demo_stage_heatmap, 3, axis=2) demo_stage_heatmaps.append(demo_stage_heatmap) demo_gt_heatmap = gt_heatmap_np[0, :, :, 0:FLAGS.num_of_joints].reshape( (FLAGS.heatmap_size, FLAGS.heatmap_size, FLAGS.num_of_joints)) demo_gt_heatmap = cv2.resize(demo_gt_heatmap, (FLAGS.input_size, FLAGS.input_size)) demo_gt_heatmap = np.amax(demo_gt_heatmap, axis=2) demo_gt_heatmap = np.reshape(demo_gt_heatmap, (FLAGS.input_size, FLAGS.input_size, 1)) demo_gt_heatmap = np.repeat(demo_gt_heatmap, 3, axis=2) if FLAGS.stages > 4: upper_img = np.concatenate((demo_stage_heatmaps[0], demo_stage_heatmaps[1], demo_stage_heatmaps[2]), axis=1) blend_img = 0.5 * demo_gt_heatmap + 0.5 * demo_img lower_img = np.concatenate((demo_stage_heatmaps[FLAGS.stages - 1], demo_gt_heatmap, blend_img), axis=1) demo_img = np.concatenate((upper_img, lower_img), axis=0) cv2.imshow('current heatmap', (demo_img * 255).astype(np.uint8)) cv2.waitKey(1000) else: upper_img = np.concatenate((demo_stage_heatmaps[FLAGS.stages - 1], demo_gt_heatmap, demo_img), axis=1) cv2.imshow('current heatmap', (upper_img * 255).astype(np.uint8)) cv2.waitKey(1000) print('##========Iter {:>6d}========##'.format(global_step)) print('Current learning rate: {:.8f}'.format(current_lr)) for stage_num in range(FLAGS.stages): print('Stage {} loss: {:>.3f}'.format(stage_num + 1, stage_losses_np[stage_num])) print('Total loss: {:>.3f}\n\n'.format(total_loss_np)) # Save models if global_step % 5000 == 1: save_path_str = 'models/' + FLAGS.saved_model_name saver.save(sess=sess, save_path=save_path_str, global_step=global_step) print('\nModel checkpoint saved...\n') # Finish training if global_step == FLAGS.training_iterations: break coord.request_stop() coord.join(threads) print('Training done.')
def main(argv): tf_device = '/gpu:0' with tf.device(tf_device): """Build graph """ if FLAGS.color_channel == 'RGB': input_data = tf.placeholder( dtype=tf.float32, shape=[None, FLAGS.input_size, FLAGS.input_size, 3], name='input_image') else: input_data = tf.placeholder( dtype=tf.float32, shape=[None, FLAGS.input_size, FLAGS.input_size, 1], name='input_image') center_map = tf.placeholder( dtype=tf.float32, shape=[None, FLAGS.input_size, FLAGS.input_size, 1], name='center_map') model = cpm_hand.CPM_Model(FLAGS.stages, FLAGS.joints + 1) model.build_model(input_data, center_map, 1) saver = tf.train.Saver() """Create session and restore weights """ sess = tf.Session() sess.run(tf.global_variables_initializer()) if FLAGS.model_path.endswith('pkl'): model.load_weights_from_file(FLAGS.model_path, sess, False) else: saver.restore(sess, FLAGS.model_path) test_center_map = cpm_utils.gaussian_img(FLAGS.input_size, FLAGS.input_size, FLAGS.input_size / 2, FLAGS.input_size / 2, FLAGS.cmap_radius) test_center_map = np.reshape(test_center_map, [1, FLAGS.input_size, FLAGS.input_size, 1]) # Check weights for variable in tf.trainable_variables(): with tf.variable_scope('', reuse=True): var = tf.get_variable(variable.name.split(':0')[0]) print(variable.name, np.mean(sess.run(var))) if not FLAGS.DEMO_TYPE.endswith(('png', 'jpg')): cam = cv2.VideoCapture(FLAGS.cam_num) # Create kalman filters if FLAGS.KALMAN_ON: kalman_filter_array = [ cv2.KalmanFilter(4, 2) for _ in range(FLAGS.joints) ] for _, joint_kalman_filter in enumerate(kalman_filter_array): joint_kalman_filter.transitionMatrix = np.array( [[1, 0, 1, 0], [0, 1, 0, 1], [0, 0, 1, 0], [0, 0, 0, 1]], np.float32) joint_kalman_filter.measurementMatrix = np.array( [[1, 0, 0, 0], [0, 1, 0, 0]], np.float32) joint_kalman_filter.processNoiseCov = np.array( [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]], np.float32) * FLAGS.kalman_noise else: kalman_filter_array = None with tf.device(tf_device): while True: t1 = time.time() if FLAGS.DEMO_TYPE.endswith(('png', 'jpg')): test_img = cpm_utils.read_image(FLAGS.DEMO_TYPE, [], FLAGS.input_size, 'IMAGE') else: test_img = cpm_utils.read_image([], cam, FLAGS.input_size, 'WEBCAM') test_img_resize = cv2.resize(test_img, (FLAGS.input_size, FLAGS.input_size)) print('img read time %f' % (time.time() - t1)) if FLAGS.color_channel == 'GRAY': test_img_resize = np.dot(test_img_resize[..., :3], [0.299, 0.587, 0.114]).reshape( (FLAGS.input_size, FLAGS.input_size, 1)) cv2.imshow('color', test_img.astype(np.uint8)) cv2.imshow('gray', test_img_resize.astype(np.uint8)) cv2.waitKey(1) test_img_input = test_img_resize / 256.0 - 0.5 test_img_input = np.expand_dims(test_img_input, axis=0) if FLAGS.DEMO_TYPE.endswith(('png', 'jpg')): # Inference t1 = time.time() predict_heatmap, stage_heatmap_np = sess.run( [ model.current_heatmap, model.stage_heatmap, ], feed_dict={ 'input_image:0': test_img_input, 'center_map:0': test_center_map }) # Show visualized image demo_img = visualize_result(test_img, FLAGS, stage_heatmap_np, kalman_filter_array) print('fps: %.2f' % (1 / (time.time() - t1))) imageio.imwrite("asdf.jpg", demo_img)