factory.py 15.5 KB
Newer Older
Yeqing Li's avatar
Yeqing Li committed
1
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
2
3
4
5
6
7
8
9
10
11
12
13
#
# 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.
Yeqing Li's avatar
Yeqing Li committed
14

Abdullah Rashwan's avatar
Abdullah Rashwan committed
15
16
17
"""Factory methods to build models."""

# Import libraries
Abdullah Rashwan's avatar
Abdullah Rashwan committed
18

Abdullah Rashwan's avatar
Abdullah Rashwan committed
19
20
21
22
23
import tensorflow as tf

from official.vision.beta.configs import image_classification as classification_cfg
from official.vision.beta.configs import maskrcnn as maskrcnn_cfg
from official.vision.beta.configs import retinanet as retinanet_cfg
Abdullah Rashwan's avatar
Abdullah Rashwan committed
24
from official.vision.beta.configs import semantic_segmentation as segmentation_cfg
Yeqing Li's avatar
Yeqing Li committed
25
from official.vision.beta.modeling import backbones
Abdullah Rashwan's avatar
Abdullah Rashwan committed
26
from official.vision.beta.modeling import classification_model
27
from official.vision.beta.modeling import decoders
Abdullah Rashwan's avatar
Abdullah Rashwan committed
28
29
from official.vision.beta.modeling import maskrcnn_model
from official.vision.beta.modeling import retinanet_model
Abdullah Rashwan's avatar
Abdullah Rashwan committed
30
from official.vision.beta.modeling import segmentation_model
Abdullah Rashwan's avatar
Abdullah Rashwan committed
31
32
from official.vision.beta.modeling.heads import dense_prediction_heads
from official.vision.beta.modeling.heads import instance_heads
Abdullah Rashwan's avatar
Abdullah Rashwan committed
33
from official.vision.beta.modeling.heads import segmentation_heads
Abdullah Rashwan's avatar
Abdullah Rashwan committed
34
35
36
37
38
39
40
41
42
43
from official.vision.beta.modeling.layers import detection_generator
from official.vision.beta.modeling.layers import mask_sampler
from official.vision.beta.modeling.layers import roi_aligner
from official.vision.beta.modeling.layers import roi_generator
from official.vision.beta.modeling.layers import roi_sampler


def build_classification_model(
    input_specs: tf.keras.layers.InputSpec,
    model_config: classification_cfg.ImageClassificationModel,
Pengchong Jin's avatar
Pengchong Jin committed
44
    l2_regularizer: tf.keras.regularizers.Regularizer = None,
Rebecca Chen's avatar
Rebecca Chen committed
45
    skip_logits_layer: bool = False) -> tf.keras.Model:  # pytype: disable=annotation-type-mismatch  # typed-keras
Abdullah Rashwan's avatar
Abdullah Rashwan committed
46
  """Builds the classification model."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
47
  norm_activation_config = model_config.norm_activation
Yeqing Li's avatar
Yeqing Li committed
48
  backbone = backbones.factory.build_backbone(
Abdullah Rashwan's avatar
Abdullah Rashwan committed
49
      input_specs=input_specs,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
50
51
      backbone_config=model_config.backbone,
      norm_activation_config=norm_activation_config,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
52
53
54
55
56
57
58
      l2_regularizer=l2_regularizer)

  model = classification_model.ClassificationModel(
      backbone=backbone,
      num_classes=model_config.num_classes,
      input_specs=input_specs,
      dropout_rate=model_config.dropout_rate,
59
      kernel_initializer=model_config.kernel_initializer,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
60
61
62
63
      kernel_regularizer=l2_regularizer,
      add_head_batch_norm=model_config.add_head_batch_norm,
      use_sync_bn=norm_activation_config.use_sync_bn,
      norm_momentum=norm_activation_config.norm_momentum,
Pengchong Jin's avatar
Pengchong Jin committed
64
65
      norm_epsilon=norm_activation_config.norm_epsilon,
      skip_logits_layer=skip_logits_layer)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
66
67
68
  return model


Fan Yang's avatar
Fan Yang committed
69
70
71
def build_maskrcnn(
    input_specs: tf.keras.layers.InputSpec,
    model_config: maskrcnn_cfg.MaskRCNN,
Rebecca Chen's avatar
Rebecca Chen committed
72
    l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:  # pytype: disable=annotation-type-mismatch  # typed-keras
Abdullah Rashwan's avatar
Abdullah Rashwan committed
73
  """Builds Mask R-CNN model."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
