def main(_): # check the running platform if platform.uname()[1] != 'dragonx-H97N-WIFI': print("Now it knows it's in a remote cluster") FLAGS.system = "remote" FLAGS.data_path = "/home/lxiaol9/ARC/EASTRNN/data/ICDAR/train" FLAGS.vis_path = "/home/lxiaol9/ARC/EASTRNN/vis/LSTM/" FLAGS.save_path = "/work/cascades/lxiaol9/ARC/EAST/checkpoints/LSTM_east/" + now.strftime("%Y%m%d-%H%M%S") FLAGS.video_path = "/home/lxiaol9/ARC/EASTRNN/data/ICDAR2013/train/" FLAGS.checkpoints_path = FLAGS.save_path FLAGS.pretrained_model_path = "/work/cascades/lxiaol9/ARC/EAST/checkpoints/east/20180921-173054/" print("############## Step1: The environment path has been set up ###############") gpus = [x.name for x in device_lib.list_local_devices() if x.device_type == "GPU"] if (FLAGS.num_gpus > len(gpus)): raise ValueError("Your machine has only %d gpus " "which is less than the requested --num_gpus=%d." % (len(gpus), FLAGS.num_gpus)) config = get_config(FLAGS) config.batch_size = 8 config.num_layers = 3 config.num_steps = 10))
flags.DEFINE_boolean("source", False, "whether load data from source") flags.DEFINE_boolean("dis_plt", False, "whether using pyplot real-time display ") flags.DEFINE_integer('save_checkpoint_steps', 1000, '') flags.DEFINE_integer('save_summary_steps', 100, '') flags.DEFINE_string( 'pretrained_model_path', '/media/dragonx/DataStorage/ARC/EASTRNN/weights/EAST/resnet_v1_50.ckpt', '') flags.DEFINE_string('geometry', 'RBOX', 'set for bb') FLAGS = flags.FLAGS save_path = '/media/dragonx/DataStorage/ARC/EASTRNN/checkpoints/LSTM/' #train_input = DetectorInputMul(save_path, 1, 2, 0) print("data has been loaded") config = get_config(FLAGS) # Global initializer for Variables in the model gpu_options = tf.GPUOptions(allow_growth=True) #global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0), trainable=False) # log: May 3rd, we need to adapt the model input, with config with tf.name_scope("Train"): # use placeholder to stand for input and targets initializer = tf.random_normal_initializer() x_train = tf.placeholder(tf.float32, shape=[None, config.num_steps, None, None, 3]) m = ArrayModel(True, config, x_train, reuse_variables=None, initializer=initializer) with tf.name_scope("Val"):
def main(_): # to increase the code robustness if platform.uname()[1] != 'dragonx-H97N-WIFI': print("Now it knows it's in a remote cluster") FLAGS.system = "remote" FLAGS.data_path = "/home/lxiaol9/ARC/EASTRNN/data/GAP_process/" FLAGS.vis_path = "/home/lxiaol9/ARC/EASTRNN/vis/LSTM/" FLAGS.save_path = "/home/lxiaol9/ARC/EASTRNN/checkpoints/LSTM/" + now.strftime( "%Y%m%d-%H%M%S") FLAGS.video_path = "/home/lxiaol9/ARC/EASTRNN/data/ICDAR/train/" FLAGS.checkpoints_path = "/home/lxiaol9/ARC/EASTRNN/weights/EAST/east_icdar2015_resnet_v1_50_rbox/" FLAGS.pretrained_model_path = "/home/lxiaol9/ARC/EASTRNN/weights/EAST/resnet_v1_50.ckpt" if not FLAGS.data_path: raise ValueError("Must set --") print( "############## Step1: The environment path has been set up ###############" ) if not FLAGS.data_path: raise ValueError("Must set --") gpus = [ x.name for x in device_lib.list_local_devices() if x.device_type == "GPU" ] if FLAGS.num_gpus > len(gpus): raise ValueError("Your machine has only %d gpus " "which is less than the requested --num_gpus=%d." % (len(gpus), FLAGS.num_gpus)) config = get_config(FLAGS) config.batch_size = 1 with tf.Graph().as_default(): # Global initializer for Variables in the model initializer = tf.random_normal_initializer() # Construct the model graph with tf.name_scope("Train"): initializer = tf.random_normal_initializer() # use placeholder to stand for input and targets x_train = tf.placeholder( tf.float32, shape=[None, config.num_steps, None, None, 3]) model = ArrayModel(True, config, x_train, reuse_variables=None, initializer=initializer) # ======================== initialize from the saved weights ============================# if platform.uname()[1] != 'dragonx-H97N-WIFI': checkpoint_path = "/home/lxiaol9/ARC/EASTRNN/checkpoints/LSTM/" else: checkpoint_path = "/media/dragonx/DataStorage/ARC/EASTRNN/checkpoints/ARC/checkpoints/LSTM/" if not os.path.exists(checkpoint_path): raise RuntimeError( 'Checkpoint `{}` not found'.format(checkpoint_path)) saver = tf.train.Saver() # restore the model from weights session = tf.Session(config=tf.ConfigProto(allow_soft_placement=False)) model_path = os.path.join(checkpoint_path, '20180818-170857-11100') logger.info('Restore from {}'.format(model_path)) saver.restore(session, model_path) print( "##############Step 2: Weight restoring successfully ################" ) ################### load all data into memory ################### if FLAGS.source is True: datapath = FLAGS.data_path test_input = HeatInputMul(datapath, 1, 12, 2) input = DetectorInputMul( datapath, 1, 12, 1) # datapath, video_start, video_end, dimension else: datapath = FLAGS.data_path test_input = HeatInputMul(datapath, 1, 12, 0) input = DetectorInputMul( datapath, 1, 12, 1) # we will use input.targets[videos, frames, vect] print( "##############Step 3: Heatmap, GT data is ready now################" ) ################### choose video and frame to test ############### i = test_input.video_name.index('Video_37_2_3') iters = 0.0 state = session.run(model.initial_state) # tensors dict fetches = { #"cost": m.cost, "final_state": model.final_state, # "loss": m.loss, "heat_map_pred": model.heatmap_predict } # frames sequence: 0, 1, 2;3, 4, 5;6, 7, 8 for step in range(int((test_input.cnt_frame[i] - 1) / 3)): feed_dict = {} data = np.zeros([ config.batch_size, config.num_steps, config.shape[0], config.shape[1], 3 ], dtype=np.float32) heat_maps = np.zeros([ config.batch_size, config.num_steps, int(config.shape[0] / 4), int(config.shape[0] / 4), 1 ], dtype=np.float32) # randomly choosing starting frame frame_set = [] video_file = FLAGS.video_path + test_input.video_name[i] + '.mp4' gt_path = FLAGS.data_path + test_input.video_name[i] + '/gt/' # frame number to choose j = step * 3 cap = cv2.VideoCapture(video_file) for m in range(config.num_steps): cap.set(1, (j + m)) ret, frame = cap.read() data[0, m, :, :, :] = cv2.resize( frame, (config.shape[0], config.shape[1])) heat_maps[0, m, :, :, 0] = cv2.resize( np.squeeze( np.load(gt_path + 'frame' + '{0:03d}'.format(j + m + 1) + '.npy')), (128, 128)) # print('choosing starting frame %d, with num_steps is %d' % (j, config.num_steps)) cap.release() frame_set.append(j) feed_dict[model.input_data] = data feed_dict[model.input_heat_maps] = heat_maps for i, (c, h) in enumerate(model.initial_state): feed_dict[c] = state[i].c feed_dict[h] = state[i].h vals = session.run(fetches, feed_dict=feed_dict) state = vals["final_state"] heat_map_pred = vals["heat_map_pred"] print(len(heat_map_pred)) # iters += config.num_steps npyname1 = FLAGS.vis_path + 'Video_37_2_3/' + 'frame' + format( j, '03d') npyname2 = FLAGS.vis_path + 'Video_37_2_3/' + 'frame' + format( j + 1, '03d') npyname3 = FLAGS.vis_path + 'Video_37_2_3/' + 'frame' + format( j + 2, '03d') np.save(npyname1, heat_map_pred[0]) np.save(npyname2, heat_map_pred[1]) np.save(npyname3, heat_map_pred[2]) print("saving frame at %d" % (j))
def main(): checkpoint_path = '/media/dragonx/DataStorage/ARC/EAST/checkpoints/LSTM_east/' idname1 = '20180924-191410' idname2 = '20180924-191410-5001' test_data_path = '/media/dragonx/DataLight/ICDAR2013/test/' save_path = '/media/dragonx/DataLight/ICDAR2013/test_results_lstm/' filename = '/media/dragonx/DataLight/ICDAR2013/test/Video_6_3_2.mp4' idx = 0 # initial frame number config = get_config(FLAGS) config.batch_size = 1 config.num_layers = 3 config.num_steps = 10 #>>>>>>>>>>>>>>>>>>>>>>Sort test video>>>>>>>>>>>>>>>>>>>>>>>>>>># video_set = [] for root, dirs, files in os.walk(test_data_path): for file in files: if file.endswith('.mp4'): video_set.append(os.path.splitext(file)[0]) index = range(0, 1) # parser for running outside # parser = argparse.ArgumentParser() # parser.add_argument('--checkpoint-path', default=checkpoint_path) # args = parser.parse_args() if not os.path.exists(checkpoint_path): raise RuntimeError('Checkpoint `{}` not found'.format(checkpoint_path)) logger.info('loading model') #>>>>>>>>>>>>>>>>>>>>>>> Loading Model >>>>>>>>>>>>>>>>>>>>>>>>># gpu_options = tf.GPUOptions(allow_growth=True) input_images = tf.placeholder(tf.float32, shape=[None, None, None, 3], name='input_images') # global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0), trainable=False) # Global initializer for Variables in the model # log: May 3rd, we need to adapt the model input, with config # with tf.name_scope("Train"): # # use placeholder to stand for input and targets # initializer = tf.random_normal_initializer() # x_train = tf.placeholder(tf.float32, shape=[None, config.num_steps, None, None, 3]) # m = ArrayModel(True, config, x_train, reuse_variables=None, initializer=initializer) with tf.name_scope("Val"): # use placeholder to stand for input and targets initializer = tf.random_normal_initializer() x_val = tf.placeholder(tf.float32, shape=[None, config.num_steps, None, None, 3]) model = ArrayModel(False, config, x_val, reuse_variables=None, initializer=initializer) var_total = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) print(var_total) #>>>>>>>>>>>>>>>>>>>>>>>> restore the model from weights>>>>>>>># soft_placement = False # var_list1 = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='feature_fusion') # var_list2 = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='resnet_v1_50') # var_list3 = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='multi_rnn_cell') # var_list4 = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='pred_module') # var_list = var_list1 + var_list2 + var_list3 + var_list4 # saver = tf.train.Saver({v.op.name: v for v in var_list}) saver = tf.train.Saver() config_proto = tf.ConfigProto(allow_soft_placement=soft_placement) # with sv.managed_session(config=config_proto) as session: # if FLAGS.restore: # print('continue training from previous checkpoint') # # ckpt = tf.train.latest_checkpoint(FLAGS.checkpoints_path) # ckpt = checkpoint_path + idname1 + '/' + idname2 # sv.saver.restore(session, ckpt) model_path = checkpoint_path + idname1 + '/' + idname2 sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) logger.info('Restore from {}'.format(model_path)) saver.restore(sess, model_path) #>>>>>>>>>>>>>>>>>>>>>>Start evaluation>>>>>>>>>>>>>>>>>>>>>>>>># P_test = [] R_test = [] f1_test = [] for k in index: P_video = [] R_video = [] f1_video = [] video_save = save_path + video_set[k] + idname1 + '_' + idname2 + '.avi' t_start = time.time() # sort up all the paths xml_solo_path = test_data_path + video_set[k] raw_video_path = test_data_path + video_set[k] + '.mp4' cap = cv2.VideoCapture(raw_video_path) frame_width = int(cap.get(3)) frame_height = int(cap.get(4)) cnt_frame = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) out = cv2.VideoWriter(video_save, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 10, (frame_width, frame_height)) # 1. load both polys and tags; 2. generate geo maps(the format of polys and tags need to match) polys_array_list, tags_array_list, id_list_list, frame_num = load_annotations_solo(xml_solo_path, \ 1, cnt_frame, frame_width, frame_height) #>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>loop over frames in the time steps >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> for i in range(int(cnt_frame / config.num_steps)): data_seq = np.zeros((1, config.num_steps, 512, 512, 3), dtype=np.float32) data_original = np.zeros( (1, config.num_steps, frame_height, frame_width, 3), dtype=np.float32) for j in range(config.num_steps): ret, frame = cap.read() # im_resized = cv2.resize(frame, (int(512), int(512))) im_resized = frame[0:512, 0:512, :] data_original[0, j, :, :, :] = frame data_seq[0, j, :, :, :] = im_resized #>>>>>>>>>>>>>>>>>>>>>>>>>Now it's time to run the model>>>>>>>>>>>>>>>>>>>>>>>>>> state = sess.run(model.initial_state) # tensors dict to run fetches = { "score_map": model.score_map_set, "geometry_map": model.geometry_set } feed_dict = {} feed_dict[model.input_data] = data_seq for i, (c, h) in enumerate(model.initial_state): feed_dict[c] = state[i].c feed_dict[h] = state[i].h timer = collections.OrderedDict([('net', 0), ('restore', 0), ('nms', 0)]) start = time.time() vals = sess.run(fetches, feed_dict=feed_dict) timer['net'] = time.time() - start #>>>>>>>>>>>>>>>>>>>>>>>>Okay!!!We could evalute the results now>>>>>>>>>>>>>>>>>>> for j in range(config.num_steps): rtparams = collections.OrderedDict() rtparams['start_time'] = datetime.datetime.now().isoformat() rtparams['image_size'] = '{}x{}'.format( frame_width, frame_height) # im_resized, (ratio_h, ratio_w) = resize_image(img) ratio_h, ratio_w = 512 / frame_height, 512 / frame_width rtparams['working_size'] = '{}x{}'.format(512, 512) # results refinement via NMS score = vals["score_map"][j] geometry = vals["geometry_map"][j] boxes, timer = detect(score_map=score, geo_map=geometry, timer=timer) logger.info( 'net {:.0f}ms, restore {:.0f}ms, nms {:.0f}ms'.format( timer['net'] * 1000, timer['restore'] * 1000, timer['nms'] * 1000)) if boxes is not None: scores = boxes[:, 8].reshape(-1) boxes = boxes[:, :8].reshape((-1, 4, 2)) boxes[:, :, 0] /= ratio_w boxes[:, :, 1] /= ratio_h duration = time.time() - start timer['overall'] = duration logger.info('[timing] {}'.format(duration)) text_lines = [] if boxes is not None: text_lines = [] for box, score in zip(boxes, scores): box = sort_poly(box.astype(np.int32)) if np.linalg.norm(box[0] - box[1]) < 5 or np.linalg.norm( box[3] - box[0]) < 5: continue tl = collections.OrderedDict( zip([ 'x0', 'y0', 'x1', 'y1', 'x2', 'y2', 'x3', 'y3' ], map(float, box.flatten()))) tl['score'] = float(score) text_lines.append(tl) pred = { 'text_lines': text_lines, 'rtparams': rtparams, 'timing': timer, } text_polys, text_tags = polys_array_list[ i * 10 + j], tags_array_list[i * 10 + j] text_polys, text_tags = check_and_validate_polys( text_polys, text_tags, (frame_height, frame_width)) # out.write(new_img) #>>>>>>>>>>>>>>>>>>>>>>>>Evaluation>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> targets = text_polys precision, recall, f1 = eval_single_frame(targets, pred) P_video.append(precision) R_video.append(recall) f1_video.append(f1) img = data_original[0, j, :, :, :] new_img = draw_illu(img.copy(), pred) new_img1 = draw_illu_gt(new_img.copy(), targets, precision, recall, f1) out.write(new_img1) # using for pre-testing if j == 0 and FLAGS.vis: fig1 = plt.figure(figsize=(20, 10)) fig1.add_subplot(1, 2, 1) plt.imshow(new_img) plt.title("Text Detection with fine-tuned EAST") fig1.add_subplot(1, 2, 2) plt.imshow(new_img1) plt.title('Text Detection Results Comparison') plt.show() if cv2.waitKey(25) & 0xFF == ord('q'): break # time.sleep(.100) else: break # evaluation on ret and gt P_test.append(np.array(P_video, dtype=np.float32)) R_test.append(np.array(R_video, dtype=np.float32)) f1_test.append(np.array(f1_video, dtype=np.float32)) print(P_video) print(R_video) print(f1_video) print("testing results are P:{}, R:{}, F1:{} on ".format( sum(P_video) / cnt_frame, sum(R_video) / cnt_frame, sum(f1_video) / cnt_frame) + video_set[k]) cap.release() out.release() # results refinement via NMS cv2.destroyAllWindows() print('here is the precision') for item in P_test: print(np.mean(item)) print('here is the recall') for item in R_test: print(np.mean(item)) print('here is the f-score') for item in f1_test: print(np.mean(item)) print(video_set)
def main(argv=None): m_cfg = sys_cfg() config = get_config(FLAGS) config.batch_size = FLAGS.batch_size_per_gpu * FLAGS.num_gpus config.num_layers = 3 config.num_steps = 5 # eval_config = get_config(FLAGS) eval_config.batch_size = 2 eval_config.num_layers = 3 eval_config.num_steps = 5 #============================ I. Model options ==============================# #>>>>>>>>>>>>>>>for PWCnet module network nn_opts = deepcopy(_DEFAULT_PWCNET_VAL_OPTIONS) if FLAGS.flownet_type is 'small': nn_opts['use_dense_cx'] = False nn_opts['use_res_cx'] = False nn_opts['pyr_lvls'] = 6 nn_opts['flow_pred_lvl'] = 2 nn_opts[ 'ckpt_path'] = '/work/cascades/lxiaol9/ARC/PWC/checkpoints/pwcnet-sm-6-2-multisteps-chairsthingsmix/pwcnet.ckpt-592000' # Model to eval else: nn_opts['use_dense_cx'] = True nn_opts['use_res_cx'] = True nn_opts['pyr_lvls'] = 6 nn_opts['flow_pred_lvl'] = 2 nn_opts[ 'ckpt_path'] = '/work/cascades/lxiaol9/ARC/PWC/checkpoints/pwcnet-lg-6-2-multisteps-chairsthingsmix/pwcnet.ckpt-595000' nn_opts['verbose'] = True nn_opts['batch_size'] = 32 # This is Batch_size per GPU(16*4/2/2 = 16) nn_opts[ 'use_tf_data'] = False # Don't use tf.data reader for this simple task nn_opts['gpu_devices'] = ['/device:GPU:0', '/device:GPU:1'] # nn_opts['controller'] = '/device:CPU:0' # Evaluate on CPU or GPU? nn_opts['adapt_info'] = (1, 436, 1024, 2) nn_opts['x_shape'] = [2, 512, 512, 3] # image pairs input shape [2, H, W, 3] nn_opts['y_shape'] = [512, 512, 2] # u,v flows output shape [H, W, 2] #>>>>>>>>>>>>>>>> For EAST module network east_opts = { 'verbose': True, 'ckpt_path': FLAGS.pretrained_model_path, 'batch_size': 40, 'batch_size_per_gpu': 20, 'gpu_devices': ['/device:GPU:0', '/device:GPU:1'], # controller device to put the model's variables on (usually, /cpu:0 or /gpu:0 -> try both!) 'controller': '/device:CPU:0', 'x_dtype': tf.float32, # image pairs input type 'x_shape': [512, 512, 3], # image pairs input shape [2, H, W, 3] 'y_score_shape': [128, 128, 1], # u,v flows output type 'y_geometry_shape': [128, 128, 5], # u,v flows output shape [H, W, 2] 'x_mask_shape': [128, 128, 1] } os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list #=============================== II. building graph for east + agg =================================# # 1.1 Input placeholders batch_size = FLAGS.batch_size_per_gpu * FLAGS.num_gpus len_seq = FLAGS.num_steps # input_images = tf.placeholder(tf.float32, shape=[batch_size*len_seq, 512, 512, 3], name='input_images') input_feat_maps = tf.placeholder(tf.float32, shape=[batch_size, len_seq, 128, 128, 32], name='input_feature_maps') input_flow_maps = tf.placeholder( tf.float32, shape=[batch_size, len_seq - 1, 128, 128, 2], name='input_flow_maps') input_score_maps = tf.placeholder(tf.float32, shape=[batch_size, len_seq, 128, 128, 1], name='input_score_maps') if FLAGS.geometry == 'RBOX': input_geo_maps = tf.placeholder( tf.float32, shape=[batch_size, len_seq, 128, 128, 5], name='input_geo_maps') else: input_geo_maps = tf.placeholder( tf.float32, shape=[batch_size, len_seq, 128, 128, 8], name='input_geo_maps') input_training_masks = tf.placeholder( tf.float32, shape=[batch_size, len_seq, 128, 128, 1], name='input_training_masks') # 1.2 lr & opt global_step = tf.get_variable('global_step', [], initializer=tf.constant_initializer(0), trainable=False) learning_rate = tf.train.exponential_decay(FLAGS.learning_rate, global_step, decay_steps=500, decay_rate=0.8, staircase=True) opt = tf.train.AdamOptimizer(learning_rate) # 1.3 add summary tf.summary.scalar('learning_rate', learning_rate) # tf.summary.image('input_images', input_images[2:20:5, :, :, :]) # 1.4 build graph in tf # input_images_split = tf.split(input_images, FLAGS.num_gpus) input_flow_maps_split = tf.split(input_flow_maps, FLAGS.num_gpus) input_feature_split = tf.split(input_feat_maps, FLAGS.num_gpus) input_score_maps_split = tf.split(input_score_maps, FLAGS.num_gpus) input_geo_maps_split = tf.split(input_geo_maps, FLAGS.num_gpus) input_training_masks_split = tf.split(input_training_masks, FLAGS.num_gpus) tower_grads = [] reuse_variables = None tvars = [] gpus = list(range(len(FLAGS.gpu_list.split(',')))) for i, gpu_id in enumerate(gpus): with tf.device('/gpu:%d' % gpu_id): with tf.name_scope('model_%d' % gpu_id) as scope: iis = input_feature_split[i] ifms = input_flow_maps_split[i] isms = input_score_maps_split[i] igms = input_geo_maps_split[i] itms = input_training_masks_split[i] # model changed to recurrent one, we only need the recurrent loss returned total_loss, model_loss = model_gru_agg.tower_loss( iis, ifms, isms, igms, itms, gpu_id=gpu_id, config=config, reuse_variables=reuse_variables) batch_norm_updates_op = tf.group( *tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope)) reuse_variables = True # tvar1 = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='tiny_embed') # tvar2 = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='pred_module') # tvars = tvar1 + tvar2 # , var_list=tvars grads = opt.compute_gradients(total_loss) tower_grads.append(grads) # 1.5 gradient parsering grads = average_gradients(tower_grads) apply_gradient_op = opt.apply_gradients(grads, global_step=global_step) # 1.6 get training operations summary_op = tf.summary.merge_all() # variable_averages = tf.train.ExponentialMovingAverage( # FLAGS.moving_average_decay, global_step) # variables_averages_op = variable_averages.apply(tf.trainable_variables()) with tf.control_dependencies([apply_gradient_op, batch_norm_updates_op]): train_op = tf.no_op(name='train_op') # 1.8 Saver & Session & Restore saver = tf.train.Saver(tf.global_variables()) # sv = tf.train.Supervisor() summary_writer = tf.summary.FileWriter(FLAGS.checkpoint_path, tf.get_default_graph()) init = tf.global_variables_initializer() g = tf.get_default_graph() with g.as_default(): config1 = tf.ConfigProto() config1.gpu_options.allow_growth = True config1.allow_soft_placement = True sess1 = tf.Session(config=config1) if FLAGS.restore: print('continue training from previous checkpoint') ckpt = FLAGS.prev_checkpoint_path saver.restore(sess1, ckpt) else: sess1.run(init) var_list1 = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='multi_rnn_cell') # var_list2 = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='pred_module') var_list_part1 = var_list1 saver_alter1 = tf.train.Saver( {v.op.name: v for v in var_list_part1}) # # var_list3 = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='tiny_embed') # # var_list4 = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='pred_module') # # var_list_part2 = var_list3 + var_list4 # # saver_alter2 = tf.train.Saver({v.op.name: v for v in var_list_part2}) print('continue training from previous weights') ckpt1 = FLAGS.prev_checkpoint_path print('Restore from {}'.format(ckpt1)) saver_alter1.restore(sess1, ckpt1) # # print('continue training from previous Flow weights') # # ckpt2 = FLAGS.prev_checkpoint_path # # print('Restore from {}'.format(ckpt2)) # # saver_alter2.restore(sess1, ckpt2) #============================= III. Other necessary componets before training =============================# print("Step 1: AGG model has been reconstructed") GPUtil.showUtilization() # >>>>>>>>>>>>>>>>>>>>>>>>>>>>> EAST model >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> # east_net = model_flow_east.EAST(mode='test', options=east_opts) print("Step 2: EAST model has been reconstructed") GPUtil.showUtilization() # >>>>>>>>>>>>>>>>>>>>>>>>>>>>> PWCnet model >>>>>>>>>>>>>>>>>>>>>>>>>>>>>># nn = ModelPWCNet(mode='test', options=nn_opts) print("Step 3: PWC model has been reconstructed") GPUtil.showUtilization() train_data_generator = icdar_smart.get_batch_seq( num_workers=FLAGS.num_readers, config=config, is_training=True) # val_data_generator = icdar.get_batch_seq(num_workers=FLAGS.num_readers, config=eval_config, is_training=False) start = time.time() #============================= IV. Training over Steps(!!!)================================================# print("Now we're starting training!!!") if not tf.gfile.Exists(FLAGS.checkpoint_path): tf.gfile.MkDir(FLAGS.checkpoint_path) else: if not FLAGS.restore: tf.gfile.DeleteRecursively(FLAGS.checkpoint_path) tf.gfile.MkDir(FLAGS.checkpoint_path) for step in range(FLAGS.max_steps): #>>>>>>>>>>>>> data if FLAGS.mode == "debug": data = [] data.append( np.ones((config.batch_size, FLAGS.num_steps, 512, 512, 3), dtype=np.float32)) data.append( np.ones((batch_size, len_seq, 128, 128, 1), dtype=np.float32)) data.append( np.ones((batch_size, len_seq, 128, 128, 5), dtype=np.float32)) data.append( np.ones((batch_size, len_seq, 128, 128, 1), dtype=np.float32)) else: data = next(train_data_generator) if step < 10: print("Data ready!!!") plt.figure(dpi=300) ax = plt.subplot(121) plt.imshow(data[1][0][0, :, :, 0]) plt.title("score map") ax = plt.subplot(122) plt.imshow(data[3][0][0, :, :, 0] * 255) plt.title("training mask") print("saving figure") plt.savefig("/home/lxiaol9/debug/running/" + str(step) + ".png") east_feed = np.reshape(data[0], [-1, 512, 512, 3]) target_frame = np.reshape( np.array(data[0])[:, 0:4, :, :, :], [-1, 512, 512, 3]) source_frame = np.reshape( np.array(data[0])[:, 1:5, :, :, :], [-1, 512, 512, 3]) flow_feed = np.concatenate((source_frame[:, np.newaxis, :, :, :], target_frame[:, np.newaxis, :, :, :]), axis=1) flow_maps_stack = [] # >>>>>>>>>>>>>>>>>>>>>>>>>>> feature extraction with EAST >>>>>>>>>>>>>>>>>>>>>>>> # rounds = int(east_feed.shape[0] / east_opts['batch_size']) feature_stack = [] flow_maps_stack = [] for r in range(rounds): feature_stack.append( east_net.sess.run( [east_net.y_hat_test_tnsr], feed_dict={ east_net.x_tnsr: east_feed[r * east_opts['batch_size']:(r + 1) * east_opts['batch_size'], :, :, :] })[0][0]) feature_maps = np.concatenate(feature_stack, axis=0) feature_maps_reshape = np.reshape(feature_maps, [-1, config.num_steps, 128, 128, 32]) #>>>>>>>>>>>>>>> flow estimation with PWCnet # x: [batch_size,2,H,W,3] uint8; x_adapt: [batch_size,2,H,W,3] float32 x_adapt, x_adapt_info = nn.adapt_x(flow_feed) if x_adapt_info is not None: y_adapt_info = (x_adapt_info[0], x_adapt_info[2], x_adapt_info[3], 2) else: y_adapt_info = None mini_batch = nn_opts['batch_size'] * nn.num_gpus rounds = int(flow_feed.shape[0] / mini_batch) for r in range(rounds): feed_dict = { nn.x_tnsr: x_adapt[r * mini_batch:(r + 1) * mini_batch, :, :, :, :] } y_hat = nn.sess.run(nn.y_hat_test_tnsr, feed_dict=feed_dict) if FLAGS.mode == "debug": print( "Step 5: now finish running one round of PWCnet for flow estimation" ) GPUtil.showUtilization() y_hats, _ = nn.postproc_y_hat_test( y_hat, y_adapt_info) # suppose to be [batch, height, width, 2] flow_maps_stack.append(y_hats[:, 1::4, 1::4, :] / 4) flow_maps = np.concatenate(flow_maps_stack, axis=0) # print("flow maps has shape ", flow_maps.shape[:]) flow_maps = np.reshape(flow_maps, [-1, FLAGS.num_steps - 1, 128, 128, 2]) #>>>>>>>>>>>>>>> running training session with g.as_default(): ml, tl, _ = sess1.run([model_loss, total_loss, train_op], \ feed_dict={input_feat_maps: feature_maps_reshape, input_score_maps: data[1], input_geo_maps: data[2], input_training_masks: data[3], input_flow_maps: flow_maps }) if FLAGS.mode == "debug": print("Step 6: running one round on training!!!") GPUtil.showUtilization() if np.isnan(tl): print('Loss diverged, stop training') break if step % 10 == 0: avg_time_per_step = (time.time() - start) / 10 avg_examples_per_second = (10 * FLAGS.batch_size_per_gpu * len(gpus)) / (time.time() - start) start = time.time() print( 'Step {:06d}, model loss {:.4f}, total loss {:.4f}, {:.2f} seconds/step, {:.2f} examples/second' .format(step, ml, tl, avg_time_per_step, avg_examples_per_second)) if step % FLAGS.save_checkpoint_steps == 0: saver.save(sess1, FLAGS.checkpoint_path + 'model.ckpt', global_step=global_step) if step % FLAGS.save_summary_steps == 0: _, tl, summary_str = sess1.run( [train_op, total_loss, summary_op], feed_dict={ input_feat_maps: feature_maps_reshape, input_score_maps: data[1], input_geo_maps: data[2], input_training_masks: data[3], input_flow_maps: flow_maps }) summary_writer.add_summary(summary_str, global_step=step)
def main(_): # check the running platform if platform.uname()[1] != 'dragonx-H97N-WIFI': print("Now it knows it's in a remote cluster") FLAGS.system = "remote" FLAGS.data_path = "/home/lxiaol9/ARC/EASTRNN/data/ICDAR/train" FLAGS.vis_path = "/home/lxiaol9/ARC/EASTRNN/vis/LSTM/" FLAGS.save_path = "/work/cascades/lxiaol9/ARC/EAST/checkpoints/LSTM_east/" + now.strftime( "%Y%m%d-%H%M%S") FLAGS.video_path = "/home/lxiaol9/ARC/EASTRNN/data/ICDAR2015/train/" FLAGS.checkpoints_path = FLAGS.save_path FLAGS.pretrained_model_path = "/work/cascades/lxiaol9/ARC/EAST/checkpoints/east/20180921-135717/" print( "############## Step1: The environment path has been set up ###############" ) gpus = [ x.name for x in device_lib.list_local_devices() if x.device_type == "GPU" ] if (FLAGS.num_gpus > len(gpus)): raise ValueError("Your machine has only %d gpus " "which is less than the requested --num_gpus=%d." % (len(gpus), FLAGS.num_gpus)) # Model parameters using config config = get_config(FLAGS) config.batch_size = 1 config.num_layers = 3 config.num_steps = 10 eval_config = get_config(FLAGS) eval_config.batch_size = 1 eval_config.num_layers = 3 eval_config.num_steps = 10 with tf.Graph().as_default(): # Global initializer for Variables in the model # log: May 3rd, we need to adapt the model input, with config with tf.name_scope("Train"): # use placeholder to stand for input and targets initializer = tf.random_normal_initializer() x_train = tf.placeholder( tf.float32, shape=[None, config.num_steps, None, None, 3]) m = ArrayModel(True, config, x_train, reuse_variables=None, initializer=initializer) print("finished Training model generation") training_score = tf.summary.image( 'score_map', m.input_score_maps[0, :, :, :, :]) training_score_pred = tf.summary.image( 'score_map_pred', tf.stack(m.score_map_set, 0) * 255) training_cost_sum = tf.summary.scalar("Loss", m.cost) training_lr = tf.summary.scalar("Learning_Rate", m.lr) training_input = tf.summary.image('input_images', m.input_data[0, :, :, :, :]) training_loss_aabb = tf.summary.scalar('geometry_AABB', m.loss_aabb) training_loss_theta = tf.summary.scalar('geometry_theta', m.loss_theta) training_loss_cls = tf.summary.scalar('classification_dice_loss', m.loss_cls) m.summary_merged = tf.summary.merge([ training_lr, training_cost_sum, training_loss_aabb, training_loss_theta, training_loss_cls, training_input, training_score, training_score_pred ]) with tf.name_scope("Val"): # use placeholder to stand for input and targets initializer = tf.random_normal_initializer() x_val = tf.placeholder( tf.float32, shape=[None, eval_config.num_steps, None, None, 3]) mvalid = ArrayModel(True, eval_config, x_val, reuse_variables=True, initializer=initializer) val_cost_sum = tf.summary.scalar("Loss", mvalid.cost) val_score = tf.summary.image( 'score_map', mvalid.input_score_maps[0, :, :, :, :]) val_score_pred = tf.summary.image( 'score_map_pred', tf.stack(mvalid.score_map_set, 0) * 255) val_input = tf.summary.image('input_images', mvalid.input_data[0, :, :, :, :]) val_loss_aabb = tf.summary.scalar('geometry_AABB', mvalid.loss_aabb) val_loss_theta = tf.summary.scalar('geometry_theta', mvalid.loss_theta) val_loss_cls = tf.summary.scalar('classification_dice_loss', mvalid.loss_cls) mvalid.summary_merged = tf.summary.merge([ val_cost_sum, val_score, val_loss_aabb, val_loss_theta, val_loss_cls, val_score_pred, val_input ]) # Now we have got our models ready, so create a dictionFLAGSary to store those computational graph models = {"Train": m} # Module 2 print("#############Step 2: models has been built############") for name, model in models.items(): model.export_ops(name) metagraph = tf.train.export_meta_graph() soft_placement = False # we could also do coding in parallel if FLAGS.num_gpus > 1: soft_placement = True util.auto_parallel(metagraph, m) # if FLAGS.pretrained_model_path is not None: # checkpoint_path = FLAGS.pretrained_model_path # ckpt_state = tf.train.get_checkpoint_state(checkpoint_path ) # ckpt = os.path.join(checkpoint_path, os.path.basename(ckpt_state.model_checkpoint_path)) # # 1. one way is to wrap original weights and change the variable names # variables_to_restore = slim.get_variables_to_restore(exclude=['fc6', 'fc7', 'fc8']) # variable_restore_op = slim.assign_from_checkpoint(ckpt, slim.get_trainable_variables(), # ignore_missing_vars=True) # Sep-18th, Try using a different saver to only initialize part of the model var_list1 = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='feature_fusion') var_list2 = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='resnet_v1_50') var_list = var_list1 + var_list2 saver_alter = tf.train.Saver({v.op.name: v for v in var_list}) config_proto = tf.ConfigProto(allow_soft_placement=soft_placement) #########################changes made up here###################### train_data_generator = icdar.get_batch_seq( num_workers=FLAGS.num_readers, config=config, is_training=True) val_data_generator = icdar.get_batch_seq(num_workers=FLAGS.num_readers, config=eval_config, is_training=False) print( "##############Step 3: Heatmap, GT data is ready now################" ) sv = tf.train.Supervisor() with sv.managed_session(config=config_proto) as session: if FLAGS.restore: print('continue training from previous checkpoint') ckpt = tf.train.latest_checkpoint(FLAGS.checkpoints_path) sv.saver.restore(session, ckpt) elif FLAGS.partially_restore: print('continue training from previous EAST checkpoint') ckpt = FLAGS.pretrained_model_path + 'model.ckpt-56092' logger.info('Restore from {}'.format(ckpt)) saver_alter.restore(session, ckpt) else: if FLAGS.pretrained_model_path is not None: variable_restore_op(session) train_writer = tf.summary.FileWriter(FLAGS.save_path + '/train/', session.graph) val_writer = tf.summary.FileWriter(FLAGS.save_path + '/val/') print("###########Step 4 : start training. ###########") for i in range(config.max_steps): lr_decay = config.lr_decay**max(i + 1 - config.max_steps, 0.0) m.assign_lr(session, FLAGS.learning_rate * lr_decay) data_train = next(train_data_generator) # apply training along the way print("Step: %d Learning Rate: %.5f" % (i + 1, session.run(m.lr))) train_loss = run_step(session, m, data_train, config, i, eval_op=m.train_op, summary_writer=train_writer, verbose=True) print("Step: %d training loss: %.5f" % (i + 1, train_loss)) if i % FLAGS.perform_val_steps == 0: data_val = next(val_data_generator) valid_loss = run_step(session, mvalid, data_val, eval_config, i, summary_writer=val_writer, verbose=True) print("Step: %d Valid loss: %.5f" % (i + 1, valid_loss)) if (i % FLAGS.save_checkpoint_steps == 0) and FLAGS.save_path: print("Saving model to %s." % FLAGS.save_path) sv.saver.save(session, FLAGS.save_path, global_step=sv.global_step)
FLAGS.save_path = "/work/cascades/lxiaol9/ARC/EAST/checkpoints/LSTM_east/" + now.strftime("%Y%m%d-%H%M%S") FLAGS.video_path = "/home/lxiaol9/ARC/EASTRNN/data/ICDAR2013/train/" FLAGS.checkpoints_path = FLAGS.save_path FLAGS.pretrained_model_path = "/work/cascades/lxiaol9/ARC/EAST/checkpoints/east/20180921-173054/" print("############## Step1: The environment path has been set up ###############") gpus = [x.name for x in device_lib.list_local_devices() if x.device_type == "GPU"] if (FLAGS.num_gpus > len(gpus)): raise ValueError("Your machine has only %d gpus " "which is less than the requested --num_gpus=%d." % (len(gpus), FLAGS.num_gpus)) config = get_config(FLAGS) config.batch_size = 8 config.num_layers = 3 config.num_steps = 10)) # Model parameters using config config = get_config(FLAGS) config.batch_size = 8 config.num_layers = 3 config.num_steps = 10 eval_config = get_config(FLAGS) eval_config.batch_size = 8 eval_config.num_layers = 3 eval_config.num_steps = 10 with tf.Graph().as_default(): # Global initializer for Variables in the model # log: May 3rd, we need to adapt the model input, with config with tf.name_scope("Train"): # use placeholder to stand for input and targets initializer = tf.random_normal_initializer() x_train = tf.placeholder(tf.float32, shape=[None, config.num_steps, None, None, 3]) m = ArrayModel(True, config, x_train, reuse_variables=None, initializer=initializer)