Guest User

det

a guest
May 4th, 2021
40
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. # Copyright (c) Facebook, Inc. and its affiliates.
  2. import inspect
  3. import logging
  4. import numpy as np
  5. from typing import Dict, List, Optional, Tuple
  6. import torch
  7. from torch import nn
  8.  
  9. from detectron2.config import configurable
  10. from detectron2.layers import ShapeSpec, nonzero_tuple
  11. from detectron2.structures import Boxes, ImageList, Instances, pairwise_iou
  12. from detectron2.utils.events import get_event_storage
  13. from detectron2.utils.registry import Registry
  14.  
  15. from ..backbone.resnet import BottleneckBlock, ResNet
  16. from ..matcher import Matcher
  17. from ..poolers import ROIPooler
  18. from ..proposal_generator.proposal_utils import add_ground_truth_to_proposals
  19. from ..sampling import subsample_labels
  20. from .box_head import build_box_head
  21. from .fast_rcnn import FastRCNNOutputLayers
  22. from .keypoint_head import build_keypoint_head
  23. from .mask_head import build_mask_head
  24. from .mutil_path_fuse_module import build_mutil_path_fuse_module
  25.  
  26. ROI_HEADS_REGISTRY = Registry("ROI_HEADS")
  27. ROI_HEADS_REGISTRY.__doc__ = """
  28. Registry for ROI heads in a generalized R-CNN model.
  29. ROIHeads take feature maps and region proposals, and
  30. perform per-region computation.
  31.  
  32. The registered object will be called with `obj(cfg, input_shape)`.
  33. The call is expected to return an :class:`ROIHeads`.
  34. """
  35.  
  36. logger = logging.getLogger(__name__)
  37.  
  38.  
  39. def build_roi_heads(cfg, input_shape):
  40.     """
  41.    Build ROIHeads defined by `cfg.MODEL.ROI_HEADS.NAME`.
  42.    """
  43.     name = cfg.MODEL.ROI_HEADS.NAME
  44.     return ROI_HEADS_REGISTRY.get(name)(cfg, input_shape)
  45.  
  46.  
  47. def select_foreground_proposals(
  48.     proposals: List[Instances], bg_label: int
  49. ) -> Tuple[List[Instances], List[torch.Tensor]]:
  50.     """
  51.    Given a list of N Instances (for N images), each containing a `gt_classes` field,
  52.    return a list of Instances that contain only instances with `gt_classes != -1 &&
  53.    gt_classes != bg_label`.
  54.  
  55.    Args:
  56.        proposals (list[Instances]): A list of N Instances, where N is the number of
  57.            images in the batch.
  58.        bg_label: label index of background class.
  59.  
  60.    Returns:
  61.        list[Instances]: N Instances, each contains only the selected foreground instances.
  62.        list[Tensor]: N boolean vector, correspond to the selection mask of
  63.            each Instances object. True for selected instances.
  64.    """
  65.     assert isinstance(proposals, (list, tuple))
  66.     assert isinstance(proposals[0], Instances)
  67.     assert proposals[0].has("gt_classes")
  68.     fg_proposals = []
  69.     fg_selection_masks = []
  70.     for proposals_per_image in proposals:
  71.         gt_classes = proposals_per_image.gt_classes
  72.         fg_selection_mask = (gt_classes != -1) & (gt_classes != bg_label)
  73.         fg_idxs = fg_selection_mask.nonzero().squeeze(1)
  74.         fg_proposals.append(proposals_per_image[fg_idxs])
  75.         fg_selection_masks.append(fg_selection_mask)
  76.     return fg_proposals, fg_selection_masks
  77.  
  78.  
  79. def select_proposals_with_visible_keypoints(proposals: List[Instances]) -> List[Instances]:
  80.     """
  81.    Args:
  82.        proposals (list[Instances]): a list of N Instances, where N is the
  83.            number of images.
  84.  
  85.    Returns:
  86.        proposals: only contains proposals with at least one visible keypoint.
  87.  
  88.    Note that this is still slightly different from Detectron.
  89.    In Detectron, proposals for training keypoint head are re-sampled from
  90.    all the proposals with IOU>threshold & >=1 visible keypoint.
  91.  
  92.    Here, the proposals are first sampled from all proposals with
  93.    IOU>threshold, then proposals with no visible keypoint are filtered out.
  94.    This strategy seems to make no difference on Detectron and is easier to implement.
  95.    """
  96.     ret = []
  97.     all_num_fg = []
  98.     for proposals_per_image in proposals:
  99.         # If empty/unannotated image (hard negatives), skip filtering for train
  100.         if len(proposals_per_image) == 0:
  101.             ret.append(proposals_per_image)
  102.             continue
  103.         gt_keypoints = proposals_per_image.gt_keypoints.tensor
  104.         # #fg x K x 3
  105.         vis_mask = gt_keypoints[:, :, 2] >= 1
  106.         xs, ys = gt_keypoints[:, :, 0], gt_keypoints[:, :, 1]
  107.         proposal_boxes = proposals_per_image.proposal_boxes.tensor.unsqueeze(dim=1)  # #fg x 1 x 4
  108.         kp_in_box = (
  109.             (xs >= proposal_boxes[:, :, 0])
  110.             & (xs <= proposal_boxes[:, :, 2])
  111.             & (ys >= proposal_boxes[:, :, 1])
  112.             & (ys <= proposal_boxes[:, :, 3])
  113.         )
  114.         selection = (kp_in_box & vis_mask).any(dim=1)
  115.         selection_idxs = nonzero_tuple(selection)[0]
  116.         all_num_fg.append(selection_idxs.numel())
  117.         ret.append(proposals_per_image[selection_idxs])
  118.  
  119.     storage = get_event_storage()
  120.     storage.put_scalar("keypoint_head/num_fg_samples", np.mean(all_num_fg))
  121.     return ret
  122.  
  123.  
  124. class ROIHeads(torch.nn.Module):
  125.     """
  126.    ROIHeads perform all per-region computation in an R-CNN.
  127.  
  128.    It typically contains logic to
  129.  
  130.    1. (in training only) match proposals with ground truth and sample them
  131.    2. crop the regions and extract per-region features using proposals
  132.    3. make per-region predictions with different heads
  133.  
  134.    It can have many variants, implemented as subclasses of this class.
  135.    This base class contains the logic to match/sample proposals.
  136.    But it is not necessary to inherit this class if the sampling logic is not needed.
  137.    """
  138.  
  139.     @configurable
  140.     def __init__(
  141.         self,
  142.         *,
  143.         num_classes,
  144.         batch_size_per_image,
  145.         positive_fraction,
  146.         proposal_matcher,
  147.         proposal_append_gt=True
  148.     ):
  149.         """
  150.        NOTE: this interface is experimental.
  151.  
  152.        Args:
  153.            num_classes (int): number of foreground classes (i.e. background is not included)
  154.            batch_size_per_image (int): number of proposals to sample for training
  155.            positive_fraction (float): fraction of positive (foreground) proposals
  156.                to sample for training.
  157.            proposal_matcher (Matcher): matcher that matches proposals and ground truth
  158.            proposal_append_gt (bool): whether to include ground truth as proposals as well
  159.        """
  160.         super().__init__()
  161.         self.batch_size_per_image = batch_size_per_image
  162.         self.positive_fraction = positive_fraction
  163.         self.num_classes = num_classes
  164.         self.proposal_matcher = proposal_matcher
  165.         self.proposal_append_gt = proposal_append_gt
  166.  
  167.     @classmethod
  168.     def from_config(cls, cfg):
  169.         return {
  170.             "batch_size_per_image": cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE,
  171.             "positive_fraction": cfg.MODEL.ROI_HEADS.POSITIVE_FRACTION,
  172.             "num_classes": cfg.MODEL.ROI_HEADS.NUM_CLASSES,
  173.             "proposal_append_gt": cfg.MODEL.ROI_HEADS.PROPOSAL_APPEND_GT,
  174.             # Matcher to assign box proposals to gt boxes
  175.             "proposal_matcher": Matcher(
  176.                 cfg.MODEL.ROI_HEADS.IOU_THRESHOLDS,
  177.                 cfg.MODEL.ROI_HEADS.IOU_LABELS,
  178.                 allow_low_quality_matches=False,
  179.             ),
  180.         }
  181.  
  182.     def _sample_proposals(
  183.         self, matched_idxs: torch.Tensor, matched_labels: torch.Tensor, gt_classes: torch.Tensor
  184.     ) -> Tuple[torch.Tensor, torch.Tensor]:
  185.         """
  186.        Based on the matching between N proposals and M groundtruth,
  187.        sample the proposals and set their classification labels.
  188.  
  189.        Args:
  190.            matched_idxs (Tensor): a vector of length N, each is the best-matched
  191.                gt index in [0, M) for each proposal.
  192.            matched_labels (Tensor): a vector of length N, the matcher's label
  193.                (one of cfg.MODEL.ROI_HEADS.IOU_LABELS) for each proposal.
  194.            gt_classes (Tensor): a vector of length M.
  195.  
  196.        Returns:
  197.            Tensor: a vector of indices of sampled proposals. Each is in [0, N).
  198.            Tensor: a vector of the same length, the classification label for
  199.                each sampled proposal. Each sample is labeled as either a category in
  200.                [0, num_classes) or the background (num_classes).
  201.        """
  202.         has_gt = gt_classes.numel() > 0
  203.         # Get the corresponding GT for each proposal
  204.         if has_gt:
  205.             gt_classes = gt_classes[matched_idxs]
  206.             # Label unmatched proposals (0 label from matcher) as background (label=num_classes)
  207.             gt_classes[matched_labels == 0] = self.num_classes
  208.             # Label ignore proposals (-1 label)
  209.             gt_classes[matched_labels == -1] = -1
  210.         else:
  211.             gt_classes = torch.zeros_like(matched_idxs) + self.num_classes
  212.  
  213.         sampled_fg_idxs, sampled_bg_idxs = subsample_labels(
  214.             gt_classes, self.batch_size_per_image, self.positive_fraction, self.num_classes
  215.         )
  216.  
  217.         sampled_idxs = torch.cat([sampled_fg_idxs, sampled_bg_idxs], dim=0)
  218.         return sampled_idxs, gt_classes[sampled_idxs]
  219.  
  220.     @torch.no_grad()
  221.     def label_and_sample_proposals(
  222.         self, proposals: List[Instances], targets: List[Instances]
  223.     ) -> List[Instances]:
  224.         """
  225.        Prepare some proposals to be used to train the ROI heads.
  226.        It performs box matching between `proposals` and `targets`, and assigns
  227.        training labels to the proposals.
  228.        It returns ``self.batch_size_per_image`` random samples from proposals and groundtruth
  229.        boxes, with a fraction of positives that is no larger than
  230.        ``self.positive_fraction``.
  231.  
  232.        Args:
  233.            See :meth:`ROIHeads.forward`
  234.  
  235.        Returns:
  236.            list[Instances]:
  237.                length `N` list of `Instances`s containing the proposals
  238.                sampled for training. Each `Instances` has the following fields:
  239.  
  240.                - proposal_boxes: the proposal boxes
  241.                - gt_boxes: the ground-truth box that the proposal is assigned to
  242.                  (this is only meaningful if the proposal has a label > 0; if label = 0
  243.                  then the ground-truth box is random)
  244.  
  245.                Other fields such as "gt_classes", "gt_masks", that's included in `targets`.
  246.        """
  247.         gt_boxes = [x.gt_boxes for x in targets]
  248.         # Augment proposals with ground-truth boxes.
  249.         # In the case of learned proposals (e.g., RPN), when training starts
  250.         # the proposals will be low quality due to random initialization.
  251.         # It's possible that none of these initial
  252.         # proposals have high enough overlap with the gt objects to be used
  253.         # as positive examples for the second stage components (box head,
  254.         # cls head, mask head). Adding the gt boxes to the set of proposals
  255.         # ensures that the second stage components will have some positive
  256.         # examples from the start of training. For RPN, this augmentation improves
  257.         # convergence and empirically improves box AP on COCO by about 0.5
  258.         # points (under one tested configuration).
  259.         if self.proposal_append_gt:
  260.             proposals = add_ground_truth_to_proposals(gt_boxes, proposals)
  261.  
  262.         proposals_with_gt = []
  263.  
  264.         num_fg_samples = []
  265.         num_bg_samples = []
  266.         for proposals_per_image, targets_per_image in zip(proposals, targets):
  267.             has_gt = len(targets_per_image) > 0
  268.             match_quality_matrix = pairwise_iou(
  269.                 targets_per_image.gt_boxes, proposals_per_image.proposal_boxes
  270.             )
  271.             matched_idxs, matched_labels = self.proposal_matcher(match_quality_matrix)
  272.             sampled_idxs, gt_classes = self._sample_proposals(
  273.                 matched_idxs, matched_labels, targets_per_image.gt_classes
  274.             )
  275.  
  276.             # Set target attributes of the sampled proposals:
  277.             proposals_per_image = proposals_per_image[sampled_idxs]
  278.             proposals_per_image.gt_classes = gt_classes
  279.  
  280.             if has_gt:
  281.                 sampled_targets = matched_idxs[sampled_idxs]
  282.                 # We index all the attributes of targets that start with "gt_"
  283.                 # and have not been added to proposals yet (="gt_classes").
  284.                 # NOTE: here the indexing waste some compute, because heads
  285.                 # like masks, keypoints, etc, will filter the proposals again,
  286.                 # (by foreground/background, or number of keypoints in the image, etc)
  287.                 # so we essentially index the data twice.
  288.                 for (trg_name, trg_value) in targets_per_image.get_fields().items():
  289.                     if trg_name.startswith("gt_") and not proposals_per_image.has(trg_name):
  290.                         proposals_per_image.set(trg_name, trg_value[sampled_targets])
  291.             # If no GT is given in the image, we don't know what a dummy gt value can be.
  292.             # Therefore the returned proposals won't have any gt_* fields, except for a
  293.             # gt_classes full of background label.
  294.  
  295.             num_bg_samples.append((gt_classes == self.num_classes).sum().item())
  296.             num_fg_samples.append(gt_classes.numel() - num_bg_samples[-1])
  297.             proposals_with_gt.append(proposals_per_image)
  298.  
  299.         # Log the number of fg/bg samples that are selected for training ROI heads
  300.         storage = get_event_storage()
  301.         storage.put_scalar("roi_head/num_fg_samples", np.mean(num_fg_samples))
  302.         storage.put_scalar("roi_head/num_bg_samples", np.mean(num_bg_samples))
  303.  
  304.         return proposals_with_gt
  305.  
  306.     def forward(
  307.         self,
  308.         images: ImageList,
  309.         features: Dict[str, torch.Tensor],
  310.         proposals: List[Instances],
  311.         targets: Optional[List[Instances]] = None,
  312.     ) -> Tuple[List[Instances], Dict[str, torch.Tensor]]:
  313.         """
  314.        Args:
  315.            images (ImageList):
  316.            features (dict[str,Tensor]): input data as a mapping from feature
  317.                map name to tensor. Axis 0 represents the number of images `N` in
  318.                the input data; axes 1-3 are channels, height, and width, which may
  319.                vary between feature maps (e.g., if a feature pyramid is used).
  320.            proposals (list[Instances]): length `N` list of `Instances`. The i-th
  321.                `Instances` contains object proposals for the i-th input image,
  322.                with fields "proposal_boxes" and "objectness_logits".
  323.            targets (list[Instances], optional): length `N` list of `Instances`. The i-th
  324.                `Instances` contains the ground-truth per-instance annotations
  325.                for the i-th input image.  Specify `targets` during training only.
  326.                It may have the following fields:
  327.  
  328.                - gt_boxes: the bounding box of each instance.
  329.                - gt_classes: the label for each instance with a category ranging in [0, #class].
  330.                - gt_masks: PolygonMasks or BitMasks, the ground-truth masks of each instance.
  331.                - gt_keypoints: NxKx3, the groud-truth keypoints for each instance.
  332.  
  333.        Returns:
  334.            list[Instances]: length `N` list of `Instances` containing the
  335.            detected instances. Returned during inference only; may be [] during training.
  336.  
  337.            dict[str->Tensor]:
  338.            mapping from a named loss to a tensor storing the loss. Used during training only.
  339.        """
  340.         raise NotImplementedError()
  341.  
  342.  
  343. @ROI_HEADS_REGISTRY.register()
  344. class Res5ROIHeads(ROIHeads):
  345.     """
  346.    The ROIHeads in a typical "C4" R-CNN model, where
  347.    the box and mask head share the cropping and
  348.    the per-region feature computation by a Res5 block.
  349.    See :paper:`ResNet` Appendix A.
  350.    """
  351.  
  352.     @configurable
  353.     def __init__(
  354.         self,
  355.         *,
  356.         in_features: List[str],
  357.         pooler: ROIPooler,
  358.         res5: nn.Module,
  359.         box_predictor: nn.Module,
  360.         mask_head: Optional[nn.Module] = None,
  361.         **kwargs
  362.     ):
  363.         """
  364.        NOTE: this interface is experimental.
  365.  
  366.        Args:
  367.            in_features (list[str]): list of backbone feature map names to use for
  368.                feature extraction
  369.            pooler (ROIPooler): pooler to extra region features from backbone
  370.            res5 (nn.Sequential): a CNN to compute per-region features, to be used by
  371.                ``box_predictor`` and ``mask_head``. Typically this is a "res5"
  372.                block from a ResNet.
  373.            box_predictor (nn.Module): make box predictions from the feature.
  374.                Should have the same interface as :class:`FastRCNNOutputLayers`.
  375.            mask_head (nn.Module): transform features to make mask predictions
  376.        """
  377.         super().__init__(**kwargs)
  378.         self.in_features = in_features
  379.         self.pooler = pooler
  380.         if isinstance(res5, (list, tuple)):
  381.             res5 = nn.Sequential(*res5)
  382.         self.res5 = res5
  383.         self.box_predictor = box_predictor
  384.         self.mask_on = mask_head is not None
  385.         if self.mask_on:
  386.             self.mask_head = mask_head
  387.  
  388.     @classmethod
  389.     def from_config(cls, cfg, input_shape):
  390.         # fmt: off
  391.         ret = super().from_config(cfg)
  392.         in_features = ret["in_features"] = cfg.MODEL.ROI_HEADS.IN_FEATURES
  393.         pooler_resolution = cfg.MODEL.ROI_BOX_HEAD.POOLER_RESOLUTION
  394.         pooler_type       = cfg.MODEL.ROI_BOX_HEAD.POOLER_TYPE
  395.         pooler_scales     = (1.0 / input_shape[in_features[0]].stride, )
  396.         sampling_ratio    = cfg.MODEL.ROI_BOX_HEAD.POOLER_SAMPLING_RATIO
  397.         mask_on           = cfg.MODEL.MASK_ON
  398.         # fmt: on
  399.         assert not cfg.MODEL.KEYPOINT_ON
  400.         assert len(in_features) == 1
  401.  
  402.         ret["pooler"] = ROIPooler(
  403.             output_size=pooler_resolution,
  404.             scales=pooler_scales,
  405.             sampling_ratio=sampling_ratio,
  406.             pooler_type=pooler_type,
  407.         )
  408.  
  409.         # Compatbility with old moco code. Might be useful.
  410.         # See notes in StandardROIHeads.from_config
  411.         if not inspect.ismethod(cls._build_res5_block):
  412.             logger.warning(
  413.                 "The behavior of _build_res5_block may change. "
  414.                 "Please do not depend on private methods."
  415.             )
  416.             cls._build_res5_block = classmethod(cls._build_res5_block)
  417.  
  418.         ret["res5"], out_channels = cls._build_res5_block(cfg)
  419.         ret["box_predictor"] = FastRCNNOutputLayers(
  420.             cfg, ShapeSpec(channels=out_channels, height=1, width=1)
  421.         )
  422.  
  423.         if mask_on:
  424.             ret["mask_head"] = build_mask_head(
  425.                 cfg,
  426.                 ShapeSpec(channels=out_channels, width=pooler_resolution, height=pooler_resolution),
  427.             )
  428.         return ret
  429.  
  430.     @classmethod
  431.     def _build_res5_block(cls, cfg):
  432.         # fmt: off
  433.         stage_channel_factor = 2 ** 3  # res5 is 8x res2
  434.         num_groups           = cfg.MODEL.RESNETS.NUM_GROUPS
  435.         width_per_group      = cfg.MODEL.RESNETS.WIDTH_PER_GROUP
  436.         bottleneck_channels  = num_groups * width_per_group * stage_channel_factor
  437.         out_channels         = cfg.MODEL.RESNETS.RES2_OUT_CHANNELS * stage_channel_factor
  438.         stride_in_1x1        = cfg.MODEL.RESNETS.STRIDE_IN_1X1
  439.         norm                 = cfg.MODEL.RESNETS.NORM
  440.         assert not cfg.MODEL.RESNETS.DEFORM_ON_PER_STAGE[-1], \
  441.             "Deformable conv is not yet supported in res5 head."
  442.         # fmt: on
  443.  
  444.         blocks = ResNet.make_stage(
  445.             BottleneckBlock,
  446.             3,
  447.             stride_per_block=[2, 1, 1],
  448.             in_channels=out_channels // 2,
  449.             bottleneck_channels=bottleneck_channels,
  450.             out_channels=out_channels,
  451.             num_groups=num_groups,
  452.             norm=norm,
  453.             stride_in_1x1=stride_in_1x1,
  454.         )
  455.         return nn.Sequential(*blocks), out_channels
  456.  
  457.     def _shared_roi_transform(self, features, boxes):
  458.         x = self.pooler(features, boxes)
  459.         return self.res5(x)
  460.  
  461.     def forward(self, images, features, proposals, targets=None):
  462.         """
  463.        See :meth:`ROIHeads.forward`.
  464.        """
  465.         del images
  466.  
  467.         if self.training:
  468.             assert targets
  469.             proposals = self.label_and_sample_proposals(proposals, targets)
  470.         del targets
  471.  
  472.         proposal_boxes = [x.proposal_boxes for x in proposals]
  473.         box_features = self._shared_roi_transform(
  474.             [features[f] for f in self.in_features], proposal_boxes
  475.         )
  476.         predictions = self.box_predictor(box_features.mean(dim=[2, 3]))
  477.  
  478.         if self.training:
  479.             del features
  480.             losses = self.box_predictor.losses(predictions, proposals)
  481.             if self.mask_on:
  482.                 proposals, fg_selection_masks = select_foreground_proposals(
  483.                     proposals, self.num_classes
  484.                 )
  485.                 # Since the ROI feature transform is shared between boxes and masks,
  486.                 # we don't need to recompute features. The mask loss is only defined
  487.                 # on foreground proposals, so we need to select out the foreground
  488.                 # features.
  489.                 mask_features = box_features[torch.cat(fg_selection_masks, dim=0)]
  490.                 del box_features
  491.                 losses.update(self.mask_head(mask_features, proposals))
  492.             return [], losses
  493.         else:
  494.             pred_instances, _ = self.box_predictor.inference(predictions, proposals)
  495.             pred_instances = self.forward_with_given_boxes(features, pred_instances)
  496.             return pred_instances, {}
  497.  
  498.     def forward_with_given_boxes(self, features, instances):
  499.         """
  500.        Use the given boxes in `instances` to produce other (non-box) per-ROI outputs.
  501.  
  502.        Args:
  503.            features: same as in `forward()`
  504.            instances (list[Instances]): instances to predict other outputs. Expect the keys
  505.                "pred_boxes" and "pred_classes" to exist.
  506.  
  507.        Returns:
  508.            instances (Instances):
  509.                the same `Instances` object, with extra
  510.                fields such as `pred_masks` or `pred_keypoints`.
  511.        """
  512.         assert not self.training
  513.         assert instances[0].has("pred_boxes") and instances[0].has("pred_classes")
  514.  
  515.         if self.mask_on:
  516.             features = [features[f] for f in self.in_features]
  517.             x = self._shared_roi_transform(features, [x.pred_boxes for x in instances])
  518.             return self.mask_head(x, instances)
  519.         else:
  520.             return instances
  521.  
  522.  
  523. @ROI_HEADS_REGISTRY.register()
  524. class StandardROIHeads(ROIHeads):
  525.     """
  526.    It's "standard" in a sense that there is no ROI transform sharing
  527.    or feature sharing between tasks.
  528.    Each head independently processes the input features by each head's
  529.    own pooler and head.
  530.  
  531.    This class is used by most models, such as FPN and C5.
  532.    To implement more models, you can subclass it and implement a different
  533.    :meth:`forward()` or a head.
  534.    """
  535.  
  536.     @configurable
  537.     def __init__(
  538.         self,
  539.         *,
  540.         box_in_features: List[str],
  541.         box_pooler: ROIPooler,
  542.         box_head: nn.Module,
  543.         box_predictor: nn.Module,
  544.         mask_in_features: Optional[List[str]] = None,
  545.         mask_pooler: Optional[ROIPooler] = None,
  546.         mask_head: Optional[nn.Module] = None,
  547.         keypoint_in_features: Optional[List[str]] = None,
  548.         keypoint_pooler: Optional[ROIPooler] = None,
  549.         keypoint_head: Optional[nn.Module] = None,
  550.         train_on_pred_boxes: bool = False,
  551.         **kwargs
  552.     ):
  553.         """
  554.        NOTE: this interface is experimental.
  555.  
  556.        Args:
  557.            box_in_features (list[str]): list of feature names to use for the box head.
  558.            box_pooler (ROIPooler): pooler to extra region features for box head
  559.            box_head (nn.Module): transform features to make box predictions
  560.            box_predictor (nn.Module): make box predictions from the feature.
  561.                Should have the same interface as :class:`FastRCNNOutputLayers`.
  562.            mask_in_features (list[str]): list of feature names to use for the mask
  563.                pooler or mask head. None if not using mask head.
  564.            mask_pooler (ROIPooler): pooler to extract region features from image features.
  565.                The mask head will then take region features to make predictions.
  566.                If None, the mask head will directly take the dict of image features
  567.                defined by `mask_in_features`
  568.            mask_head (nn.Module): transform features to make mask predictions
  569.            keypoint_in_features, keypoint_pooler, keypoint_head: similar to ``mask_*``.
  570.            train_on_pred_boxes (bool): whether to use proposal boxes or
  571.                predicted boxes from the box head to train other heads.
  572.        """
  573.         super().__init__(**kwargs)
  574.         # keep self.in_features for backward compatibility
  575.         self.in_features = self.box_in_features = box_in_features
  576.         self.box_pooler = box_pooler
  577.         self.box_head = box_head
  578.         self.box_predictor = box_predictor
  579.  
  580.         self.mask_on = mask_in_features is not None
  581.         if self.mask_on:
  582.             self.mask_in_features = mask_in_features
  583.             self.mask_pooler = mask_pooler
  584.             self.mask_head = mask_head
  585.  
  586.         self.keypoint_on = keypoint_in_features is not None
  587.         if self.keypoint_on:
  588.             self.keypoint_in_features = keypoint_in_features
  589.             self.keypoint_pooler = keypoint_pooler
  590.             self.keypoint_head = keypoint_head
  591.  
  592.         self.train_on_pred_boxes = train_on_pred_boxes
  593.         self.mutil_path_fuse_on = cfg.MODEL.TEXTFUSENET_MUTIL_PATH_FUSE_ON
  594.         if self.mutil_path_fuse_on:
  595.             self._init_seg_head(cfg)
  596.             self._init_mutil_path_fuse_module(cfg)
  597.  
  598.     @classmethod
  599.     def from_config(cls, cfg, input_shape):
  600.         ret = super().from_config(cfg)
  601.         ret["train_on_pred_boxes"] = cfg.MODEL.ROI_BOX_HEAD.TRAIN_ON_PRED_BOXES
  602.         # Subclasses that have not been updated to use from_config style construction
  603.         # may have overridden _init_*_head methods. In this case, those overridden methods
  604.         # will not be classmethods and we need to avoid trying to call them here.
  605.         # We test for this with ismethod which only returns True for bound methods of cls.
  606.         # Such subclasses will need to handle calling their overridden _init_*_head methods.
  607.         if inspect.ismethod(cls._init_box_head):
  608.             ret.update(cls._init_box_head(cfg, input_shape))
  609.         if inspect.ismethod(cls._init_mask_head):
  610.             ret.update(cls._init_mask_head(cfg, input_shape))
  611.         if inspect.ismethod(cls._init_keypoint_head):
  612.             ret.update(cls._init_keypoint_head(cfg, input_shape))
  613.         return ret
  614.  
  615.     @classmethod
  616.     def _init_box_head(cls, cfg, input_shape):
  617.         # fmt: off
  618.         in_features       = cfg.MODEL.ROI_HEADS.IN_FEATURES
  619.         pooler_resolution = cfg.MODEL.ROI_BOX_HEAD.POOLER_RESOLUTION
  620.         pooler_scales     = tuple(1.0 / input_shape[k].stride for k in in_features)
  621.         sampling_ratio    = cfg.MODEL.ROI_BOX_HEAD.POOLER_SAMPLING_RATIO
  622.         pooler_type       = cfg.MODEL.ROI_BOX_HEAD.POOLER_TYPE
  623.         # fmt: on
  624.  
  625.         # If StandardROIHeads is applied on multiple feature maps (as in FPN),
  626.         # then we share the same predictors and therefore the channel counts must be the same
  627.         in_channels = [input_shape[f].channels for f in in_features]
  628.         # Check all channel counts are equal
  629.         assert len(set(in_channels)) == 1, in_channels
  630.         in_channels = in_channels[0]
  631.  
  632.         box_pooler = ROIPooler(
  633.             output_size=pooler_resolution,
  634.             scales=pooler_scales,
  635.             sampling_ratio=sampling_ratio,
  636.             pooler_type=pooler_type,
  637.         )
  638.         # Here we split "box head" and "box predictor", which is mainly due to historical reasons.
  639.         # They are used together so the "box predictor" layers should be part of the "box head".
  640.         # New subclasses of ROIHeads do not need "box predictor"s.
  641.         box_head = build_box_head(
  642.             cfg, ShapeSpec(channels=in_channels, height=pooler_resolution, width=pooler_resolution)
  643.         )
  644.         box_predictor = FastRCNNOutputLayers(cfg, box_head.output_shape)
  645.         return {
  646.             "box_in_features": in_features,
  647.             "box_pooler": box_pooler,
  648.             "box_head": box_head,
  649.             "box_predictor": box_predictor,
  650.         }
  651.  
  652.     @classmethod
  653.     def _init_mask_head(cls, cfg, input_shape):
  654.         if not cfg.MODEL.MASK_ON:
  655.             return {}
  656.         # fmt: off
  657.         in_features       = cfg.MODEL.ROI_HEADS.IN_FEATURES
  658.         pooler_resolution = cfg.MODEL.ROI_MASK_HEAD.POOLER_RESOLUTION
  659.         pooler_scales     = tuple(1.0 / input_shape[k].stride for k in in_features)
  660.         sampling_ratio    = cfg.MODEL.ROI_MASK_HEAD.POOLER_SAMPLING_RATIO
  661.         pooler_type       = cfg.MODEL.ROI_MASK_HEAD.POOLER_TYPE
  662.         # fmt: on
  663.  
  664.         in_channels = [input_shape[f].channels for f in in_features][0]
  665.  
  666.         ret = {"mask_in_features": in_features}
  667.         ret["mask_pooler"] = (
  668.             ROIPooler(
  669.                 output_size=pooler_resolution,
  670.                 scales=pooler_scales,
  671.                 sampling_ratio=sampling_ratio,
  672.                 pooler_type=pooler_type,
  673.             )
  674.             if pooler_type
  675.             else None
  676.         )
  677.         if pooler_type:
  678.             shape = ShapeSpec(
  679.                 channels=in_channels, width=pooler_resolution, height=pooler_resolution
  680.             )
  681.         else:
  682.             shape = {f: input_shape[f] for f in in_features}
  683.         ret["mask_head"] = build_mask_head(cfg, shape)
  684.         return ret
  685.  
  686.     @classmethod
  687.     def _init_keypoint_head(cls, cfg, input_shape):
  688.         if not cfg.MODEL.KEYPOINT_ON:
  689.             return {}
  690.         # fmt: off
  691.         in_features       = cfg.MODEL.ROI_HEADS.IN_FEATURES
  692.         pooler_resolution = cfg.MODEL.ROI_KEYPOINT_HEAD.POOLER_RESOLUTION
  693.         pooler_scales     = tuple(1.0 / input_shape[k].stride for k in in_features)  # noqa
  694.         sampling_ratio    = cfg.MODEL.ROI_KEYPOINT_HEAD.POOLER_SAMPLING_RATIO
  695.         pooler_type       = cfg.MODEL.ROI_KEYPOINT_HEAD.POOLER_TYPE
  696.         # fmt: on
  697.  
  698.         in_channels = [input_shape[f].channels for f in in_features][0]
  699.  
  700.         ret = {"keypoint_in_features": in_features}
  701.         ret["keypoint_pooler"] = (
  702.             ROIPooler(
  703.                 output_size=pooler_resolution,
  704.                 scales=pooler_scales,
  705.                 sampling_ratio=sampling_ratio,
  706.                 pooler_type=pooler_type,
  707.             )
  708.             if pooler_type
  709.             else None
  710.         )
  711.         if pooler_type:
  712.             shape = ShapeSpec(
  713.                 channels=in_channels, width=pooler_resolution, height=pooler_resolution
  714.             )
  715.         else:
  716.             shape = {f: input_shape[f] for f in in_features}
  717.         ret["keypoint_head"] = build_keypoint_head(cfg, shape)
  718.         return ret
  719.        
  720.        
  721.     def _init_seg_head(self, cfg):
  722.         self.fpn_features_fused_level = cfg.MODEL.TEXTFUSENET_SEG_HEAD.FPN_FEATURES_FUSED_LEVEL
  723.         self.seg_head = build_seg_head(cfg)
  724.  
  725.     def _init_mutil_path_fuse_module(self, cfg):
  726.         self.mutil_path_fuse_module = build_mutil_path_fuse_module(cfg)
  727.  
  728.  
  729.     def forward(
  730.         self,
  731.         images: ImageList,
  732.         features: Dict[str, torch.Tensor],
  733.         proposals: List[Instances],
  734.         targets: Optional[List[Instances]] = None,
  735.     ) -> Tuple[List[Instances], Dict[str, torch.Tensor]]:
  736.         """
  737.        See :class:`ROIHeads.forward`.
  738.        """
  739.         del images
  740.         if self.training:
  741.             assert targets, "'targets' argument is required during training"
  742.             proposals = self.label_and_sample_proposals(proposals, targets)
  743.         del targets
  744.  
  745.         if self.training:
  746.             losses = self._forward_box(features, proposals)
  747.             # Usually the original proposals used by the box head are used by the mask, keypoint
  748.             # heads. But when `self.train_on_pred_boxes is True`, proposals will contain boxes
  749.             # predicted by the box head.
  750.             losses.update(self._forward_mask(features, proposals))
  751.             losses.update(self._forward_keypoint(features, proposals))
  752.             return proposals, losses
  753.         else:
  754.             pred_instances = self._forward_box(features, proposals)
  755.             # During inference cascaded prediction is used: the mask and keypoints heads are only
  756.             # applied to the top scoring box detections.
  757.             pred_instances = self.forward_with_given_boxes(features, pred_instances)
  758.             return pred_instances, {}
  759.  
  760.     def forward_with_given_boxes(
  761.         self, features: Dict[str, torch.Tensor], instances: List[Instances]
  762.     ) -> List[Instances]:
  763.         """
  764.        Use the given boxes in `instances` to produce other (non-box) per-ROI outputs.
  765.  
  766.        This is useful for downstream tasks where a box is known, but need to obtain
  767.        other attributes (outputs of other heads).
  768.        Test-time augmentation also uses this.
  769.  
  770.        Args:
  771.            features: same as in `forward()`
  772.            instances (list[Instances]): instances to predict other outputs. Expect the keys
  773.                "pred_boxes" and "pred_classes" to exist.
  774.  
  775.        Returns:
  776.            list[Instances]:
  777.                the same `Instances` objects, with extra
  778.                fields such as `pred_masks` or `pred_keypoints`.
  779.        """
  780.         assert not self.training
  781.         assert instances[0].has("pred_boxes") and instances[0].has("pred_classes")
  782.  
  783.         instances = self._forward_mask(features, instances)
  784.         instances = self._forward_keypoint(features, instances)
  785.         return instances
  786.  
  787.     def _forward_box(self, features: Dict[str, torch.Tensor], proposals: List[Instances]):
  788.         """
  789.        Forward logic of the box prediction branch. If `self.train_on_pred_boxes is True`,
  790.            the function puts predicted boxes in the `proposal_boxes` field of `proposals` argument.
  791.  
  792.        Args:
  793.            features (dict[str, Tensor]): mapping from feature map names to tensor.
  794.                Same as in :meth:`ROIHeads.forward`.
  795.            proposals (list[Instances]): the per-image object proposals with
  796.                their matching ground truth.
  797.                Each has fields "proposal_boxes", and "objectness_logits",
  798.                "gt_classes", "gt_boxes".
  799.  
  800.        Returns:
  801.            In training, a dict of losses.
  802.            In inference, a list of `Instances`, the predicted instances.
  803.        """
  804.         features = [features[f] for f in self.box_in_features]
  805.         box_features = self.box_pooler(features, [x.proposal_boxes for x in proposals])
  806.         box_features = self.box_head(box_features)
  807.         predictions = self.box_predictor(box_features)
  808.         del box_features
  809.  
  810.         if self.training:
  811.             losses = self.box_predictor.losses(predictions, proposals)
  812.             # proposals is modified in-place below, so losses must be computed first.
  813.             if self.train_on_pred_boxes:
  814.                 with torch.no_grad():
  815.                     pred_boxes = self.box_predictor.predict_boxes_for_gt_classes(
  816.                         predictions, proposals
  817.                     )
  818.                     for proposals_per_image, pred_boxes_per_image in zip(proposals, pred_boxes):
  819.                         proposals_per_image.proposal_boxes = Boxes(pred_boxes_per_image)
  820.             return losses
  821.         else:
  822.             pred_instances, _ = self.box_predictor.inference(predictions, proposals)
  823.             return pred_instances
  824.  
  825.  
  826.        
  827.     def _forward_mask(self, features, instances,targets=None):
  828.         """
  829.        Forward logic of the mask prediction branch.
  830.        Args:
  831.            features (list[Tensor]): #level input features for mask prediction
  832.            instances (list[Instances]): the per-image instances to train/predict masks.
  833.                In training, they can be the proposals.
  834.                In inference, they can be the predicted boxes.
  835.        Returns:
  836.            In training, a dict of losses.
  837.            In inference, update `instances` with new fields "pred_masks" and return it.
  838.        """
  839.         if not self.mask_on:
  840.             return {} if self.training else instances
  841.  
  842.         if self.training:
  843.             # The loss is only defined on positive proposals.
  844.             proposals, _ = select_foreground_proposals(instances, self.num_classes)
  845.             proposal_boxes = [x.proposal_boxes for x in proposals]
  846.             mask_features = self.mask_pooler(features, proposal_boxes)
  847.  
  848.             ###############################################  mutil_path_fuse  ###################################################
  849.             if self.mutil_path_fuse_on:
  850.                 image_shape = [proposals[0].image_size[1], proposals[0].image_size[0]]
  851.                 seg_logits, global_context = self.seg_head(features, self.fpn_features_fused_level, proposal_boxes, image_shape)
  852.                 mask_features = self.mutil_path_fuse_module(mask_features, global_context, proposals)
  853.             #####################################################################################################################
  854.  
  855.             mask_logits = self.mask_head(mask_features)
  856.  
  857.             if self.mutil_path_fuse_on:
  858.                 seg_loss = build_seg_head_loss()
  859.                 return {
  860.                         "loss_mask": mask_rcnn_loss(mask_logits, proposals),
  861.                         "loss_seg": seg_loss(seg_logits, targets),
  862.                 }
  863.             else:
  864.                 return {"loss_mask": mask_rcnn_loss(mask_logits, proposals)}
  865.         else:
  866.             pred_boxes = [x.pred_boxes for x in instances]
  867.             mask_features = self.mask_pooler(features, pred_boxes)
  868.  
  869.             ###############################################  mutil_path_fuse  ###################################################
  870.             if self.mutil_path_fuse_on:
  871.                 image_shape = [instances[0].image_size[1], instances[0].image_size[0]]
  872.                 seg_logits,global_context = self.seg_head(features, self.fpn_features_fused_level, pred_boxes, image_shape)
  873.                 mask_features = self.mutil_path_fuse_module(mask_features, global_context, instances)
  874.             ######################################################################################################################
  875.  
  876.             mask_logits = self.mask_head(mask_features)
  877.             mask_rcnn_inference(mask_logits, instances)
  878.             return instances
  879.            
  880.  
  881.     def _forward_keypoint(self, features: Dict[str, torch.Tensor], instances: List[Instances]):
  882.         """
  883.        Forward logic of the keypoint prediction branch.
  884.  
  885.        Args:
  886.            features (dict[str, Tensor]): mapping from feature map names to tensor.
  887.                Same as in :meth:`ROIHeads.forward`.
  888.            instances (list[Instances]): the per-image instances to train/predict keypoints.
  889.                In training, they can be the proposals.
  890.                In inference, they can be the boxes predicted by R-CNN box head.
  891.  
  892.        Returns:
  893.            In training, a dict of losses.
  894.            In inference, update `instances` with new fields "pred_keypoints" and return it.
  895.        """
  896.         if not self.keypoint_on:
  897.             return {} if self.training else instances
  898.  
  899.         if self.training:
  900.             # head is only trained on positive proposals with >=1 visible keypoints.
  901.             instances, _ = select_foreground_proposals(instances, self.num_classes)
  902.             instances = select_proposals_with_visible_keypoints(instances)
  903.  
  904.         if self.keypoint_pooler is not None:
  905.             features = [features[f] for f in self.keypoint_in_features]
  906.             boxes = [x.proposal_boxes if self.training else x.pred_boxes for x in instances]
  907.             features = self.keypoint_pooler(features, boxes)
  908.         else:
  909.             features = {f: features[f] for f in self.keypoint_in_features}
  910.         return self.keypoint_head(features, instances)
RAW Paste Data