74
  norm_activation_config = model_config.norm_activation
Yeqing Li's avatar
Yeqing Li committed
75
  backbone = backbones.factory.build_backbone(
Abdullah Rashwan's avatar
Abdullah Rashwan committed
76
      input_specs=input_specs,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
77
78
      backbone_config=model_config.backbone,
      norm_activation_config=norm_activation_config,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
79
      l2_regularizer=l2_regularizer)
Cristina Vasconcelos's avatar
Cristina Vasconcelos committed
80
  backbone_features = backbone(tf.keras.Input(input_specs.shape[1:]))
Abdullah Rashwan's avatar
Abdullah Rashwan committed
81

82
  decoder = decoders.factory.build_decoder(
Abdullah Rashwan's avatar
Abdullah Rashwan committed
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
      input_specs=backbone.output_specs,
      model_config=model_config,
      l2_regularizer=l2_regularizer)

  rpn_head_config = model_config.rpn_head
  roi_generator_config = model_config.roi_generator
  roi_sampler_config = model_config.roi_sampler
  roi_aligner_config = model_config.roi_aligner
  detection_head_config = model_config.detection_head
  generator_config = model_config.detection_generator
  num_anchors_per_location = (
      len(model_config.anchor.aspect_ratios) * model_config.anchor.num_scales)

  rpn_head = dense_prediction_heads.RPNHead(
      min_level=model_config.min_level,
      max_level=model_config.max_level,
      num_anchors_per_location=num_anchors_per_location,
      num_convs=rpn_head_config.num_convs,
      num_filters=rpn_head_config.num_filters,
      use_separable_conv=rpn_head_config.use_separable_conv,
      activation=norm_activation_config.activation,
      use_sync_bn=norm_activation_config.use_sync_bn,
      norm_momentum=norm_activation_config.norm_momentum,
      norm_epsilon=norm_activation_config.norm_epsilon,
      kernel_regularizer=l2_regularizer)

  detection_head = instance_heads.DetectionHead(
      num_classes=model_config.num_classes,
      num_convs=detection_head_config.num_convs,
      num_filters=detection_head_config.num_filters,
      use_separable_conv=detection_head_config.use_separable_conv,
      num_fcs=detection_head_config.num_fcs,
      fc_dims=detection_head_config.fc_dims,
Xianzhi Du's avatar
Xianzhi Du committed
116
      class_agnostic_bbox_pred=detection_head_config.class_agnostic_bbox_pred,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
117
118
119
120
      activation=norm_activation_config.activation,
      use_sync_bn=norm_activation_config.use_sync_bn,
      norm_momentum=norm_activation_config.norm_momentum,
      norm_epsilon=norm_activation_config.norm_epsilon,
Xianzhi Du's avatar
Xianzhi Du committed
121
122
      kernel_regularizer=l2_regularizer,
      name='detection_head')
Cristina Vasconcelos's avatar
Cristina Vasconcelos committed
123

Cristina Vasconcelos's avatar
Cristina Vasconcelos committed
124
  # Builds decoder and region proposal network:
Cristina Vasconcelos's avatar
Cristina Vasconcelos committed
125
126
127
128
  if decoder:
    decoder_features = decoder(backbone_features)
    rpn_head(decoder_features)

Xianzhi Du's avatar
Xianzhi Du committed
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
  if roi_sampler_config.cascade_iou_thresholds:
    detection_head_cascade = [detection_head]
    for cascade_num in range(len(roi_sampler_config.cascade_iou_thresholds)):
      detection_head = instance_heads.DetectionHead(
          num_classes=model_config.num_classes,
          num_convs=detection_head_config.num_convs,
          num_filters=detection_head_config.num_filters,
          use_separable_conv=detection_head_config.use_separable_conv,
          num_fcs=detection_head_config.num_fcs,
          fc_dims=detection_head_config.fc_dims,
          class_agnostic_bbox_pred=detection_head_config
          .class_agnostic_bbox_pred,
          activation=norm_activation_config.activation,
          use_sync_bn=norm_activation_config.use_sync_bn,
          norm_momentum=norm_activation_config.norm_momentum,
          norm_epsilon=norm_activation_config.norm_epsilon,
          kernel_regularizer=l2_regularizer,
          name='detection_head_{}'.format(cascade_num + 1))
