def train(): """The main function that runs training""" ## data image, ih, iw, gt_boxes, gt_masks, num_instances, img_id = \ datasets.get_dataset(FLAGS.dataset_name, FLAGS.dataset_split_name, FLAGS.dataset_dir, FLAGS.im_batch, is_training=True) data_queue = tf.RandomShuffleQueue( capacity=32, min_after_dequeue=16, dtypes=(image.dtype, ih.dtype, iw.dtype, gt_boxes.dtype, gt_masks.dtype, num_instances.dtype, img_id.dtype)) enqueue_op = data_queue.enqueue( (image, ih, iw, gt_boxes, gt_masks, num_instances, img_id)) data_queue_runner = tf.train.QueueRunner(data_queue, [enqueue_op] * 4) tf.add_to_collection(tf.GraphKeys.QUEUE_RUNNERS, data_queue_runner) (image, ih, iw, gt_boxes, gt_masks, num_instances, img_id) = data_queue.dequeue() im_shape = tf.shape(image) image = tf.reshape(image, (im_shape[0], im_shape[1], im_shape[2], 3)) ## network logits, end_points, pyramid_map = network.get_network( FLAGS.network, image, weight_decay=FLAGS.weight_decay, is_training=True) outputs = pyramid_network.build(end_points, im_shape[1], im_shape[2], pyramid_map, num_classes=81, base_anchors=9, is_training=True, gt_boxes=gt_boxes, gt_masks=gt_masks, loss_weights=[0.2, 0.2, 1.0, 0.2, 1.0]) total_loss = outputs['total_loss'] losses = outputs['losses'] batch_info = outputs['batch_info'] regular_loss = tf.add_n( tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) input_image = end_points['input'] final_box = outputs['final_boxes']['box'] final_cls = outputs['final_boxes']['cls'] final_prob = outputs['final_boxes']['prob'] final_gt_cls = outputs['final_boxes']['gt_cls'] # Get the computation result from output final_mask = outputs['mask']['mask'] print('final_mask', final_mask.shape) gt = outputs['gt'] # replace the draw_bbox # draw_mask(step, # np.uint8((np.array(input_imagenp[0])/2.0+0.5)*255.0), # name='est', # bbox=final_boxnp, # mask=final_masknp, # label=final_clsnp, # prob=final_probnp, # gt_label=np.argmax(np.asarray(final_gt_clsnp),axis=1), # ) # draw_mask(step, # np.uint8((np.array(input_imagenp[0])/2.0+0.5)*255.0), # name='gt', # bbox=gtnp[:,0:4], # mask=gt_masksnp, # label=np.asarray(gtnp[:,4], dtype=np.uint8), # ) ############################# tmp_0 = outputs['losses'] tmp_1 = outputs['losses'] tmp_2 = outputs['losses'] tmp_3 = outputs['losses'] tmp_4 = outputs['losses'] # tmp_0 = outputs['tmp_0'] # tmp_1 = outputs['tmp_1'] # tmp_2 = outputs['tmp_2'] tmp_3 = outputs['tmp_3'] tmp_4 = outputs['tmp_4'] ############################ ## solvers global_step = slim.create_global_step() update_op = solve(global_step) cropped_rois = tf.get_collection('__CROPPED__')[0] transposed = tf.get_collection('__TRANSPOSED__')[0] #os.environ["CUDA_VISIBLE_DEVICES"]="1" config = tf.ConfigProto(allow_soft_placement=True) #config.gpu_options.allow_growth = True config.gpu_options.per_process_gpu_memory_fraction = 0.95 ##gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.45) ##sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) sess = tf.Session(config=config) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) sess.run(init_op) summary_op = tf.summary.merge_all() logdir = os.path.join(FLAGS.train_dir, strftime('%Y%m%d%H%M%S', gmtime())) if not os.path.exists(logdir): os.makedirs(logdir) summary_writer = tf.summary.FileWriter(logdir, graph=sess.graph) ## restore restore(sess) ## main loop coord = tf.train.Coordinator() threads = [] # print (tf.get_collection(tf.GraphKeys.QUEUE_RUNNERS)) for qr in tf.get_collection(tf.GraphKeys.QUEUE_RUNNERS): threads.extend( qr.create_threads(sess, coord=coord, daemon=True, start=True)) tf.train.start_queue_runners(sess=sess, coord=coord) saver = tf.train.Saver(max_to_keep=20) for step in range(FLAGS.max_iters): start_time = time.time() #rGet the mask results s_, tot_loss, reg_lossnp, img_id_str, \ rpn_box_loss, rpn_cls_loss, refined_box_loss, refined_cls_loss, mask_loss, \ gt_boxesnp, \ rpn_batch_pos, rpn_batch, refine_batch_pos, refine_batch, mask_batch_pos, mask_batch, \ input_imagenp, final_masknp, final_boxnp, final_clsnp, final_probnp, final_gt_clsnp, gtnp,gt_masksnp, tmp_0np, tmp_1np, tmp_2np, tmp_3np, tmp_4np= \ sess.run([update_op, total_loss, regular_loss, img_id] + losses + [gt_boxes] + batch_info + [input_image] + [final_mask] + [final_box] + [final_cls] + [final_prob] + [final_gt_cls] + [gt] + [gt_masks] + [tmp_0] + [tmp_1] + [tmp_2] + [tmp_3] + [tmp_4]) duration_time = time.time() - start_time if step % 1 == 0: print( """iter %d: image-id:%07d, time:%.3f(sec), regular_loss: %.6f, """ """total-loss %.4f(%.4f, %.4f, %.6f, %.4f, %.4f), """ """instances: %d, """ """batch:(%d|%d, %d|%d, %d|%d)""" % (step, img_id_str, duration_time, reg_lossnp, tot_loss, rpn_box_loss, rpn_cls_loss, refined_box_loss, refined_cls_loss, mask_loss, gt_boxesnp.shape[0], rpn_batch_pos, rpn_batch, refine_batch_pos, refine_batch, mask_batch_pos, mask_batch)) draw_bbox( step, np.uint8((np.array(input_imagenp[0]) / 2.0 + 0.5) * 255.0), name='est', bbox=final_boxnp, label=final_clsnp, prob=final_probnp, gt_label=np.argmax(np.asarray(final_gt_clsnp), axis=1), ) draw_bbox( step, np.uint8((np.array(input_imagenp[0]) / 2.0 + 0.5) * 255.0), name='gt', bbox=gtnp[:, 0:4], label=np.asarray(gtnp[:, 4], dtype=np.uint8), ) print("labels") # print (cat_id_to_cls_name(np.unique(np.argmax(np.asarray(final_gt_clsnp),axis=1)))[1:]) # print (cat_id_to_cls_name(np.unique(np.asarray(gt_boxesnp, dtype=np.uint8)[:,4]))) print( cat_id_to_cls_name( np.unique(np.argmax(np.asarray(tmp_3np), axis=1)))[1:]) #print (cat_id_to_cls_name(np.unique(np.argmax(np.asarray(gt_boxesnp)[:,4],axis=1)))) print("classes") print( cat_id_to_cls_name( np.unique(np.argmax(np.array(tmp_4np), axis=1)))) # print (np.asanyarray(tmp_3np)) #print ("ordered rois") #print (np.asarray(tmp_0np)[0]) #print ("pyramid_feature") #print () #print(np.unique(np.argmax(np.array(final_probnp),axis=1))) #for var, val in zip(tmp_2, tmp_2np): # print(var.name) #print(np.argmax(np.array(tmp_0np),axis=1)) if np.isnan(tot_loss) or np.isinf(tot_loss): print(gt_boxesnp) raise if step % 100 == 0: summary_str = sess.run(summary_op) summary_writer.add_summary(summary_str, step) summary_writer.flush() if (step % 10000 == 0 or step + 1 == FLAGS.max_iters) and step != 0: checkpoint_path = os.path.join( FLAGS.train_dir, FLAGS.dataset_name + '_' + FLAGS.network + '_model.ckpt') saver.save(sess, checkpoint_path, global_step=step) if coord.should_stop(): coord.request_stop() coord.join(threads)
def test(): """The main function that runs training""" ## data image, original_image_height, original_image_width, image_height, image_width, gt_boxes, gt_masks, num_instances, image_id = \ datasets.get_dataset(FLAGS.dataset_name, FLAGS.dataset_split_name_test, FLAGS.dataset_dir, FLAGS.im_batch, is_training=False) im_shape = tf.shape(image) image = tf.reshape(image, (im_shape[0], im_shape[1], im_shape[2], 3)) ## network logits, end_points, pyramid_map = network.get_network(FLAGS.network, image, weight_decay=0.0, batch_norm_decay=0.0, is_training=True) outputs = pyramid_network.build(end_points, im_shape[1], im_shape[2], pyramid_map, num_classes=81, base_anchors=3, is_training=False, gt_boxes=None, gt_masks=None, loss_weights=[0.0, 0.0, 0.0, 0.0, 0.0]) input_image = end_points['input'] testing_mask_rois = outputs['mask_ordered_rois'] testing_mask_final_mask = outputs['mask_final_mask'] testing_mask_final_clses = outputs['mask_final_clses'] testing_mask_final_scores = outputs['mask_final_scores'] ## solvers global_step = slim.create_global_step() gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8) sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) # init_op = tf.group( # tf.global_variables_initializer(), # tf.local_variables_initializer() # ) # sess.run(init_op) # summary_op = tf.summary.merge_all() logdir = os.path.join(FLAGS.train_dir, strftime('%Y%m%d%H%M%S', gmtime())) if not os.path.exists(logdir): os.makedirs(logdir) summary_writer = tf.summary.FileWriter(logdir, graph=sess.graph) ## restore restore(sess) tf.train.start_queue_runners(sess=sess) ## main loop # for step in range(FLAGS.max_iters): for step in range(82783):#range(40503): start_time = time.time() image_id_str, original_image_heightnp, original_image_widthnp, image_heightnp, image_widthnp, \ gt_boxesnp, gt_masksnp,\ input_imagenp,\ testing_mask_roisnp, testing_mask_final_masknp, testing_mask_final_clsesnp, testing_mask_final_scoresnp = \ sess.run([image_id] + [original_image_height] + [original_image_width] + [image_height] + [image_width] +\ [gt_boxes] + [gt_masks] +\ [input_image] + \ [testing_mask_rois] + [testing_mask_final_mask] + [testing_mask_final_clses] + [testing_mask_final_scores]) duration_time = time.time() - start_time if step % 1 == 0: print ( """iter %d: image-id:%07d, time:%.3f(sec), """ """instances: %d, """ % (step, image_id_str, duration_time, gt_boxesnp.shape[0])) if step % 1 == 0: draw_bbox(step, np.uint8((np.array(input_imagenp[0])/2.0+0.5)*255.0), name='test_est', bbox=testing_mask_roisnp, label=testing_mask_final_clsesnp, prob=testing_mask_final_scoresnp, mask=testing_mask_final_masknp, vis_th=0.5) draw_bbox(step, np.uint8((np.array(input_imagenp[0])/2.0+0.5)*255.0), name='test_gt', bbox=gt_boxesnp[:,0:4], label=gt_boxesnp[:,4].astype(np.int32), prob=np.ones((gt_boxesnp.shape[0],81), dtype=np.float32),) print ("predict") # LOG (cat_id_to_cls_name(np.unique(np.argmax(np.array(training_rcnn_clsesnp),axis=1)))) print (cat_id_to_cls_name(testing_mask_final_clsesnp)) print (np.max(np.array(testing_mask_final_scoresnp),axis=1)) _collectData(image_id_str, testing_mask_final_clsesnp, testing_mask_roisnp, testing_mask_final_scoresnp, original_image_heightnp, original_image_widthnp, image_heightnp, image_widthnp, testing_mask_final_masknp)
def train(): """The main function that runs training""" ## data image, original_image_height, original_image_width, image_height, image_width, gt_boxes, gt_masks, num_instances, image_id = \ datasets.get_dataset(FLAGS.dataset_name, FLAGS.dataset_split_name, FLAGS.dataset_dir, FLAGS.im_batch, is_training=True) ## queuing data data_queue = tf.RandomShuffleQueue(capacity=32, min_after_dequeue=16, dtypes=( image.dtype, original_image_height.dtype, original_image_width.dtype, image_height.dtype, image_width.dtype, gt_boxes.dtype, gt_masks.dtype, num_instances.dtype, image_id.dtype)) enqueue_op = data_queue.enqueue((image, original_image_height, original_image_width, image_height, image_width, gt_boxes, gt_masks, num_instances, image_id)) data_queue_runner = tf.train.QueueRunner(data_queue, [enqueue_op] * 4) tf.add_to_collection(tf.GraphKeys.QUEUE_RUNNERS, data_queue_runner) (image, original_image_height, original_image_width, image_height, image_width, gt_boxes, gt_masks, num_instances, image_id) = data_queue.dequeue() im_shape = tf.shape(image) image = tf.reshape(image, (im_shape[0], im_shape[1], im_shape[2], 3)) ## network logits, end_points, pyramid_map = network.get_network(FLAGS.network, image, weight_decay=FLAGS.weight_decay, batch_norm_decay=FLAGS.batch_norm_decay, is_training=True) outputs = pyramid_network.build(end_points, image_height, image_width, pyramid_map, num_classes=81, base_anchors=3,#9#15 is_training=True, gt_boxes=gt_boxes, gt_masks=gt_masks, loss_weights=[1.0, 1.0, 10.0, 1.0, 10.0]) # loss_weights=[10.0, 1.0, 0.0, 0.0, 0.0]) # loss_weights=[100.0, 100.0, 1000.0, 10.0, 100.0]) # loss_weights=[0.2, 0.2, 1.0, 0.2, 1.0]) # loss_weights=[0.1, 0.01, 10.0, 0.1, 1.0]) total_loss = outputs['total_loss'] losses = outputs['losses'] batch_info = outputs['batch_info'] regular_loss = tf.add_n(tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)) input_image = end_points['input'] training_rcnn_rois = outputs['training_rcnn_rois'] training_rcnn_clses = outputs['training_rcnn_clses'] training_rcnn_clses_target = outputs['training_rcnn_clses_target'] training_rcnn_scores = outputs['training_rcnn_scores'] training_mask_rois = outputs['training_mask_rois'] training_mask_clses_target = outputs['training_mask_clses_target'] training_mask_final_mask = outputs['training_mask_final_mask'] training_mask_final_mask_target = outputs['training_mask_final_mask_target'] tmp_0 = outputs['rpn']['P2']['shape'] tmp_1 = outputs['rpn']['P3']['shape'] tmp_2 = outputs['rpn']['P4']['shape'] tmp_3 = outputs['rpn']['P5']['shape'] ## solvers global_step = slim.create_global_step() update_op = solve(global_step) cropped_rois = tf.get_collection('__CROPPED__')[0] transposed = tf.get_collection('__TRANSPOSED__')[0] gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.95) #gpu_options = tf.GPUOptions(allow_growth=True) sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) #sess = tf.InteractiveSession(config=tf.ConfigProto(gpu_options=gpu_options)) init_op = tf.group( tf.global_variables_initializer(), tf.local_variables_initializer() ) sess.run(init_op) summary_op = tf.summary.merge_all() logdir = os.path.join(FLAGS.train_dir, strftime('%Y%m%d%H%M%S', gmtime())) if not os.path.exists(logdir): os.makedirs(logdir) summary_writer = tf.summary.FileWriter(logdir, graph=sess.graph) ## restore restore(sess) ## coord settings coord = tf.train.Coordinator() threads = [] for qr in tf.get_collection(tf.GraphKeys.QUEUE_RUNNERS): threads.extend(qr.create_threads(sess, coord=coord, daemon=True, start=True)) tf.train.start_queue_runners(sess=sess, coord=coord) ## saver init saver = tf.train.Saver(max_to_keep=20) ## finalize the graph for checking memory leak sess.graph.finalize() ## main loop for step in range(FLAGS.max_iters): start_time = time.time() s_, tot_loss, reg_lossnp, image_id_str, \ rpn_box_loss, rpn_cls_loss, rcnn_box_loss, rcnn_cls_loss, mask_loss, \ gt_boxesnp, tmp_0np, tmp_1np, tmp_2np, tmp_3np, \ rpn_batch_pos, rpn_batch, rcnn_batch_pos, rcnn_batch, mask_batch_pos, mask_batch, \ input_imagenp, \ training_rcnn_roisnp, training_rcnn_clsesnp, training_rcnn_clses_targetnp, training_rcnn_scoresnp, training_mask_roisnp, training_mask_clses_targetnp, training_mask_final_masknp, training_mask_final_mask_targetnp = \ sess.run([update_op, total_loss, regular_loss, image_id] + losses + [gt_boxes] + [tmp_0] + [tmp_1] + [tmp_2] +[tmp_3] + batch_info + [input_image] + [training_rcnn_rois] + [training_rcnn_clses] + [training_rcnn_clses_target] + [training_rcnn_scores] + [training_mask_rois] + [training_mask_clses_target] + [training_mask_final_mask] + [training_mask_final_mask_target]) duration_time = time.time() - start_time if step % 1 == 0: LOG ( """iter %d: image-id:%07d, time:%.3f(sec), regular_loss: %.6f, """ """total-loss %.4f(%.4f, %.4f, %.6f, %.4f, %.4f), """ """instances: %d, """ """batch:(%d|%d, %d|%d, %d|%d)""" % (step, image_id_str, duration_time, reg_lossnp, tot_loss, rpn_box_loss, rpn_cls_loss, rcnn_box_loss, rcnn_cls_loss, mask_loss, gt_boxesnp.shape[0], rpn_batch_pos, rpn_batch, rcnn_batch_pos, rcnn_batch, mask_batch_pos, mask_batch)) LOG ("target") LOG (cat_id_to_cls_name(np.unique(np.argmax(np.asarray(training_rcnn_clses_targetnp),axis=1)))) LOG ("predict") LOG (cat_id_to_cls_name(np.unique(np.argmax(np.array(training_rcnn_clsesnp),axis=1)))) LOG (tmp_0np) LOG (tmp_1np) LOG (tmp_2np) LOG (tmp_3np) if step % 50 == 0: draw_bbox(step, np.uint8((np.array(input_imagenp[0])/2.0+0.5)*255.0), name='train_est', bbox=training_rcnn_roisnp, label=np.argmax(np.array(training_rcnn_scoresnp),axis=1), prob=training_rcnn_scoresnp, # bbox=training_mask_roisnp, # label=training_mask_clses_targetnp, # prob=np.zeros((training_mask_final_masknp.shape[0],81), dtype=np.float32)+1.0, # mask=training_mask_final_masknp, vis_all=True) draw_bbox(step, np.uint8((np.array(input_imagenp[0])/2.0+0.5)*255.0), name='train_gt', bbox=training_rcnn_roisnp, label=np.argmax(np.array(training_rcnn_clses_targetnp),axis=1), prob=np.zeros((training_rcnn_clsesnp.shape[0],81), dtype=np.float32)+1.0, # bbox=training_mask_roisnp, # label=training_mask_clses_targetnp, # prob=np.zeros((training_mask_final_masknp.shape[0],81), dtype=np.float32)+1.0, # mask=training_mask_final_mask_targetnp, vis_all=True) if np.isnan(tot_loss) or np.isinf(tot_loss): LOG (gt_boxesnp) raise if step % 100 == 0: summary_str = sess.run(summary_op) summary_writer.add_summary(summary_str, step) summary_writer.flush() if (step % 500 == 0 or step + 1 == FLAGS.max_iters) and step != 0: checkpoint_path = os.path.join(FLAGS.train_dir, FLAGS.dataset_name + '_' + FLAGS.network + '_model.ckpt') saver.save(sess, checkpoint_path, global_step=step) if coord.should_stop(): coord.request_stop() coord.join(threads) gc.collect()
def test(): """The main function that runs training""" ## data image, ih, iw, gt_boxes, gt_masks, num_instances, img_id = \ datasets.get_dataset(FLAGS.dataset_name, FLAGS.dataset_split_name, FLAGS.dataset_dir, FLAGS.im_batch, is_training=False) im_shape = tf.shape(image) image = tf.reshape(image, (im_shape[0], im_shape[1], im_shape[2], 3)) ## network logits, end_points, pyramid_map = network.get_network( FLAGS.network, image, weight_decay=FLAGS.weight_decay, is_training=False) outputs = pyramid_network.build(end_points, im_shape[1], im_shape[2], pyramid_map, num_classes=81, base_anchors=15, is_training=False, gt_boxes=None, gt_masks=None, loss_weights=[0.0, 0.0, 0.0, 0.0, 0.0]) input_image = end_points['input'] testing_mask_rois = outputs['mask_ordered_rois'] testing_mask_final_mask = outputs['mask_final_mask'] testing_mask_final_clses = outputs['mask_final_clses'] testing_mask_final_scores = outputs['mask_final_scores'] ############################# tmp_0 = outputs['tmp_0'] tmp_1 = outputs['tmp_1'] tmp_2 = outputs['tmp_2'] tmp_3 = outputs['tmp_3'] tmp_4 = outputs['tmp_4'] tmp_5 = outputs['tmp_5'] ############################ ## solvers global_step = slim.create_global_step() #update_op = solve(global_step) cropped_rois = tf.get_collection('__CROPPED__')[0] transposed = tf.get_collection('__TRANSPOSED__')[0] gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8) sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) sess.run(init_op) summary_op = tf.summary.merge_all() logdir = os.path.join(FLAGS.train_dir, strftime('%Y%m%d%H%M%S', gmtime())) if not os.path.exists(logdir): os.makedirs(logdir) summary_writer = tf.summary.FileWriter(logdir, graph=sess.graph) ## restore restore(sess) ## main loop coord = tf.train.Coordinator() threads = [] # print (tf.get_collection(tf.GraphKeys.QUEUE_RUNNERS)) for qr in tf.get_collection(tf.GraphKeys.QUEUE_RUNNERS): threads.extend( qr.create_threads(sess, coord=coord, daemon=True, start=True)) tf.train.start_queue_runners(sess=sess, coord=coord) saver = tf.train.Saver(max_to_keep=20) for step in range(FLAGS.max_iters): start_time = time.time() img_id_str, \ gt_boxesnp, \ input_imagenp, tmp_0np, tmp_1np, tmp_2np, tmp_3np, tmp_4np, tmp_5np, \ testing_mask_roisnp, testing_mask_final_masknp, testing_mask_final_clsesnp, testing_mask_final_scoresnp = \ sess.run([img_id] + \ [gt_boxes] + \ [input_image] + [tmp_0] + [tmp_1] + [tmp_2] + [tmp_3] + [tmp_4] + [tmp_5] + \ [testing_mask_rois] + [testing_mask_final_mask] + [testing_mask_final_clses] + [testing_mask_final_scores]) duration_time = time.time() - start_time if step % 1 == 0: print("""iter %d: image-id:%07d, time:%.3f(sec), """ """instances: %d, """ % (step, img_id_str, duration_time, gt_boxesnp.shape[0])) if step % 1 == 0: draw_bbox( step, np.uint8((np.array(input_imagenp[0]) / 2.0 + 0.5) * 255.0), name='test_est', bbox=testing_mask_roisnp, label=testing_mask_final_clsesnp, prob=testing_mask_final_scoresnp, mask=testing_mask_final_masknp, )
def forward_test_single_image(): if not os.path.exists(save_dir_bbox): os.makedirs(save_dir_bbox) if not os.path.exists(save_dir_mask): os.makedirs(save_dir_mask) file_pathname = testdata_base_dir + '*.' + file_pattern image_paths = glob.glob(file_pathname) #with .jpg/.png image_names = glob.glob(file_pathname) #no .jpg/.png for i in range(len(image_paths)): image_names[i] = image_paths[i][len(testdata_base_dir):-4] print(image_paths) print(image_names) TEST_image = tf.placeholder(tf.float32, shape=[1, None, None, 3]) im_shape = tf.shape(TEST_image) ## network logits, end_points, pyramid_map = network.get_network( FLAGS.network, TEST_image, weight_decay=FLAGS.weight_decay, is_training=True) outputs = pyramid_network.build(end_points, im_shape[1], im_shape[2], pyramid_map, num_classes=81, base_anchors=9, is_training=True, gt_boxes=None, gt_masks=None, loss_weights=[0.2, 0.2, 1.0, 0.2, 1.0]) input_image = end_points['input'] print("input_image.shape", input_image.shape) final_box = outputs['final_boxes']['box'] final_cls = outputs['final_boxes']['cls'] final_prob = outputs['final_boxes']['prob'] final_mask = outputs['mask']['mask'] config = tf.ConfigProto() config.gpu_options.allow_growth = True # config.gpu_options.per_process_gpu_memory_fraction = 0.5 sess = tf.Session(config=config) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) sess.run(init_op) ## restore trained model restore(sess) ## main loop coord = tf.train.Coordinator() threads = [] # print (tf.get_collection(tf.GraphKeys.QUEUE_RUNNERS)) for qr in tf.get_collection(tf.GraphKeys.QUEUE_RUNNERS): threads.extend( qr.create_threads(sess, coord=coord, daemon=True, start=True)) tf.train.start_queue_runners(sess=sess, coord=coord) for i in range(len(image_paths)): ## read test image test_img = Image.open(image_paths[i]) test_img = test_img.convert("RGB") test_img_i = np.array(test_img, dtype=np.uint8) test_img = np.array(test_img, dtype=np.float32) test_img = test_img[np.newaxis, ...] print("test_img.shape", test_img.shape) print("test_img_i.shape", test_img_i.shape) # start_time = time.time() input_imagenp, final_boxnp, final_clsnp, final_probnp, \ final_masknp= \ sess.run([input_image] + [final_box] + [final_cls] + [final_prob] + [final_mask], feed_dict={TEST_image:test_img}) # duration_time = time.time() - start_time draw_bbox( test_img_i, type='est', bbox=final_boxnp, label=final_clsnp, prob=final_probnp, gt_label=None, save_dir=save_dir_bbox, save_name=image_names[i], ) print("final_masknp.shape\n", final_masknp.shape) draw_mask( test_img_i, type='est', bbox=final_boxnp, mask=final_masknp, label=final_clsnp, prob=final_probnp, gt_label=None, save_dir=save_dir_mask, save_name=image_names[i], ) if coord.should_stop(): coord.request_stop() coord.join(threads)