text: # Copyright (c) OpenMMLab. All rights reserved.
  # Copyright (c) 2019 Western Digital Corporation or its affiliates.
  import torch
  
  from ..builder import DETECTORS
  from .single_stage import SingleStageDetector
  from mmdet.core import bbox2result
  
  @DETECTORS.register_module()
  class YOLOV4(SingleStageDetector):
  
      def __init__(self,
                   backbone,
                   neck,
                   bbox_head,
                   train_cfg=None,
                   test_cfg=None,
                   pretrained=None,
                   init_cfg=None):
          super(YOLOV4, self).__init__(backbone, neck, bbox_head, train_cfg,
                                       test_cfg, pretrained, init_cfg)
  
      def onnx_export(self, img, img_metas):
          """Test function for exporting to ONNX, without test time augmentation.
  
          Args:
              img (torch.Tensor): input images.
              img_metas (list[dict]): List of image information.
  
          Returns:
              tuple[Tensor, Tensor]: dets of shape [N, num_det, 5]
                  and class labels of shape [N, num_det].
          """
          x = self.extract_feat(img)
          outs = self.bbox_head.forward(x)
          # get shape as tensor
          img_shape = torch._shape_as_tensor(img)[2:]
          img_metas[0]['img_shape_for_onnx'] = img_shape
  
          det_bboxes, det_labels = self.bbox_head.onnx_export(*outs, img_metas)
  
          return det_bboxes, det_labels
  
  
      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
  
      def forward_dummy(self, img):
          """Used for computing network flops.
  
          See `mmdetection/tools/analysis_tools/get_flops.py`
          """
          x = self.extract_feat(img)
          outs = self.bbox_head(x)
          return outs
  
      def forward_train(self,
                        img,
                        img_metas,
                        gt_bboxes,
                        gt_labels,
                        gt_bboxes_ignore=None):
          """
          Args:
              img (Tensor): Input images of shape (N, C, H, W).
                  Typically these should be mean centered and std scaled.
              img_metas (list[dict]): A List of image info dict where each dict
                  has: 'img_shape', 'scale_factor', 'flip', and may also contain
                  'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
                  For details on the values of these keys see
                  :class:`mmdet.datasets.pipelines.Collect`.
              gt_bboxes (list[Tensor]): Each item are the truth boxes for each
                  image in [tl_x, tl_y, br_x, br_y] format.
              gt_labels (list[Tensor]): Class indices corresponding to each box
              gt_bboxes_ignore (None | list[Tensor]): Specify which bounding
                  boxes can be ignored when computing the loss.
  
          Returns:
              dict[str, Tensor]: A dictionary of loss components.
          """
          # img = img.data[0]
          # img_metas = img_metas.data[0]
          # gt_bboxes = gt_bboxes.data[0]
          # gt_labels = gt_labels.data[0]
          super(SingleStageDetector, self).forward_train(img, img_metas)
          x = self.extract_feat(img)
          losses = self.bbox_head.forward_train(x, img_metas, gt_bboxes,
                                                gt_labels, gt_bboxes_ignore)
          return losses
  
      def simple_test(self, img, img_metas, rescale=False):
          """Test function without test-time augmentation.
  
          Args:
              img (torch.Tensor): Images with shape (N, C, H, W).
              img_metas (list[dict]): List of image information.
              rescale (bool, optional): Whether to rescale the results.
                  Defaults to False.
  
          Returns:
              list[list[np.ndarray]]: BBox results of each image and classes.
                  The outer list corresponds to each image. The inner list
                  corresponds to each class.
          """
          feat = self.extract_feat(img)
          results_list = self.bbox_head.simple_test(
              feat, img_metas, rescale=rescale)
          bbox_results = [
              bbox2result(det_bboxes, det_labels, self.bbox_head.num_classes)
              for det_bboxes, det_labels in results_list
          ]
          return bbox_results
  
      def aug_test(self, imgs, img_metas, rescale=False):
          """Test function with test time augmentation.
  
          Args:
              imgs (list[Tensor]): the outer list indicates test-time
                  augmentations and inner Tensor should have a shape NxCxHxW,
                  which contains all images in the batch.
              img_metas (list[list[dict]]): the outer list indicates test-time
                  augs (multiscale, flip, etc.) and the inner list indicates
                  images in a batch. each dict has image information.
              rescale (bool, optional): Whether to rescale the results.
                  Defaults to False.
  
          Returns:
              list[list[np.ndarray]]: BBox results of each image and classes.
                  The outer list corresponds to each image. The inner list
                  corresponds to each class.
          """
          assert hasattr(self.bbox_head, 'aug_test'), \
              f'{self.bbox_head.__class__.__name__}' \
              ' does not support test-time augmentation'
  
          feats = self.extract_feats(imgs)
          results_list = self.bbox_head.aug_test(
              feats, img_metas, rescale=rescale)
          bbox_results = [
              bbox2result(det_bboxes, det_labels, self.bbox_head.num_classes)
              for det_bboxes, det_labels in results_list
          ]
          return bbox_results
  
      def onnx_export(self, img, img_metas, with_nms=True):
          """Test function without test time augmentation.
  
          Args:
              img (torch.Tensor): input images.
              img_metas (list[dict]): List of image information.
  
          Returns:
              tuple[Tensor, Tensor]: dets of shape [N, num_det, 5]
                  and class labels of shape [N, num_det].
          """
          x = self.extract_feat(img)
          outs = self.bbox_head(x)
          # get origin input shape to support onnx dynamic shape
  
          # get shape as tensor
          img_shape = torch._shape_as_tensor(img)[2:]
          img_metas[0]['img_shape_for_onnx'] = img_shape
          # get pad input shape to support onnx dynamic shape for exporting
          # `CornerNet` and `CentripetalNet`, which 'pad_shape' is used
          # for inference
          img_metas[0]['pad_shape_for_onnx'] = img_shape
  
          if len(outs) == 2:
              # add dummy score_factor
              outs = (*outs, None)
          # TODO Can we change to `get_bboxes` when `onnx_export` fail
          det_bboxes, det_labels = self.bbox_head.onnx_export(
              *outs, img_metas, with_nms=with_nms)
  
          return det_bboxes, det_labels
  
      def train_step(self, data, optimizer):
          """The iteration step during training.
  
          This method defines an iteration step during training, except for the
          back propagation and optimizer updating, which are done in an optimizer
          hook. Note that in some complicated cases or models, the whole process
          including back propagation and optimizer updating is also defined in
          this method, such as GAN.
  
          Args:
              data (dict): The output of dataloader.
              optimizer (:obj:`torch.optim.Optimizer` | dict): The optimizer of
                  runner is passed to ``train_step()``. This argument is unused
                  and reserved.
  
          Returns:
              dict: It should contain at least 3 keys: ``loss``, ``log_vars``, \
                  ``num_samples``.
  
                  - ``loss`` is a tensor for back propagation, which can be a
                    weighted sum of multiple losses.
                  - ``log_vars`` contains all the variables to be sent to the
                    logger.
                  - ``num_samples`` indicates the batch size (when the model is
                    DDP, it means the batch size on each GPU), which is used for
                    averaging the logs.
          """
          losses = self(**data)
          loss, log_vars = self._parse_losses(losses)
  
          outputs = dict(
              loss=loss, log_vars=log_vars, num_samples=len(data['img_metas']))
  
          return outputs
  
      def forward_test(self, imgs, img_metas, **kwargs):
          """
          Args:
              imgs (List[Tensor]): the outer list indicates test-time
                  augmentations and inner Tensor should have a shape NxCxHxW,
                  which contains all images in the batch.
              img_metas (List[List[dict]]): the outer list indicates test-time
                  augs (multiscale, flip, etc.) and the inner list indicates
                  images in a batch.
          """
          # img_metas = [x.data for x in img_metas]
          # img_metas = [x[0] for x in img_metas]
          for var, name in [(imgs, 'imgs'), (img_metas, 'img_metas')]:
              if not isinstance(var, list):
                  raise TypeError(f'{name} must be a list, but got {type(var)}')
  
          num_augs = len(imgs)
          if num_augs != len(img_metas):
              raise ValueError(f'num of augmentations ({len(imgs)}) '
                               f'!= num of image meta ({len(img_metas)})')
  
          # NOTE the batched image size information may be useful, e.g.
          # in DETR, this is needed for the construction of masks, which is
          # then used for the transformer_head.
          for img, img_meta in zip(imgs, img_metas):
              batch_size = len(img_meta)
              for img_id in range(batch_size):
                  img_meta[img_id]['batch_input_shape'] = tuple(img.size()[-2:])
  
          if num_augs == 1:
              # proposals (List[List[Tensor]]): the outer list indicates
              # test-time augs (multiscale, flip, etc.) and the inner list
              # indicates images in a batch.
              # The Tensor should have a shape Px4, where P is the number of
              # proposals.
              if 'proposals' in kwargs:
                  kwargs['proposals'] = kwargs['proposals'][0]
              return self.simple_test(imgs[0], img_metas[0], **kwargs)
          else:
              assert imgs[0].size(0) == 1, 'aug test does not support ' \
                                           'inference with batch size ' \
                                           f'{imgs[0].size(0)}'
              # TODO: support test augmentation for predefined proposals
              assert 'proposals' not in kwargs
              return self.aug_test(imgs, img_metas, **kwargs)
  
      def val_step(self, data, optimizer=None):
          """The iteration step during validation.
  
          This method shares the same signature as :func:`train_step`, but used
          during val epochs. Note that the evaluation after training epochs is
          not implemented with this method, but an evaluation hook.
          """
          losses = self(**data)
          loss, log_vars = self._parse_losses(losses)
  
          outputs = dict(
              loss=loss, log_vars=log_vars, num_samples=len(data['img_metas']))
  
          return outputs