Cristina Vasconcelos's avatar
Cristina Vasconcelos committed
147

Xianzhi Du's avatar
Xianzhi Du committed
148
149
      detection_head_cascade.append(detection_head)
    detection_head = detection_head_cascade
Abdullah Rashwan's avatar
Abdullah Rashwan committed
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166

  roi_generator_obj = roi_generator.MultilevelROIGenerator(
      pre_nms_top_k=roi_generator_config.pre_nms_top_k,
      pre_nms_score_threshold=roi_generator_config.pre_nms_score_threshold,
      pre_nms_min_size_threshold=(
          roi_generator_config.pre_nms_min_size_threshold),
      nms_iou_threshold=roi_generator_config.nms_iou_threshold,
      num_proposals=roi_generator_config.num_proposals,
      test_pre_nms_top_k=roi_generator_config.test_pre_nms_top_k,
      test_pre_nms_score_threshold=(
          roi_generator_config.test_pre_nms_score_threshold),
      test_pre_nms_min_size_threshold=(
          roi_generator_config.test_pre_nms_min_size_threshold),
      test_nms_iou_threshold=roi_generator_config.test_nms_iou_threshold,
      test_num_proposals=roi_generator_config.test_num_proposals,
      use_batched_nms=roi_generator_config.use_batched_nms)

Xianzhi Du's avatar
Xianzhi Du committed
167
  roi_sampler_cascade = []
Abdullah Rashwan's avatar
Abdullah Rashwan committed
168
169
170
171
172
173
174
175
176
  roi_sampler_obj = roi_sampler.ROISampler(
      mix_gt_boxes=roi_sampler_config.mix_gt_boxes,
      num_sampled_rois=roi_sampler_config.num_sampled_rois,
      foreground_fraction=roi_sampler_config.foreground_fraction,
      foreground_iou_threshold=roi_sampler_config.foreground_iou_threshold,
      background_iou_high_threshold=(
          roi_sampler_config.background_iou_high_threshold),
      background_iou_low_threshold=(
          roi_sampler_config.background_iou_low_threshold))
Xianzhi Du's avatar
Xianzhi Du committed
177
178
179
180
181
182
183
184
185
186
187
188
  roi_sampler_cascade.append(roi_sampler_obj)
  # Initialize addtional roi simplers for cascade heads.
  if roi_sampler_config.cascade_iou_thresholds:
    for iou in roi_sampler_config.cascade_iou_thresholds:
      roi_sampler_obj = roi_sampler.ROISampler(
          mix_gt_boxes=False,
          num_sampled_rois=roi_sampler_config.num_sampled_rois,
          foreground_iou_threshold=iou,
          background_iou_high_threshold=iou,
          background_iou_low_threshold=0.0,
          skip_subsampling=True)
      roi_sampler_cascade.append(roi_sampler_obj)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
189
190
191
192
193
194

  roi_aligner_obj = roi_aligner.MultilevelROIAligner(
      crop_size=roi_aligner_config.crop_size,
      sample_offset=roi_aligner_config.sample_offset)

  detection_generator_obj = detection_generator.DetectionGenerator(
Fan Yang's avatar
Fan Yang committed
195
      apply_nms=generator_config.apply_nms,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
196
197
198
199
      pre_nms_top_k=generator_config.pre_nms_top_k,
      pre_nms_score_threshold=generator_config.pre_nms_score_threshold,
      nms_iou_threshold=generator_config.nms_iou_threshold,
      max_num_detections=generator_config.max_num_detections,
Xianzhi Du's avatar
Xianzhi Du committed
200
      nms_version=generator_config.nms_version,
Xianzhi Du's avatar
Xianzhi Du committed
201
202
      use_cpu_nms=generator_config.use_cpu_nms,
      soft_nms_sigma=generator_config.soft_nms_sigma)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
