model_builder.py 23.1 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""A function to build a DetectionModel from configuration."""
17

18
import functools
19

20
21
22
23
24
25
26
27
28
from object_detection.builders import anchor_generator_builder
from object_detection.builders import box_coder_builder
from object_detection.builders import box_predictor_builder
from object_detection.builders import hyperparams_builder
from object_detection.builders import image_resizer_builder
from object_detection.builders import losses_builder
from object_detection.builders import matcher_builder
from object_detection.builders import post_processing_builder
from object_detection.builders import region_similarity_calculator_builder as sim_calc
29
from object_detection.core import balanced_positive_negative_sampler as sampler
30
from object_detection.core import post_processing
31
from object_detection.core import target_assigner
32
33
34
35
from object_detection.meta_architectures import faster_rcnn_meta_arch
from object_detection.meta_architectures import rfcn_meta_arch
from object_detection.meta_architectures import ssd_meta_arch
from object_detection.models import faster_rcnn_inception_resnet_v2_feature_extractor as frcnn_inc_res
36
37
from object_detection.models import faster_rcnn_inception_v2_feature_extractor as frcnn_inc_v2
from object_detection.models import faster_rcnn_nas_feature_extractor as frcnn_nas
38
from object_detection.models import faster_rcnn_pnas_feature_extractor as frcnn_pnas
39
from object_detection.models import faster_rcnn_resnet_v1_feature_extractor as frcnn_resnet_v1
40
from object_detection.models import ssd_resnet_v1_fpn_feature_extractor as ssd_resnet_v1_fpn
41
from object_detection.models import ssd_resnet_v1_ppn_feature_extractor as ssd_resnet_v1_ppn
42
from object_detection.models.embedded_ssd_mobilenet_v1_feature_extractor import EmbeddedSSDMobileNetV1FeatureExtractor
43
from object_detection.models.ssd_inception_v2_feature_extractor import SSDInceptionV2FeatureExtractor
44
from object_detection.models.ssd_inception_v3_feature_extractor import SSDInceptionV3FeatureExtractor
45
from object_detection.models.ssd_mobilenet_v1_feature_extractor import SSDMobileNetV1FeatureExtractor
46
47
from object_detection.models.ssd_mobilenet_v1_fpn_feature_extractor import SSDMobileNetV1FpnFeatureExtractor
from object_detection.models.ssd_mobilenet_v1_ppn_feature_extractor import SSDMobileNetV1PpnFeatureExtractor
48
from object_detection.models.ssd_mobilenet_v2_feature_extractor import SSDMobileNetV2FeatureExtractor
49
from object_detection.models.ssd_mobilenet_v2_fpn_feature_extractor import SSDMobileNetV2FpnFeatureExtractor
50
from object_detection.predictors import rfcn_box_predictor
51
from object_detection.protos import model_pb2
52
from object_detection.utils import ops
53
54
55
56
57
# BEGIN GOOGLE-INTERNAL
# TODO(lzc): move ssd_mask_meta_arch to third party when it has decent
# performance relative to a comparable Mask R-CNN model (b/112561592).
from google3.image.understanding.object_detection.meta_architectures import ssd_mask_meta_arch
# END GOOGLE-INTERNAL
58
59
60
61

# A map of names to SSD feature extractors.
SSD_FEATURE_EXTRACTOR_CLASS_MAP = {
    'ssd_inception_v2': SSDInceptionV2FeatureExtractor,
62
    'ssd_inception_v3': SSDInceptionV3FeatureExtractor,
63
    'ssd_mobilenet_v1': SSDMobileNetV1FeatureExtractor,
64
65
    'ssd_mobilenet_v1_fpn': SSDMobileNetV1FpnFeatureExtractor,
    'ssd_mobilenet_v1_ppn': SSDMobileNetV1PpnFeatureExtractor,
66
    'ssd_mobilenet_v2': SSDMobileNetV2FeatureExtractor,
67
    'ssd_mobilenet_v2_fpn': SSDMobileNetV2FpnFeatureExtractor,
68
69
70
    'ssd_resnet50_v1_fpn': ssd_resnet_v1_fpn.SSDResnet50V1FpnFeatureExtractor,
    'ssd_resnet101_v1_fpn': ssd_resnet_v1_fpn.SSDResnet101V1FpnFeatureExtractor,
    'ssd_resnet152_v1_fpn': ssd_resnet_v1_fpn.SSDResnet152V1FpnFeatureExtractor,
71
72
73
74
75
    'ssd_resnet50_v1_ppn': ssd_resnet_v1_ppn.SSDResnet50V1PpnFeatureExtractor,
    'ssd_resnet101_v1_ppn':
        ssd_resnet_v1_ppn.SSDResnet101V1PpnFeatureExtractor,
    'ssd_resnet152_v1_ppn':
        ssd_resnet_v1_ppn.SSDResnet152V1PpnFeatureExtractor,
76
    'embedded_ssd_mobilenet_v1': EmbeddedSSDMobileNetV1FeatureExtractor,
77
78
79
80
}

# A map of names to Faster R-CNN feature extractors.
FASTER_RCNN_FEATURE_EXTRACTOR_CLASS_MAP = {
Vivek Rathod's avatar
Vivek Rathod committed
81
82
    'faster_rcnn_nas':
    frcnn_nas.FasterRCNNNASFeatureExtractor,
83
84
    'faster_rcnn_pnas':
    frcnn_pnas.FasterRCNNPNASFeatureExtractor,
85
86
87
88
    'faster_rcnn_inception_resnet_v2':
    frcnn_inc_res.FasterRCNNInceptionResnetV2FeatureExtractor,
    'faster_rcnn_inception_v2':
    frcnn_inc_v2.FasterRCNNInceptionV2FeatureExtractor,
89
90
91
92
93
94
95
96
97
    'faster_rcnn_resnet50':
    frcnn_resnet_v1.FasterRCNNResnet50FeatureExtractor,
    'faster_rcnn_resnet101':
    frcnn_resnet_v1.FasterRCNNResnet101FeatureExtractor,
    'faster_rcnn_resnet152':
    frcnn_resnet_v1.FasterRCNNResnet152FeatureExtractor,
}


98
99
def build(model_config, is_training, add_summaries=True,
          add_background_class=True):
100
101
102
103
104
105
  """Builds a DetectionModel based on the model config.

  Args:
    model_config: A model.proto object containing the config for the desired
      DetectionModel.
    is_training: True if this model is being built for training purposes.
