coreml-community/ControlNet-v1-1-Annotators-cpu

1# Copyright (c) Facebook, Inc. and its affiliates.2from typing import List3import torch4 5from annotator.oneformer.detectron2.layers import nonzero_tuple6 7 8# TODO: the name is too general9class Matcher(object):10    """11    This class assigns to each predicted "element" (e.g., a box) a ground-truth12    element. Each predicted element will have exactly zero or one matches; each13    ground-truth element may be matched to zero or more predicted elements.14 15    The matching is determined by the MxN match_quality_matrix, that characterizes16    how well each (ground-truth, prediction)-pair match each other. For example,17    if the elements are boxes, this matrix may contain box intersection-over-union18    overlap values.19 20    The matcher returns (a) a vector of length N containing the index of the21    ground-truth element m in [0, M) that matches to prediction n in [0, N).22    (b) a vector of length N containing the labels for each prediction.23    """24 25    def __init__(26        self, thresholds: List[float], labels: List[int], allow_low_quality_matches: bool = False27    ):28        """29        Args:30            thresholds (list): a list of thresholds used to stratify predictions31                into levels.32            labels (list): a list of values to label predictions belonging at33                each level. A label can be one of {-1, 0, 1} signifying34                {ignore, negative class, positive class}, respectively.35            allow_low_quality_matches (bool): if True, produce additional matches36                for predictions with maximum match quality lower than high_threshold.37                See set_low_quality_matches_ for more details.38 39            For example,40                thresholds = [0.3, 0.5]41                labels = [0, -1, 1]42                All predictions with iou < 0.3 will be marked with 0 and43                thus will be considered as false positives while training.44                All predictions with 0.3 <= iou < 0.5 will be marked with -1 and45                thus will be ignored.46                All predictions with 0.5 <= iou will be marked with 1 and47                thus will be considered as true positives.48        """49        # Add -inf and +inf to first and last position in thresholds50        thresholds = thresholds[:]51        assert thresholds[0] > 052        thresholds.insert(0, -float("inf"))53        thresholds.append(float("inf"))54        # Currently torchscript does not support all + generator55        assert all([low <= high for (low, high) in zip(thresholds[:-1], thresholds[1:])])56        assert all([l in [-1, 0, 1] for l in labels])57        assert len(labels) == len(thresholds) - 158        self.thresholds = thresholds59        self.labels = labels60        self.allow_low_quality_matches = allow_low_quality_matches61 62    def __call__(self, match_quality_matrix):63        """64        Args:65            match_quality_matrix (Tensor[float]): an MxN tensor, containing the66                pairwise quality between M ground-truth elements and N predicted67                elements. All elements must be >= 0 (due to the us of `torch.nonzero`68                for selecting indices in :meth:`set_low_quality_matches_`).69 70        Returns:71            matches (Tensor[int64]): a vector of length N, where matches[i] is a matched72                ground-truth index in [0, M)73            match_labels (Tensor[int8]): a vector of length N, where pred_labels[i] indicates74                whether a prediction is a true or false positive or ignored75        """76        assert match_quality_matrix.dim() == 277        if match_quality_matrix.numel() == 0:78            default_matches = match_quality_matrix.new_full(79                (match_quality_matrix.size(1),), 0, dtype=torch.int6480            )81            # When no gt boxes exist, we define IOU = 0 and therefore set labels82            # to `self.labels[0]`, which usually defaults to background class 083            # To choose to ignore instead, can make labels=[-1,0,-1,1] + set appropriate thresholds84            default_match_labels = match_quality_matrix.new_full(85                (match_quality_matrix.size(1),), self.labels[0], dtype=torch.int886            )87            return default_matches, default_match_labels88 89        assert torch.all(match_quality_matrix >= 0)90 91        # match_quality_matrix is M (gt) x N (predicted)92        # Max over gt elements (dim 0) to find best gt candidate for each prediction93        matched_vals, matches = match_quality_matrix.max(dim=0)94 95        match_labels = matches.new_full(matches.size(), 1, dtype=torch.int8)96 97        for (l, low, high) in zip(self.labels, self.thresholds[:-1], self.thresholds[1:]):98            low_high = (matched_vals >= low) & (matched_vals < high)99            match_labels[low_high] = l100 101        if self.allow_low_quality_matches:102            self.set_low_quality_matches_(match_labels, match_quality_matrix)103 104        return matches, match_labels105 106    def set_low_quality_matches_(self, match_labels, match_quality_matrix):107        """108        Produce additional matches for predictions that have only low-quality matches.109        Specifically, for each ground-truth G find the set of predictions that have110        maximum overlap with it (including ties); for each prediction in that set, if111        it is unmatched, then match it to the ground-truth G.112 113        This function implements the RPN assignment case (i) in Sec. 3.1.2 of114        :paper:`Faster R-CNN`.115        """116        # For each gt, find the prediction with which it has highest quality117        highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1)118        # Find the highest quality match available, even if it is low, including ties.119        # Note that the matches qualities must be positive due to the use of120        # `torch.nonzero`.121        _, pred_inds_with_highest_quality = nonzero_tuple(122            match_quality_matrix == highest_quality_foreach_gt[:, None]123        )124        # If an anchor was labeled positive only due to a low-quality match125        # with gt_A, but it has larger overlap with gt_B, it's matched index will still be gt_B.126        # This follows the implementation in Detectron, and is found to have no significant impact.127        match_labels[pred_inds_with_highest_quality] = 1128