Spaces:
Runtime error
Runtime error
Last commit not found
| # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. | |
| import torch | |
| # from maskrcnn_benchmark.layers import smooth_l1_loss | |
| from maskrcnn_benchmark.modeling.matcher import Matcher | |
| from maskrcnn_benchmark.modeling.utils import cat | |
| from maskrcnn_benchmark.structures.boxlist_ops import boxlist_iou | |
| from torch.nn import functional as F | |
| def project_masks_on_boxes(segmentation_masks, proposals, discretization_size): | |
| """ | |
| Given segmentation masks and the bounding boxes corresponding | |
| to the location of the masks in the image, this function | |
| crops and resizes the masks in the position defined by the | |
| boxes. This prepares the masks for them to be fed to the | |
| loss computation as the targets. | |
| Arguments: | |
| segmentation_masks: an instance of SegmentationMask | |
| proposals: an instance of BoxList | |
| """ | |
| masks = [] | |
| M = discretization_size | |
| device = proposals.bbox.device | |
| proposals = proposals.convert("xyxy") | |
| assert segmentation_masks.size == proposals.size, "{}, {}".format( | |
| segmentation_masks, proposals | |
| ) | |
| # TODO put the proposals on the CPU, as the representation for the | |
| # masks is not efficient GPU-wise (possibly several small tensors for | |
| # representing a single instance mask) | |
| proposals = proposals.bbox.to(torch.device("cpu")) | |
| for segmentation_mask, proposal in zip(segmentation_masks, proposals): | |
| # crop the masks, resize them to the desired resolution and | |
| # then convert them to the tensor representation, | |
| # instead of the list representation that was used | |
| cropped_mask = segmentation_mask.crop(proposal) | |
| scaled_mask = cropped_mask.resize((M, M)) | |
| mask = scaled_mask.convert(mode="mask") | |
| masks.append(mask) | |
| if len(masks) == 0: | |
| return torch.empty(0, dtype=torch.float32, device=device) | |
| return torch.stack(masks, dim=0).to(device, dtype=torch.float32) | |
| class MaskRCNNLossComputation(object): | |
| def __init__(self, proposal_matcher, discretization_size): | |
| """ | |
| Arguments: | |
| proposal_matcher (Matcher) | |
| discretization_size (int) | |
| """ | |
| self.proposal_matcher = proposal_matcher | |
| self.discretization_size = discretization_size | |
| def match_targets_to_proposals(self, proposal, target): | |
| match_quality_matrix = boxlist_iou(target, proposal) | |
| matched_idxs = self.proposal_matcher(match_quality_matrix) | |
| # Mask RCNN needs "labels" and "masks "fields for creating the targets | |
| target = target.copy_with_fields(["labels", "masks"]) | |
| # get the targets corresponding GT for each proposal | |
| # NB: need to clamp the indices because we can have a single | |
| # GT in the image, and matched_idxs can be -2, which goes | |
| # out of bounds | |
| matched_targets = target[matched_idxs.clamp(min=0)] | |
| matched_targets.add_field("matched_idxs", matched_idxs) | |
| return matched_targets | |
| def prepare_targets(self, proposals, targets): | |
| labels = [] | |
| masks = [] | |
| for proposals_per_image, targets_per_image in zip(proposals, targets): | |
| matched_targets = self.match_targets_to_proposals( | |
| proposals_per_image, targets_per_image | |
| ) | |
| matched_idxs = matched_targets.get_field("matched_idxs") | |
| labels_per_image = matched_targets.get_field("labels") | |
| labels_per_image = labels_per_image.to(dtype=torch.int64) | |
| # this can probably be removed, but is left here for clarity | |
| # and completeness | |
| neg_inds = matched_idxs == Matcher.BELOW_LOW_THRESHOLD | |
| labels_per_image[neg_inds] = 0 | |
| # mask scores are only computed on positive samples | |
| positive_inds = torch.nonzero(labels_per_image > 0).squeeze(1) | |
| segmentation_masks = matched_targets.get_field("masks") | |
| segmentation_masks = segmentation_masks[positive_inds] | |
| positive_proposals = proposals_per_image[positive_inds] | |
| masks_per_image = project_masks_on_boxes( | |
| segmentation_masks, positive_proposals, self.discretization_size | |
| ) | |
| labels.append(labels_per_image) | |
| masks.append(masks_per_image) | |
| return labels, masks | |
| def __call__(self, proposals, mask_logits, targets): | |
| """ | |
| Arguments: | |
| proposals (list[BoxList]) | |
| mask_logits (Tensor) | |
| targets (list[BoxList]) | |
| Return: | |
| mask_loss (Tensor): scalar tensor containing the loss | |
| """ | |
| labels, mask_targets = self.prepare_targets(proposals, targets) | |
| labels = cat(labels, dim=0) | |
| mask_targets = cat(mask_targets, dim=0) | |
| positive_inds = torch.nonzero(labels > 0).squeeze(1) | |
| labels_pos = labels[positive_inds] | |
| # torch.mean (in binary_cross_entropy_with_logits) doesn't | |
| # accept empty tensors, so handle it separately | |
| if mask_targets.numel() == 0: | |
| return mask_logits.sum() * 0 | |
| mask_loss = F.binary_cross_entropy_with_logits( | |
| mask_logits[positive_inds, labels_pos], mask_targets | |
| ) | |
| return mask_loss | |
| class CharMaskRCNNLossComputation(object): | |
| def __init__(self, use_weighted_loss=False): | |
| """ | |
| Arguments: | |
| proposal_matcher (Matcher) | |
| discretization_size (int) | |
| """ | |
| self.use_weighted_loss = use_weighted_loss | |
| def __call__( | |
| self, | |
| proposals, | |
| mask_logits, | |
| char_mask_logits, | |
| mask_targets, | |
| char_mask_targets, | |
| char_mask_weights, | |
| ): | |
| """ | |
| Arguments: | |
| proposals (list[BoxList]) | |
| mask_logits (Tensor) | |
| targets (list[BoxList]) | |
| Return: | |
| mask_loss (Tensor): scalar tensor containing the loss | |
| """ | |
| mask_targets = cat(mask_targets, dim=0) | |
| char_mask_targets = cat(char_mask_targets, dim=0) | |
| char_mask_weights = cat(char_mask_weights, dim=0) | |
| char_mask_weights = char_mask_weights.mean(dim=0) | |
| # torch.mean (in binary_cross_entropy_with_logits) doesn't | |
| # accept empty tensors, so handle it separately | |
| if mask_targets.numel() == 0 or char_mask_targets.numel() == 0: | |
| return mask_logits.sum() * 0, char_mask_targets.sum() * 0 | |
| mask_loss = F.binary_cross_entropy_with_logits( | |
| mask_logits.squeeze(dim=1), mask_targets | |
| ) | |
| if self.use_weighted_loss: | |
| char_mask_loss = F.cross_entropy( | |
| char_mask_logits, char_mask_targets, char_mask_weights, ignore_index=-1 | |
| ) | |
| else: | |
| char_mask_loss = F.cross_entropy( | |
| char_mask_logits, char_mask_targets, ignore_index=-1 | |
| ) | |
| return mask_loss, char_mask_loss | |
| class SeqMaskRCNNLossComputation(object): | |
| def __init__(self): | |
| """ | |
| Arguments: | |
| proposal_matcher (Matcher) | |
| discretization_size (int) | |
| """ | |
| def __call__( | |
| self, | |
| proposals, | |
| mask_logits, | |
| mask_targets, | |
| ): | |
| """ | |
| Arguments: | |
| proposals (list[BoxList]) | |
| mask_logits (Tensor) | |
| targets (list[BoxList]) | |
| Return: | |
| mask_loss (Tensor): scalar tensor containing the loss | |
| """ | |
| mask_targets = cat(mask_targets, dim=0) | |
| # torch.mean (in binary_cross_entropy_with_logits) doesn't | |
| # accept empty tensors, so handle it separately | |
| if mask_targets.numel() == 0: | |
| return mask_logits.sum() * 0 | |
| mask_loss = F.binary_cross_entropy_with_logits( | |
| mask_logits.squeeze(dim=1), mask_targets | |
| ) | |
| return mask_loss | |
| def make_roi_mask_loss_evaluator(cfg): | |
| matcher = Matcher( | |
| cfg.MODEL.ROI_HEADS.FG_IOU_THRESHOLD, | |
| cfg.MODEL.ROI_HEADS.BG_IOU_THRESHOLD, | |
| allow_low_quality_matches=False, | |
| ) | |
| if cfg.MODEL.CHAR_MASK_ON: | |
| loss_evaluator = CharMaskRCNNLossComputation( | |
| use_weighted_loss=cfg.MODEL.ROI_MASK_HEAD.USE_WEIGHTED_CHAR_MASK | |
| ) | |
| else: | |
| if cfg.SEQUENCE.SEQ_ON: | |
| loss_evaluator = SeqMaskRCNNLossComputation() | |
| else: | |
| loss_evaluator = MaskRCNNLossComputation( | |
| matcher, cfg.MODEL.ROI_MASK_HEAD.RESOLUTION | |
| ) | |
| return loss_evaluator | |