106
    add_summaries: Whether to add tensorflow summaries in the model graph.
107
108
109
110
    add_background_class: Whether to add an implicit background class to one-hot
      encodings of groundtruth labels. Set to false if using groundtruth labels
      with an explicit background class or using multiclass scores instead of
      truth in the case of distillation. Ignored in the case of faster_rcnn.
111
112
113
114
115
116
117
118
119
120
  Returns:
    DetectionModel based on the config.

  Raises:
    ValueError: On invalid meta architecture or model.
  """
  if not isinstance(model_config, model_pb2.DetectionModel):
    raise ValueError('model_config not of type model_pb2.DetectionModel.')
  meta_architecture = model_config.WhichOneof('model')
  if meta_architecture == 'ssd':
121
122
    return _build_ssd_model(model_config.ssd, is_training, add_summaries,
                            add_background_class)
123
  if meta_architecture == 'faster_rcnn':
124
125
    return _build_faster_rcnn_model(model_config.faster_rcnn, is_training,
                                    add_summaries)
126
127
128
129
  raise ValueError('Unknown meta architecture: {}'.format(meta_architecture))


def _build_ssd_feature_extractor(feature_extractor_config, is_training,
130
                                 reuse_weights=None):
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
  """Builds a ssd_meta_arch.SSDFeatureExtractor based on config.

  Args:
    feature_extractor_config: A SSDFeatureExtractor proto config from ssd.proto.
    is_training: True if this feature extractor is being built for training.
    reuse_weights: if the feature extractor should reuse weights.

  Returns:
    ssd_meta_arch.SSDFeatureExtractor based on config.

  Raises:
    ValueError: On invalid feature extractor type.
  """
  feature_type = feature_extractor_config.type
  depth_multiplier = feature_extractor_config.depth_multiplier
  min_depth = feature_extractor_config.min_depth
147
  pad_to_multiple = feature_extractor_config.pad_to_multiple
148
  use_explicit_padding = feature_extractor_config.use_explicit_padding
149
  use_depthwise = feature_extractor_config.use_depthwise
150
151
  conv_hyperparams = hyperparams_builder.build(
      feature_extractor_config.conv_hyperparams, is_training)
152
153
  override_base_feature_extractor_hyperparams = (
      feature_extractor_config.override_base_feature_extractor_hyperparams)
154
155
156
157
158

  if feature_type not in SSD_FEATURE_EXTRACTOR_CLASS_MAP:
    raise ValueError('Unknown ssd feature_extractor: {}'.format(feature_type))

  feature_extractor_class = SSD_FEATURE_EXTRACTOR_CLASS_MAP[feature_type]
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
  kwargs = {
      'is_training':
          is_training,
      'depth_multiplier':
          depth_multiplier,
      'min_depth':
          min_depth,
      'pad_to_multiple':
          pad_to_multiple,
      'conv_hyperparams_fn':
          conv_hyperparams,
      'reuse_weights':
          reuse_weights,
      'use_explicit_padding':
          use_explicit_padding,
      'use_depthwise':
          use_depthwise,
      'override_base_feature_extractor_hyperparams':
          override_base_feature_extractor_hyperparams
  }

  if feature_extractor_config.HasField('fpn'):
    kwargs.update({
182
183
184
185
186
187
        'fpn_min_level':
            feature_extractor_config.fpn.min_level,
        'fpn_max_level':
            feature_extractor_config.fpn.max_level,
        'additional_layer_depth':
            feature_extractor_config.fpn.additional_layer_depth,
188
189
190
    })

  return feature_extractor_class(**kwargs)
191
192


193
194
def _build_ssd_model(ssd_config, is_training, add_summaries,
                     add_background_class=True):
195
196
197
198
199
200
  """Builds an SSD detection model based on the model config.

  Args:
    ssd_config: A ssd.proto object containing the config for the desired
      SSDMetaArch.
    is_training: True if this model is being built for training purposes.
