def read_and_prepocess_single_img(filename_queue, shortside_len, is_training): img_name, img, gtboxes_and_label, num_objects = read_single_example_and_decode(filename_queue) img = tf.cast(img, tf.float32) if is_training: # prob is 0.3: convert to gray img = image_preprocess_multi_gpu_aug.random_rgb2gray(img_tensor=img, gtboxes_and_label=gtboxes_and_label) # rotate with 0.5 prob. and if rotate, if will random choose a theta from : tf.range(-90, 90+16, delta=15) img, gtboxes_and_label = image_preprocess_multi_gpu_aug.random_rotate_img(img_tensor=img, gtboxes_and_label=gtboxes_and_label) img, gtboxes_and_label, img_h, img_w = image_preprocess_multi_gpu_aug.short_side_resize(img_tensor=img, gtboxes_and_label=gtboxes_and_label, target_shortside_len=shortside_len, max_len=cfgs.IMG_MAX_LENGTH) img, gtboxes_and_label = image_preprocess_multi_gpu_aug.random_flip_left_right(img_tensor=img, gtboxes_and_label=gtboxes_and_label) img, gtboxes_and_label = image_preprocess_multi_gpu_aug.random_flip_up_dowm(img_tensor=img, gtboxes_and_label=gtboxes_and_label) else: img, gtboxes_and_label, img_h, img_w = image_preprocess_multi_gpu_aug.short_side_resize(img_tensor=img, gtboxes_and_label=gtboxes_and_label, target_shortside_len=shortside_len) gtboxes_and_label = get_horizen_minAreaRectangle(gtboxes_and_label) if cfgs.NET_NAME in ['resnet101_v1d', 'resnet50_v1d']: img = img / 255. - tf.constant([cfgs.PIXEL_MEAN_]) else: img = img - tf.constant([cfgs.PIXEL_MEAN]) # sub pixel mean at last return img_name, img, gtboxes_and_label, num_objects, img_h, img_w
def build_fastrcnn(self, feature_to_cropped, rois, img_shape): with tf.variable_scope('Fast-RCNN'): # 5. ROI Pooling with tf.variable_scope('rois_pooling'): rois = boxes_utils.get_horizen_minAreaRectangle(rois, False) pooled_features = self.roi_pooling( feature_maps=feature_to_cropped, rois=rois, img_shape=img_shape) # xmin, ymin, xmax, ymax = tf.unstack(rois, axis=1) # # h = tf.maximum(ymax - ymin, 0) # w = tf.maximum(xmax - xmin, 0) # x_c = (xmax + xmin) // 2 # y_c = (ymax + ymin) // 2 # theta = tf.ones_like(h) * -90 # rois = tf.transpose(tf.stack([x_c, y_c, h, w, theta])) # 6. inferecne rois in Fast-RCNN to obtain fc_flatten features if self.base_network_name.startswith('resnet'): fc_flatten = resnet.restnet_head(input=pooled_features, is_training=self.is_training, scope=self.base_network_name) elif self.base_network_name.startswith('MobilenetV2'): fc_flatten = mobilenet_v2.mobilenetv2_head( inputs=pooled_features, is_training=self.is_training) else: raise NotImplementedError('only support resnet and mobilenet') # 7. cls and reg in Fast-RCNN with slim.arg_scope([slim.fully_connected], weights_regularizer=slim.l2_regularizer( cfgs.WEIGHT_DECAY)): cls_score = slim.fully_connected( fc_flatten, num_outputs=cfgs.CLASS_NUM + 1, weights_initializer=cfgs.INITIALIZER, activation_fn=None, trainable=self.is_training, scope='cls_fc') bbox_pred = slim.fully_connected( fc_flatten, num_outputs=(cfgs.CLASS_NUM + 1) * 5, weights_initializer=cfgs.BBOX_INITIALIZER, activation_fn=None, trainable=self.is_training, scope='reg_fc') # for convient. It also produce (cls_num +1) bboxes cls_score = tf.reshape(cls_score, [-1, cfgs.CLASS_NUM + 1]) bbox_pred = tf.reshape(bbox_pred, [-1, 5 * (cfgs.CLASS_NUM + 1)]) return bbox_pred, cls_score
def test(img_num): with tf.Graph().as_default(): img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \ next_batch(dataset_name=cfgs.DATASET_NAME, batch_size=cfgs.BATCH_SIZE, shortside_len=cfgs.SHORT_SIDE_LEN, is_training=False) gtboxes_and_label, head = get_head( tf.squeeze(gtboxes_and_label_batch, 0)) gtboxes_and_label = tf.py_func(back_forward_convert, inp=[gtboxes_and_label], Tout=tf.float32) gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 6]) head_quadrant = tf.py_func(get_head_quadrant, inp=[head, gtboxes_and_label], Tout=tf.float32) head_quadrant = tf.reshape(head_quadrant, [-1, 1]) gtboxes_and_label_minAreaRectangle = get_horizen_minAreaRectangle( gtboxes_and_label) gtboxes_and_label_minAreaRectangle = tf.reshape( gtboxes_and_label_minAreaRectangle, [-1, 5]) # *********************************************************************************************** # * share net * # *********************************************************************************************** _, share_net = get_network_byname(net_name=cfgs.NET_NAME, inputs=img_batch, num_classes=None, is_training=True, output_stride=None, global_pool=False, spatial_squeeze=False) # *********************************************************************************************** # * RPN * # *********************************************************************************************** rpn = build_rpn.RPN( net_name=cfgs.NET_NAME, inputs=img_batch, gtboxes_and_label=None, is_training=False, share_head=cfgs.SHARE_HEAD, share_net=share_net, stride=cfgs.STRIDE, anchor_ratios=cfgs.ANCHOR_RATIOS, anchor_scales=cfgs.ANCHOR_SCALES, scale_factors=cfgs.SCALE_FACTORS, base_anchor_size_list=cfgs. BASE_ANCHOR_SIZE_LIST, # P2, P3, P4, P5, P6 level=cfgs.LEVEL, top_k_nms=cfgs.RPN_TOP_K_NMS, rpn_nms_iou_threshold=cfgs.RPN_NMS_IOU_THRESHOLD, max_proposals_num=cfgs.MAX_PROPOSAL_NUM, rpn_iou_positive_threshold=cfgs.RPN_IOU_POSITIVE_THRESHOLD, rpn_iou_negative_threshold=cfgs.RPN_IOU_NEGATIVE_THRESHOLD, rpn_mini_batch_size=cfgs.RPN_MINIBATCH_SIZE, rpn_positives_ratio=cfgs.RPN_POSITIVE_RATE, remove_outside_anchors=False, # whether remove anchors outside rpn_weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME]) # rpn predict proposals rpn_proposals_boxes, rpn_proposals_scores = rpn.rpn_proposals( ) # rpn_score shape: [300, ] # *********************************************************************************************** # * Fast RCNN * # *********************************************************************************************** fast_rcnn = build_fast_rcnn.FastRCNN( feature_pyramid=rpn.feature_pyramid, rpn_proposals_boxes=rpn_proposals_boxes, rpn_proposals_scores=rpn_proposals_scores, img_shape=tf.shape(img_batch), img_batch=img_batch, roi_size=cfgs.ROI_SIZE, roi_pool_kernel_size=cfgs.ROI_POOL_KERNEL_SIZE, scale_factors=cfgs.SCALE_FACTORS, gtboxes_and_label=None, gtboxes_and_label_minAreaRectangle= gtboxes_and_label_minAreaRectangle, fast_rcnn_nms_iou_threshold=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD, fast_rcnn_maximum_boxes_per_img=100, fast_rcnn_nms_max_boxes_per_class=cfgs. FAST_RCNN_NMS_MAX_BOXES_PER_CLASS, show_detections_score_threshold=cfgs.FINAL_SCORE_THRESHOLD, # show detections which score >= 0.6 num_classes=cfgs.CLASS_NUM, fast_rcnn_minibatch_size=cfgs.FAST_RCNN_MINIBATCH_SIZE, fast_rcnn_positives_ratio=cfgs.FAST_RCNN_POSITIVE_RATE, fast_rcnn_positives_iou_threshold=cfgs. FAST_RCNN_IOU_POSITIVE_THRESHOLD, # iou>0.5 is positive, iou<0.5 is negative use_dropout=cfgs.USE_DROPOUT, weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME], is_training=False, level=cfgs.LEVEL, head_quadrant=head_quadrant) fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category, \ fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, fast_rcnn_head_quadrant, \ num_of_objects_rotate, detection_category_rotate = fast_rcnn.fast_rcnn_predict() # train init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = restore_model.get_restorer() config = tf.ConfigProto() # config.gpu_options.per_process_gpu_memory_fraction = 0.5 config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess, coord) for i in range(img_num): start = time.time() _img_name_batch, _img_batch, _gtboxes_and_label, _gtboxes_and_label_minAreaRectangle, _head_quadrant, \ _fast_rcnn_decode_boxes, _fast_rcnn_score, _detection_category, _fast_rcnn_decode_boxes_rotate, \ _fast_rcnn_score_rotate, _fast_rcnn_head_quadrant, _detection_category_rotate \ = sess.run([img_name_batch, img_batch, gtboxes_and_label, gtboxes_and_label_minAreaRectangle, head_quadrant, fast_rcnn_decode_boxes, fast_rcnn_score, detection_category, fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, fast_rcnn_head_quadrant, detection_category_rotate]) end = time.time() _img_batch = np.squeeze(_img_batch, axis=0) _img_batch_fpn_horizonal = help_utils.draw_box_cv( _img_batch, boxes=_fast_rcnn_decode_boxes, labels=_detection_category, scores=_fast_rcnn_score) _img_batch_fpn_rotate = help_utils.draw_rotate_box_cv( _img_batch, boxes=_fast_rcnn_decode_boxes_rotate, labels=_detection_category_rotate, scores=_fast_rcnn_score_rotate, head=np.argmax(_fast_rcnn_head_quadrant, axis=1)) mkdir(cfgs.TEST_SAVE_PATH) cv2.imwrite( cfgs.TEST_SAVE_PATH + '/{}_horizontal_fpn.jpg'.format(str(_img_name_batch[0])), _img_batch_fpn_horizonal) cv2.imwrite( cfgs.TEST_SAVE_PATH + '/{}_rotate_fpn.jpg'.format(str(_img_name_batch[0])), _img_batch_fpn_rotate) temp_label_horizontal = np.reshape(_gtboxes_and_label[:, -1:], [ -1, ]).astype(np.int64) temp_label_rotate = np.reshape(_gtboxes_and_label[:, -1:], [ -1, ]).astype(np.int64) _img_batch_gt_horizontal = help_utils.draw_box_cv( _img_batch, boxes=_gtboxes_and_label_minAreaRectangle[:, :-1], labels=temp_label_horizontal, scores=None) _img_batch_gt_rotate = help_utils.draw_rotate_box_cv( _img_batch, boxes=_gtboxes_and_label[:, :-1], labels=temp_label_rotate, scores=None, head=np.reshape(_head_quadrant, [ -1, ])) cv2.imwrite( cfgs.TEST_SAVE_PATH + '/{}_horizontal_gt.jpg'.format(str(_img_name_batch[0])), _img_batch_gt_horizontal) cv2.imwrite( cfgs.TEST_SAVE_PATH + '/{}_rotate_gt.jpg'.format(str(_img_name_batch[0])), _img_batch_gt_rotate) view_bar( '{} image cost {}s'.format(str(_img_name_batch[0]), (end - start)), i + 1, img_num) coord.request_stop() coord.join(threads)
def eval_ship(img_num, mode): with tf.Graph().as_default(): img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \ next_batch(dataset_name=cfgs.DATASET_NAME, batch_size=cfgs.BATCH_SIZE, shortside_len=cfgs.SHORT_SIDE_LEN, is_training=False) gtboxes_and_label = tf.py_func( back_forward_convert, inp=[tf.squeeze(gtboxes_and_label_batch, 0)], Tout=tf.float32) gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 6]) gtboxes_and_label_minAreaRectangle = get_horizen_minAreaRectangle( gtboxes_and_label) gtboxes_and_label_minAreaRectangle = tf.reshape( gtboxes_and_label_minAreaRectangle, [-1, 5]) # *********************************************************************************************** # * share net * # *********************************************************************************************** _, share_net = get_network_byname(net_name=cfgs.NET_NAME, inputs=img_batch, num_classes=None, is_training=True, output_stride=None, global_pool=False, spatial_squeeze=False) # *********************************************************************************************** # * RPN * # *********************************************************************************************** rpn = build_rpn.RPN( net_name=cfgs.NET_NAME, inputs=img_batch, gtboxes_and_label=None, is_training=False, share_head=False, share_net=share_net, stride=cfgs.STRIDE, anchor_ratios=cfgs.ANCHOR_RATIOS, anchor_scales=cfgs.ANCHOR_SCALES, scale_factors=cfgs.SCALE_FACTORS, base_anchor_size_list=cfgs. BASE_ANCHOR_SIZE_LIST, # P2, P3, P4, P5, P6 level=cfgs.LEVEL, top_k_nms=cfgs.RPN_TOP_K_NMS, rpn_nms_iou_threshold=cfgs.RPN_NMS_IOU_THRESHOLD, max_proposals_num=cfgs.MAX_PROPOSAL_NUM, rpn_iou_positive_threshold=cfgs.RPN_IOU_POSITIVE_THRESHOLD, rpn_iou_negative_threshold=cfgs.RPN_IOU_NEGATIVE_THRESHOLD, rpn_mini_batch_size=cfgs.RPN_MINIBATCH_SIZE, rpn_positives_ratio=cfgs.RPN_POSITIVE_RATE, remove_outside_anchors=False, # whether remove anchors outside rpn_weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME]) # rpn predict proposals rpn_proposals_boxes, rpn_proposals_scores = rpn.rpn_proposals( ) # rpn_score shape: [300, ] # *********************************************************************************************** # * Fast RCNN * # *********************************************************************************************** fast_rcnn = build_fast_rcnn1.FastRCNN( feature_pyramid=rpn.feature_pyramid, rpn_proposals_boxes=rpn_proposals_boxes, rpn_proposals_scores=rpn_proposals_scores, img_shape=tf.shape(img_batch), roi_size=cfgs.ROI_SIZE, roi_pool_kernel_size=cfgs.ROI_POOL_KERNEL_SIZE, scale_factors=cfgs.SCALE_FACTORS, gtboxes_and_label=None, gtboxes_and_label_minAreaRectangle= gtboxes_and_label_minAreaRectangle, fast_rcnn_nms_iou_threshold=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD, fast_rcnn_maximum_boxes_per_img=100, fast_rcnn_nms_max_boxes_per_class=cfgs. FAST_RCNN_NMS_MAX_BOXES_PER_CLASS, show_detections_score_threshold=cfgs.FINAL_SCORE_THRESHOLD, # show detections which score >= 0.6 num_classes=cfgs.CLASS_NUM, fast_rcnn_minibatch_size=cfgs.FAST_RCNN_MINIBATCH_SIZE, fast_rcnn_positives_ratio=cfgs.FAST_RCNN_POSITIVE_RATE, fast_rcnn_positives_iou_threshold=cfgs. FAST_RCNN_IOU_POSITIVE_THRESHOLD, # iou>0.5 is positive, iou<0.5 is negative use_dropout=cfgs.USE_DROPOUT, weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME], is_training=False, level=cfgs.LEVEL) fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category, \ fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, num_of_objects_rotate, detection_category_rotate = \ fast_rcnn.fast_rcnn_predict() if mode == 0: fast_rcnn_decode_boxes_rotate = get_horizen_minAreaRectangle( fast_rcnn_decode_boxes_rotate, False) # train init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = restore_model.get_restorer() with tf.Session() as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess, coord) gtboxes_horizontal_dict = {} predict_horizontal_dict = {} gtboxes_rotate_dict = {} predict_rotate_dict = {} for i in range(img_num): start = time.time() _img_name_batch, _img_batch, _gtboxes_and_label, _gtboxes_and_label_minAreaRectangle, \ _fast_rcnn_decode_boxes, _fast_rcnn_score, _detection_category, _fast_rcnn_decode_boxes_rotate, \ _fast_rcnn_score_rotate, _detection_category_rotate \ = sess.run([img_name_batch, img_batch, gtboxes_and_label, gtboxes_and_label_minAreaRectangle, fast_rcnn_decode_boxes, fast_rcnn_score, detection_category, fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, detection_category_rotate]) end = time.time() # gtboxes convert dict gtboxes_horizontal_dict[str(_img_name_batch[0])] = [] predict_horizontal_dict[str(_img_name_batch[0])] = [] gtboxes_rotate_dict[str(_img_name_batch[0])] = [] predict_rotate_dict[str(_img_name_batch[0])] = [] gtbox_horizontal_list, predict_horizontal_list = \ make_dict_packle(_gtboxes_and_label_minAreaRectangle, _fast_rcnn_decode_boxes, _fast_rcnn_score, _detection_category) if mode == 0: gtbox_rotate_list, predict_rotate_list = \ make_dict_packle(_gtboxes_and_label_minAreaRectangle, _fast_rcnn_decode_boxes_rotate, _fast_rcnn_score_rotate, _detection_category_rotate) else: gtbox_rotate_list, predict_rotate_list = \ make_dict_packle(_gtboxes_and_label, _fast_rcnn_decode_boxes_rotate, _fast_rcnn_score_rotate, _detection_category_rotate) gtboxes_horizontal_dict[str( _img_name_batch[0])].extend(gtbox_horizontal_list) predict_horizontal_dict[str( _img_name_batch[0])].extend(predict_horizontal_list) gtboxes_rotate_dict[str( _img_name_batch[0])].extend(gtbox_rotate_list) predict_rotate_dict[str( _img_name_batch[0])].extend(predict_rotate_list) view_bar( '{} image cost {}s'.format(str(_img_name_batch[0]), (end - start)), i + 1, img_num) fw1 = open('gtboxes_horizontal_dict.pkl', 'w') fw2 = open('predict_horizontal_dict.pkl', 'w') fw3 = open('gtboxes_rotate_dict.pkl', 'w') fw4 = open('predict_rotate_dict.pkl', 'w') pickle.dump(gtboxes_horizontal_dict, fw1) pickle.dump(predict_horizontal_dict, fw2) pickle.dump(gtboxes_rotate_dict, fw3) pickle.dump(predict_rotate_dict, fw4) fw1.close() fw2.close() fw3.close() fw4.close() coord.request_stop() coord.join(threads)
def train(): faster_rcnn = build_whole_network.DetectionNetwork( base_network_name=cfgs.NET_NAME, is_training=True) with tf.name_scope('get_batch'): img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \ next_batch(dataset_name=cfgs.DATASET_NAME, # 'pascal', 'coco' batch_size=cfgs.BATCH_SIZE, shortside_len=cfgs.IMG_SHORT_SIDE_LEN, is_training=True) mask_batch = tf.py_func( image_preprocess.get_mask, [tf.squeeze(img_batch, 0), tf.squeeze(gtboxes_and_label_batch, 0)], [tf.float32]) gtboxes_and_label = tf.py_func( back_forward_convert, inp=[tf.squeeze(gtboxes_and_label_batch, 0)], Tout=tf.float32) gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 6]) gtboxes_and_label_AreaRectangle = get_horizen_minAreaRectangle( gtboxes_and_label) gtboxes_and_label_AreaRectangle = tf.reshape( gtboxes_and_label_AreaRectangle, [-1, 5]) with tf.name_scope('draw_gtboxes'): gtboxes_in_img = draw_box_with_color( img_batch, tf.reshape(gtboxes_and_label_AreaRectangle, [-1, 5])[:, :-1], text=tf.shape(gtboxes_and_label_AreaRectangle)[0]) gtboxes_rotate_in_img = draw_box_with_color_rotate( img_batch, tf.reshape(gtboxes_and_label, [-1, 6])[:, :-1], text=tf.shape(gtboxes_and_label)[0]) biases_regularizer = tf.no_regularizer weights_regularizer = tf.contrib.layers.l2_regularizer(cfgs.WEIGHT_DECAY) # list as many types of layers as possible, even if they are not used now with slim.arg_scope([ slim.conv2d, slim.conv2d_in_plane, slim.conv2d_transpose, slim.separable_conv2d, slim.fully_connected ], weights_regularizer=weights_regularizer, biases_regularizer=biases_regularizer, biases_initializer=tf.constant_initializer(0.0)): final_boxes_h, final_scores_h, final_category_h, \ final_boxes_r, final_scores_r, final_category_r, loss_dict = faster_rcnn.build_whole_detection_network( input_img_batch=img_batch, gtboxes_r_batch=gtboxes_and_label, gtboxes_h_batch=gtboxes_and_label_AreaRectangle, mask_batch=mask_batch[0]) dets_in_img = draw_boxes_with_categories_and_scores( img_batch=img_batch, boxes=final_boxes_h, labels=final_category_h, scores=final_scores_h) dets_rotate_in_img = draw_boxes_with_categories_and_scores_rotate( img_batch=img_batch, boxes=final_boxes_r, labels=final_category_r, scores=final_scores_r) # ----------------------------------------------------------------------------------------------------build loss weight_decay_loss = tf.add_n(slim.losses.get_regularization_losses()) rpn_location_loss = loss_dict['rpn_loc_loss'] rpn_cls_loss = loss_dict['rpn_cls_loss'] rpn_total_loss = rpn_location_loss + rpn_cls_loss fastrcnn_cls_loss_h = loss_dict['fastrcnn_cls_loss_h'] fastrcnn_loc_loss_h = loss_dict['fastrcnn_loc_loss_h'] fastrcnn_cls_loss_r = loss_dict['fastrcnn_cls_loss_r'] fastrcnn_loc_loss_r = loss_dict['fastrcnn_loc_loss_r'] fastrcnn_total_loss = fastrcnn_cls_loss_h + fastrcnn_loc_loss_h + fastrcnn_cls_loss_r + fastrcnn_loc_loss_r if cfgs.ADD_ATTENTION: attention_loss = loss_dict['attention_loss'] total_loss = rpn_total_loss + fastrcnn_total_loss + attention_loss + weight_decay_loss else: attention_loss = tf.convert_to_tensor(0) total_loss = rpn_total_loss + fastrcnn_total_loss + weight_decay_loss # ---------------------------------------------------------------------------------------------------add summary tf.summary.scalar('RPN_LOSS/cls_loss', rpn_cls_loss) tf.summary.scalar('RPN_LOSS/location_loss', rpn_location_loss) tf.summary.scalar('RPN_LOSS/rpn_total_loss', rpn_total_loss) tf.summary.scalar('FAST_LOSS/fastrcnn_cls_loss_h', fastrcnn_cls_loss_h) tf.summary.scalar('FAST_LOSS/fastrcnn_location_loss_h', fastrcnn_loc_loss_h) tf.summary.scalar('FAST_LOSS/fastrcnn_cls_loss_r', fastrcnn_cls_loss_r) tf.summary.scalar('FAST_LOSS/fastrcnn_location_loss_r', fastrcnn_loc_loss_r) tf.summary.scalar('FAST_LOSS/fastrcnn_total_loss', fastrcnn_total_loss) if cfgs.ADD_ATTENTION: tf.summary.scalar('ATTENTION_LOSS/attention_loss', attention_loss) tf.summary.scalar('LOSS/total_loss', total_loss) tf.summary.scalar('LOSS/regular_weights', weight_decay_loss) tf.summary.image('img/gtboxes', gtboxes_in_img) tf.summary.image('img/gtboxes_rotate', gtboxes_rotate_in_img) tf.summary.image('img/mask', mask_batch) tf.summary.image('img/dets', dets_in_img) tf.summary.image('img/dets_rotate', dets_rotate_in_img) # ---------------------------------------------------------------------------------------------add summary global_step = slim.get_or_create_global_step() lr = tf.train.piecewise_constant( global_step, boundaries=[ np.int64(cfgs.DECAY_STEP[0]), np.int64(cfgs.DECAY_STEP[1]) ], values=[cfgs.LR, cfgs.LR / 10., cfgs.LR / 100.]) tf.summary.scalar('lr', lr) optimizer = tf.train.MomentumOptimizer(lr, momentum=cfgs.MOMENTUM) # ---------------------------------------------------------------------------------------------compute gradients gradients = faster_rcnn.get_gradients(optimizer, total_loss) # enlarge_gradients for bias if cfgs.MUTILPY_BIAS_GRADIENT: gradients = faster_rcnn.enlarge_gradients_for_bias(gradients) if cfgs.GRADIENT_CLIPPING_BY_NORM: with tf.name_scope('clip_gradients'): gradients = slim.learning.clip_gradient_norms( gradients, cfgs.GRADIENT_CLIPPING_BY_NORM) # ---------------------------------------------------------------------------------------------compute gradients # train_op train_op = optimizer.apply_gradients(grads_and_vars=gradients, global_step=global_step) summary_op = tf.summary.merge_all() init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = faster_rcnn.get_restorer() saver = tf.train.Saver(max_to_keep=10) config = tf.ConfigProto() config.gpu_options.allow_growth = True # with tf.Session(config=config) as sess: # with tf.Session() as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess, coord) summary_path = os.path.join(cfgs.SUMMARY_PATH, cfgs.VERSION) tools.mkdir(summary_path) summary_writer = tf.summary.FileWriter(summary_path, graph=sess.graph) for step in range(cfgs.MAX_ITERATION): training_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())) if step % cfgs.SHOW_TRAIN_INFO_INTE != 0 and step % cfgs.SMRY_ITER != 0: _, global_stepnp = sess.run([train_op, global_step]) else: if step % cfgs.SHOW_TRAIN_INFO_INTE == 0 and step % cfgs.SMRY_ITER != 0: start = time.time() _global_step, _img_name_batch, _rpn_location_loss, _rpn_classification_loss, \ _rpn_total_loss, _fast_rcnn_location_loss, _fast_rcnn_classification_loss, \ _fast_rcnn_location_rotate_loss, _fast_rcnn_classification_rotate_loss, \ _fast_rcnn_total_loss, _attention_loss, _total_loss, _ = \ sess.run([global_step, img_name_batch, rpn_location_loss, rpn_cls_loss, rpn_total_loss, fastrcnn_loc_loss_h, fastrcnn_cls_loss_h, fastrcnn_loc_loss_r, fastrcnn_cls_loss_r, fastrcnn_total_loss, attention_loss, total_loss, train_op]) end = time.time() print(""" {}: step{} image_name:{} |\t rpn_loc_loss:{} |\t rpn_cla_loss:{} |\t rpn_total_loss:{} | fast_rcnn_loc_loss:{} |\t fast_rcnn_cla_loss:{} |\t fast_rcnn_loc_rotate_loss:{} |\t fast_rcnn_cla_rotate_loss:{} |\t fast_rcnn_total_loss:{} |\t attention_loss:{} |\t total_loss:{} |\t pre_cost_time:{}s""" \ .format(training_time, _global_step, str(_img_name_batch[0]), _rpn_location_loss, _rpn_classification_loss, _rpn_total_loss, _fast_rcnn_location_loss, _fast_rcnn_classification_loss, _fast_rcnn_location_rotate_loss, _fast_rcnn_classification_rotate_loss, _fast_rcnn_total_loss, _attention_loss, _total_loss, (end - start))) else: if step % cfgs.SMRY_ITER == 0: _, global_stepnp, summary_str = sess.run( [train_op, global_step, summary_op]) summary_writer.add_summary(summary_str, global_stepnp) summary_writer.flush() if (step > 0 and step % cfgs.SAVE_WEIGHTS_INTE == 0) or (step == cfgs.MAX_ITERATION - 1): save_dir = os.path.join(cfgs.TRAINED_CKPT, cfgs.VERSION) if not os.path.exists(save_dir): os.makedirs(save_dir) save_ckpt = os.path.join( save_dir, 'voc_' + str(global_stepnp) + 'model.ckpt') saver.save(sess, save_ckpt) print(' weights had been saved') coord.request_stop() coord.join(threads)
def train(): with tf.Graph().as_default(): with tf.name_scope('get_batch'): img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \ next_batch(dataset_name=cfgs.DATASET_NAME, batch_size=cfgs.BATCH_SIZE, shortside_len=cfgs.SHORT_SIDE_LEN, is_training=True) gtboxes_and_label = tf.py_func( back_forward_convert, inp=[tf.squeeze(gtboxes_and_label_batch, 0)], Tout=tf.float32) gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 6]) gtboxes_and_label_minAreaRectangle = get_horizen_minAreaRectangle( gtboxes_and_label) gtboxes_and_label_minAreaRectangle = tf.reshape( gtboxes_and_label_minAreaRectangle, [-1, 5]) with tf.name_scope('draw_gtboxes'): gtboxes_in_img = draw_box_with_color( img_batch, tf.reshape(gtboxes_and_label_minAreaRectangle, [-1, 5])[:, :-1], text=tf.shape(gtboxes_and_label_minAreaRectangle)[0]) # *********************************************************************************************** # * share net * # *********************************************************************************************** _, share_net = get_network_byname(net_name=cfgs.NET_NAME, inputs=img_batch, num_classes=None, is_training=True, output_stride=None, global_pool=False, spatial_squeeze=False) # *********************************************************************************************** # * rpn * # *********************************************************************************************** rpn = build_rpn.RPN( net_name=cfgs.NET_NAME, inputs=img_batch, gtboxes_and_label=gtboxes_and_label_minAreaRectangle, is_training=True, share_head=cfgs.SHARE_HEAD, share_net=share_net, stride=cfgs.STRIDE, anchor_ratios=cfgs.ANCHOR_RATIOS, anchor_scales=cfgs.ANCHOR_SCALES, scale_factors=cfgs.SCALE_FACTORS, base_anchor_size_list=cfgs. BASE_ANCHOR_SIZE_LIST, # P2, P3, P4, P5, P6 level=cfgs.LEVEL, top_k_nms=cfgs.RPN_TOP_K_NMS, rpn_nms_iou_threshold=cfgs.RPN_NMS_IOU_THRESHOLD, max_proposals_num=cfgs.MAX_PROPOSAL_NUM, rpn_iou_positive_threshold=cfgs.RPN_IOU_POSITIVE_THRESHOLD, rpn_iou_negative_threshold=cfgs. RPN_IOU_NEGATIVE_THRESHOLD, # iou>=0.7 is positive box, iou< 0.3 is negative rpn_mini_batch_size=cfgs.RPN_MINIBATCH_SIZE, rpn_positives_ratio=cfgs.RPN_POSITIVE_RATE, remove_outside_anchors=False, # whether remove anchors outside rpn_weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME]) rpn_proposals_boxes, rpn_proposals_scores = rpn.rpn_proposals( ) # rpn_score shape: [300, ] rpn_location_loss, rpn_classification_loss = rpn.rpn_losses() rpn_total_loss = rpn_classification_loss + rpn_location_loss with tf.name_scope('draw_proposals'): # score > 0.5 is object rpn_object_boxes_indices = tf.reshape( tf.where(tf.greater(rpn_proposals_scores, 0.5)), [-1]) rpn_object_boxes = tf.gather(rpn_proposals_boxes, rpn_object_boxes_indices) rpn_proposals_objcet_boxes_in_img = draw_box_with_color( img_batch, rpn_object_boxes, text=tf.shape(rpn_object_boxes)[0]) rpn_proposals_boxes_in_img = draw_box_with_color( img_batch, rpn_proposals_boxes, text=tf.shape(rpn_proposals_boxes)[0]) # *********************************************************************************************** # * Fast RCNN * # *********************************************************************************************** fast_rcnn = build_fast_rcnn.FastRCNN( feature_pyramid=rpn.feature_pyramid, rpn_proposals_boxes=rpn_proposals_boxes, rpn_proposals_scores=rpn_proposals_scores, img_shape=tf.shape(img_batch), roi_size=cfgs.ROI_SIZE, roi_pool_kernel_size=cfgs.ROI_POOL_KERNEL_SIZE, scale_factors=cfgs.SCALE_FACTORS, gtboxes_and_label=gtboxes_and_label, gtboxes_and_label_minAreaRectangle= gtboxes_and_label_minAreaRectangle, fast_rcnn_nms_iou_threshold=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD, fast_rcnn_maximum_boxes_per_img=100, fast_rcnn_nms_max_boxes_per_class=cfgs. FAST_RCNN_NMS_MAX_BOXES_PER_CLASS, show_detections_score_threshold=cfgs. FINAL_SCORE_THRESHOLD, # show detections which score >= 0.6 num_classes=cfgs.CLASS_NUM, fast_rcnn_minibatch_size=cfgs.FAST_RCNN_MINIBATCH_SIZE, fast_rcnn_positives_ratio=cfgs.FAST_RCNN_POSITIVE_RATE, fast_rcnn_positives_iou_threshold=cfgs. FAST_RCNN_IOU_POSITIVE_THRESHOLD, # iou>0.5 is positive, iou<0.5 is negative use_dropout=cfgs.USE_DROPOUT, weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME], is_training=True, level=cfgs.LEVEL) fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category = \ fast_rcnn.fast_rcnn_predict() fast_rcnn_location_loss, fast_rcnn_classification_loss = fast_rcnn.fast_rcnn_loss( ) fast_rcnn_total_loss = fast_rcnn_location_loss + fast_rcnn_classification_loss with tf.name_scope('draw_boxes_with_categories'): fast_rcnn_predict_boxes_in_imgs = draw_boxes_with_categories( img_batch=img_batch, boxes=fast_rcnn_decode_boxes, labels=detection_category, scores=fast_rcnn_score) # train total_loss = slim.losses.get_total_loss() global_step = slim.get_or_create_global_step() lr = tf.train.piecewise_constant( global_step, boundaries=[np.int64(20000), np.int64(40000)], values=[cfgs.LR, cfgs.LR / 10, cfgs.LR / 100]) tf.summary.scalar('lr', lr) optimizer = tf.train.MomentumOptimizer(lr, momentum=cfgs.MOMENTUM) train_op = slim.learning.create_train_op( total_loss, optimizer, global_step) # rpn_total_loss, # train_op = optimizer.minimize(second_classification_loss, global_step) # *********************************************************************************************** # * Summary * # *********************************************************************************************** # ground truth and predict tf.summary.image('img/gtboxes', gtboxes_in_img) tf.summary.image('img/faster_rcnn_predict', fast_rcnn_predict_boxes_in_imgs) # rpn loss and image tf.summary.scalar('rpn/rpn_location_loss', rpn_location_loss) tf.summary.scalar('rpn/rpn_classification_loss', rpn_classification_loss) tf.summary.scalar('rpn/rpn_total_loss', rpn_total_loss) tf.summary.scalar('fast_rcnn/fast_rcnn_location_loss', fast_rcnn_location_loss) tf.summary.scalar('fast_rcnn/fast_rcnn_classification_loss', fast_rcnn_classification_loss) tf.summary.scalar('fast_rcnn/fast_rcnn_total_loss', fast_rcnn_total_loss) tf.summary.scalar('loss/total_loss', total_loss) tf.summary.image('rpn/rpn_all_boxes', rpn_proposals_boxes_in_img) tf.summary.image('rpn/rpn_object_boxes', rpn_proposals_objcet_boxes_in_img) # learning_rate tf.summary.scalar('learning_rate', lr) summary_op = tf.summary.merge_all() init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = restore_model.get_restorer() saver = tf.train.Saver(max_to_keep=10) config = tf.ConfigProto() # config.gpu_options.per_process_gpu_memory_fraction = 0.5 config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess, coord) summary_path = os.path.join(FLAGS.summary_path, cfgs.VERSION) mkdir(summary_path) summary_writer = tf.summary.FileWriter(summary_path, graph=sess.graph) for step in range(cfgs.MAX_ITERATION): training_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())) start = time.time() _global_step, _img_name_batch, _rpn_location_loss, _rpn_classification_loss, \ _rpn_total_loss, _fast_rcnn_location_loss, _fast_rcnn_classification_loss, \ _fast_rcnn_total_loss, _total_loss, _ = \ sess.run([global_step, img_name_batch, rpn_location_loss, rpn_classification_loss, rpn_total_loss, fast_rcnn_location_loss, fast_rcnn_classification_loss, fast_rcnn_total_loss, total_loss, train_op]) end = time.time() if step % 10 == 0: print(""" {}: step{} image_name:{} |\t rpn_loc_loss:{} |\t rpn_cla_loss:{} |\t rpn_total_loss:{} | fast_rcnn_loc_loss:{} |\t fast_rcnn_cla_loss:{} |\t fast_rcnn_total_loss:{} | total_loss:{} |\t pre_cost_time:{}s""" \ .format(training_time, _global_step, str(_img_name_batch[0]), _rpn_location_loss, _rpn_classification_loss, _rpn_total_loss, _fast_rcnn_location_loss, _fast_rcnn_classification_loss, _fast_rcnn_total_loss, _total_loss, (end - start))) if step % 50 == 0: summary_str = sess.run(summary_op) summary_writer.add_summary(summary_str, _global_step) summary_writer.flush() if (step > 0 and step % 1000 == 0) or (step == cfgs.MAX_ITERATION - 1): save_dir = os.path.join(FLAGS.trained_checkpoint, cfgs.VERSION) mkdir(save_dir) save_ckpt = os.path.join( save_dir, 'voc_' + str(_global_step) + 'model.ckpt') saver.save(sess, save_ckpt) print(' weights had been saved') coord.request_stop() coord.join(threads)
def build_whole_detection_network(self, input_img_batch, gtboxes_h_batch, gtboxes_r_batch): if self.is_training: # ensure shape is [M, 5] gtboxes_r_batch = tf.reshape(gtboxes_r_batch, [-1, 6]) gtboxes_r_batch = tf.cast(gtboxes_r_batch, tf.float32) gtboxes_batch = tf.reshape(gtboxes_h_batch, [-1, 5]) gtboxes_batch = tf.cast(gtboxes_batch, tf.float32) img_shape = tf.shape(input_img_batch) # 1. build base network feature_to_cropped = self.build_base_network(input_img_batch) # 2. build rpn with tf.variable_scope('build_rpn', regularizer=slim.l2_regularizer( cfgs.WEIGHT_DECAY)): rpn_conv3x3 = slim.conv2d(feature_to_cropped, 512, [3, 3], trainable=self.is_training, weights_initializer=cfgs.INITIALIZER, activation_fn=tf.nn.relu, scope='rpn_conv/3x3') rpn_cls_score = slim.conv2d(rpn_conv3x3, self.num_anchors_per_location * 2, [1, 1], stride=1, trainable=self.is_training, weights_initializer=cfgs.INITIALIZER, activation_fn=None, scope='rpn_cls_score') rpn_box_pred = slim.conv2d( rpn_conv3x3, self.num_anchors_per_location * 4, [1, 1], stride=1, trainable=self.is_training, weights_initializer=cfgs.BBOX_INITIALIZER, activation_fn=None, scope='rpn_bbox_pred') rpn_box_pred = tf.reshape(rpn_box_pred, [-1, 4]) rpn_cls_score = tf.reshape(rpn_cls_score, [-1, 2]) rpn_cls_prob = slim.softmax(rpn_cls_score, scope='rpn_cls_prob') # 3. generate_anchors featuremap_height, featuremap_width = tf.shape( feature_to_cropped)[1], tf.shape(feature_to_cropped)[2] featuremap_height = tf.cast(featuremap_height, tf.float32) featuremap_width = tf.cast(featuremap_width, tf.float32) anchors = anchor_utils.make_anchors( base_anchor_size=cfgs.BASE_ANCHOR_SIZE_LIST[0], anchor_scales=cfgs.ANCHOR_SCALES, anchor_ratios=cfgs.ANCHOR_RATIOS, featuremap_height=featuremap_height, featuremap_width=featuremap_width, stride=cfgs.ANCHOR_STRIDE, name="make_anchors_forRPN") # with tf.variable_scope('make_anchors'): # anchors = anchor_utils.make_anchors(height=featuremap_height, # width=featuremap_width, # feat_stride=cfgs.ANCHOR_STRIDE[0], # anchor_scales=cfgs.ANCHOR_SCALES, # anchor_ratios=cfgs.ANCHOR_RATIOS, base_size=16 # ) # 4. postprocess rpn proposals. such as: decode, clip, NMS with tf.variable_scope('postprocess_RPN'): # rpn_cls_prob = tf.reshape(rpn_cls_score, [-1, 2]) # rpn_cls_prob = slim.softmax(rpn_cls_prob, scope='rpn_cls_prob') # rpn_box_pred = tf.reshape(rpn_box_pred, [-1, 4]) rois, roi_scores = postprocess_rpn_proposals( rpn_bbox_pred=rpn_box_pred, rpn_cls_prob=rpn_cls_prob, img_shape=img_shape, anchors=anchors, is_training=self.is_training) # rois shape [-1, 4] # +++++++++++++++++++++++++++++++++++++add img smry+++++++++++++++++++++++++++++++++++++++++++++++++++++++ if self.is_training: rois_in_img = show_box_in_tensor.draw_boxes_with_scores( img_batch=input_img_batch, boxes=rois, scores=roi_scores) tf.summary.image('all_rpn_rois', rois_in_img) score_gre_05 = tf.reshape( tf.where(tf.greater_equal(roi_scores, 0.5)), [-1]) score_gre_05_rois = tf.gather(rois, score_gre_05) score_gre_05_score = tf.gather(roi_scores, score_gre_05) score_gre_05_in_img = show_box_in_tensor.draw_boxes_with_scores( img_batch=input_img_batch, boxes=score_gre_05_rois, scores=score_gre_05_score) tf.summary.image('score_greater_05_rois', score_gre_05_in_img) # ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ if self.is_training: with tf.variable_scope('sample_anchors_minibatch'): rpn_labels, rpn_bbox_targets = \ tf.py_func( anchor_target_layer, [gtboxes_batch, img_shape, anchors], [tf.float32, tf.float32]) rpn_bbox_targets = tf.reshape(rpn_bbox_targets, [-1, 4]) rpn_labels = tf.to_int32(rpn_labels, name="to_int32") rpn_labels = tf.reshape(rpn_labels, [-1]) self.add_anchor_img_smry(input_img_batch, anchors, rpn_labels) # --------------------------------------add smry---------------------------------------------------------------- rpn_cls_category = tf.argmax(rpn_cls_prob, axis=1) kept_rpppn = tf.reshape(tf.where(tf.not_equal(rpn_labels, -1)), [-1]) rpn_cls_category = tf.gather(rpn_cls_category, kept_rpppn) acc = tf.reduce_mean( tf.to_float( tf.equal(rpn_cls_category, tf.to_int64(tf.gather(rpn_labels, kept_rpppn))))) tf.summary.scalar('ACC/rpn_accuracy', acc) with tf.control_dependencies([rpn_labels]): with tf.variable_scope('sample_RCNN_minibatch_stage1'): stage1_rois, stage1_labels, stage1_bbox_targets = \ tf.py_func(proposal_target_layer_3, [rois, rois, gtboxes_batch, gtboxes_r_batch, cfgs.FAST_RCNN_IOU_POSITIVE_THRESHOLD[0]], [tf.float32, tf.float32, tf.float32]) stage1_rois = tf.reshape(stage1_rois, [-1, 4]) stage1_labels = tf.to_int32(stage1_labels) stage1_labels = tf.reshape(stage1_labels, [-1]) stage1_bbox_targets = tf.reshape( stage1_bbox_targets, [-1, 5 * (cfgs.CLASS_NUM + 1)]) self.add_roi_batch_img_smry(input_img_batch, stage1_rois, stage1_labels, 'stage1') #stage1_bbox_targets_h = boxes_utils.get_horizen_minAreaRectangle(stage1_bbox_targets, False) else: stage1_rois = rois # -------------------------------------------------------------------------------------------------------------# # Fast-RCNN-before1 # # -------------------------------------------------------------------------------------------------------------# # 5. build Fast-RCNN-before1 # rois = tf.Print(rois, [tf.shape(rois)], 'rois shape', summarize=10) stage1_bbox_pred_fliter, stage1_bbox_pred, stage1_cls_score = self.build_fastrcnn( feature_to_cropped=feature_to_cropped, rois=stage1_rois, img_shape=img_shape, scope='stage1') # bbox_pred shape: [-1, 4*(cls_num+1)]. # cls_score shape: [-1, cls_num+1] stage1_cls_prob = slim.softmax(stage1_cls_score, 'stage1_cls_prob') stage1_cls_category = tf.argmax(stage1_cls_prob, axis=1) # ----------------------------------------------add smry------------------------------------------------------- if self.is_training: stage1_fast_acc = tf.reduce_mean( tf.to_float( tf.equal(stage1_cls_category, tf.to_int64(stage1_labels)))) tf.summary.scalar('ACC/stage1_fast_acc', stage1_fast_acc) # postprocess_fastrcnn_before1 # return x,y,w,h,theta stage1_bbox = self.postprocess_cascade( rois=stage1_rois, bbox_ppred=stage1_bbox_pred_fliter, scope='stage1', five=False) #stage1_bbox_h = boxes_utils.get_horizen_minAreaRectangle(stage1_bbox, with_label=False) if self.is_training: overlaps = iou_rotate.iou_rotate_calculate(stage1_bbox, gtboxes_r_batch[:, :-1], use_gpu=True, gpu_id=0) if self.is_training: with tf.control_dependencies([stage1_bbox]): with tf.variable_scope('sample_RCNN_minibatch_stage2'): stage2_rois, stage2_labels, stage2_bbox_targets = \ tf.py_func(proposal_target_layer_r, [stage1_bbox,gtboxes_r_batch, overlaps, cfgs.FAST_RCNN_IOU_POSITIVE_THRESHOLD[1]], [tf.float32, tf.float32,tf.float32]) stage2_rois = tf.reshape(stage2_rois, [-1, 5]) # 斜 stage2_labels = tf.to_int32(stage2_labels) stage2_labels = tf.reshape(stage2_labels, [-1]) stage2_bbox_targets = tf.reshape( stage2_bbox_targets, [-1, 5 * (cfgs.CLASS_NUM + 1)]) self.add_roi_batch_img_smry_rotate(input_img_batch, stage2_rois, stage2_labels, 'stage2') else: stage2_rois = stage1_bbox # -------------------------------------------------------------------------------------------------------------# # Fast-RCNN-before2 # # -------------------------------------------------------------------------------------------------------------# # 6. build Fast-RCNN-before2 # rois = tf.Print(rois, [tf.shape(rois)], 'rois shape', summarize=10) # stage2_rois = tf.stop_gradient(stage2_rois) stage2_rois_h = boxes_utils.get_horizen_minAreaRectangle( stage2_rois, with_label=False) ##斜变正 stage2_bbox_pred_fliter, stage2_bbox_pred, stage2_cls_score = self.build_fastrcnn( feature_to_cropped=feature_to_cropped, rois=stage2_rois_h, img_shape=img_shape, scope='stage2') # bbox_pred shape: [-1, 4*(cls_num+1)]. # cls_score shape: [-1, cls_num+1] stage2_cls_prob = slim.softmax(stage2_cls_score, 'stage2_cls_prob') stage2_cls_category = tf.argmax(stage2_cls_prob, axis=1) # ----------------------------------------------add smry------------------------------------------------------- if self.is_training: stage2_fast_acc = tf.reduce_mean( tf.to_float( tf.equal(stage2_cls_category, tf.to_int64(stage2_labels)))) tf.summary.scalar('ACC/stage2_fast_acc', stage2_fast_acc) # postprocess_fastrcnn_before2 stage2_bbox = self.postprocess_cascade( rois=stage2_rois, bbox_ppred=stage2_bbox_pred_fliter, scope='stage2') #stage2_bbox_h = boxes_utils.get_horizen_minAreaRectangle(stage2_bbox, with_label=False) if self.is_training: overlaps = iou_rotate.iou_rotate_calculate(stage2_bbox, gtboxes_r_batch[:, :-1], use_gpu=True, gpu_id=0) if self.is_training: with tf.control_dependencies([stage2_bbox]): with tf.variable_scope('sample_RCNN_minibatch_stage3'): stage3_rois, stage3_labels, stage3_bbox_targets = \ tf.py_func(proposal_target_layer_r, [stage2_bbox, gtboxes_r_batch, overlaps, cfgs.FAST_RCNN_IOU_POSITIVE_THRESHOLD[2]], [tf.float32, tf.float32, tf.float32]) stage3_rois = tf.reshape(stage3_rois, [-1, 5]) stage3_labels = tf.to_int32(stage3_labels) stage3_labels = tf.reshape(stage3_labels, [-1]) stage3_bbox_targets = tf.reshape( stage3_bbox_targets, [-1, 5 * (cfgs.CLASS_NUM + 1)]) self.add_roi_batch_img_smry_rotate(input_img_batch, stage3_rois, stage3_labels, 'stage3') else: stage3_rois = stage2_bbox # -------------------------------------------------------------------------------------------------------------# # Fast-RCNN # # -------------------------------------------------------------------------------------------------------------# # 7. build Fast-RCNN # rois = tf.Print(rois, [tf.shape(rois)], 'rois shape', summarize=10) # stage3_rois = tf.stop_gradient(stage3_rois) stage3_rois_h = boxes_utils.get_horizen_minAreaRectangle( stage3_rois, with_label=False) ##斜变正 stage3_bbox_pred, stage3_cls_score = self.build_fastrcnn( feature_to_cropped=feature_to_cropped, rois=stage3_rois_h, img_shape=img_shape, scope='stage3') # bbox_pred shape: [-1, 4*(cls_num+1)]. # cls_score shape: [-1, cls_num+1] stage3_cls_prob = slim.softmax(stage3_cls_score, 'stage3_cls_prob') stage3_cls_category = tf.argmax(stage3_cls_prob, axis=1) # ----------------------------------------------add smry------------------------------------------------------- if self.is_training: fast_acc = tf.reduce_mean( tf.to_float( tf.equal(stage3_cls_category, tf.to_int64(stage3_labels)))) tf.summary.scalar('ACC/fast_acc', fast_acc) # postprocess_fastrcnn if not self.is_training: with slim.arg_scope([ slim.conv2d, slim.conv2d_in_plane, slim.conv2d_transpose, slim.separable_conv2d, slim.fully_connected ], reuse=True): _, _, final_scores_stage2 = self.build_fastrcnn( feature_to_cropped=feature_to_cropped, rois=stage3_rois_h, img_shape=img_shape, scope='stage2') final_scores_stage2 = slim.softmax(final_scores_stage2, 'final_scores_stage2') _, _, final_scores_stage1 = self.build_fastrcnn( feature_to_cropped=feature_to_cropped, rois=stage3_rois_h, img_shape=img_shape, scope='stage1') final_scores_stage1 = slim.softmax(final_scores_stage1, 'final_scores_stage1') # choose which stage to export cls_prob = tf.add(final_scores_stage2, final_scores_stage1) cls_prob = tf.add(cls_prob, stage3_cls_prob) / 3 return self.postprocess_fastrcnn_r(rois=stage3_rois, bbox_ppred=stage3_bbox_pred, scores=cls_prob, img_shape=img_shape, scope='stage3') else: ''' when trian. We need build Loss ''' loss_dict = self.build_loss( rpn_box_pred=rpn_box_pred, rpn_bbox_targets=rpn_bbox_targets, rpn_cls_score=rpn_cls_score, rpn_labels=rpn_labels, bbox_pred=stage3_bbox_pred, bbox_targets=stage3_bbox_targets, stage2_bbox_pred=stage2_bbox_pred, stage2_bbox_targets=stage2_bbox_targets, stage1_bbox_pred=stage1_bbox_pred, stage1_bbox_targets=stage1_bbox_targets, cls_score=stage3_cls_score, labels=stage3_labels, stage2_cls_score=stage2_cls_score, stage2_labels=stage2_labels, stage1_cls_score=stage1_cls_score, stage1_labels=stage1_labels) final_bbox, final_scores, final_category = self.postprocess_fastrcnn_r( rois=stage3_rois, bbox_ppred=stage3_bbox_pred, scores=stage3_cls_prob, img_shape=img_shape, scope='stage3') return final_bbox, final_scores, final_category, loss_dict
def eval_text(img_num, mode): with tf.Graph().as_default(): faster_rcnn = build_whole_network.DetectionNetwork( base_network_name=cfgs.NET_NAME, is_training=False) # 通过readtfrecord来读取文件 img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \ next_batch(dataset_name=cfgs.DATASET_NAME, batch_size=cfgs.BATCH_SIZE, shortside_len=cfgs.IMG_SHORT_SIDE_LEN, is_training=False) gtboxes_and_label = tf.py_func( back_forward_convert, inp=[tf.squeeze(gtboxes_and_label_batch, 0)], Tout=tf.float32) gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 6]) gtboxes_and_label_minAreaRectangle = get_horizen_minAreaRectangle( gtboxes_and_label) gtboxes_and_label_minAreaRectangle = tf.reshape( gtboxes_and_label_minAreaRectangle, [-1, 5]) final_boxes_h, final_scores_h, final_category_h, \ final_boxes_r, final_scores_r, final_category_r= faster_rcnn.build_whole_detection_network( input_img_batch=img_batch, gtboxes_r_batch=gtboxes_and_label, gtboxes_h_batch=gtboxes_and_label_minAreaRectangle) if mode == 0: final_boxes_r = get_horizen_minAreaRectangle(final_boxes_r, False) # train init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = faster_rcnn.get_restorer() config = tf.ConfigProto() # config.gpu_options.per_process_gpu_memory_fraction = 0.5 config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess, coord) gtboxes_horizontal_dict = {} predict_horizontal_dict = {} gtboxes_rotate_dict = {} predict_rotate_dict = {} for i in range(img_num): start = time.time() _img_name_batch, _img_batch, _gtboxes_and_label, _gtboxes_and_label_minAreaRectangle, \ _final_boxes_h, _final_scores_h, _final_category_h, _final_boxes_r, \ _final_scores_r, _final_category_r \ = sess.run([img_name_batch, img_batch, gtboxes_and_label, gtboxes_and_label_minAreaRectangle, final_boxes_h, final_scores_h, final_category_h, final_boxes_r, final_scores_r, final_category_r]) end = time.time() # det_detections_h = draw_box_in_img.draw_box_cv(np.squeeze(_img_batch, 0), # boxes=_final_boxes_h, # labels=_final_category_h, # scores=final_scores_h) # det_detections_r = draw_box_in_img.draw_rotate_box_cv(np.squeeze(_img_batch, 0), # boxes=_final_boxes_r, # labels=_final_category_r, # scores= _final_scores_r) # save_dir = os.path.join('/home/sws/code/Tfrecord/', cfgs.VERSION) # tools.mkdir(save_dir) # cv2.imwrite(save_dir + '/' + _img_name_batch[0] + '_h.jpg', # det_detections_h) # cv2.imwrite(save_dir + '/' + _img_name_batch[0] + '_r.jpg', # det_detections_r) # gtboxes convert dict gtboxes_horizontal_dict[str(_img_name_batch[0])] = [] predict_horizontal_dict[str(_img_name_batch[0])] = [] gtboxes_rotate_dict[str(_img_name_batch[0])] = [] predict_rotate_dict[str(_img_name_batch[0])] = [] gtbox_horizontal_list, predict_horizontal_list = \ make_dict_packle(_gtboxes_and_label_minAreaRectangle, _final_boxes_h, _final_scores_h, _final_category_h) # if mode == 0: # gtbox_rotate_list, predict_rotate_list = \ # make_dict_packle(_gtboxes_and_label_minAreaRectangle, _fast_rcnn_decode_boxes_rotate, # _fast_rcnn_score_rotate, _detection_category_rotate) if mode == 0: gtbox_rotate_list, predict_rotate_list = \ make_dict_packle(_gtboxes_and_label_minAreaRectangle, _final_boxes_r, _final_scores_r, _final_category_r) else: gtbox_rotate_list, predict_rotate_list = \ make_dict_packle(_gtboxes_and_label, _final_boxes_r, _final_scores_r, _final_category_r) gtboxes_horizontal_dict[str( _img_name_batch[0])].extend(gtbox_horizontal_list) predict_horizontal_dict[str( _img_name_batch[0])].extend(predict_horizontal_list) gtboxes_rotate_dict[str( _img_name_batch[0])].extend(gtbox_rotate_list) predict_rotate_dict[str( _img_name_batch[0])].extend(predict_rotate_list) view_bar( '{} image cost {}s'.format(str(_img_name_batch[0]), (end - start)), i + 1, img_num) print('\n') fw1 = open('gtboxes_horizontal_dict.pkl', 'w') fw2 = open('predict_horizontal_dict.pkl', 'w') fw3 = open('gtboxes_rotate_dict.pkl', 'w') fw4 = open('predict_rotate_dict.pkl', 'w') pickle.dump(gtboxes_horizontal_dict, fw1) pickle.dump(predict_horizontal_dict, fw2) pickle.dump(gtboxes_rotate_dict, fw3) pickle.dump(predict_rotate_dict, fw4) fw1.close() fw2.close() fw3.close() fw4.close() coord.request_stop() coord.join(threads)
def test(img_num): with tf.Graph().as_default(): img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \ next_batch(dataset_name=cfgs.DATASET_NAME, batch_size=cfgs.BATCH_SIZE, shortside_len=cfgs.SHORT_SIDE_LEN, is_training=False) gtboxes_and_label = tf.py_func( back_forward_convert, inp=[tf.squeeze(gtboxes_and_label_batch, 0)], Tout=tf.float32) gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 6]) gtboxes_and_label_minAreaRectangle = get_horizen_minAreaRectangle( gtboxes_and_label) gtboxes_and_label_minAreaRectangle = tf.reshape( gtboxes_and_label_minAreaRectangle, [-1, 5]) # *********************************************************************************************** # * share net * # *********************************************************************************************** _, share_net = get_network_byname(net_name=cfgs.NET_NAME, inputs=img_batch, num_classes=None, is_training=True, output_stride=None, global_pool=False, spatial_squeeze=False) # *********************************************************************************************** # * RPN * # *********************************************************************************************** rpn = build_rpn.RPN( net_name=cfgs.NET_NAME, inputs=img_batch, gtboxes_and_label=None, is_training=False, share_head=cfgs.SHARE_HEAD, share_net=share_net, stride=cfgs.STRIDE, anchor_ratios=cfgs.ANCHOR_RATIOS, anchor_scales=cfgs.ANCHOR_SCALES, scale_factors=cfgs.SCALE_FACTORS, base_anchor_size_list=cfgs. BASE_ANCHOR_SIZE_LIST, # P2, P3, P4, P5, P6 level=cfgs.LEVEL, top_k_nms=cfgs.RPN_TOP_K_NMS, rpn_nms_iou_threshold=cfgs.RPN_NMS_IOU_THRESHOLD, max_proposals_num=cfgs.MAX_PROPOSAL_NUM, rpn_iou_positive_threshold=cfgs.RPN_IOU_POSITIVE_THRESHOLD, rpn_iou_negative_threshold=cfgs.RPN_IOU_NEGATIVE_THRESHOLD, rpn_mini_batch_size=cfgs.RPN_MINIBATCH_SIZE, rpn_positives_ratio=cfgs.RPN_POSITIVE_RATE, remove_outside_anchors=False, # whether remove anchors outside rpn_weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME]) # rpn predict proposals rpn_proposals_boxes, rpn_proposals_scores = rpn.rpn_proposals( ) # rpn_score shape: [300, ] # *********************************************************************************************** # * Fast RCNN * # *********************************************************************************************** fast_rcnn = build_fast_rcnn1.FastRCNN( feature_pyramid=rpn.feature_pyramid, rpn_proposals_boxes=rpn_proposals_boxes, rpn_proposals_scores=rpn_proposals_scores, img_shape=tf.shape(img_batch), roi_size=cfgs.ROI_SIZE, roi_pool_kernel_size=cfgs.ROI_POOL_KERNEL_SIZE, scale_factors=cfgs.SCALE_FACTORS, gtboxes_and_label=None, gtboxes_and_label_minAreaRectangle= gtboxes_and_label_minAreaRectangle, fast_rcnn_nms_iou_threshold=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD, fast_rcnn_maximum_boxes_per_img=100, fast_rcnn_nms_max_boxes_per_class=cfgs. FAST_RCNN_NMS_MAX_BOXES_PER_CLASS, show_detections_score_threshold=cfgs.FINAL_SCORE_THRESHOLD, # show detections which score >= 0.6 num_classes=cfgs.CLASS_NUM, fast_rcnn_minibatch_size=cfgs.FAST_RCNN_MINIBATCH_SIZE, fast_rcnn_positives_ratio=cfgs.FAST_RCNN_POSITIVE_RATE, fast_rcnn_positives_iou_threshold=cfgs. FAST_RCNN_IOU_POSITIVE_THRESHOLD, # iou>0.5 is positive, iou<0.5 is negative use_dropout=cfgs.USE_DROPOUT, weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME], is_training=False, level=cfgs.LEVEL) fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category, \ fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, num_of_objects_rotate, detection_category_rotate,_, _, _, _ = \ fast_rcnn.fast_rcnn_predict() # train init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = restore_model.get_restorer() config = tf.ConfigProto() # config.gpu_options.per_process_gpu_memory_fraction = 0.5 config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess, coord) TrueDetection = 0 FalseAlarm = 0 TotalObjects = 0 TotalDetections = 0 for i in range(img_num): start = time.time() _img_name_batch, _img_batch, _gtboxes_and_label, _gtboxes_and_label_minAreaRectangle, \ _fast_rcnn_decode_boxes, _fast_rcnn_score, _detection_category, _fast_rcnn_decode_boxes_rotate, \ _fast_rcnn_score_rotate, _detection_category_rotate \ = sess.run([img_name_batch, img_batch, gtboxes_and_label, gtboxes_and_label_minAreaRectangle, fast_rcnn_decode_boxes, fast_rcnn_score, detection_category, fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, detection_category_rotate]) end = time.time() _img_batch = np.squeeze(_img_batch, axis=0) _img_batch_fpn_horizonal = help_utils.draw_box_cv( _img_batch, boxes=_fast_rcnn_decode_boxes, labels=_detection_category, scores=_fast_rcnn_score) _img_batch_fpn_rotate = help_utils.draw_rotate_box_cv( _img_batch, boxes=_fast_rcnn_decode_boxes_rotate, labels=_detection_category_rotate, scores=_fast_rcnn_score_rotate) mkdir(cfgs.TEST_SAVE_PATH) cv2.imwrite( cfgs.TEST_SAVE_PATH + str(_img_name_batch[0]) + '_horizontal_fpn.jpg', _img_batch_fpn_horizonal) cv2.imwrite( cfgs.TEST_SAVE_PATH + str(_img_name_batch[0]) + '_rotate_fpn.jpg', _img_batch_fpn_rotate) temp_label_horizontal = np.reshape(_gtboxes_and_label[:, -1:], [ -1, ]).astype(np.int64) temp_label_rotate = np.reshape(_gtboxes_and_label[:, -1:], [ -1, ]).astype(np.int64) _img_batch_gt_horizontal = help_utils.draw_box_cv( _img_batch, boxes=_gtboxes_and_label_minAreaRectangle[:, :-1], labels=temp_label_horizontal, scores=None) _img_batch_gt_rotate = help_utils.draw_rotate_box_cv( _img_batch, boxes=_gtboxes_and_label[:, :-1], labels=temp_label_rotate, scores=None) cv2.imwrite( cfgs.TEST_SAVE_PATH + str(_img_name_batch[0]) + '_horizontal_gt.jpg', _img_batch_gt_horizontal) cv2.imwrite( cfgs.TEST_SAVE_PATH + str(_img_name_batch[0]) + '_rotate_gt.jpg', _img_batch_gt_rotate) Current_TrueDetection = 0 lock_list = [] for m in range(len(_gtboxes_and_label)): for n, box in enumerate(_fast_rcnn_decode_boxes_rotate): if n in lock_list: continue y_c, x_c, h, w, theta = box[0], box[1], box[2], box[ 3], box[4] r1 = ((x_c, y_c), (w, h), theta) gt_box = _gtboxes_and_label[m] gt_yc, gt_xc, gt_h, gt_w, gt_theta = gt_box[0], gt_box[ 1], gt_box[2], gt_box[3], gt_box[4] r2 = ((gt_xc, gt_yc), (gt_w, gt_h), gt_theta) int_pts = cv2.rotatedRectangleIntersection(r1, r2)[0] if int_pts > 0.5: Current_TrueDetection += 1 lock_list.append(n) TrueDetection += Current_TrueDetection if (len(_fast_rcnn_decode_boxes_rotate) > 0): FalseAlarm += len( _fast_rcnn_score_rotate) - Current_TrueDetection TotalDetections += len(_fast_rcnn_score_rotate) TotalObjects += len(_gtboxes_and_label) view_bar( '{} image cost {}s'.format(str(_img_name_batch[0]), (end - start)), i + 1, img_num) Pd = float(TrueDetection) / float(TotalDetections) Pf = float(FalseAlarm) / float(TotalDetections) F1 = 2 * float(Pd * (1 - Pf)) / float(Pd + (1 - Pf)) Op = float(TrueDetection) / float(TotalObjects) print("\nACC: " + str(Pd)) print("FalseAlarm: " + str(Pf)) print("F1: " + str(F1)) print("Op: " + str(TrueDetection) + "/" + str(TotalObjects) + "=" + str(Op)) coord.request_stop() coord.join(threads)
def eval_ship(img_num): with tf.Graph().as_default(): img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \ next_batch(dataset_name=cfgs.DATASET_NAME, batch_size=cfgs.BATCH_SIZE, shortside_len=cfgs.SHORT_SIDE_LEN, is_training=True) gtboxes_and_label, head = get_head( tf.squeeze(gtboxes_and_label_batch, 0)) gtboxes_and_label = tf.py_func(back_forward_convert, inp=[gtboxes_and_label], Tout=tf.float32) gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 6]) head_quadrant = tf.py_func(get_head_quadrant, inp=[head, gtboxes_and_label], Tout=tf.float32) head_quadrant = tf.reshape(head_quadrant, [-1, 1]) gtboxes_and_label_minAreaRectangle = get_horizen_minAreaRectangle( gtboxes_and_label) gtboxes_and_label_minAreaRectangle = tf.reshape( gtboxes_and_label_minAreaRectangle, [-1, 5]) # *********************************************************************************************** # * share net * # *********************************************************************************************** _, share_net = get_network_byname(net_name=cfgs.NET_NAME, inputs=img_batch, num_classes=None, is_training=True, output_stride=None, global_pool=False, spatial_squeeze=False) # *********************************************************************************************** # * RPN * # *********************************************************************************************** rpn = build_rpn.RPN( net_name=cfgs.NET_NAME, inputs=img_batch, gtboxes_and_label=None, is_training=False, share_head=cfgs.SHARE_HEAD, share_net=share_net, stride=cfgs.STRIDE, anchor_ratios=cfgs.ANCHOR_RATIOS, anchor_scales=cfgs.ANCHOR_SCALES, scale_factors=cfgs.SCALE_FACTORS, base_anchor_size_list=cfgs. BASE_ANCHOR_SIZE_LIST, # P2, P3, P4, P5, P6 level=cfgs.LEVEL, top_k_nms=cfgs.RPN_TOP_K_NMS, rpn_nms_iou_threshold=cfgs.RPN_NMS_IOU_THRESHOLD, max_proposals_num=cfgs.MAX_PROPOSAL_NUM, rpn_iou_positive_threshold=cfgs.RPN_IOU_POSITIVE_THRESHOLD, rpn_iou_negative_threshold=cfgs.RPN_IOU_NEGATIVE_THRESHOLD, rpn_mini_batch_size=cfgs.RPN_MINIBATCH_SIZE, rpn_positives_ratio=cfgs.RPN_POSITIVE_RATE, remove_outside_anchors=False, # whether remove anchors outside rpn_weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME]) # rpn predict proposals rpn_proposals_boxes, rpn_proposals_scores = rpn.rpn_proposals( ) # rpn_score shape: [300, ] # *********************************************************************************************** # * Fast RCNN * # *********************************************************************************************** fast_rcnn = build_fast_rcnn.FastRCNN( feature_pyramid=rpn.feature_pyramid, rpn_proposals_boxes=rpn_proposals_boxes, rpn_proposals_scores=rpn_proposals_scores, img_shape=tf.shape(img_batch), img_batch=img_batch, roi_size=cfgs.ROI_SIZE, roi_pool_kernel_size=cfgs.ROI_POOL_KERNEL_SIZE, scale_factors=cfgs.SCALE_FACTORS, gtboxes_and_label=None, gtboxes_and_label_minAreaRectangle= gtboxes_and_label_minAreaRectangle, fast_rcnn_nms_iou_threshold=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD, fast_rcnn_maximum_boxes_per_img=100, fast_rcnn_nms_max_boxes_per_class=cfgs. FAST_RCNN_NMS_MAX_BOXES_PER_CLASS, show_detections_score_threshold=cfgs.FINAL_SCORE_THRESHOLD, # show detections which score >= 0.6 num_classes=cfgs.CLASS_NUM, fast_rcnn_minibatch_size=cfgs.FAST_RCNN_MINIBATCH_SIZE, fast_rcnn_positives_ratio=cfgs.FAST_RCNN_POSITIVE_RATE, fast_rcnn_positives_iou_threshold=cfgs. FAST_RCNN_IOU_POSITIVE_THRESHOLD, # iou>0.5 is positive, iou<0.5 is negative use_dropout=cfgs.USE_DROPOUT, weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME], is_training=False, level=cfgs.LEVEL, head_quadrant=head_quadrant) fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category, \ fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, fast_rcnn_head_quadrant, \ num_of_objects_rotate, detection_category_rotate = fast_rcnn.fast_rcnn_predict() # train init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = restore_model.get_restorer() config = tf.ConfigProto() # config.gpu_options.per_process_gpu_memory_fraction = 0.5 config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess, coord) gtboxes_horizontal_dict = {} gtboxes_rotate_dict = {} all_boxes_h = [] all_boxes_r = [] all_img_names = [] for i in range(img_num): start = time.time() _img_name_batch, _img_batch, _gtboxes_and_label, _gtboxes_and_label_minAreaRectangle, \ _fast_rcnn_decode_boxes, _fast_rcnn_score, _detection_category, _fast_rcnn_decode_boxes_rotate, \ _fast_rcnn_score_rotate, _detection_category_rotate \ = sess.run([img_name_batch, img_batch, gtboxes_and_label, gtboxes_and_label_minAreaRectangle, fast_rcnn_decode_boxes, fast_rcnn_score, detection_category, fast_rcnn_decode_boxes_rotate, fast_rcnn_score_rotate, detection_category_rotate]) end = time.time() # gtboxes convert dict gtboxes_horizontal_dict[str(_img_name_batch[0])] = [] gtboxes_rotate_dict[str(_img_name_batch[0])] = [] gtbox_horizontal_list, gtbox_rotate_list = make_dict_packle( _gtboxes_and_label, _gtboxes_and_label_minAreaRectangle) xmin, ymin, xmax, ymax = _fast_rcnn_decode_boxes[:, 1], _fast_rcnn_decode_boxes[:, 0], _fast_rcnn_decode_boxes[:, 3], _fast_rcnn_decode_boxes[:, 2] x_c, y_c, w, h, theta = _fast_rcnn_decode_boxes_rotate[:, 1], _fast_rcnn_decode_boxes_rotate[:, 0], _fast_rcnn_decode_boxes_rotate[:, 3], \ _fast_rcnn_decode_boxes_rotate[:, 2], _fast_rcnn_decode_boxes_rotate[:, 4] boxes_h = np.transpose(np.stack([xmin, ymin, xmax, ymax])) boxes_r = np.transpose(np.stack([x_c, y_c, w, h, theta])) dets_h = np.hstack((_detection_category.reshape(-1, 1), _fast_rcnn_score.reshape(-1, 1), boxes_h)) dets_r = np.hstack( (_detection_category_rotate.reshape(-1, 1), _fast_rcnn_score_rotate.reshape(-1, 1), boxes_r)) all_boxes_h.append(dets_h) all_boxes_r.append(dets_r) all_img_names.append(str(_img_name_batch[0])) gtboxes_horizontal_dict[str( _img_name_batch[0])].extend(gtbox_horizontal_list) gtboxes_rotate_dict[str( _img_name_batch[0])].extend(gtbox_rotate_list) print(str(_img_name_batch[0])) view_bar( '{} image cost {}s'.format(str(_img_name_batch[0]), (end - start)), i + 1, img_num) write_voc_results_file(all_boxes_h, all_img_names, cfgs.EVALUATE_R_DIR, 0) write_voc_results_file(all_boxes_r, all_img_names, cfgs.EVALUATE_R_DIR, 1) fw1 = open('gtboxes_horizontal_dict.pkl', 'wb') fw2 = open('gtboxes_rotate_dict.pkl', 'wb') pickle.dump(gtboxes_horizontal_dict, fw1) pickle.dump(gtboxes_rotate_dict, fw2) fw1.close() fw2.close() coord.request_stop() coord.join(threads)
def test_rotate(): with tf.Graph().as_default(): with tf.name_scope('get_batch'): img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \ next_batch(dataset_name=cfgs.DATASET_NAME, batch_size=cfgs.BATCH_SIZE, shortside_len=cfgs.SHORT_SIDE_LEN, is_training=True) gtboxes_and_label, head = get_head( tf.squeeze(gtboxes_and_label_batch, 0)) gtboxes_and_label = tf.py_func(back_forward_convert, inp=[gtboxes_and_label], Tout=tf.float32) gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 6]) head_quadrant = tf.py_func(get_head_quadrant, inp=[head, gtboxes_and_label], Tout=tf.float32) head_quadrant = tf.reshape(head_quadrant, [-1, 1]) gtboxes_and_label_minAreaRectangle = get_horizen_minAreaRectangle( gtboxes_and_label) gtboxes_and_label_minAreaRectangle = tf.reshape( gtboxes_and_label_minAreaRectangle, [-1, 5]) with tf.name_scope('draw_gtboxes'): gtboxes_in_img = draw_box_with_color( img_batch, tf.reshape(gtboxes_and_label_minAreaRectangle, [-1, 5])[:, :-1], text=tf.shape(gtboxes_and_label_minAreaRectangle)[0]) gtboxes_rotate_in_img = draw_box_with_color_rotate( img_batch, tf.reshape(gtboxes_and_label, [-1, 6])[:, :-1], text=tf.shape(gtboxes_and_label)[0], head=head_quadrant) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) config = tf.ConfigProto() # config.gpu_options.per_process_gpu_memory_fraction = 0.5 config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess, coord) for i in range(650): img_gtboxes, img_gtboxes_rotate, img_name = sess.run( [gtboxes_in_img, gtboxes_rotate_in_img, img_name_batch]) img_gtboxes = np.squeeze(img_gtboxes, axis=0) img_gtboxes_rotate = np.squeeze(img_gtboxes_rotate, axis=0) print(i) cv2.imwrite( cfgs.INFERENCE_SAVE_PATH + '/{}_horizontal_fpn.jpg'.format(str(img_name[0])), img_gtboxes) cv2.imwrite( cfgs.INFERENCE_SAVE_PATH + '/{}_rotate_fpn.jpg'.format(str(img_name[0])), img_gtboxes_rotate) coord.request_stop() coord.join(threads)
def eval_dict_convert(img_num, mode): with tf.Graph().as_default(): # img = tf.placeholder(shape=[None, None, 3], dtype=tf.uint8) img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch = \ next_batch(dataset_name=cfgs.DATASET_NAME, batch_size=cfgs.BATCH_SIZE, shortside_len=cfgs.SHORT_SIDE_LEN, is_training=False) gtboxes_and_label = tf.py_func( back_forward_convert, inp=[tf.squeeze(gtboxes_and_label_batch, 0)], Tout=tf.float32) gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 6]) gtboxes_and_label_minAreaRectangle = get_horizen_minAreaRectangle( gtboxes_and_label) gtboxes_and_label_minAreaRectangle = tf.reshape( gtboxes_and_label_minAreaRectangle, [-1, 5]) # *********************************************************************************************** # * share net * # *********************************************************************************************** _, share_net = get_network_byname(net_name=cfgs.NET_NAME, inputs=img_batch, num_classes=None, is_training=True, output_stride=None, global_pool=False, spatial_squeeze=False) # *********************************************************************************************** # * RPN * # *********************************************************************************************** rpn = build_rpn.RPN( net_name=cfgs.NET_NAME, inputs=img_batch, gtboxes_and_label=gtboxes_and_label, is_training=False, share_head=cfgs.SHARE_HEAD, share_net=share_net, anchor_ratios=cfgs.ANCHOR_RATIOS, anchor_scales=cfgs.ANCHOR_SCALES, anchor_angles=cfgs.ANCHOR_ANGLES, scale_factors=cfgs.SCALE_FACTORS, base_anchor_size_list=cfgs. BASE_ANCHOR_SIZE_LIST, # P2, P3, P4, P5, P6 level=cfgs.LEVEL, anchor_stride=cfgs.ANCHOR_STRIDE, pool_stride=cfgs.POOL_STRIDE, top_k_nms=cfgs.RPN_TOP_K_NMS, kernel_size=cfgs.KERNEL_SIZE, use_angles_condition=False, anchor_angle_threshold=cfgs.RPN_ANCHOR_ANGLES_THRESHOLD, nms_angle_threshold=cfgs.RPN_NMS_ANGLES_THRESHOLD, rpn_nms_iou_threshold=cfgs.RPN_NMS_IOU_THRESHOLD, max_proposals_num=cfgs.MAX_PROPOSAL_NUM, rpn_iou_positive_threshold=cfgs.RPN_IOU_POSITIVE_THRESHOLD, rpn_iou_negative_threshold=cfgs.RPN_IOU_NEGATIVE_THRESHOLD, # iou>=0.7 is positive box, iou< 0.3 is negative rpn_mini_batch_size=cfgs.RPN_MINIBATCH_SIZE, rpn_positives_ratio=cfgs.RPN_POSITIVE_RATE, remove_outside_anchors=cfgs. IS_FILTER_OUTSIDE_BOXES, # whether remove anchors outside rpn_weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME], scope='') rpn_proposals_boxes, rpn_proposals_scores = rpn.rpn_proposals( ) # rpn_score shape: [300, ] _, _, rpn_predict_boxes, rpn_predict_scores = rpn.rpn_losses() # *********************************************************************************************** # * Fast RCNN * # *********************************************************************************************** fast_rcnn = build_fast_rcnn.FastRCNN( img_batch=img_batch, feature_pyramid=rpn.feature_pyramid, rpn_proposals_boxes=rpn_proposals_boxes, rpn_proposals_scores=rpn_proposals_scores, stop_gradient_for_proposals=False, img_shape=tf.shape(img_batch), roi_size=cfgs.ROI_SIZE, roi_pool_kernel_size=cfgs.ROI_POOL_KERNEL_SIZE, scale_factors=cfgs.SCALE_FACTORS, gtboxes_and_label=None, fast_rcnn_nms_iou_threshold=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD, top_k_nms=cfgs.FAST_RCNN_TOP_K_NMS, nms_angle_threshold=cfgs.FAST_RCNN_NMS_ANGLES_THRESHOLD, use_angle_condition=False, level=cfgs.LEVEL, fast_rcnn_maximum_boxes_per_img=100, fast_rcnn_nms_max_boxes_per_class=cfgs. FAST_RCNN_NMS_MAX_BOXES_PER_CLASS, show_detections_score_threshold=cfgs.FINAL_SCORE_THRESHOLD, # show detections which score >= 0.6 num_classes=cfgs.CLASS_NUM, fast_rcnn_minibatch_size=cfgs.FAST_RCNN_MINIBATCH_SIZE, fast_rcnn_positives_ratio=cfgs.FAST_RCNN_POSITIVE_RATE, fast_rcnn_positives_iou_threshold=cfgs. FAST_RCNN_IOU_POSITIVE_THRESHOLD, boxes_angle_threshold=cfgs.FAST_RCNN_BOXES_ANGLES_THRESHOLD, use_dropout=cfgs.USE_DROPOUT, weight_decay=cfgs.WEIGHT_DECAY[cfgs.NET_NAME], is_training=False) fast_rcnn_decode_boxes, fast_rcnn_score, num_of_objects, detection_category = \ fast_rcnn.fast_rcnn_predict() ############################################################################################## if cfgs.NEED_AUXILIARY: predict_boxes = tf.concat( [fast_rcnn_decode_boxes, rpn_predict_boxes], axis=0) predict_scores = tf.concat( [fast_rcnn_score, rpn_predict_scores - 0.2], axis=0) rpn_predict_label = tf.ones([ tf.shape(rpn_predict_scores)[0], ], tf.int64) labels = tf.concat([detection_category, rpn_predict_label], axis=0) # valid_indices = nms_rotate.nms_rotate(decode_boxes=predict_boxes, # scores=predict_scores, # iou_threshold=0.15, # max_output_size=30, # use_angle_condition=False, # angle_threshold=15, # use_gpu=True) valid_indices = tf.py_func(nms_rotate.nms_rotate_cpu, inp=[ predict_boxes, predict_scores, tf.constant(0.15, tf.float32), tf.constant(30, tf.float32) ], Tout=tf.int64) fast_rcnn_decode_boxes = tf.gather(predict_boxes, valid_indices) fast_rcnn_score = tf.gather(predict_scores, valid_indices) detection_category = tf.gather(labels, valid_indices) ############################################################################################## if mode == 0: fast_rcnn_decode_boxes = get_horizen_minAreaRectangle( fast_rcnn_decode_boxes, False) # train init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) restorer, restore_ckpt = restore_model.get_restorer() config = tf.ConfigProto() # config.gpu_options.per_process_gpu_memory_fraction = 0.5 config.gpu_options.allow_growth = True with tf.Session(config=config) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess, coord) gtboxes_dict = {} predict_dict = {} for i in range(img_num): start = time.time() _img_name_batch, _img_batch, _gtboxes_and_label, _fast_rcnn_decode_boxes, \ _gtboxes_and_label_minAreaRectangle, _fast_rcnn_score, _detection_category \ = sess.run([img_name_batch, img_batch, gtboxes_and_label, fast_rcnn_decode_boxes, gtboxes_and_label_minAreaRectangle, fast_rcnn_score, detection_category]) end = time.time() # gtboxes convert dict gtboxes_dict[str(_img_name_batch[0])] = [] predict_dict[str(_img_name_batch[0])] = [] # for j, box in enumerate(_gtboxes_and_label): # bbox_dict = {} # bbox_dict['bbox'] = np.array(_gtboxes_and_label[j, :-1], np.float64) # bbox_dict['name'] = LABEl_NAME_MAP[int(_gtboxes_and_label[j, -1])] # gtbox_dict[str(_img_name_batch[0])].append(bbox_dict) # # for label in NAME_LABEL_MAP.keys(): # if label == 'back_ground': # continue # else: # temp_dict = {} # temp_dict['name'] = label # # ind = np.where(_detection_category == NAME_LABEL_MAP[label])[0] # temp_boxes = _fast_rcnn_decode_boxes[ind] # temp_score = np.reshape(_fast_rcnn_score[ind], [-1, 1]) # temp_dict['bbox'] = np.array(np.concatenate([temp_boxes, temp_score], axis=1), np.float64) # predict_dict[str(_img_name_batch[0])].append(temp_dict) if mode == 0: gtboxes_list, predict_list = \ make_dict_packle(_gtboxes_and_label_minAreaRectangle, _fast_rcnn_decode_boxes, _fast_rcnn_score, _detection_category) else: gtboxes_list, predict_list = \ make_dict_packle(_gtboxes_and_label, _fast_rcnn_decode_boxes, _fast_rcnn_score, _detection_category) gtboxes_dict[str(_img_name_batch[0])].extend(gtboxes_list) predict_dict[str(_img_name_batch[0])].extend(predict_list) view_bar( '{} image cost {}s'.format(str(_img_name_batch[0]), (end - start)), i + 1, img_num) fw1 = open('gtboxes_dict.pkl', 'w') fw2 = open('predict_dict.pkl', 'w') pickle.dump(gtboxes_dict, fw1) pickle.dump(predict_dict, fw2) fw1.close() fw2.close() coord.request_stop() coord.join(threads)
def train(): with tf.Graph().as_default(), tf.device('/cpu:0'): global_step = slim.get_or_create_global_step() lr = tf.train.piecewise_constant(global_step, boundaries=[np.int64(cfgs.DECAY_STEP[0]), np.int64(cfgs.DECAY_STEP[1]), np.int64(cfgs.DECAY_STEP[2])], values=[cfgs.LR, cfgs.LR / 10., cfgs.LR / 100., cfgs.LR / 1000.]) tf.summary.scalar('lr', lr) optimizer = tf.train.MomentumOptimizer(lr, momentum=cfgs.MOMENTUM) faster_rcnn = build_whole_network.DetectionNetwork(base_network_name=cfgs.NET_NAME, is_training=True) with tf.name_scope('get_batch'): num_gpu = len(cfgs.GPU_GROUP.strip().split(',')) img_name_batch, img_batch, gtboxes_and_label_batch, num_objects_batch, img_h_batch, img_w_batch = \ next_batch(dataset_name=cfgs.DATASET_NAME, # 'pascal', 'coco' batch_size=cfgs.BATCH_SIZE * num_gpu, shortside_len=cfgs.IMG_SHORT_SIDE_LEN, is_training=True) # gtboxes_and_label = tf.reshape(gtboxes_and_label_batch, [-1, 5]) # if cfgs.NET_NAME in ['resnet101_v1d', 'resnet50_v1d']: # img_batch = img_batch / tf.constant([cfgs.PIXEL_STD]) # data processing inputs_list = [] for i in range(num_gpu): # img_name = img_name_batch[i] img = tf.expand_dims(img_batch[i], axis=0) gtboxes_and_label = tf.cast(tf.reshape(gtboxes_and_label_batch[i], [-1, 9]), tf.float32) num_objects = num_objects_batch[i] num_objects = tf.cast(tf.reshape(num_objects, [-1, ]), tf.float32) img_h = img_h_batch[i] img_w = img_w_batch[i] # img_h = tf.cast(tf.reshape(img_h, [-1, ]), tf.float32) # img_w = tf.cast(tf.reshape(img_w, [-1, ]), tf.float32) inputs_list.append([img, gtboxes_and_label, num_objects, img_h, img_w]) tower_grads = [] biases_regularizer = tf.no_regularizer weights_regularizer = tf.contrib.layers.l2_regularizer(cfgs.WEIGHT_DECAY) total_loss_dict = { 'rpn_cls_loss': tf.constant(0., tf.float32), 'rpn_loc_loss': tf.constant(0., tf.float32), 'fastrcnn_cls_loss_h': tf.constant(0., tf.float32), 'fastrcnn_loc_loss_h': tf.constant(0., tf.float32), 'fastrcnn_cls_loss_r': tf.constant(0., tf.float32), 'fastrcnn_loc_loss_r': tf.constant(0., tf.float32), 'total_losses': tf.constant(0., tf.float32), } with tf.variable_scope(tf.get_variable_scope()): for i in range(num_gpu): with tf.device('/gpu:%d' % i): with tf.name_scope('tower_%d' % i): with slim.arg_scope( [slim.model_variable, slim.variable], device='/device:CPU:0'): with slim.arg_scope([slim.conv2d, slim.conv2d_in_plane, slim.conv2d_transpose, slim.separable_conv2d, slim.fully_connected], weights_regularizer=weights_regularizer, biases_regularizer=biases_regularizer, biases_initializer=tf.constant_initializer(0.0)): gtboxes_and_label = tf.py_func(get_gtboxes_and_label, inp=[inputs_list[i][1], inputs_list[i][2]], Tout=tf.float32) gtboxes_and_label = tf.py_func(back_forward_convert, inp=[gtboxes_and_label], Tout=tf.float32) gtboxes_and_label = tf.reshape(gtboxes_and_label, [-1, 6]) gtboxes_and_label_AreaRectangle = get_horizen_minAreaRectangle(gtboxes_and_label) gtboxes_and_label_AreaRectangle = tf.reshape(gtboxes_and_label_AreaRectangle, [-1, 5]) img = inputs_list[i][0] img_shape = inputs_list[i][-2:] img = tf.image.crop_to_bounding_box(image=img, offset_height=0, offset_width=0, target_height=tf.cast(img_shape[0], tf.int32), target_width=tf.cast(img_shape[1], tf.int32)) final_boxes_h, final_scores_h, final_category_h, \ final_boxes_r, final_scores_r, final_category_r, loss_dict = faster_rcnn.build_whole_detection_network( input_img_batch=img, gtboxes_r_batch=gtboxes_and_label, gtboxes_h_batch=gtboxes_and_label_AreaRectangle) if cfgs.ADD_BOX_IN_TENSORBOARD: dets_in_img = draw_boxes_with_categories_and_scores(img_batch=img, boxes=final_boxes_h, labels=final_category_h, scores=final_scores_h) dets_rotate_in_img = draw_boxes_with_categories_and_scores_rotate(img_batch=img, boxes=final_boxes_r, labels=final_category_r, scores=final_scores_r) tf.summary.image('Compare/final_detection_rotate_gpu:%d' % i, dets_rotate_in_img) total_losses = 0.0 for k in loss_dict.keys(): total_losses += loss_dict[k] total_loss_dict[k] += loss_dict[k] / num_gpu total_losses = total_losses / num_gpu total_loss_dict['total_losses'] += total_losses if i == num_gpu - 1: regularization_losses = tf.get_collection( tf.GraphKeys.REGULARIZATION_LOSSES) # weight_decay_loss = tf.add_n(slim.losses.get_regularization_losses()) total_losses = total_losses + tf.add_n(regularization_losses) tf.get_variable_scope().reuse_variables() grads = optimizer.compute_gradients(total_losses) # average loss tower_grads.append(grads) for k in total_loss_dict.keys(): tf.summary.scalar('{}/{}'.format(k.split('_')[0], k), total_loss_dict[k]) if len(tower_grads) > 1: grads = sum_gradients(tower_grads) else: grads = tower_grads[0] apply_gradient_op = optimizer.apply_gradients(grads, global_step=global_step) variable_averages = tf.train.ExponentialMovingAverage(0.9999, global_step) variables_averages_op = variable_averages.apply(tf.trainable_variables()) train_op = tf.group(apply_gradient_op, variables_averages_op) # train_op = optimizer.apply_gradients(final_gvs, global_step=global_step) summary_op = tf.summary.merge_all() restorer, restore_ckpt = faster_rcnn.get_restorer() saver = tf.train.Saver(max_to_keep=10) init_op = tf.group( tf.global_variables_initializer(), tf.local_variables_initializer() ) tfconfig = tf.ConfigProto( allow_soft_placement=True, log_device_placement=False) #tfconfig.gpu_options.allow_growth = True tfconfig.gpu_options.per_process_gpu_memory_fraction = 0.7 with tf.Session(config=tfconfig) as sess: sess.run(init_op) if not restorer is None: restorer.restore(sess, restore_ckpt) print('restore model') coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(coord=coord, sess=sess) summary_path = os.path.join(cfgs.SUMMARY_PATH, cfgs.VERSION) tools.mkdir(summary_path) summary_writer = tf.summary.FileWriter(summary_path, graph=sess.graph) for step in range(cfgs.MAX_ITERATION // num_gpu): training_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())) if step % cfgs.SHOW_TRAIN_INFO_INTE != 0 and step % cfgs.SMRY_ITER != 0: _, global_stepnp = sess.run([train_op, global_step]) else: if step % cfgs.SHOW_TRAIN_INFO_INTE == 0 and step % cfgs.SMRY_ITER != 0: start = time.time() _, global_stepnp, total_loss_dict_ = \ sess.run([train_op, global_step, total_loss_dict]) end = time.time() print('***'*20) print("""%s: global_step:%d current_step:%d""" % (training_time, (global_stepnp-1)*num_gpu, step*num_gpu)) print("""per_cost_time:%.3fs""" % ((end - start) / num_gpu)) loss_str = '' for k in total_loss_dict_.keys(): loss_str += '%s:%.3f\n' % (k, total_loss_dict_[k]) print(loss_str) else: if step % cfgs.SMRY_ITER == 0: _, global_stepnp, summary_str = sess.run([train_op, global_step, summary_op]) summary_writer.add_summary(summary_str, (global_stepnp-1)*num_gpu) summary_writer.flush() if (step > 0 and step % (cfgs.SAVE_WEIGHTS_INTE // num_gpu) == 0) or (step >= cfgs.MAX_ITERATION // num_gpu - 1): save_dir = os.path.join(cfgs.TRAINED_CKPT, cfgs.VERSION) if not os.path.exists(save_dir): os.mkdir(save_dir) save_ckpt = os.path.join(save_dir, 'coco_' + str((global_stepnp-1)*num_gpu) + 'model.ckpt') saver.save(sess, save_ckpt) print(' weights had been saved') coord.request_stop() coord.join(threads)