(features, labels, mode, params)
Our model_fn for ResNet to be used with our Estimator.
Our model_fn for ResNet to be used with our Estimator.
Our model_fn for ResNet to be used with our Estimator .
def lighr_head_model_fn(features, labels, mode, params):
    """Our model_fn for ResNet to be used with our Estimator."""
    num_anchors_list = labels['num_anchors_list']
    num_feature_layers = len(num_anchors_list)

    shape = labels['targets'][-1]
    if mode != tf.estimator.ModeKeys.TRAIN:
        org_image = labels['targets'][-2]
        isdifficult = labels['targets'][-3]
        bbox_img = labels['targets'][-4]
        gbboxes_raw = labels['targets'][-5]
        glabels_raw = labels['targets'][-6]

    glabels = labels['targets'][:num_feature_layers][0]
    gtargets = labels['targets'][num_feature_layers : 2 * num_feature_layers][0]
    gscores = labels['targets'][2 * num_feature_layers : 3 * num_feature_layers][0]

    #features = tf.ones([4,480,480,3]) * 0.5
    with tf.variable_scope(params['model_scope'], default_name = None, values = [features], reuse=tf.AUTO_REUSE):
        rpn_feat_map, backbone_feat = xception_body.XceptionBody(features, params['num_classes'], is_training=(mode == tf.estimator.ModeKeys.TRAIN), data_format=params['data_format'])
        #rpn_feat_map = tf.Print(rpn_feat_map,[tf.shape(rpn_feat_map), rpn_feat_map,backbone_feat])
        rpn_cls_score, rpn_bbox_pred = xception_body.get_rpn(rpn_feat_map, num_anchors_list[0], (mode == tf.estimator.ModeKeys.TRAIN), params['data_format'], 'rpn_head')

        large_sep_feature = xception_body.large_sep_kernel(backbone_feat, 256, 10 * 7 * 7, (mode == tf.estimator.ModeKeys.TRAIN), params['data_format'], 'large_sep_feature')

        if params['data_format'] == 'channels_first':
            rpn_cls_score = tf.transpose(rpn_cls_score, [0, 2, 3, 1])
            rpn_bbox_pred = tf.transpose(rpn_bbox_pred, [0, 2, 3, 1])

        rpn_cls_score = tf.reshape(rpn_cls_score, [-1, 2])
        rpn_object_score = tf.nn.softmax(rpn_cls_score)[:, -1]

        rpn_object_score = tf.reshape(rpn_object_score, [1, -1])
        rpn_location_pred = tf.reshape(rpn_bbox_pred, [1, -1, 4])

        rpn_bboxes_pred = labels['rpn_decode_fn'](rpn_location_pred)

        proposals_bboxes = xception_body.get_proposals(rpn_object_score, rpn_bboxes_pred, None, params['rpn_pre_nms_top_n'], params['rpn_post_nms_top_n'], params['rpn_nms_thres'], params['rpn_min_size'], (mode == tf.estimator.ModeKeys.TRAIN), params['data_format'])
        #proposals_targets = tf.Print(proposals_targets, [proposals_targets], message='proposals_targets0:')

        cls_score, bboxes_reg = xception_body.get_head(large_sep_feature, lambda input_, bboxes_, grid_width_, grid_height_ : ps_roi_align(input_, bboxes_, grid_width_, grid_height_, pool_method), 7, 7, None, proposals_bboxes, params['num_classes'], (mode == tf.estimator.ModeKeys.TRAIN), False, 0, params['data_format'], 'final_head')

        head_bboxes_pred = labels['head_decode_fn'](proposals_bboxes, bboxes_reg)

        head_cls_score = tf.reshape(cls_score, [-1, params['num_classes']])
        head_cls_score = tf.nn.softmax(head_cls_score)
        head_bboxes_pred = tf.reshape(head_bboxes_pred, [-1, 4])

        shape = tf.squeeze(shape, axis = 0)
        glabels = tf.squeeze(glabels, axis = 0)
        gtargets = tf.squeeze(gtargets, axis = 0)
        gscores = tf.squeeze(gscores, axis = 0)
        if mode != tf.estimator.ModeKeys.TRAIN:
            org_image = tf.squeeze(org_image, axis = 0)
            isdifficult = tf.squeeze(isdifficult, axis = 0)
            gbboxes_raw = tf.squeeze(gbboxes_raw, axis = 0)
            glabels_raw = tf.squeeze(glabels_raw, axis = 0)
            bbox_img = tf.squeeze(bbox_img, axis = 0)

        eval_ops, save_image_op = bboxes_eval(org_image, shape, bbox_img, cls_score, head_bboxes_pred, glabels_raw, gbboxes_raw, isdifficult, params['num_classes'])
        _ = tf.identity(save_image_op, name='save_image_with_bboxes_op')

    with tf.control_dependencies([save_image_op]):
        weight_decay_loss = params['weight_decay'] * tf.add_n([tf.nn.l2_loss(v) for v in tf.trainable_variables() if 'batch_normalization' not in v.name])

    predictions = {
        'classes': tf.argmax(head_cls_score, axis=-1),
        'probabilities': tf.reduce_max(head_cls_score, axis=-1),
        'bboxes_predict': head_bboxes_pred,
        'saved_image_index': save_image_op }

    summary_hook = tf.train.SummarySaverHook(
                        save_secs=FLAGS.save_summary_steps,
                        output_dir=FLAGS.model_dir,
                        summary_op=tf.summary.merge_all())

    if mode == tf.estimator.ModeKeys.EVAL:
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions,
                        evaluation_hooks = [summary_hook],
                        loss=weight_decay_loss, eval_metric_ops=eval_ops)
    else:
        raise ValueError('This script only support predict mode!')
