Spaces:
Sleeping
Sleeping
| # Copyright 2023 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. | |
| """Model architecture factory.""" | |
| from __future__ import absolute_import | |
| from __future__ import division | |
| from __future__ import print_function | |
| from official.legacy.detection.modeling.architecture import fpn | |
| from official.legacy.detection.modeling.architecture import heads | |
| from official.legacy.detection.modeling.architecture import identity | |
| from official.legacy.detection.modeling.architecture import nn_ops | |
| from official.legacy.detection.modeling.architecture import resnet | |
| from official.legacy.detection.modeling.architecture import spinenet | |
| def norm_activation_generator(params): | |
| return nn_ops.norm_activation_builder( | |
| momentum=params.batch_norm_momentum, | |
| epsilon=params.batch_norm_epsilon, | |
| trainable=params.batch_norm_trainable, | |
| activation=params.activation) | |
| def backbone_generator(params): | |
| """Generator function for various backbone models.""" | |
| if params.architecture.backbone == 'resnet': | |
| resnet_params = params.resnet | |
| backbone_fn = resnet.Resnet( | |
| resnet_depth=resnet_params.resnet_depth, | |
| activation=params.norm_activation.activation, | |
| norm_activation=norm_activation_generator( | |
| params.norm_activation)) | |
| elif params.architecture.backbone == 'spinenet': | |
| spinenet_params = params.spinenet | |
| backbone_fn = spinenet.SpineNetBuilder(model_id=spinenet_params.model_id) | |
| else: | |
| raise ValueError('Backbone model `{}` is not supported.' | |
| .format(params.architecture.backbone)) | |
| return backbone_fn | |
| def multilevel_features_generator(params): | |
| """Generator function for various FPN models.""" | |
| if params.architecture.multilevel_features == 'fpn': | |
| fpn_params = params.fpn | |
| fpn_fn = fpn.Fpn( | |
| min_level=params.architecture.min_level, | |
| max_level=params.architecture.max_level, | |
| fpn_feat_dims=fpn_params.fpn_feat_dims, | |
| use_separable_conv=fpn_params.use_separable_conv, | |
| activation=params.norm_activation.activation, | |
| use_batch_norm=fpn_params.use_batch_norm, | |
| norm_activation=norm_activation_generator( | |
| params.norm_activation)) | |
| elif params.architecture.multilevel_features == 'identity': | |
| fpn_fn = identity.Identity() | |
| else: | |
| raise ValueError('The multi-level feature model `{}` is not supported.' | |
| .format(params.architecture.multilevel_features)) | |
| return fpn_fn | |
| def retinanet_head_generator(params): | |
| """Generator function for RetinaNet head architecture.""" | |
| head_params = params.retinanet_head | |
| anchors_per_location = params.anchor.num_scales * len( | |
| params.anchor.aspect_ratios) | |
| return heads.RetinanetHead( | |
| params.architecture.min_level, | |
| params.architecture.max_level, | |
| params.architecture.num_classes, | |
| anchors_per_location, | |
| head_params.num_convs, | |
| head_params.num_filters, | |
| head_params.use_separable_conv, | |
| norm_activation=norm_activation_generator(params.norm_activation)) | |
| def rpn_head_generator(params): | |
| """Generator function for RPN head architecture.""" | |
| head_params = params.rpn_head | |
| anchors_per_location = params.anchor.num_scales * len( | |
| params.anchor.aspect_ratios) | |
| return heads.RpnHead( | |
| params.architecture.min_level, | |
| params.architecture.max_level, | |
| anchors_per_location, | |
| head_params.num_convs, | |
| head_params.num_filters, | |
| head_params.use_separable_conv, | |
| params.norm_activation.activation, | |
| head_params.use_batch_norm, | |
| norm_activation=norm_activation_generator(params.norm_activation)) | |
| def oln_rpn_head_generator(params): | |
| """Generator function for OLN-proposal (OLN-RPN) head architecture.""" | |
| head_params = params.rpn_head | |
| anchors_per_location = params.anchor.num_scales * len( | |
| params.anchor.aspect_ratios) | |
| return heads.OlnRpnHead( | |
| params.architecture.min_level, | |
| params.architecture.max_level, | |
| anchors_per_location, | |
| head_params.num_convs, | |
| head_params.num_filters, | |
| head_params.use_separable_conv, | |
| params.norm_activation.activation, | |
| head_params.use_batch_norm, | |
| norm_activation=norm_activation_generator(params.norm_activation)) | |
| def fast_rcnn_head_generator(params): | |
| """Generator function for Fast R-CNN head architecture.""" | |
| head_params = params.frcnn_head | |
| return heads.FastrcnnHead( | |
| params.architecture.num_classes, | |
| head_params.num_convs, | |
| head_params.num_filters, | |
| head_params.use_separable_conv, | |
| head_params.num_fcs, | |
| head_params.fc_dims, | |
| params.norm_activation.activation, | |
| head_params.use_batch_norm, | |
| norm_activation=norm_activation_generator(params.norm_activation)) | |
| def oln_box_score_head_generator(params): | |
| """Generator function for Scoring Fast R-CNN head architecture.""" | |
| head_params = params.frcnn_head | |
| return heads.OlnBoxScoreHead( | |
| params.architecture.num_classes, | |
| head_params.num_convs, | |
| head_params.num_filters, | |
| head_params.use_separable_conv, | |
| head_params.num_fcs, | |
| head_params.fc_dims, | |
| params.norm_activation.activation, | |
| head_params.use_batch_norm, | |
| norm_activation=norm_activation_generator(params.norm_activation)) | |
| def mask_rcnn_head_generator(params): | |
| """Generator function for Mask R-CNN head architecture.""" | |
| head_params = params.mrcnn_head | |
| return heads.MaskrcnnHead( | |
| params.architecture.num_classes, | |
| params.architecture.mask_target_size, | |
| head_params.num_convs, | |
| head_params.num_filters, | |
| head_params.use_separable_conv, | |
| params.norm_activation.activation, | |
| head_params.use_batch_norm, | |
| norm_activation=norm_activation_generator(params.norm_activation)) | |
| def oln_mask_score_head_generator(params): | |
| """Generator function for Scoring Mask R-CNN head architecture.""" | |
| head_params = params.mrcnn_head | |
| return heads.OlnMaskScoreHead( | |
| params.architecture.num_classes, | |
| params.architecture.mask_target_size, | |
| head_params.num_convs, | |
| head_params.num_filters, | |
| head_params.use_separable_conv, | |
| params.norm_activation.activation, | |
| head_params.use_batch_norm, | |
| norm_activation=norm_activation_generator(params.norm_activation)) | |
| def shapeprior_head_generator(params): | |
| """Generator function for shape prior head architecture.""" | |
| head_params = params.shapemask_head | |
| return heads.ShapemaskPriorHead( | |
| params.architecture.num_classes, | |
| head_params.num_downsample_channels, | |
| head_params.mask_crop_size, | |
| head_params.use_category_for_mask, | |
| head_params.shape_prior_path) | |
| def coarsemask_head_generator(params): | |
| """Generator function for ShapeMask coarse mask head architecture.""" | |
| head_params = params.shapemask_head | |
| return heads.ShapemaskCoarsemaskHead( | |
| params.architecture.num_classes, | |
| head_params.num_downsample_channels, | |
| head_params.mask_crop_size, | |
| head_params.use_category_for_mask, | |
| head_params.num_convs, | |
| norm_activation=norm_activation_generator(params.norm_activation)) | |
| def finemask_head_generator(params): | |
| """Generator function for Shapemask fine mask head architecture.""" | |
| head_params = params.shapemask_head | |
| return heads.ShapemaskFinemaskHead( | |
| params.architecture.num_classes, | |
| head_params.num_downsample_channels, | |
| head_params.mask_crop_size, | |
| head_params.use_category_for_mask, | |
| head_params.num_convs, | |
| head_params.upsample_factor) | |