201
    add_summaries: Whether to add tf summaries in the model.
202
203
204
205
    add_background_class: Whether to add an implicit background class to one-hot
      encodings of groundtruth labels. Set to false if using groundtruth labels
      with an explicit background class or using multiclass scores instead of
      truth in the case of distillation.
206
207
  Returns:
    SSDMetaArch based on the config.
208

209
210
211
212
213
214
215
  Raises:
    ValueError: If ssd_config.type is not recognized (i.e. not registered in
      model_class_map).
  """
  num_classes = ssd_config.num_classes

  # Feature extractor
216
  feature_extractor = _build_ssd_feature_extractor(
217
      feature_extractor_config=ssd_config.feature_extractor,
218
      is_training=is_training)
219
220
221
222
223

  box_coder = box_coder_builder.build(ssd_config.box_coder)
  matcher = matcher_builder.build(ssd_config.matcher)
  region_similarity_calculator = sim_calc.build(
      ssd_config.similarity_calculator)
224
  encode_background_as_zeros = ssd_config.encode_background_as_zeros
225
  negative_class_weight = ssd_config.negative_class_weight
226
227
228
229
230
231
232
233
234
  ssd_box_predictor = box_predictor_builder.build(hyperparams_builder.build,
                                                  ssd_config.box_predictor,
                                                  is_training, num_classes)
  anchor_generator = anchor_generator_builder.build(
      ssd_config.anchor_generator)
  image_resizer_fn = image_resizer_builder.build(ssd_config.image_resizer)
  non_max_suppression_fn, score_conversion_fn = post_processing_builder.build(
      ssd_config.post_processing)
  (classification_loss, localization_loss, classification_weight,
235
236
   localization_weight, hard_example_miner,
   random_example_sampler) = losses_builder.build(ssd_config.loss)
237
  normalize_loss_by_num_matches = ssd_config.normalize_loss_by_num_matches
238
  normalize_loc_loss_by_codesize = ssd_config.normalize_loc_loss_by_codesize
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
  weight_regression_loss_by_score = (ssd_config.weight_regression_loss_by_score)

  target_assigner_instance = target_assigner.TargetAssigner(
      region_similarity_calculator,
      matcher,
      box_coder,
      negative_class_weight=negative_class_weight,
      weight_regression_loss_by_score=weight_regression_loss_by_score)

  expected_classification_loss_under_sampling = None
  if ssd_config.use_expected_classification_loss_under_sampling:
    expected_classification_loss_under_sampling = functools.partial(
        ops.expected_classification_loss_under_sampling,
        minimum_negative_sampling=ssd_config.minimum_negative_sampling,
        desired_negative_sampling_ratio=ssd_config.
        desired_negative_sampling_ratio)
255

256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
  ssd_meta_arch_fn = ssd_meta_arch.SSDMetaArch
  # BEGIN GOOGLE-INTERNAL
  # TODO(lzc): move ssd_mask_meta_arch to third party when it has decent
  # performance relative to a comparable Mask R-CNN model (b/112561592).
  predictor_config = ssd_config.box_predictor
  predict_instance_masks = False
  if predictor_config.WhichOneof(
      'box_predictor_oneof') == 'convolutional_box_predictor':
    predict_instance_masks = (
        predictor_config.convolutional_box_predictor.HasField('mask_head'))
  elif predictor_config.WhichOneof(
      'box_predictor_oneof') == 'weight_shared_convolutional_box_predictor':
    predict_instance_masks = (
        predictor_config.weight_shared_convolutional_box_predictor.HasField(
            'mask_head'))
  if predict_instance_masks:
    ssd_meta_arch_fn = ssd_mask_meta_arch.SSDMaskMetaArch
  # END GOOGLE-INTERNAL

  return ssd_meta_arch_fn(
      is_training=is_training,
      anchor_generator=anchor_generator,
      box_predictor=ssd_box_predictor,
      box_coder=box_coder,
      feature_extractor=feature_extractor,
      encode_background_as_zeros=encode_background_as_zeros,
      image_resizer_fn=image_resizer_fn,
      non_max_suppression_fn=non_max_suppression_fn,
      score_conversion_fn=score_conversion_fn,
      classification_loss=classification_loss,
      localization_loss=localization_loss,
      classification_loss_weight=classification_weight,
      localization_loss_weight=localization_weight,
      normalize_loss_by_num_matches=normalize_loss_by_num_matches,
      hard_example_miner=hard_example_miner,
291
      target_assigner_instance=target_assigner_instance,
292
      add_summaries=add_summaries,
293
294
      normalize_loc_loss_by_codesize=normalize_loc_loss_by_codesize,
      freeze_batchnorm=ssd_config.freeze_batchnorm,
295
      inplace_batchnorm_update=ssd_config.inplace_batchnorm_update,
296
      add_background_class=add_background_class,
297
298
299
      random_example_sampler=random_example_sampler,
      expected_classification_loss_under_sampling=
      expected_classification_loss_under_sampling)
300
301
302


def _build_faster_rcnn_feature_extractor(
303
304
    feature_extractor_config, is_training, reuse_weights=None,
    inplace_batchnorm_update=False):
305
306
307
308
309
310
311
  """Builds a faster_rcnn_meta_arch.FasterRCNNFeatureExtractor based on config.

  Args:
    feature_extractor_config: A FasterRcnnFeatureExtractor proto config from
      faster_rcnn.proto.
    is_training: True if this feature extractor is being built for training.
    reuse_weights: if the feature extractor should reuse weights.
