text: import torch
  import torch.nn as nn
  
  from mmdet.core import bbox2result, bbox2roi, build_assigner, build_sampler, triplet_bbox2roi, embedding2result
  from .. import builder
  from ..registry import DETECTORS
  from .base import BaseDetector
  from .test_mixins import BBoxTestMixin, MaskTestMixin, RPNTestMixin
  
  
  @DETECTORS.register_module
  class TwoStageDetector_Triplet(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,
                   embedding_head=None,
                   mask_roi_extractor=None,
                   mask_head=None,
                   train_cfg=None,
                   test_cfg=None,
                   pretrained=None):
          super(TwoStageDetector_Triplet, 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:
              self.bbox_roi_extractor = builder.build_roi_extractor(
                  bbox_roi_extractor)
              self.bbox_head = builder.build_head(bbox_head)
  
          # embedding_head used to output embedded roi feature
          if embedding_head is not None:
              self.embedding_head = builder.build_head(embedding_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
  
          self.init_weights(pretrained=pretrained)
  
      @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_Triplet, self).init_weights(pretrained)
          self.backbone.init_weights(pretrained=pretrained)
          if self.with_neck:
              if isinstance(self.neck, nn.Sequential):
                  for m in self.neck:
                      m.init_weights()
              else:
                  self.neck.init_weights()
          if self.with_shared_head:
              self.shared_head.init_weights(pretrained=pretrained)
          if self.with_rpn:
              self.rpn_head.init_weights()
          if self.with_bbox:
              self.bbox_roi_extractor.init_weights()
              self.bbox_head.init_weights()
          if self.with_mask:
              self.mask_head.init_weights()
              if not self.share_roi_extractor:
                  self.mask_roi_extractor.init_weights()
  
      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 `mmedetection/tools/get_flops.py`
          """
          outs = ()
          # backbone
          x = self.extract_feat(img)
          # rpn
          if self.with_rpn:
              rpn_outs = self.rpn_head(x)
              outs = outs + (rpn_outs,)
          proposals = torch.randn(1000, 4).cuda()
          # bbox head
          rois = bbox2roi([proposals])
          if self.with_bbox:
              bbox_feats = self.bbox_roi_extractor(
                  x[:self.bbox_roi_extractor.num_inputs], rois)
              if self.with_shared_head:
                  bbox_feats = self.shared_head(bbox_feats)
              cls_score, bbox_pred = self.bbox_head(bbox_feats)
              outs = outs + (cls_score, bbox_pred)
          # mask head
          if self.with_mask:
              mask_rois = rois[:100]
              mask_feats = self.mask_roi_extractor(
                  x[:self.mask_roi_extractor.num_inputs], mask_rois)
              if self.with_shared_head:
                  mask_feats = self.shared_head(mask_feats)
              mask_pred = self.mask_head(mask_feats)
              outs = outs + (mask_pred,)
          return outs
  
      # def forward_train_single(self,
      #                          img,
      #                          img_meta,
      #                          gt_bboxes,
      #                          gt_labels,
      #                          gt_bboxes_ignore=None,
      #                          gt_masks=None,
      #                          proposals=None):
      #     """
      #     Args:
      #         img (Tensor): of shape (N, C, H, W) encoding input images.
      #             Typically these should be mean centered and std scaled.
      #
      #         img_meta (list[dict]): 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
      #             `mmdet/datasets/pipelines/formatting.py: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.
      #
      #         gt_masks (None | Tensor) : true segmentation masks for each box
      #             used if the architecture supports a segmentation task.
      #
      #         proposals : override rpn proposals with custom proposals. Use when
      #             `with_rpn` is False.
      #
      #     Returns:
      #         dict[str, Tensor]: a dictionary of loss components
      #     """
      #     x = self.extract_feat(img)
      #
      #     losses = dict()
      #
      #     # RPN forward and loss
      #     if self.with_rpn:
      #         rpn_outs = self.rpn_head(x)
      #         rpn_loss_inputs = rpn_outs + (gt_bboxes, img_meta,
      #                                       self.train_cfg.rpn)
      #         rpn_losses = self.rpn_head.loss(
      #             *rpn_loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore)
      #         losses.update(rpn_losses)
      #
      #         proposal_cfg = self.train_cfg.get('rpn_proposal',
      #                                           self.test_cfg.rpn)
      #         proposal_inputs = rpn_outs + (img_meta, proposal_cfg)
      #         proposal_list = self.rpn_head.get_bboxes(*proposal_inputs)
      #     else:
      #         proposal_list = proposals
      #
      #     # assign gts and sample proposals
      #     if self.with_bbox or self.with_mask:
      #         bbox_assigner = build_assigner(self.train_cfg.rcnn.assigner)
      #         bbox_sampler = build_sampler(
      #             self.train_cfg.rcnn.sampler, context=self)
      #         num_imgs = img.size(0)
      #         if gt_bboxes_ignore is None:
      #             gt_bboxes_ignore = [None for _ in range(num_imgs)]
      #         sampling_results = []
      #         for i in range(num_imgs):
      #             assign_result = bbox_assigner.assign(proposal_list[i],
      #                                                  gt_bboxes[i],
      #                                                  gt_bboxes_ignore[i],
      #                                                  gt_labels[i])
      #             sampling_result = bbox_sampler.sample(
      #                 assign_result,
      #                 proposal_list[i],
      #                 gt_bboxes[i],
      #                 gt_labels[i],
      #                 feats=[lvl_feat[i][None] for lvl_feat in x])
      #             sampling_results.append(sampling_result)
      #
      #     # bbox head forward and loss
      #     if self.with_bbox:
      #         rois = bbox2roi([res.bboxes for res in sampling_results])
      #         # TODO: a more flexible way to decide which feature maps to use
      #         bbox_feats = self.bbox_roi_extractor(
      #             x[:self.bbox_roi_extractor.num_inputs], rois)
      #         if self.with_shared_head:
      #             bbox_feats = self.shared_head(bbox_feats)
      #         cls_score, bbox_pred = self.bbox_head(bbox_feats)
      #
      #         bbox_targets = self.bbox_head.get_target(sampling_results,
      #                                                  gt_bboxes, gt_labels,
      #                                                  self.train_cfg.rcnn)
      #         loss_bbox = self.bbox_head.loss(cls_score, bbox_pred,
      #                                         *bbox_targets)
      #         losses.update(loss_bbox)
      #
      #     # mask head forward and loss
      #     if self.with_mask:
      #         if not self.share_roi_extractor:
      #             pos_rois = bbox2roi(
      #                 [res.pos_bboxes for res in sampling_results])
      #             mask_feats = self.mask_roi_extractor(
      #                 x[:self.mask_roi_extractor.num_inputs], pos_rois)
      #             if self.with_shared_head:
      #                 mask_feats = self.shared_head(mask_feats)
      #         else:
      #             pos_inds = []
      #             device = bbox_feats.device
      #             for res in sampling_results:
      #                 pos_inds.append(
      #                     torch.ones(
      #                         res.pos_bboxes.shape[0],
      #                         device=device,
      #                         dtype=torch.uint8))
      #                 pos_inds.append(
      #                     torch.zeros(
      #                         res.neg_bboxes.shape[0],
      #                         device=device,
      #                         dtype=torch.uint8))
      #             pos_inds = torch.cat(pos_inds)
      #             mask_feats = bbox_feats[pos_inds]
      #
      #         if mask_feats.shape[0] > 0:
      #             mask_pred = self.mask_head(mask_feats)
      #             mask_targets = self.mask_head.get_target(
      #                 sampling_results, gt_masks, self.train_cfg.rcnn)
      #             pos_labels = torch.cat(
      #                 [res.pos_gt_labels for res in sampling_results])
      #             loss_mask = self.mask_head.loss(mask_pred, mask_targets,
      #                                             pos_labels)
      #             losses.update(loss_mask)
      #
      #     return losses
  
      def forward_train_triplet(self,
                                img_,
                                img_meta_,
                                gt_bboxes_,
                                gt_labels_,
                                gt_instances_,
                                gt_viewpoints_=None,
                                gt_display_=None,
                                gt_bboxes_ignore_=None,
                                gt_masks_=None,
                                proposals_=None):
          """
          triplet forward function.
          :param img_:
          :param img_meta_:
          :param gt_bboxes_:
          :param gt_labels_:
          :param gt_instances_:
          :param gt_viewpoints_:
          :param gt_display_:
          :param gt_bboxes_ignore_:
          :param gt_masks_:
          :param proposals_:
          :return:
          """
          losses = dict()
          triplet_assign_result = []
          triplet_proposals = []
          triplet_gt_bboxes = []
          triplet_gt_instances = []
          triplet_feats = []
          batch_size = img_[0].size(0)
  
          for i in range(3):
              img = img_[i]
              img_meta = img_meta_[i]
              gt_bboxes = gt_bboxes_[i]
              gt_labels = gt_labels_[i]
              gt_instances = gt_instances_[i]
              gt_viewpoints = gt_viewpoints_[i] if (gt_viewpoints_ is not None) else None
              gt_display = gt_display_[i] if (gt_display_ is not None) else None
              gt_bboxes_ignore = gt_bboxes_ignore_[i] if (gt_bboxes_ignore_ is not None) else None
              gt_masks = gt_masks_[i] if (gt_bboxes_ignore_ is not None) else None
              proposals = proposals_[i] if (gt_bboxes_ignore_ is not None) else None
  
              x = self.extract_feat(img)
  
              # RPN forward and loss
              if self.with_rpn:
                  rpn_outs = self.rpn_head(x)
                  rpn_loss_inputs = rpn_outs + (gt_bboxes, img_meta,
                                                self.train_cfg.rpn)
                  rpn_losses = self.rpn_head.loss(
                      *rpn_loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore)
                  for k, v in rpn_losses.items():
                      if k not in losses.keys():
                          losses.setdefault(k, v)
                      else:
                          for j in range(len(v)):
                              losses[k][j] += v[j]
                  # losses.update(rpn_losses)
  
                  proposal_cfg = self.train_cfg.get('rpn_proposal',
                                                    self.test_cfg.rpn)
                  proposal_inputs = rpn_outs + (img_meta, proposal_cfg)
                  proposal_list = self.rpn_head.get_bboxes(*proposal_inputs)
              else:
                  proposal_list = proposals
  
              # assign gts and sample proposals
              if self.with_bbox or self.with_mask:
                  bbox_assigner = build_assigner(self.train_cfg.rcnn.assigner)
                  bbox_sampler = build_sampler(
                      self.train_cfg.rcnn.sampler, context=self)
                  num_imgs = img.size(0)
                  if gt_bboxes_ignore is None:
                      gt_bboxes_ignore = [None for _ in range(num_imgs)]
                  sampling_results = []
                  for i in range(num_imgs):
                      assign_result = bbox_assigner.assign(proposal_list[i],
                                                           gt_bboxes[i],
                                                           gt_bboxes_ignore[i],
                                                           gt_labels[i])
                      sampling_result = bbox_sampler.sample(
                          assign_result,
                          proposal_list[i],
                          gt_bboxes[i],
                          gt_labels[i],
                          feats=[lvl_feat[i][None] for lvl_feat in x])
                      sampling_results.append(sampling_result)
  
              # bbox head forward and loss
              if self.with_bbox:
                  rois = bbox2roi([res.bboxes for res in sampling_results])
                  # TODO: a more flexible way to decide which feature maps to use
                  bbox_feats = self.bbox_roi_extractor(
                      x[:self.bbox_roi_extractor.num_inputs], rois)
                  if self.with_shared_head:
                      bbox_feats = self.shared_head(bbox_feats)
                  cls_score, bbox_pred = self.bbox_head(bbox_feats)
  
                  bbox_targets = self.bbox_head.get_target(sampling_results,
                                                           gt_bboxes, gt_labels,
                                                           self.train_cfg.rcnn)
                  loss_bbox = self.bbox_head.loss(cls_score, bbox_pred,
                                                  *bbox_targets)
                  for k, v in loss_bbox.items():
                      if k not in losses.keys():
                          losses.setdefault(k, v)
                      else:
                          losses[k] += v
  
                  # losses.update(loss_bbox)
  
              # assign instance id
              instance_assigner = build_assigner(self.train_cfg.triplet.assigner)
  
              num_imgs = img.size(0)
              if gt_bboxes_ignore is None:
                  gt_bboxes_ignore = [None for _ in range(num_imgs)]
  
              assign_result_lst = []
              for i in range(num_imgs):
                  instance_assign_result = instance_assigner.assign(proposal_list[i],
                                                                    gt_bboxes[i],
                                                                    gt_bboxes_ignore[i],
                                                                    gt_instances[i])
                  assign_result_lst.append(instance_assign_result)
              triplet_assign_result.append(assign_result_lst)
              triplet_proposals.append(proposal_list)
              triplet_gt_bboxes.append(gt_bboxes)
              triplet_gt_instances.append(gt_instances)
              triplet_feats.append(x)
  
          instance_sampler = build_sampler(
              self.train_cfg.triplet.sampler)
  
          # sampling triplet bboxes pair
          triplet_sampling_result = instance_sampler.sample(triplet_assign_result,
                                                            triplet_proposals,
                                                            triplet_gt_bboxes,
                                                            triplet_gt_instances,
                                                            batch_size)
  
          triplet_rois = triplet_bbox2roi(triplet_sampling_result)
          anchor = []
          neg = []
          positive = []
          for t_roi in triplet_rois:
              roi_feat_lst = []
              for idx, roi in enumerate(t_roi):
                  roi_feat = self.bbox_roi_extractor(triplet_feats[idx][:self.bbox_roi_extractor.num_inputs], roi.unsqueeze(0))
                  roi_feat_lst.append(roi_feat)
              anchor.append(roi_feat_lst[0])
              neg.append(roi_feat_lst[1])
              positive.append(roi_feat_lst[2])
  
          anchor = torch.cat(anchor)
          neg = torch.cat(neg)
          positive = torch.cat(positive)
  
          # triplet shape: (3,), triplet_list should contain results of a batch
          dan, dap = self.embedding_head(anchor, neg, positive)
          triplet_loss = self.embedding_head.loss(dan, dap, margin=0.2, loss_weight=3)
  
          losses.update(triplet_loss)
          losses['acc'] /= 3
          return losses
  
      def forward_train(self,
                        img,
                        img_meta,
                        gt_bboxes,
                        gt_labels,
                        gt_instances,
                        gt_bboxes_ignore=None,
                        gt_masks=None,
                        proposals=None):
          return self.forward_train_triplet(img,
                                            img_meta,
                                            gt_bboxes,
                                            gt_labels,
                                            gt_instances,
                                            gt_bboxes_ignore_=gt_bboxes_ignore,
                                            gt_masks_=gt_masks,
                                            proposals_=proposals)
  
      async def async_simple_test(self,
                                  img,
                                  img_meta,
                                  proposals=None,
                                  rescale=False):
          """Async test without augmentation."""
          assert self.with_bbox, "Bbox head must be implemented."
          x = self.extract_feat(img)
  
          if proposals is None:
              proposal_list = await self.async_test_rpn(x, img_meta,
                                                        self.test_cfg.rpn)
          else:
              proposal_list = proposals
  
          det_bboxes, det_labels = await self.async_test_bboxes(
              x, img_meta, proposal_list, self.test_cfg.rcnn, rescale=rescale)
          bbox_results = bbox2result(det_bboxes, det_labels,
                                     self.bbox_head.num_classes)
  
          if not self.with_mask:
              return bbox_results
          else:
              segm_results = await self.async_test_mask(
                  x,
                  img_meta,
                  det_bboxes,
                  det_labels,
                  rescale=rescale,
                  mask_test_cfg=self.test_cfg.get('mask'))
              return bbox_results, segm_results
  
      def simple_test(self, img, img_meta, proposals=None, rescale=False):
          """Test without augmentation."""
          assert self.with_bbox, "Bbox head must be implemented."
  
          x = self.extract_feat(img)
  
          if proposals is None:
              proposal_list = self.simple_test_rpn(x, img_meta,
                                                   self.test_cfg.rpn)
          else:
              proposal_list = proposals
  
          if not self.with_embedding:
              det_bboxes, det_labels = self.simple_test_bboxes(
                  x, img_meta, proposal_list, self.test_cfg.rcnn, rescale=rescale)
              bbox_results = bbox2result(det_bboxes, det_labels,
                                         self.bbox_head.num_classes)
          else:
              det_bboxes, det_labels, det_embedding = self.simple_test_bboxes_embedded(
                  x, img_meta, proposal_list, self.test_cfg.rcnn, rescale=rescale)
              bbox_results = bbox2result(det_bboxes, det_labels,
                                         self.bbox_head.num_classes)
              embedding_results = embedding2result(det_embedding, det_labels, self.bbox_head.num_classes)
  
          if self.with_mask:
              segm_results = self.simple_test_mask(
                  x, img_meta, det_bboxes, det_labels, rescale=rescale)
  
          if self.with_mask and self.with_embedding:
              return bbox_results, segm_results, embedding_results
          elif self.with_mask:
              return bbox_results, segm_results
          elif self.with_embedding:
              return bbox_results, embedding_results
          else:
              return bbox_results
  
      def aug_test(self, imgs, img_metas, rescale=False):
          """Test with augmentations.
  
          If rescale is False, then returned bboxes and masks will fit the scale
          of imgs[0].
          """
          # recompute feats to save memory
          proposal_list = self.aug_test_rpn(
              self.extract_feats(imgs), img_metas, self.test_cfg.rpn)
          det_bboxes, det_labels = self.aug_test_bboxes(
              self.extract_feats(imgs), img_metas, proposal_list,
              self.test_cfg.rcnn)
  
          if rescale:
              _det_bboxes = det_bboxes
          else:
              _det_bboxes = det_bboxes.clone()
              _det_bboxes[:, :4] *= img_metas[0][0]['scale_factor']
          bbox_results = bbox2result(_det_bboxes, det_labels,
                                     self.bbox_head.num_classes)
  
          # det_bboxes always keep the original scale
          if self.with_mask:
              segm_results = self.aug_test_mask(
                  self.extract_feats(imgs), img_metas, det_bboxes, det_labels)
              return bbox_results, segm_results
          else:
              return bbox_results
  