def lighr_head_model_fn ( features , labels , mode , params ) : num_anchors_list = labels [ 'num_anchors_list' ] num_feature_layers = len ( num_anchors_list ) shape = labels [ 'targets' ] [ - 1 ] if mode != tf . estimator . ModeKeys . TRAIN : org_image = labels [ 'targets' ] [ - 2 ] isdifficult = labels [ 'targets' ] [ - 3 ] bbox_img = labels [ 'targets' ] [ - 4 ] gbboxes_raw = labels [ 'targets' ] [ - 5 ] glabels_raw = labels [ 'targets' ] [ - 6 ] glabels = labels [ 'targets' ] [ : num_feature_layers ] [ 0 ] gtargets = labels [ 'targets' ] [ num_feature_layers : 2 * num_feature_layers ] [ 0 ] gscores = labels [ 'targets' ] [ 2 * num_feature_layers : 3 * num_feature_layers ] [ 0 ] #features = tf.ones([4,480,480,3]) * 0.5 with tf . variable_scope ( params [ 'model_scope' ] , default_name = None , values = [ features ] , reuse = tf . AUTO_REUSE ) : rpn_feat_map , backbone_feat = xception_body . XceptionBody ( features , params [ 'num_classes' ] , is_training = ( mode == tf . estimator . ModeKeys . TRAIN ) , data_format = params [ 'data_format' ] ) #rpn_feat_map = tf.Print(rpn_feat_map,[tf.shape(rpn_feat_map), rpn_feat_map,backbone_feat]) rpn_cls_score , rpn_bbox_pred = xception_body . get_rpn ( rpn_feat_map , num_anchors_list [ 0 ] , ( mode == tf . estimator . ModeKeys . TRAIN ) , params [ 'data_format' ] , 'rpn_head' ) large_sep_feature = xception_body . large_sep_kernel ( backbone_feat , 256 , 10 * 7 * 7 , ( mode == tf . estimator . ModeKeys . TRAIN ) , params [ 'data_format' ] , 'large_sep_feature' ) if params [ 'data_format' ] == 'channels_first' : rpn_cls_score = tf . transpose ( rpn_cls_score , [ 0 , 2 , 3 , 1 ] ) rpn_bbox_pred = tf . transpose ( rpn_bbox_pred , [ 0 , 2 , 3 , 1 ] ) rpn_cls_score = tf . reshape ( rpn_cls_score , [ - 1 , 2 ] ) rpn_object_score = tf . nn . softmax ( rpn_cls_score ) [ : , - 1 ] rpn_object_score = tf . reshape ( rpn_object_score , [ 1 , - 1 ] ) rpn_location_pred = tf . reshape ( rpn_bbox_pred , [ 1 , - 1 , 4 ] ) rpn_bboxes_pred = labels [ 'rpn_decode_fn' ] ( rpn_location_pred ) proposals_bboxes = xception_body . get_proposals ( rpn_object_score , rpn_bboxes_pred , None , params [ 'rpn_pre_nms_top_n' ] , params [ 'rpn_post_nms_top_n' ] , params [ 'rpn_nms_thres' ] , params [ 'rpn_min_size' ] , ( mode == tf . estimator . ModeKeys . TRAIN ) , params [ 'data_format' ] ) #proposals_targets = tf.Print(proposals_targets, [proposals_targets], message='proposals_targets0:') cls_score , bboxes_reg = xception_body . get_head ( large_sep_feature , lambda input_ , bboxes_ , grid_width_ , grid_height_ : ps_roi_align ( input_ , bboxes_ , grid_width_ , grid_height_ , pool_method ) , 7 , 7 , None , proposals_bboxes , params [ 'num_classes' ] , ( mode == tf . estimator . ModeKeys . TRAIN ) , False , 0 , params [ 'data_format' ] , 'final_head' ) head_bboxes_pred = labels [ 'head_decode_fn' ] ( proposals_bboxes , bboxes_reg ) head_cls_score = tf . reshape ( cls_score , [ - 1 , params [ 'num_classes' ] ] ) head_cls_score = tf . nn . softmax ( head_cls_score ) head_bboxes_pred = tf . reshape ( head_bboxes_pred , [ - 1 , 4 ] ) shape = tf . squeeze ( shape , axis = 0 ) glabels = tf . squeeze ( glabels , axis = 0 ) gtargets = tf . squeeze ( gtargets , axis = 0 ) gscores = tf . squeeze ( gscores , axis = 0 ) if mode != tf . estimator . ModeKeys . TRAIN : org_image = tf . squeeze ( org_image , axis = 0 ) isdifficult = tf . squeeze ( isdifficult , axis = 0 ) gbboxes_raw = tf . squeeze ( gbboxes_raw , axis = 0 ) glabels_raw = tf . squeeze ( glabels_raw , axis = 0 ) bbox_img = tf . squeeze ( bbox_img , axis = 0 ) eval_ops , save_image_op = bboxes_eval ( org_image , shape , bbox_img , cls_score , head_bboxes_pred , glabels_raw , gbboxes_raw , isdifficult , params [ 'num_classes' ] ) _ = tf . identity ( save_image_op , name = 'save_image_with_bboxes_op' ) with tf . control_dependencies ( [ save_image_op ] ) : weight_decay_loss = params [ 'weight_decay' ] * tf . add_n ( [ tf . nn . l2_loss ( v ) for v in tf . trainable_variables ( ) if 'batch_normalization' not in v . name ] ) predictions = { 'classes' : tf . argmax ( head_cls_score , axis = - 1 ) , 'probabilities' : tf . reduce_max ( head_cls_score , axis = - 1 ) , 'bboxes_predict' : head_bboxes_pred , 'saved_image_index' : save_image_op } summary_hook = tf . train . SummarySaverHook ( save_secs = FLAGS . save_summary_steps , output_dir = FLAGS . model_dir , summary_op = tf . summary . merge_all ( ) ) if mode == tf . estimator . ModeKeys . EVAL : return tf . estimator . EstimatorSpec ( mode = mode , predictions = predictions , evaluation_hooks = [ summary_hook ] , loss = weight_decay_loss , eval_metric_ops = eval_ops ) else : raise ValueError ( 'This script only support predict mode!' )