312
313
314
315
316
    inplace_batchnorm_update: Whether to update batch_norm inplace during
      training. This is required for batch norm to work correctly on TPUs. When
      this is false, user must add a control dependency on
      tf.GraphKeys.UPDATE_OPS for train/loss op in order to update the batch
      norm moving average parameters.
317
318
319
320
321
322
323

  Returns:
    faster_rcnn_meta_arch.FasterRCNNFeatureExtractor based on config.

  Raises:
    ValueError: On invalid feature extractor type.
  """
324
325
  if inplace_batchnorm_update:
    raise ValueError('inplace batchnorm updates not supported.')
326
327
328
  feature_type = feature_extractor_config.type
  first_stage_features_stride = (
      feature_extractor_config.first_stage_features_stride)
329
  batch_norm_trainable = feature_extractor_config.batch_norm_trainable
330
331
332
333
334
335
336

  if feature_type not in FASTER_RCNN_FEATURE_EXTRACTOR_CLASS_MAP:
    raise ValueError('Unknown Faster R-CNN feature_extractor: {}'.format(
        feature_type))
  feature_extractor_class = FASTER_RCNN_FEATURE_EXTRACTOR_CLASS_MAP[
      feature_type]
  return feature_extractor_class(
337
338
      is_training, first_stage_features_stride,
      batch_norm_trainable, reuse_weights)
339
340


341
def _build_faster_rcnn_model(frcnn_config, is_training, add_summaries):
342
343
344
345
346
347
348
  """Builds a Faster R-CNN or R-FCN detection model based on the model config.

  Builds R-FCN model if the second_stage_box_predictor in the config is of type
  `rfcn_box_predictor` else builds a Faster R-CNN model.

  Args:
    frcnn_config: A faster_rcnn.proto object containing the config for the