203
204
205
206
207
208
209
210
211
212
213

  if model_config.include_mask:
    mask_head = instance_heads.MaskHead(
        num_classes=model_config.num_classes,
        upsample_factor=model_config.mask_head.upsample_factor,
        num_convs=model_config.mask_head.num_convs,
        num_filters=model_config.mask_head.num_filters,
        use_separable_conv=model_config.mask_head.use_separable_conv,
        activation=model_config.norm_activation.activation,
        norm_momentum=model_config.norm_activation.norm_momentum,
        norm_epsilon=model_config.norm_activation.norm_epsilon,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
214
215
        kernel_regularizer=l2_regularizer,
        class_agnostic=model_config.mask_head.class_agnostic)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236

    mask_sampler_obj = mask_sampler.MaskSampler(
        mask_target_size=(
            model_config.mask_roi_aligner.crop_size *
            model_config.mask_head.upsample_factor),
        num_sampled_masks=model_config.mask_sampler.num_sampled_masks)

    mask_roi_aligner_obj = roi_aligner.MultilevelROIAligner(
        crop_size=model_config.mask_roi_aligner.crop_size,
        sample_offset=model_config.mask_roi_aligner.sample_offset)
  else:
    mask_head = None
    mask_sampler_obj = None
    mask_roi_aligner_obj = None

  model = maskrcnn_model.MaskRCNNModel(
      backbone=backbone,
      decoder=decoder,
      rpn_head=rpn_head,
      detection_head=detection_head,
      roi_generator=roi_generator_obj,
Xianzhi Du's avatar
Xianzhi Du committed
237
      roi_sampler=roi_sampler_cascade,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
238
239
240
241
      roi_aligner=roi_aligner_obj,
      detection_generator=detection_generator_obj,
      mask_head=mask_head,
      mask_sampler=mask_sampler_obj,
Xianzhi Du's avatar
Xianzhi Du committed
242
243
      mask_roi_aligner=mask_roi_aligner_obj,
      class_agnostic_bbox_pred=detection_head_config.class_agnostic_bbox_pred,
244
245
246
247
248
249
      cascade_class_ensemble=detection_head_config.cascade_class_ensemble,
      min_level=model_config.min_level,
      max_level=model_config.max_level,
      num_scales=model_config.anchor.num_scales,
      aspect_ratios=model_config.anchor.aspect_ratios,
      anchor_size=model_config.anchor.anchor_size)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
250
251
252
  return model


