text: # -*- coding: utf-8 -*-
  # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
  # File:
  
  import torch
  
  from detectron2.modeling import ROI_HEADS_REGISTRY, StandardROIHeads
  from detectron2.modeling.poolers import ROIPooler
  
  # Is this function good to expose as an API?
  from detectron2.modeling.roi_heads.roi_heads import select_foreground_proposals
  
  from .densepose_head import (
      build_densepose_data_filter,
      build_densepose_head,
      build_densepose_losses,
      build_densepose_predictor,
      densepose_inference,
  )
  
  
  @ROI_HEADS_REGISTRY.register()
  class DensePoseROIHeads(StandardROIHeads):
      """
      A Standard ROIHeads which contains an addition of DensePose head.
      """
  
      def __init__(self, cfg, input_shape):
          super().__init__(cfg, input_shape)
          self._init_densepose_head(cfg)
  
      def _init_densepose_head(self, cfg):
          # fmt: off
          self.densepose_on          = cfg.MODEL.DENSEPOSE_ON
          if not self.densepose_on:
              return
          self.densepose_data_filter = build_densepose_data_filter(cfg)
          dp_pooler_resolution       = cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_RESOLUTION
          dp_pooler_scales           = tuple(1.0 / self.feature_strides[k] for k in self.in_features)
          dp_pooler_sampling_ratio   = cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_SAMPLING_RATIO
          dp_pooler_type             = cfg.MODEL.ROI_DENSEPOSE_HEAD.POOLER_TYPE
          # fmt: on
          in_channels = [self.feature_channels[f] for f in self.in_features][0]
          self.densepose_pooler = ROIPooler(
              output_size=dp_pooler_resolution,
              scales=dp_pooler_scales,
              sampling_ratio=dp_pooler_sampling_ratio,
              pooler_type=dp_pooler_type,
          )
          self.densepose_head = build_densepose_head(cfg, in_channels)
          self.densepose_predictor = build_densepose_predictor(
              cfg, self.densepose_head.n_out_channels
          )
          self.densepose_losses = build_densepose_losses(cfg)
  
      def _forward_densepose(self, features, instances):
          """
          Forward logic of the densepose prediction branch.
  
          Args:
              features (list[Tensor]): #level input features for densepose prediction
              instances (list[Instances]): the per-image instances to train/predict densepose.
                  In training, they can be the proposals.
                  In inference, they can be the predicted boxes.
  
          Returns:
              In training, a dict of losses.
              In inference, update `instances` with new fields "densepose" and return it.
          """
          if not self.densepose_on:
              return {} if self.training else instances
  
          if self.training:
              proposals, _ = select_foreground_proposals(instances, self.num_classes)
              proposals_dp = self.densepose_data_filter(proposals)
              if len(proposals_dp) > 0:
                  proposal_boxes = [x.proposal_boxes for x in proposals_dp]
                  features_dp = self.densepose_pooler(features, proposal_boxes)
                  densepose_head_outputs = self.densepose_head(features_dp)
                  densepose_outputs, _ = self.densepose_predictor(densepose_head_outputs)
                  densepose_loss_dict = self.densepose_losses(proposals_dp, densepose_outputs)
                  return densepose_loss_dict
          else:
              pred_boxes = [x.pred_boxes for x in instances]
              features_dp = self.densepose_pooler(features, pred_boxes)
              if len(features_dp) > 0:
                  densepose_head_outputs = self.densepose_head(features_dp)
                  densepose_outputs, _ = self.densepose_predictor(densepose_head_outputs)
              else:
                  # If no detection occured instances
                  # set densepose_outputs to empty tensors
                  empty_tensor = torch.zeros(size=(0, 0, 0, 0), device=features_dp.device)
                  densepose_outputs = tuple([empty_tensor] * 4)
  
              densepose_inference(densepose_outputs, instances)
              return instances
  
      def forward(self, images, features, proposals, targets=None):
          features_list = [features[f] for f in self.in_features]
  
          instances, losses = super().forward(images, features, proposals, targets)
          del targets, images
  
          if self.training:
              losses.update(self._forward_densepose(features_list, instances))
          else:
              instances = self._forward_densepose(features_list, instances)
          return instances, losses
  