Mjolnir65/FasterRCNN

1import math2from collections import OrderedDict3from typing import Dict, List, Optional, Tuple4 5import torch6from torch import nn, Tensor7from torch.nn import functional as F8from torchvision.ops import complete_box_iou_loss, distance_box_iou_loss, FrozenBatchNorm2d, generalized_box_iou_loss9 10 11class BalancedPositiveNegativeSampler:12    """13    This class samples batches, ensuring that they contain a fixed proportion of positives14    """15 16    def __init__(self, batch_size_per_image: int, positive_fraction: float) -> None:17        """18        Args:19            batch_size_per_image (int): number of elements to be selected per image20            positive_fraction (float): percentage of positive elements per batch21        """22        self.batch_size_per_image = batch_size_per_image23        self.positive_fraction = positive_fraction24 25    def __call__(self, matched_idxs: List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]:26        """27        Args:28            matched_idxs: list of tensors containing -1, 0 or positive values.29                Each tensor corresponds to a specific image.30                -1 values are ignored, 0 are considered as negatives and > 0 as31                positives.32 33        Returns:34            pos_idx (list[tensor])35            neg_idx (list[tensor])36 37        Returns two lists of binary masks for each image.38        The first list contains the positive elements that were selected,39        and the second list the negative example.40        """41        pos_idx = []42        neg_idx = []43        for matched_idxs_per_image in matched_idxs:44            positive = torch.where(matched_idxs_per_image >= 1)[0]45            negative = torch.where(matched_idxs_per_image == 0)[0]46 47            num_pos = int(self.batch_size_per_image * self.positive_fraction)48            # protect against not enough positive examples49            num_pos = min(positive.numel(), num_pos)50            num_neg = self.batch_size_per_image - num_pos51            # protect against not enough negative examples52            num_neg = min(negative.numel(), num_neg)53 54            # randomly select positive and negative examples55            perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos]56            perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg]57 58            pos_idx_per_image = positive[perm1]59            neg_idx_per_image = negative[perm2]60 61            # create binary mask from indices62            pos_idx_per_image_mask = torch.zeros_like(matched_idxs_per_image, dtype=torch.uint8)63            neg_idx_per_image_mask = torch.zeros_like(matched_idxs_per_image, dtype=torch.uint8)64 65            pos_idx_per_image_mask[pos_idx_per_image] = 166            neg_idx_per_image_mask[neg_idx_per_image] = 167 68            pos_idx.append(pos_idx_per_image_mask)69            neg_idx.append(neg_idx_per_image_mask)70 71        return pos_idx, neg_idx72 73 74@torch.jit._script_if_tracing75def encode_boxes(reference_boxes: Tensor, proposals: Tensor, weights: Tensor) -> Tensor:76    """77    Encode a set of proposals with respect to some78    reference boxes79 80    Args:81        reference_boxes (Tensor): reference boxes82        proposals (Tensor): boxes to be encoded83        weights (Tensor[4]): the weights for ``(x, y, w, h)``84    """85 86    # perform some unpacking to make it JIT-fusion friendly87    wx = weights[0]88    wy = weights[1]89    ww = weights[2]90    wh = weights[3]91 92    proposals_x1 = proposals[:, 0].unsqueeze(1)93    proposals_y1 = proposals[:, 1].unsqueeze(1)94    proposals_x2 = proposals[:, 2].unsqueeze(1)95    proposals_y2 = proposals[:, 3].unsqueeze(1)96 97    reference_boxes_x1 = reference_boxes[:, 0].unsqueeze(1)98    reference_boxes_y1 = reference_boxes[:, 1].unsqueeze(1)99    reference_boxes_x2 = reference_boxes[:, 2].unsqueeze(1)100    reference_boxes_y2 = reference_boxes[:, 3].unsqueeze(1)101 102    # implementation starts here103    ex_widths = proposals_x2 - proposals_x1104    ex_heights = proposals_y2 - proposals_y1105    ex_ctr_x = proposals_x1 + 0.5 * ex_widths106    ex_ctr_y = proposals_y1 + 0.5 * ex_heights107 108    gt_widths = reference_boxes_x2 - reference_boxes_x1109    gt_heights = reference_boxes_y2 - reference_boxes_y1110    gt_ctr_x = reference_boxes_x1 + 0.5 * gt_widths111    gt_ctr_y = reference_boxes_y1 + 0.5 * gt_heights112 113    targets_dx = wx * (gt_ctr_x - ex_ctr_x) / ex_widths114    targets_dy = wy * (gt_ctr_y - ex_ctr_y) / ex_heights115    targets_dw = ww * torch.log(gt_widths / ex_widths)116    targets_dh = wh * torch.log(gt_heights / ex_heights)117 118    targets = torch.cat((targets_dx, targets_dy, targets_dw, targets_dh), dim=1)119    return targets120 121 122class BoxCoder:123    """124    This class encodes and decodes a set of bounding boxes into125    the representation used for training the regressors.126    """127 128    def __init__(129        self, weights: Tuple[float, float, float, float], bbox_xform_clip: float = math.log(1000.0 / 16)130    ) -> None:131        """132        Args:133            weights (4-element tuple)134            bbox_xform_clip (float)135        """136        self.weights = weights137        self.bbox_xform_clip = bbox_xform_clip138 139    def encode(self, reference_boxes: List[Tensor], proposals: List[Tensor]) -> List[Tensor]:140        boxes_per_image = [len(b) for b in reference_boxes]141        reference_boxes = torch.cat(reference_boxes, dim=0)142        proposals = torch.cat(proposals, dim=0)143        targets = self.encode_single(reference_boxes, proposals)144        return targets.split(boxes_per_image, 0)145 146    def encode_single(self, reference_boxes: Tensor, proposals: Tensor) -> Tensor:147        """148        Encode a set of proposals with respect to some149        reference boxes150 151        Args:152            reference_boxes (Tensor): reference boxes153            proposals (Tensor): boxes to be encoded154        """155        dtype = reference_boxes.dtype156        device = reference_boxes.device157        weights = torch.as_tensor(self.weights, dtype=dtype, device=device)158        targets = encode_boxes(reference_boxes, proposals, weights)159 160        return targets161 162    def decode(self, rel_codes: Tensor, boxes: List[Tensor]) -> Tensor:163        torch._assert(164            isinstance(boxes, (list, tuple)),165            "This function expects boxes of type list or tuple.",166        )167        torch._assert(168            isinstance(rel_codes, torch.Tensor),169            "This function expects rel_codes of type torch.Tensor.",170        )171        boxes_per_image = [b.size(0) for b in boxes]172        concat_boxes = torch.cat(boxes, dim=0)173        box_sum = 0174        for val in boxes_per_image:175            box_sum += val176        if box_sum > 0:177            rel_codes = rel_codes.reshape(box_sum, -1)178        pred_boxes = self.decode_single(rel_codes, concat_boxes)179        if box_sum > 0:180            pred_boxes = pred_boxes.reshape(box_sum, -1, 4)181        return pred_boxes182 183    def decode_single(self, rel_codes: Tensor, boxes: Tensor) -> Tensor:184        """185        From a set of original boxes and encoded relative box offsets,186        get the decoded boxes.187 188        Args:189            rel_codes (Tensor): encoded boxes190            boxes (Tensor): reference boxes.191        """192 193        boxes = boxes.to(rel_codes.dtype)194 195        widths = boxes[:, 2] - boxes[:, 0]196        heights = boxes[:, 3] - boxes[:, 1]197        ctr_x = boxes[:, 0] + 0.5 * widths198        ctr_y = boxes[:, 1] + 0.5 * heights199 200        wx, wy, ww, wh = self.weights201        dx = rel_codes[:, 0::4] / wx202        dy = rel_codes[:, 1::4] / wy203        dw = rel_codes[:, 2::4] / ww204        dh = rel_codes[:, 3::4] / wh205 206        # Prevent sending too large values into torch.exp()207        dw = torch.clamp(dw, max=self.bbox_xform_clip)208        dh = torch.clamp(dh, max=self.bbox_xform_clip)209 210        pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]211        pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]212        pred_w = torch.exp(dw) * widths[:, None]213        pred_h = torch.exp(dh) * heights[:, None]214 215        # Distance from center to box's corner.216        c_to_c_h = torch.tensor(0.5, dtype=pred_ctr_y.dtype, device=pred_h.device) * pred_h217        c_to_c_w = torch.tensor(0.5, dtype=pred_ctr_x.dtype, device=pred_w.device) * pred_w218 219        pred_boxes1 = pred_ctr_x - c_to_c_w220        pred_boxes2 = pred_ctr_y - c_to_c_h221        pred_boxes3 = pred_ctr_x + c_to_c_w222        pred_boxes4 = pred_ctr_y + c_to_c_h223        pred_boxes = torch.stack((pred_boxes1, pred_boxes2, pred_boxes3, pred_boxes4), dim=2).flatten(1)224        return pred_boxes225 226 227class BoxLinearCoder:228    """229    The linear box-to-box transform defined in FCOS. The transformation is parameterized230    by the distance from the center of (square) src box to 4 edges of the target box.231    """232 233    def __init__(self, normalize_by_size: bool = True) -> None:234        """235        Args:236            normalize_by_size (bool): normalize deltas by the size of src (anchor) boxes.237        """238        self.normalize_by_size = normalize_by_size239 240    def encode(self, reference_boxes: Tensor, proposals: Tensor) -> Tensor:241        """242        Encode a set of proposals with respect to some reference boxes243 244        Args:245            reference_boxes (Tensor): reference boxes246            proposals (Tensor): boxes to be encoded247 248        Returns:249            Tensor: the encoded relative box offsets that can be used to250            decode the boxes.251 252        """253 254        # get the center of reference_boxes255        reference_boxes_ctr_x = 0.5 * (reference_boxes[..., 0] + reference_boxes[..., 2])256        reference_boxes_ctr_y = 0.5 * (reference_boxes[..., 1] + reference_boxes[..., 3])257 258        # get box regression transformation deltas259        target_l = reference_boxes_ctr_x - proposals[..., 0]260        target_t = reference_boxes_ctr_y - proposals[..., 1]261        target_r = proposals[..., 2] - reference_boxes_ctr_x262        target_b = proposals[..., 3] - reference_boxes_ctr_y263 264        targets = torch.stack((target_l, target_t, target_r, target_b), dim=-1)265 266        if self.normalize_by_size:267            reference_boxes_w = reference_boxes[..., 2] - reference_boxes[..., 0]268            reference_boxes_h = reference_boxes[..., 3] - reference_boxes[..., 1]269            reference_boxes_size = torch.stack(270                (reference_boxes_w, reference_boxes_h, reference_boxes_w, reference_boxes_h), dim=-1271            )272            targets = targets / reference_boxes_size273        return targets274 275    def decode(self, rel_codes: Tensor, boxes: Tensor) -> Tensor:276 277        """278        From a set of original boxes and encoded relative box offsets,279        get the decoded boxes.280 281        Args:282            rel_codes (Tensor): encoded boxes283            boxes (Tensor): reference boxes.284 285        Returns:286            Tensor: the predicted boxes with the encoded relative box offsets.287 288        .. note::289            This method assumes that ``rel_codes`` and ``boxes`` have same size for 0th dimension. i.e. ``len(rel_codes) == len(boxes)``.290 291        """292 293        boxes = boxes.to(dtype=rel_codes.dtype)294 295        ctr_x = 0.5 * (boxes[..., 0] + boxes[..., 2])296        ctr_y = 0.5 * (boxes[..., 1] + boxes[..., 3])297 298        if self.normalize_by_size:299            boxes_w = boxes[..., 2] - boxes[..., 0]300            boxes_h = boxes[..., 3] - boxes[..., 1]301 302            list_box_size = torch.stack((boxes_w, boxes_h, boxes_w, boxes_h), dim=-1)303            rel_codes = rel_codes * list_box_size304 305        pred_boxes1 = ctr_x - rel_codes[..., 0]306        pred_boxes2 = ctr_y - rel_codes[..., 1]307        pred_boxes3 = ctr_x + rel_codes[..., 2]308        pred_boxes4 = ctr_y + rel_codes[..., 3]309 310        pred_boxes = torch.stack((pred_boxes1, pred_boxes2, pred_boxes3, pred_boxes4), dim=-1)311        return pred_boxes312 313 314class Matcher:315    """316    This class assigns to each predicted "element" (e.g., a box) a ground-truth317    element. Each predicted element will have exactly zero or one matches; each318    ground-truth element may be assigned to zero or more predicted elements.319 320    Matching is based on the MxN match_quality_matrix, that characterizes how well321    each (ground-truth, predicted)-pair match. For example, if the elements are322    boxes, the matrix may contain box IoU overlap values.323 324    The matcher returns a tensor of size N containing the index of the ground-truth325    element m that matches to prediction n. If there is no match, a negative value326    is returned.327    """328 329    BELOW_LOW_THRESHOLD = -1330    BETWEEN_THRESHOLDS = -2331 332    __annotations__ = {333        "BELOW_LOW_THRESHOLD": int,334        "BETWEEN_THRESHOLDS": int,335    }336 337    def __init__(self, high_threshold: float, low_threshold: float, allow_low_quality_matches: bool = False) -> None:338        """339        Args:340            high_threshold (float): quality values greater than or equal to341                this value are candidate matches.342            low_threshold (float): a lower quality threshold used to stratify343                matches into three levels:344                1) matches >= high_threshold345                2) BETWEEN_THRESHOLDS matches in [low_threshold, high_threshold)346                3) BELOW_LOW_THRESHOLD matches in [0, low_threshold)347            allow_low_quality_matches (bool): if True, produce additional matches348                for predictions that have only low-quality match candidates. See349                set_low_quality_matches_ for more details.350        """351        self.BELOW_LOW_THRESHOLD = -1352        self.BETWEEN_THRESHOLDS = -2353        torch._assert(low_threshold <= high_threshold, "low_threshold should be <= high_threshold")354        self.high_threshold = high_threshold355        self.low_threshold = low_threshold356        self.allow_low_quality_matches = allow_low_quality_matches357 358    def __call__(self, match_quality_matrix: Tensor) -> Tensor:359        """360        Args:361            match_quality_matrix (Tensor[float]): an MxN tensor, containing the362            pairwise quality between M ground-truth elements and N predicted elements.363 364        Returns:365            matches (Tensor[int64]): an N tensor where N[i] is a matched gt in366            [0, M - 1] or a negative value indicating that prediction i could not367            be matched.368        """369        if match_quality_matrix.numel() == 0:370            # empty targets or proposals not supported during training371            if match_quality_matrix.shape[0] == 0:372                raise ValueError("No ground-truth boxes available for one of the images during training")373            else:374                raise ValueError("No proposal boxes available for one of the images during training")375 376        # match_quality_matrix is M (gt) x N (predicted)377        # Max over gt elements (dim 0) to find best gt candidate for each prediction378        matched_vals, matches = match_quality_matrix.max(dim=0)379        if self.allow_low_quality_matches:380            all_matches = matches.clone()381        else:382            all_matches = None  # type: ignore[assignment]383 384        # Assign candidate matches with low quality to negative (unassigned) values385        below_low_threshold = matched_vals < self.low_threshold386        between_thresholds = (matched_vals >= self.low_threshold) & (matched_vals < self.high_threshold)387        matches[below_low_threshold] = self.BELOW_LOW_THRESHOLD388        matches[between_thresholds] = self.BETWEEN_THRESHOLDS389 390        if self.allow_low_quality_matches:391            if all_matches is None:392                torch._assert(False, "all_matches should not be None")393            else:394                self.set_low_quality_matches_(matches, all_matches, match_quality_matrix)395 396        return matches397 398    def set_low_quality_matches_(self, matches: Tensor, all_matches: Tensor, match_quality_matrix: Tensor) -> None:399        """400        Produce additional matches for predictions that have only low-quality matches.401        Specifically, for each ground-truth find the set of predictions that have402        maximum overlap with it (including ties); for each prediction in that set, if403        it is unmatched, then match it to the ground-truth with which it has the highest404        quality value.405        """406        # For each gt, find the prediction with which it has the highest quality407        highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1)408        # Find the highest quality match available, even if it is low, including ties409        gt_pred_pairs_of_highest_quality = torch.where(match_quality_matrix == highest_quality_foreach_gt[:, None])410        # Example gt_pred_pairs_of_highest_quality:411        # (tensor([0, 1, 1, 2, 2, 3, 3, 4, 5, 5]),412        #  tensor([39796, 32055, 32070, 39190, 40255, 40390, 41455, 45470, 45325, 46390]))413        # Each element in the first tensor is a gt index, and each element in second tensor is a prediction index414        # Note how gt items 1, 2, 3, and 5 each have two ties415 416        pred_inds_to_update = gt_pred_pairs_of_highest_quality[1]417        matches[pred_inds_to_update] = all_matches[pred_inds_to_update]418 419 420class SSDMatcher(Matcher):421    def __init__(self, threshold: float) -> None:422        super().__init__(threshold, threshold, allow_low_quality_matches=False)423 424    def __call__(self, match_quality_matrix: Tensor) -> Tensor:425        matches = super().__call__(match_quality_matrix)426 427        # For each gt, find the prediction with which it has the highest quality428        _, highest_quality_pred_foreach_gt = match_quality_matrix.max(dim=1)429        matches[highest_quality_pred_foreach_gt] = torch.arange(430            highest_quality_pred_foreach_gt.size(0), dtype=torch.int64, device=highest_quality_pred_foreach_gt.device431        )432 433        return matches434 435 436def overwrite_eps(model: nn.Module, eps: float) -> None:437    """438    This method overwrites the default eps values of all the439    FrozenBatchNorm2d layers of the model with the provided value.440    This is necessary to address the BC-breaking change introduced441    by the bug-fix at pytorch/vision#2933. The overwrite is applied442    only when the pretrained weights are loaded to maintain compatibility443    with previous versions.444 445    Args:446        model (nn.Module): The model on which we perform the overwrite.447        eps (float): The new value of eps.448    """449    for module in model.modules():450        if isinstance(module, FrozenBatchNorm2d):451            module.eps = eps452 453 454def retrieve_out_channels(model: nn.Module, size: Tuple[int, int]) -> List[int]:455    """456    This method retrieves the number of output channels of a specific model.457 458    Args:459        model (nn.Module): The model for which we estimate the out_channels.460            It should return a single Tensor or an OrderedDict[Tensor].461        size (Tuple[int, int]): The size (wxh) of the input.462 463    Returns:464        out_channels (List[int]): A list of the output channels of the model.465    """466    in_training = model.training467    model.eval()468 469    with torch.no_grad():470        # Use dummy data to retrieve the feature map sizes to avoid hard-coding their values471        device = next(model.parameters()).device472        tmp_img = torch.zeros((1, 3, size[1], size[0]), device=device)473        features = model(tmp_img)474        if isinstance(features, torch.Tensor):475            features = OrderedDict([("0", features)])476        out_channels = [x.size(1) for x in features.values()]477 478    if in_training:479        model.train()480 481    return out_channels482 483 484@torch.jit.unused485def _fake_cast_onnx(v: Tensor) -> int:486    return v  # type: ignore[return-value]487 488 489def _topk_min(input: Tensor, orig_kval: int, axis: int) -> int:490    """491    ONNX spec requires the k-value to be less than or equal to the number of inputs along492    provided dim. Certain models use the number of elements along a particular axis instead of K493    if K exceeds the number of elements along that axis. Previously, python's min() function was494    used to determine whether to use the provided k-value or the specified dim axis value.495 496    However, in cases where the model is being exported in tracing mode, python min() is497    static causing the model to be traced incorrectly and eventually fail at the topk node.498    In order to avoid this situation, in tracing mode, torch.min() is used instead.499 500    Args:501        input (Tensor): The original input tensor.502        orig_kval (int): The provided k-value.503        axis(int): Axis along which we retrieve the input size.504 505    Returns:506        min_kval (int): Appropriately selected k-value.507    """508    if not torch.jit.is_tracing():509        return min(orig_kval, input.size(axis))510    axis_dim_val = torch._shape_as_tensor(input)[axis].unsqueeze(0)511    min_kval = torch.min(torch.cat((torch.tensor([orig_kval], dtype=axis_dim_val.dtype), axis_dim_val), 0))512    return _fake_cast_onnx(min_kval)513 514 515def _box_loss(516    type: str,517    box_coder: BoxCoder,518    anchors_per_image: Tensor,519    matched_gt_boxes_per_image: Tensor,520    bbox_regression_per_image: Tensor,521    cnf: Optional[Dict[str, float]] = None,522) -> Tensor:523    torch._assert(type in ["l1", "smooth_l1", "ciou", "diou", "giou"], f"Unsupported loss: {type}")524 525    if type == "l1":526        target_regression = box_coder.encode_single(matched_gt_boxes_per_image, anchors_per_image)527        return F.l1_loss(bbox_regression_per_image, target_regression, reduction="sum")528    elif type == "smooth_l1":529        target_regression = box_coder.encode_single(matched_gt_boxes_per_image, anchors_per_image)530        beta = cnf["beta"] if cnf is not None and "beta" in cnf else 1.0531        return F.smooth_l1_loss(bbox_regression_per_image, target_regression, reduction="sum", beta=beta)532    else:533        bbox_per_image = box_coder.decode_single(bbox_regression_per_image, anchors_per_image)534        eps = cnf["eps"] if cnf is not None and "eps" in cnf else 1e-7535        if type == "ciou":536            return complete_box_iou_loss(bbox_per_image, matched_gt_boxes_per_image, reduction="sum", eps=eps)537        if type == "diou":538            return distance_box_iou_loss(bbox_per_image, matched_gt_boxes_per_image, reduction="sum", eps=eps)539        # otherwise giou540        return generalized_box_iou_loss(bbox_per_image, matched_gt_boxes_per_image, reduction="sum", eps=eps)541