Fan Yang's avatar
Fan Yang committed
253
254
255
def build_retinanet(
    input_specs: tf.keras.layers.InputSpec,
    model_config: retinanet_cfg.RetinaNet,
Rebecca Chen's avatar
Rebecca Chen committed
256
    l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:  # pytype: disable=annotation-type-mismatch  # typed-keras
Abdullah Rashwan's avatar
Abdullah Rashwan committed
257
  """Builds RetinaNet model."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
258
  norm_activation_config = model_config.norm_activation
Yeqing Li's avatar
Yeqing Li committed
259
  backbone = backbones.factory.build_backbone(
Abdullah Rashwan's avatar
Abdullah Rashwan committed
260
      input_specs=input_specs,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
261
262
      backbone_config=model_config.backbone,
      norm_activation_config=norm_activation_config,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
263
      l2_regularizer=l2_regularizer)
Cristina Vasconcelos's avatar
Cristina Vasconcelos committed
264
  backbone_features = backbone(tf.keras.Input(input_specs.shape[1:]))
Abdullah Rashwan's avatar
Abdullah Rashwan committed
265

266
  decoder = decoders.factory.build_decoder(
Abdullah Rashwan's avatar
Abdullah Rashwan committed
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
      input_specs=backbone.output_specs,
      model_config=model_config,
      l2_regularizer=l2_regularizer)

  head_config = model_config.head
  generator_config = model_config.detection_generator
  num_anchors_per_location = (
      len(model_config.anchor.aspect_ratios) * model_config.anchor.num_scales)

  head = dense_prediction_heads.RetinaNetHead(
      min_level=model_config.min_level,
      max_level=model_config.max_level,
      num_classes=model_config.num_classes,
      num_anchors_per_location=num_anchors_per_location,
      num_convs=head_config.num_convs,
      num_filters=head_config.num_filters,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
283
284
285
      attribute_heads=[
          cfg.as_dict() for cfg in (head_config.attribute_heads or [])
      ],
Abdullah Rashwan's avatar
Abdullah Rashwan committed
286
287
288
289
290
291
292
      use_separable_conv=head_config.use_separable_conv,
      activation=norm_activation_config.activation,
      use_sync_bn=norm_activation_config.use_sync_bn,
      norm_momentum=norm_activation_config.norm_momentum,
      norm_epsilon=norm_activation_config.norm_epsilon,
      kernel_regularizer=l2_regularizer)

Cristina Vasconcelos's avatar
Cristina Vasconcelos committed
293
294
295
296
297
  # Builds decoder and head so that their trainable weights are initialized
  if decoder:
    decoder_features = decoder(backbone_features)
    _ = head(decoder_features)

Abdullah Rashwan's avatar
Abdullah Rashwan committed
298
  detection_generator_obj = detection_generator.MultilevelDetectionGenerator(
Fan Yang's avatar
Fan Yang committed
299
      apply_nms=generator_config.apply_nms,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
300
301
302
303
      pre_nms_top_k=generator_config.pre_nms_top_k,
      pre_nms_score_threshold=generator_config.pre_nms_score_threshold,
      nms_iou_threshold=generator_config.nms_iou_threshold,
      max_num_detections=generator_config.max_num_detections,
Xianzhi Du's avatar
Xianzhi Du committed
304
      nms_version=generator_config.nms_version,
Xianzhi Du's avatar
Xianzhi Du committed
305
306
      use_cpu_nms=generator_config.use_cpu_nms,
      soft_nms_sigma=generator_config.soft_nms_sigma)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
307
308

  model = retinanet_model.RetinaNetModel(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
309
310
311
312
313
314
315
316
317
      backbone,
      decoder,
      head,
      detection_generator_obj,
      min_level=model_config.min_level,
      max_level=model_config.max_level,
      num_scales=model_config.anchor.num_scales,
      aspect_ratios=model_config.anchor.aspect_ratios,
      anchor_size=model_config.anchor.anchor_size)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
318
  return model
Abdullah Rashwan's avatar
Abdullah Rashwan committed
319
320
321
322


def build_segmentation_model(
    input_specs: tf.keras.layers.InputSpec,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
323
    model_config: segmentation_cfg.SemanticSegmentationModel,
Rebecca Chen's avatar
Rebecca Chen committed
324
    l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:  # pytype: disable=annotation-type-mismatch  # typed-keras
Abdullah Rashwan's avatar
Abdullah Rashwan committed
325
  """Builds Segmentation model."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
326
  norm_activation_config = model_config.norm_activation
Abdullah Rashwan's avatar
Abdullah Rashwan committed
327
328
  backbone = backbones.factory.build_backbone(
      input_specs=input_specs,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
329
330
      backbone_config=model_config.backbone,
      norm_activation_config=norm_activation_config,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
331
332
      l2_regularizer=l2_regularizer)

333
  decoder = decoders.factory.build_decoder(
Abdullah Rashwan's avatar
Abdullah Rashwan committed
334
335
336
337
338
339
340
341
342
343
      input_specs=backbone.output_specs,
      model_config=model_config,
      l2_regularizer=l2_regularizer)

  head_config = model_config.head

  head = segmentation_heads.SegmentationHead(
      num_classes=model_config.num_classes,
      level=head_config.level,
      num_convs=head_config.num_convs,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
344
      prediction_kernel_size=head_config.prediction_kernel_size,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
345
      num_filters=head_config.num_filters,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
346
      use_depthwise_convolution=head_config.use_depthwise_convolution,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
347
      upsample_factor=head_config.upsample_factor,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
348
349
350
      feature_fusion=head_config.feature_fusion,
      low_level=head_config.low_level,
      low_level_num_filters=head_config.low_level_num_filters,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
351
352
353
354
355
356
357
358
      activation=norm_activation_config.activation,
      use_sync_bn=norm_activation_config.use_sync_bn,
      norm_momentum=norm_activation_config.norm_momentum,
      norm_epsilon=norm_activation_config.norm_epsilon,
      kernel_regularizer=l2_regularizer)

  model = segmentation_model.SegmentationModel(backbone, decoder, head)
  return model