example-apps/PDD/pcb-defect-detection/libs/networks/build_whole_network2.py

   1 # -*-coding: utf-8 -*-
   2
   3 from __future__ import absolute_import, division, print_function
   4
   5 import os
   6 import tensorflow as tf
   7 import tensorflow.contrib.slim as slim
   8 import numpy as np
   9
  10 from libs.networks import resnet
  11 from libs.networks import mobilenet_v2
  12 from libs.box_utils import encode_and_decode
  13 from libs.box_utils import boxes_utils
  14 from libs.box_utils import anchor_utils
  15 from libs.configs import cfgs
  16 from libs.losses import losses
  17 from libs.box_utils import show_box_in_tensor
  18 from libs.detection_oprations.proposal_opr import postprocess_rpn_proposals
  19 from libs.detection_oprations.anchor_target_layer_without_boxweight import anchor_target_layer
  20 from libs.detection_oprations.proposal_target_layer import proposal_target_layer
  21
  22
  23 class DetectionNetwork(object):
  24
  25     def __init__(self, base_network_name, is_training):
  26
  27         self.base_network_name = base_network_name
  28         self.is_training = is_training
  29         self.num_anchors_per_location = len(cfgs.ANCHOR_SCALES) * len(cfgs.ANCHOR_RATIOS)
  30
  31     def build_base_network(self, input_img_batch):
  32
  33         if self.base_network_name.startswith('resnet_v1'):
  34             return resnet.resnet_base(input_img_batch, scope_name=self.base_network_name, is_training=self.is_training)
  35
  36         elif self.base_network_name.startswith('MobilenetV2'):
  37             return mobilenet_v2.mobilenetv2_base(input_img_batch, is_training=self.is_training)
  38
  39         else:
  40             raise ValueError('Sry, we only support resnet or mobilenet_v2')
  41
  42     def postprocess_fastrcnn(self, rois, bbox_ppred, scores, img_shape):
  43         '''
  44
  45         :param rois:[-1, 4]
  46         :param bbox_ppred: [-1, (cfgs.Class_num+1) * 4]
  47         :param scores: [-1, cfgs.Class_num + 1]
  48         :return:
  49         '''
  50
  51         with tf.name_scope('postprocess_fastrcnn'):
  52             rois = tf.stop_gradient(rois)
  53             scores = tf.stop_gradient(scores)
  54             bbox_ppred = tf.reshape(bbox_ppred, [-1, cfgs.CLASS_NUM + 1, 4])
  55             bbox_ppred = tf.stop_gradient(bbox_ppred)
  56
  57             bbox_pred_list = tf.unstack(bbox_ppred, axis=1)
  58             score_list = tf.unstack(scores, axis=1)
  59
  60             allclasses_boxes = []
  61             allclasses_scores = []
  62             categories = []
  63             for i in range(1, cfgs.CLASS_NUM+1):
  64
  65                 # 1. decode boxes in each class
  66                 tmp_encoded_box = bbox_pred_list[i]
  67                 tmp_score = score_list[i]
  68                 tmp_decoded_boxes = encode_and_decode.decode_boxes(encoded_boxes=tmp_encoded_box,
  69                                                                    reference_boxes=rois,
  70                                                                    scale_factors=cfgs.ROI_SCALE_FACTORS)
  71                 # tmp_decoded_boxes = encode_and_decode.decode_boxes(boxes=rois,
  72                 #                                                    deltas=tmp_encoded_box,
  73                 #                                                    scale_factor=cfgs.ROI_SCALE_FACTORS)
  74
  75                 # 2. clip to img boundaries
  76                 tmp_decoded_boxes = boxes_utils.clip_boxes_to_img_boundaries(decode_boxes=tmp_decoded_boxes,
  77                                                                              img_shape=img_shape)
  78
  79                 # 3. NMS
  80                 keep = tf.image.non_max_suppression(
  81                     boxes=tmp_decoded_boxes,
  82                     scores=tmp_score,
  83                     max_output_size=cfgs.FAST_RCNN_NMS_MAX_BOXES_PER_CLASS,
  84                     iou_threshold=cfgs.FAST_RCNN_NMS_IOU_THRESHOLD)
  85
  86                 perclass_boxes = tf.gather(tmp_decoded_boxes, keep)
  87                 perclass_scores = tf.gather(tmp_score, keep)
  88
  89                 allclasses_boxes.append(perclass_boxes)
  90                 allclasses_scores.append(perclass_scores)
  91                 categories.append(tf.ones_like(perclass_scores) * i)
  92
  93             final_boxes = tf.concat(allclasses_boxes, axis=0)
  94             final_scores = tf.concat(allclasses_scores, axis=0)
  95             final_category = tf.concat(categories, axis=0)
  96
  97             if self.is_training:
  98                 '''
  99                 in training. We should show the detecitons in the tensorboard. So we add this.
 100                 '''
 101                 kept_indices = tf.reshape(tf.where(tf.greater_equal(final_scores, cfgs.SHOW_SCORE_THRSHOLD)), [-1])
 102
 103                 final_boxes = tf.gather(final_boxes, kept_indices)
 104                 final_scores = tf.gather(final_scores, kept_indices)
 105                 final_category = tf.gather(final_category, kept_indices)
 106
 107         return final_boxes, final_scores, final_category
 108
 109     def roi_pooling(self, feature_maps, rois, img_shape, scope):
 110         '''
 111         Here use roi warping as roi_pooling
 112
 113         :param featuremaps_dict: feature map to crop
 114         :param rois: shape is [-1, 4]. [x1, y1, x2, y2]
 115         :return:
 116         '''
 117
 118         with tf.variable_scope('ROI_Warping_'+scope):
 119             img_h, img_w = tf.cast(img_shape[1], tf.float32), tf.cast(img_shape[2], tf.float32)
 120             N = tf.shape(rois)[0]
 121             x1, y1, x2, y2 = tf.unstack(rois, axis=1)
 122
 123             normalized_x1 = x1 / img_w
 124             normalized_x2 = x2 / img_w
 125             normalized_y1 = y1 / img_h
 126             normalized_y2 = y2 / img_h
 127
 128             normalized_rois = tf.transpose(
 129                 tf.stack([normalized_y1, normalized_x1, normalized_y2, normalized_x2]), name='get_normalized_rois')
 130
 131             normalized_rois = tf.stop_gradient(normalized_rois)
 132
 133             cropped_roi_features = tf.image.crop_and_resize(feature_maps, normalized_rois,
 134                                                             box_ind=tf.zeros(shape=[N],
 135                                                                              dtype=tf.int32),
 136                                                             crop_size=[cfgs.ROI_SIZE, cfgs.ROI_SIZE],
 137                                                             name='CROP_AND_RESIZE'
 138                                                             )
 139             roi_features = slim.max_pool2d(cropped_roi_features,
 140                                            [cfgs.ROI_POOL_KERNEL_SIZE, cfgs.ROI_POOL_KERNEL_SIZE],
 141                                            stride=cfgs.ROI_POOL_KERNEL_SIZE)
 142
 143         return roi_features
 144
 145     def build_fastrcnn(self, P_list, rois_list, img_shape):
 146
 147         with tf.variable_scope('Fast-RCNN'):
 148             # 5. ROI Pooling
 149             with tf.variable_scope('rois_pooling'):
 150                 pooled_features_list = []
 151                 for level_name, p, rois in zip(cfgs.LEVLES, P_list, rois_list):  # exclude P6_rois
 152                     # p = tf.Print(p, [tf.shape(p)], summarize=10, message=level_name+'SHPAE***')
 153                     pooled_features = self.roi_pooling(feature_maps=p, rois=rois, img_shape=img_shape,
 154                                                        scope=level_name)
 155                     pooled_features_list.append(pooled_features)
 156
 157                 pooled_features = tf.concat(pooled_features_list, axis=0) # [minibatch_size, H, W, C]
 158
 159             # 6. inferecne rois in Fast-RCNN to obtain fc_flatten features
 160             if self.base_network_name.startswith('resnet'):
 161                 fc_flatten = resnet.restnet_head(inputs=pooled_features,
 162                                                  is_training=self.is_training,
 163                                                  scope_name=self.base_network_name)
 164             elif self.base_network_name.startswith('Mobile'):
 165                 fc_flatten = mobilenet_v2.mobilenetv2_head(inputs=pooled_features,
 166                                                            is_training=self.is_training)
 167             else:
 168                 raise NotImplementedError('only support resnet and mobilenet')
 169
 170             # 7. cls and reg in Fast-RCNN
 171             with slim.arg_scope([slim.fully_connected], weights_regularizer=slim.l2_regularizer(cfgs.WEIGHT_DECAY)):
 172
 173                 cls_score = slim.fully_connected(fc_flatten,
 174                                                  num_outputs=cfgs.CLASS_NUM+1,
 175                                                  weights_initializer=cfgs.INITIALIZER,
 176                                                  activation_fn=None, trainable=self.is_training,
 177                                                  scope='cls_fc')
 178
 179                 bbox_pred = slim.fully_connected(fc_flatten,
 180                                                  num_outputs=(cfgs.CLASS_NUM+1)*4,
 181                                                  weights_initializer=cfgs.BBOX_INITIALIZER,
 182                                                  activation_fn=None, trainable=self.is_training,
 183                                                  scope='reg_fc')
 184                 # for convient. It also produce (cls_num +1) bboxes
 185
 186                 cls_score = tf.reshape(cls_score, [-1, cfgs.CLASS_NUM+1])
 187                 bbox_pred = tf.reshape(bbox_pred, [-1, 4*(cfgs.CLASS_NUM+1)])
 188
 189         return bbox_pred, cls_score
 190
 191     def assign_levels(self, all_rois, labels=None, bbox_targets=None):
 192         '''
 193
 194         :param all_rois:
 195         :param labels:
 196         :param bbox_targets:
 197         :return:
 198         '''
 199         with tf.name_scope('assign_levels'):
 200             # all_rois = tf.Print(all_rois, [tf.shape(all_rois)], summarize=10, message='ALL_ROIS_SHAPE*****')
 201             xmin, ymin, xmax, ymax = tf.unstack(all_rois, axis=1)
 202
 203             h = tf.maximum(0., ymax - ymin)
 204             w = tf.maximum(0., xmax - xmin)
 205
 206             levels = tf.floor(4. + tf.log(tf.sqrt(w * h + 1e-8) / 224.0) / tf.log(2.))  # 4 + log_2(***)
 207             # use floor instead of round
 208
 209             min_level = int(cfgs.LEVLES[0][-1])
 210             max_level = min(5, int(cfgs.LEVLES[-1][-1]))
 211             levels = tf.maximum(levels, tf.ones_like(levels) * min_level)  # level minimum is 2
 212             levels = tf.minimum(levels, tf.ones_like(levels) * max_level)  # level maximum is 5
 213
 214             levels = tf.stop_gradient(tf.reshape(levels, [-1]))
 215
 216             def get_rois(levels, level_i, rois, labels, bbox_targets):
 217
 218                 level_i_indices = tf.reshape(tf.where(tf.equal(levels, level_i)), [-1])
 219                 tf.summary.scalar('LEVEL/LEVEL_%d_rois_NUM'%level_i, tf.shape(level_i_indices)[0])
 220                 level_i_rois = tf.gather(rois, level_i_indices)
 221
 222                 # if self.is_training:
 223                 #     level_i_rois = tf.stop_gradient(tf.concat([level_i_rois, [[0, 0, 0., 0.]]], axis=0))
 224                 #     # to avoid the num of level i rois is 0.0, which will broken the BP in tf
 225                 #
 226                 #     level_i_labels = tf.gather(labels, level_i_indices)
 227                 #     level_i_labels = tf.stop_gradient(tf.concat([level_i_labels, [0]], axis=0))
 228                 #
 229                 #     level_i_targets = tf.gather(bbox_targets, level_i_indices)
 230                 #     level_i_targets = tf.stop_gradient(tf.concat([level_i_targets,
 231                 #                                                   tf.zeros(shape=(1, 4*(cfgs.CLASS_NUM+1)), dtype=tf.float32)],
 232                 #                                                  axis=0))
 233                 #
 234                 #     return level_i_rois, level_i_labels, level_i_targets
 235                 if self.is_training:
 236                     level_i_labels = tf.gather(labels, level_i_indices)
 237                     level_i_targets = tf.gather(bbox_targets, level_i_indices)
 238                     return level_i_rois, level_i_labels, level_i_targets
 239                 else:
 240                     return level_i_rois, None, None
 241
 242             rois_list = []
 243             labels_list = []
 244             targets_list = []
 245             for i in range(min_level, max_level+1):
 246                 P_i_rois, P_i_labels, P_i_targets = get_rois(levels, level_i=i, rois=all_rois,
 247                                                              labels=labels,
 248                                                              bbox_targets=bbox_targets)
 249                 rois_list.append(P_i_rois)
 250                 labels_list.append(P_i_labels)
 251                 targets_list.append(P_i_targets)
 252
 253             if self.is_training:
 254                 all_labels = tf.concat(labels_list, axis=0)
 255                 all_targets = tf.concat(targets_list, axis=0)
 256
 257                 return rois_list, \
 258                        tf.stop_gradient(all_labels), \
 259                        tf.stop_gradient(all_targets)
 260             else:
 261                 return rois_list  # [P2_rois, P3_rois, P4_rois, P5_rois] Note: P6 do not assign rois
 262     # ---------------------------------------------------------------------------------------
 263
 264     def add_anchor_img_smry(self, img, anchors, labels):
 265
 266         positive_anchor_indices = tf.reshape(tf.where(tf.greater_equal(labels, 1)), [-1])
 267         negative_anchor_indices = tf.reshape(tf.where(tf.equal(labels, 0)), [-1])
 268
 269         positive_anchor = tf.gather(anchors, positive_anchor_indices)
 270         negative_anchor = tf.gather(anchors, negative_anchor_indices)
 271
 272         pos_in_img = show_box_in_tensor.only_draw_boxes(img_batch=img,
 273                                                         boxes=positive_anchor)
 274         neg_in_img = show_box_in_tensor.only_draw_boxes(img_batch=img,
 275                                                         boxes=negative_anchor)
 276
 277         tf.summary.image('positive_anchor', pos_in_img)
 278         tf.summary.image('negative_anchors', neg_in_img)
 279
 280     def add_roi_batch_img_smry(self, img, rois, labels):
 281         positive_roi_indices = tf.reshape(tf.where(tf.greater_equal(labels, 1)), [-1])
 282
 283         negative_roi_indices = tf.reshape(tf.where(tf.equal(labels, 0)), [-1])
 284
 285         pos_roi = tf.gather(rois, positive_roi_indices)
 286         neg_roi = tf.gather(rois, negative_roi_indices)
 287
 288
 289         pos_in_img = show_box_in_tensor.only_draw_boxes(img_batch=img,
 290                                                                boxes=pos_roi)
 291         neg_in_img = show_box_in_tensor.only_draw_boxes(img_batch=img,
 292                                                                boxes=neg_roi)
 293         tf.summary.image('pos_rois', pos_in_img)
 294         tf.summary.image('neg_rois', neg_in_img)
 295
 296     def build_loss(self, rpn_box_pred, rpn_bbox_targets, rpn_cls_score, rpn_labels,
 297                    bbox_pred, bbox_targets, cls_score, labels):
 298         '''
 299
 300         :param rpn_box_pred: [-1, 4]
 301         :param rpn_bbox_targets: [-1, 4]
 302         :param rpn_cls_score: [-1]
 303         :param rpn_labels: [-1]
 304         :param bbox_pred: [-1, 4*(cls_num+1)]
 305         :param bbox_targets: [-1, 4*(cls_num+1)]
 306         :param cls_score: [-1, cls_num+1]
 307         :param labels: [-1]
 308         :return:
 309         '''
 310         with tf.variable_scope('build_loss') as sc:
 311             with tf.variable_scope('rpn_loss'):
 312
 313                 rpn_bbox_loss = losses.smooth_l1_loss_rpn(bbox_pred=rpn_box_pred,
 314                                                           bbox_targets=rpn_bbox_targets,
 315                                                           label=rpn_labels,
 316                                                           sigma=cfgs.RPN_SIGMA)
 317                 # rpn_cls_loss:
 318                 # rpn_cls_score = tf.reshape(rpn_cls_score, [-1, 2])
 319                 # rpn_labels = tf.reshape(rpn_labels, [-1])
 320                 # ensure rpn_labels shape is [-1]
 321                 rpn_select = tf.reshape(tf.where(tf.not_equal(rpn_labels, -1)), [-1])
 322                 rpn_cls_score = tf.reshape(tf.gather(rpn_cls_score, rpn_select), [-1, 2])
 323                 rpn_labels = tf.reshape(tf.gather(rpn_labels, rpn_select), [-1])
 324                 rpn_cls_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(logits=rpn_cls_score,
 325                                                                                              labels=rpn_labels))
 326
 327                 rpn_cls_loss = rpn_cls_loss * cfgs.RPN_CLASSIFICATION_LOSS_WEIGHT
 328                 rpn_bbox_loss = rpn_bbox_loss * cfgs.RPN_LOCATION_LOSS_WEIGHT
 329
 330             with tf.variable_scope('FastRCNN_loss'):
 331                 if not cfgs.FAST_RCNN_MINIBATCH_SIZE == -1:
 332                     bbox_loss = losses.smooth_l1_loss_rcnn(bbox_pred=bbox_pred,
 333                                                            bbox_targets=bbox_targets,
 334                                                            label=labels,
 335                                                            num_classes=cfgs.CLASS_NUM + 1,
 336                                                            sigma=cfgs.FASTRCNN_SIGMA)
 337
 338                     # cls_score = tf.reshape(cls_score, [-1, cfgs.CLASS_NUM + 1])
 339                     # labels = tf.reshape(labels, [-1])
 340                     cls_loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
 341                         logits=cls_score,
 342                         labels=labels))  # beacause already sample before
 343                 else:
 344                     '''
 345                     applying OHEM here
 346                     '''
 347                     print(20 * "@@")
 348                     print("@@" + 10 * " " + "TRAIN WITH OHEM ...")
 349                     print(20 * "@@")
 350                     cls_loss = bbox_loss = losses.sum_ohem_loss(
 351                         cls_score=cls_score,
 352                         label=labels,
 353                         bbox_targets=bbox_targets,
 354                         nr_ohem_sampling=128,
 355                         nr_classes=cfgs.CLASS_NUM + 1)
 356                 cls_loss = cls_loss * cfgs.FAST_RCNN_CLASSIFICATION_LOSS_WEIGHT
 357                 bbox_loss = bbox_loss * cfgs.FAST_RCNN_LOCATION_LOSS_WEIGHT
 358             loss_dict = {
 359                 'rpn_cls_loss': rpn_cls_loss,
 360                 'rpn_loc_loss': rpn_bbox_loss,
 361                 'fastrcnn_cls_loss': cls_loss,
 362                 'fastrcnn_loc_loss': bbox_loss
 363             }
 364         return loss_dict
 365
 366
 367     def build_whole_detection_network(self, input_img_batch, gtboxes_batch):
 368
 369         if self.is_training:
 370             # ensure shape is [M, 5]
 371             gtboxes_batch = tf.reshape(gtboxes_batch, [-1, 5])
 372             gtboxes_batch = tf.cast(gtboxes_batch, tf.float32)
 373
 374         img_shape = tf.shape(input_img_batch)
 375
 376         # 1. build base network
 377         P_dict = self.build_base_network(input_img_batch)  # [P2, P3, P4, P5, P6]
 378
 379         # 2. build fpn  by build rpn for each level
 380         with tf.variable_scope("build_FPN", regularizer=slim.l2_regularizer(cfgs.WEIGHT_DECAY)):
 381             fpn_box_delta = {}
 382             fpn_cls_score = {}
 383             fpn_cls_prob = {}
 384             for key in cfgs.LEVLES:
 385                 if cfgs.SHARE_HEADS:
 386                     reuse_flag = None if key == cfgs.LEVLES[0] else True
 387                     scope_list = ['fpn_conv/3x3', 'fpn_cls_score', 'fpn_bbox_pred']
 388                 else:
 389                     reuse_flag = None
 390                     scope_list= ['fpn_conv/3x3_%s' % key, 'fpn_cls_score_%s' % key, 'fpn_bbox_pred_%s' % key]
 391                 rpn_conv3x3 = slim.conv2d(
 392                     P_dict[key], 512, [3, 3],
 393                     trainable=self.is_training, weights_initializer=cfgs.INITIALIZER, padding="SAME",
 394                     activation_fn=tf.nn.relu,
 395                     scope=scope_list[0],
 396                     reuse=reuse_flag)
 397                 rpn_cls_score = slim.conv2d(rpn_conv3x3, self.num_anchors_per_location*2, [1, 1], stride=1,
 398                                             trainable=self.is_training, weights_initializer=cfgs.INITIALIZER,
 399                                             activation_fn=None, padding="VALID",
 400                                             scope=scope_list[1],
 401                                             reuse=reuse_flag)
 402                 rpn_box_delta = slim.conv2d(rpn_conv3x3, self.num_anchors_per_location*4, [1, 1], stride=1,
 403                                             trainable=self.is_training, weights_initializer=cfgs.BBOX_INITIALIZER,
 404                                             activation_fn=None, padding="VALID",
 405                                             scope=scope_list[2],
 406                                             reuse=reuse_flag)
 407                 fpn_box_delta[key] = tf.reshape(rpn_box_delta, [-1, 4])
 408                 fpn_cls_score[key] = tf.reshape(rpn_cls_score, [-1, 2])
 409                 fpn_cls_prob[key] = slim.softmax(fpn_cls_score[key])
 410
 411         # 3. generate anchors for fpn. (by generate for each level)
 412         anchors_dict = {}
 413         anchor_list = []
 414         with tf.name_scope("generate_FPN_anchors"):
 415             for key in cfgs.LEVLES:
 416                 p_h, p_w = tf.to_float(tf.shape(P_dict[key])[1]), tf.to_float(tf.shape(P_dict[key])[2])
 417                 id_ = int(key[-1]) - int(cfgs.LEVLES[0][-1]) # such as : 2-2, 3-3
 418                 tmp_anchors = anchor_utils.make_anchors(base_anchor_size=cfgs.BASE_ANCHOR_SIZE_LIST[id_],
 419                                                         anchor_scales=cfgs.ANCHOR_SCALES,
 420                                                         anchor_ratios=cfgs.ANCHOR_RATIOS,
 421                                                         stride=cfgs.ANCHOR_STRIDE_LIST[id_],
 422                                                         featuremap_height=p_h,
 423                                                         featuremap_width=p_w,
 424                                                         name='%s_make_anchors' % key)
 425                 anchors_dict[key] = tmp_anchors
 426                 anchor_list.append(tmp_anchors)
 427         all_anchors = tf.concat(anchor_list, axis=0)
 428
 429         # 4. postprocess fpn proposals. such as: decode, clip, NMS
 430         #    Need to Note: Here we NMS for each level instead of NMS for all anchors.
 431         rois_list = []
 432         rois_scores_list = []
 433         with tf.name_scope("postproces_fpn"):
 434             for key in cfgs.LEVLES:
 435                 tmp_rois, tmp_roi_scores = postprocess_rpn_proposals(rpn_bbox_pred=fpn_box_delta[key],
 436                                                                      rpn_cls_prob=fpn_cls_prob[key],
 437                                                                      img_shape=img_shape,
 438                                                                      anchors=anchors_dict[key],
 439                                                                      is_training=self.is_training)
 440                 rois_list.append(tmp_rois)
 441                 rois_scores_list.append(tmp_roi_scores)
 442             allrois = tf.concat(rois_list, axis=0)
 443             allrois_scores = tf.concat(rois_scores_list, axis=0)
 444             fpn_topk = cfgs.FPN_TOP_K_PER_LEVEL_TRAIN if self.is_training else cfgs.FPN_TOP_K_PER_LEVEL_TEST
 445             topk = tf.minimum(fpn_topk, tf.shape(allrois)[0])
 446
 447             rois_scores, topk_indices = tf.nn.top_k(allrois_scores, k=topk)
 448
 449             rois = tf.stop_gradient(tf.gather(allrois, topk_indices))
 450             rois_scores = tf.stop_gradient(rois_scores)
 451
 452             # +++++++++++++++++++++++++++++++++++++add img smry+++++++++++++++++++++++++++++++++++++++++++++++++++++++
 453             if self.is_training:
 454                 score_gre_05 = tf.reshape(tf.where(tf.greater_equal(rois_scores, 0.5)), [-1])
 455                 score_gre_05_rois = tf.gather(rois, score_gre_05)
 456                 score_gre_05_score = tf.gather(rois_scores, score_gre_05)
 457                 score_gre_05_in_img = show_box_in_tensor.draw_boxes_with_scores(img_batch=input_img_batch,
 458                                                                                 boxes=score_gre_05_rois,
 459                                                                                 scores=score_gre_05_score)
 460                 tf.summary.image('score_greater_05_rois', score_gre_05_in_img)
 461             # ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 462
 463         # sample for fpn. We should concat all the anchors
 464         if self.is_training:
 465             with tf.variable_scope('sample_anchors_minibatch'):
 466                 fpn_labels, fpn_bbox_targets = \
 467                     tf.py_func(
 468                         anchor_target_layer,
 469                         [gtboxes_batch, img_shape, all_anchors],
 470                         [tf.float32, tf.float32])
 471                 fpn_bbox_targets = tf.reshape(fpn_bbox_targets, [-1, 4])
 472                 fpn_labels = tf.to_int32(fpn_labels, name="to_int32")
 473                 fpn_labels = tf.reshape(fpn_labels, [-1])
 474                 self.add_anchor_img_smry(input_img_batch, all_anchors, fpn_labels)
 475
 476             with tf.control_dependencies([fpn_labels]):
 477                 with tf.variable_scope('sample_RCNN_minibatch'):
 478                     rois, labels, bbox_targets = \
 479                     tf.py_func(proposal_target_layer,
 480                                [rois, gtboxes_batch],
 481                                [tf.float32, tf.float32, tf.float32])
 482                     rois = tf.reshape(rois, [-1, 4])
 483                     labels = tf.to_int32(labels)
 484                     labels = tf.reshape(labels, [-1])
 485                     bbox_targets = tf.reshape(bbox_targets, [-1, 4*(cfgs.CLASS_NUM+1)])
 486                     self.add_roi_batch_img_smry(input_img_batch, rois, labels)
 487         if self.is_training:
 488             rois_list, labels, bbox_targets = self.assign_levels(all_rois=rois,
 489                                                                  labels=labels,
 490                                                                  bbox_targets=bbox_targets)
 491         else:
 492             rois_list = self.assign_levels(all_rois=rois)  # rois_list: [P2_rois, P3_rois, P4_rois, P5_rois]
 493
 494         # -------------------------------------------------------------------------------------------------------------#
 495         #                                            Fast-RCNN                                                         #
 496         # -------------------------------------------------------------------------------------------------------------#
 497
 498         # 5. build Fast-RCNN
 499         # rois = tf.Print(rois, [tf.shape(rois)], 'rois shape', summarize=10)
 500         bbox_pred, cls_score = self.build_fastrcnn(P_list=[P_dict[key] for key in cfgs.LEVLES],
 501                                                    rois_list=rois_list,
 502                                                    img_shape=img_shape)
 503         # bbox_pred shape: [-1, 4*(cls_num+1)].
 504         # cls_score shape： [-1, cls_num+1]
 505
 506         cls_prob = slim.softmax(cls_score, 'cls_prob')
 507
 508
 509         # ----------------------------------------------add smry-------------------------------------------------------
 510         if self.is_training:
 511             cls_category = tf.argmax(cls_prob, axis=1)
 512             fast_acc = tf.reduce_mean(tf.to_float(tf.equal(cls_category, tf.to_int64(labels))))
 513             tf.summary.scalar('ACC/fast_acc', fast_acc)
 514
 515         rois = tf.concat(rois_list, axis=0, name='concat_rois')
 516         #  6. postprocess_fastrcnn
 517         if not self.is_training:
 518             return self.postprocess_fastrcnn(rois=rois, bbox_ppred=bbox_pred, scores=cls_prob, img_shape=img_shape)
 519         else:
 520             '''
 521             when trian. We need build Loss
 522             '''
 523             loss_dict = self.build_loss(rpn_box_pred=tf.concat([fpn_box_delta[key] for key in cfgs.LEVLES], axis=0),
 524                                         rpn_bbox_targets=fpn_bbox_targets,
 525                                         rpn_cls_score=tf.concat([fpn_cls_score[key] for key in cfgs.LEVLES], axis=0),
 526                                         rpn_labels=fpn_labels,
 527                                         bbox_pred=bbox_pred,
 528                                         bbox_targets=bbox_targets,
 529                                         cls_score=cls_score,
 530                                         labels=labels)
 531
 532             final_bbox, final_scores, final_category = self.postprocess_fastrcnn(rois=rois,
 533                                                                                  bbox_ppred=bbox_pred,
 534                                                                                  scores=cls_prob,
 535                                                                                  img_shape=img_shape)
 536             return final_bbox, final_scores, final_category, loss_dict
 537
 538     def get_restorer(self):
 539         checkpoint_path = tf.train.latest_checkpoint(os.path.join(cfgs.TRAINED_CKPT, cfgs.VERSION))
 540
 541         if checkpoint_path != None:
 542             restorer = tf.train.Saver()
 543             print("model restore from :", checkpoint_path)
 544         else:
 545             checkpoint_path = cfgs.PRETRAINED_CKPT
 546             print("model restore from pretrained mode, path is :", checkpoint_path)
 547
 548             model_variables = slim.get_model_variables()
 549             # for var in model_variables:
 550             #     print(var.name)
 551             # print(20*"__++__++__")
 552
 553             def name_in_ckpt_rpn(var):
 554                 return var.op.name
 555
 556             def name_in_ckpt_fastrcnn_head(var):
 557                 '''
 558                 Fast-RCNN/resnet_v1_50/block4 -->resnet_v1_50/block4
 559                 Fast-RCNN/MobilenetV2/** -- > MobilenetV2 **
 560                 :param var:
 561                 :return:
 562                 '''
 563                 return '/'.join(var.op.name.split('/')[1:])
 564             nameInCkpt_Var_dict = {}
 565             for var in model_variables:
 566                 if var.name.startswith(self.base_network_name):
 567                     var_name_in_ckpt = name_in_ckpt_rpn(var)
 568                     nameInCkpt_Var_dict[var_name_in_ckpt] = var
 569             restore_variables = nameInCkpt_Var_dict
 570             for key, item in restore_variables.items():
 571                 print("var_in_graph: ", item.name)
 572                 print("var_in_ckpt: ", key)
 573                 print(20*"___")
 574             restorer = tf.train.Saver(restore_variables)
 575             print(20 * "****")
 576             print("restore from pretrained_weighs in IMAGE_NET")
 577         return restorer, checkpoint_path
 578
 579     def get_gradients(self, optimizer, loss):
 580         '''
 581
 582         :param optimizer:
 583         :param loss:
 584         :return:
 585
 586         return vars and grads that not be fixed
 587         '''
 588
 589         # if cfgs.FIXED_BLOCKS > 0:
 590         #     trainable_vars = tf.trainable_variables()
 591         #     # trained_vars = slim.get_trainable_variables()
 592         #     start_names = [cfgs.NET_NAME + '/block%d'%i for i in range(1, cfgs.FIXED_BLOCKS+1)] + \
 593         #                   [cfgs.NET_NAME + '/conv1']
 594         #     start_names = tuple(start_names)
 595         #     trained_var_list = []
 596         #     for var in trainable_vars:
 597         #         if not var.name.startswith(start_names):
 598         #             trained_var_list.append(var)
 599         #     # slim.learning.train()
 600         #     grads = optimizer.compute_gradients(loss, var_list=trained_var_list)
 601         #     return grads
 602         # else:
 603         #     return optimizer.compute_gradients(loss)
 604         return optimizer.compute_gradients(loss)
 605
 606     def enlarge_gradients_for_bias(self, gradients):
 607
 608         final_gradients = []
 609         with tf.variable_scope("Gradient_Mult") as scope:
 610             for grad, var in gradients:
 611                 scale = 1.0
 612                 if cfgs.MUTILPY_BIAS_GRADIENT and './biases' in var.name:
 613                     scale = scale * cfgs.MUTILPY_BIAS_GRADIENT
 614                 if not np.allclose(scale, 1.0):
 615                     grad = tf.multiply(grad, scale)
 616                 final_gradients.append((grad, var))
 617         return final_gradients
 618
 619
 620
 621
 622
 623
 624
 625
 626
 627
 628
 629
 630
 631
 632
 633
 634
 635
 636
 637