349
      desired FasterRCNNMetaArch or RFCNMetaArch.
350
    is_training: True if this model is being built for training purposes.
351
    add_summaries: Whether to add tf summaries in the model.
352
353
354

  Returns:
    FasterRCNNMetaArch based on the config.
355

356
357
358
359
360
361
362
363
  Raises:
    ValueError: If frcnn_config.type is not recognized (i.e. not registered in
      model_class_map).
  """
  num_classes = frcnn_config.num_classes
  image_resizer_fn = image_resizer_builder.build(frcnn_config.image_resizer)

  feature_extractor = _build_faster_rcnn_feature_extractor(
364
365
      frcnn_config.feature_extractor, is_training,
      frcnn_config.inplace_batchnorm_update)
366

367
  number_of_stages = frcnn_config.number_of_stages
368
369
370
  first_stage_anchor_generator = anchor_generator_builder.build(
      frcnn_config.first_stage_anchor_generator)

371
372
373
374
  first_stage_target_assigner = target_assigner.create_target_assigner(
      'FasterRCNN',
      'proposal',
      use_matmul_gather=frcnn_config.use_matmul_gather_in_matcher)
375
  first_stage_atrous_rate = frcnn_config.first_stage_atrous_rate
376
  first_stage_box_predictor_arg_scope_fn = hyperparams_builder.build(
377
378
379
380
381
      frcnn_config.first_stage_box_predictor_conv_hyperparams, is_training)
  first_stage_box_predictor_kernel_size = (
      frcnn_config.first_stage_box_predictor_kernel_size)
  first_stage_box_predictor_depth = frcnn_config.first_stage_box_predictor_depth
  first_stage_minibatch_size = frcnn_config.first_stage_minibatch_size
382
383
384
  # TODO(bhattad): When eval is supported using static shapes, add separate
  # use_static_shapes_for_trainig and use_static_shapes_for_evaluation.
  use_static_shapes = frcnn_config.use_static_shapes and is_training
385
386
  first_stage_sampler = sampler.BalancedPositiveNegativeSampler(
      positive_fraction=frcnn_config.first_stage_positive_balance_fraction,
387
      is_static=frcnn_config.use_static_balanced_label_sampler and is_training)
388
  first_stage_max_proposals = frcnn_config.first_stage_max_proposals
389
390
391
392
393
394
395
396
397
398
399
400
401
402
  if (frcnn_config.first_stage_nms_iou_threshold < 0 or
      frcnn_config.first_stage_nms_iou_threshold > 1.0):
    raise ValueError('iou_threshold not in [0, 1.0].')
  if (is_training and frcnn_config.second_stage_batch_size >
      first_stage_max_proposals):
    raise ValueError('second_stage_batch_size should be no greater than '
                     'first_stage_max_proposals.')
  first_stage_non_max_suppression_fn = functools.partial(
      post_processing.batch_multiclass_non_max_suppression,
      score_thresh=frcnn_config.first_stage_nms_score_threshold,
      iou_thresh=frcnn_config.first_stage_nms_iou_threshold,
      max_size_per_class=frcnn_config.first_stage_max_proposals,
      max_total_size=frcnn_config.first_stage_max_proposals,
      use_static_shapes=use_static_shapes and is_training)
403
404
405
406
407
408
409
410
  first_stage_loc_loss_weight = (
      frcnn_config.first_stage_localization_loss_weight)
  first_stage_obj_loss_weight = frcnn_config.first_stage_objectness_loss_weight

  initial_crop_size = frcnn_config.initial_crop_size
  maxpool_kernel_size = frcnn_config.maxpool_kernel_size
  maxpool_stride = frcnn_config.maxpool_stride

411
412
413
414
  second_stage_target_assigner = target_assigner.create_target_assigner(
      'FasterRCNN',
      'detection',
      use_matmul_gather=frcnn_config.use_matmul_gather_in_matcher)
415
416
417
418
419
420
  second_stage_box_predictor = box_predictor_builder.build(
      hyperparams_builder.build,
      frcnn_config.second_stage_box_predictor,
      is_training=is_training,
      num_classes=num_classes)
  second_stage_batch_size = frcnn_config.second_stage_batch_size
421
422
  second_stage_sampler = sampler.BalancedPositiveNegativeSampler(
      positive_fraction=frcnn_config.second_stage_balance_fraction,
423
      is_static=frcnn_config.use_static_balanced_label_sampler and is_training)
424
425
426
427
  (second_stage_non_max_suppression_fn, second_stage_score_conversion_fn
  ) = post_processing_builder.build(frcnn_config.second_stage_post_processing)
  second_stage_localization_loss_weight = (
      frcnn_config.second_stage_localization_loss_weight)
428
429
430
  second_stage_classification_loss = (
      losses_builder.build_faster_rcnn_classification_loss(
          frcnn_config.second_stage_classification_loss))
431
432
  second_stage_classification_loss_weight = (
      frcnn_config.second_stage_classification_loss_weight)
433
434
  second_stage_mask_prediction_loss_weight = (
      frcnn_config.second_stage_mask_prediction_loss_weight)
435
436
437
438
439
440
441
442

  hard_example_miner = None
  if frcnn_config.HasField('hard_example_miner'):
    hard_example_miner = losses_builder.build_hard_example_miner(
        frcnn_config.hard_example_miner,
        second_stage_classification_loss_weight,
        second_stage_localization_loss_weight)

443
444
445
  crop_and_resize_fn = (
      ops.matmul_crop_and_resize if frcnn_config.use_matmul_crop_and_resize
      else ops.native_crop_and_resize)
446
447
  clip_anchors_to_image = (
      frcnn_config.clip_anchors_to_image)
448

449
450
451
452
453
  common_kwargs = {
      'is_training': is_training,
      'num_classes': num_classes,
      'image_resizer_fn': image_resizer_fn,
      'feature_extractor': feature_extractor,
454
      'number_of_stages': number_of_stages,
455
      'first_stage_anchor_generator': first_stage_anchor_generator,
456
      'first_stage_target_assigner': first_stage_target_assigner,
457
      'first_stage_atrous_rate': first_stage_atrous_rate,
458
459
      'first_stage_box_predictor_arg_scope_fn':
      first_stage_box_predictor_arg_scope_fn,
460
461
462
463
      'first_stage_box_predictor_kernel_size':
      first_stage_box_predictor_kernel_size,
      'first_stage_box_predictor_depth': first_stage_box_predictor_depth,
      'first_stage_minibatch_size': first_stage_minibatch_size,
464
      'first_stage_sampler': first_stage_sampler,
465
      'first_stage_non_max_suppression_fn': first_stage_non_max_suppression_fn,
466
467
468
      'first_stage_max_proposals': first_stage_max_proposals,
      'first_stage_localization_loss_weight': first_stage_loc_loss_weight,
      'first_stage_objectness_loss_weight': first_stage_obj_loss_weight,
469
      'second_stage_target_assigner': second_stage_target_assigner,
470
      'second_stage_batch_size': second_stage_batch_size,
471
      'second_stage_sampler': second_stage_sampler,
472
473
474
475
476
      'second_stage_non_max_suppression_fn':
      second_stage_non_max_suppression_fn,
      'second_stage_score_conversion_fn': second_stage_score_conversion_fn,
      'second_stage_localization_loss_weight':
      second_stage_localization_loss_weight,
477
478
      'second_stage_classification_loss':
      second_stage_classification_loss,
479
480
      'second_stage_classification_loss_weight':
      second_stage_classification_loss_weight,
481
      'hard_example_miner': hard_example_miner,
482
      'add_summaries': add_summaries,
483
484
485
486
      'crop_and_resize_fn': crop_and_resize_fn,
      'clip_anchors_to_image': clip_anchors_to_image,
      'use_static_shapes': use_static_shapes,
      'resize_masks': frcnn_config.resize_masks
487
  }
488

489
490
  if isinstance(second_stage_box_predictor,
                rfcn_box_predictor.RfcnBoxPredictor):
491
492
493
494
495
496
497
498
499
    return rfcn_meta_arch.RFCNMetaArch(
        second_stage_rfcn_box_predictor=second_stage_box_predictor,
        **common_kwargs)
  else:
    return faster_rcnn_meta_arch.FasterRCNNMetaArch(
        initial_crop_size=initial_crop_size,
        maxpool_kernel_size=maxpool_kernel_size,
        maxpool_stride=maxpool_stride,
        second_stage_mask_rcnn_box_predictor=second_stage_box_predictor,
500
501
        second_stage_mask_prediction_loss_weight=(
            second_stage_mask_prediction_loss_weight),
502
        **common_kwargs)