text: # Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  # SPDX-License-Identifier: Apache-2.0
  # -*- coding: utf-8 -*-
  import tensorflow as tf
  from .. import builder
  from ..registry import DETECTORS
  from .base import BaseDetector
  from .test_mixins import BBoxTestMixin, RPNTestMixin 
  
  
  @DETECTORS.register_module
  class TwoStageDetector(BaseDetector, RPNTestMixin, BBoxTestMixin): # MaskTestMixin):
      """Base class for two-stage detectors.
  
      Two-stage detectors typically consisting of a region proposal network and a
      task-specific regression head.
      """
  
      def __init__(self,
                   backbone,
                   neck=None,
                   shared_head=None,
                   rpn_head=None,
                   bbox_roi_extractor=None,
                   bbox_head=None,
                   mask_roi_extractor=None,
                   mask_head=None,
                   train_cfg=None,
                   test_cfg=None,
                   pretrained=None):
          super(TwoStageDetector, self).__init__()
          self.backbone = builder.build_backbone(backbone)
          if neck is not None:
              self.neck = builder.build_neck(neck)
  
          if shared_head is not None:
              self.shared_head = builder.build_shared_head(shared_head)
  
          if rpn_head is not None:
              self.rpn_head = builder.build_head(rpn_head)
  
          if bbox_head is not None:
              if bbox_roi_extractor is not None:
                  self.bbox_roi_extractor = builder.build_roi_extractor(
                      bbox_roi_extractor)
              self.bbox_head = builder.build_head(bbox_head)
  
          if mask_head is not None:
              if mask_roi_extractor is not None:
                  self.mask_roi_extractor = builder.build_roi_extractor(
                      mask_roi_extractor)
                  self.share_roi_extractor = False
              else:
                  self.share_roi_extractor = True
                  self.mask_roi_extractor = self.bbox_roi_extractor
              self.mask_head = builder.build_head(mask_head)
  
          self.train_cfg = train_cfg
          self.test_cfg = test_cfg
  
  
      @property
      def with_rpn(self):
          return hasattr(self, 'rpn_head') and self.rpn_head is not None
  
      def init_weights(self, pretrained=None):
          super(TwoStageDetector, self).init_weights(pretrained)
  
      def extract_feat(self, img):
          """Directly extract features from the backbone+neck
          """
          x = self.backbone(img)
          if self.with_neck:
              x = self.neck(x)
          return x
  
  
  