| 4 | |
| 5 | # useful when there are huge number of gt boxes |
| 6 | class LoopMatcher(object): |
| 7 | def __init__( |
| 8 | self, thresholds: List[float], labels: List[int], allow_low_quality_matches: bool = False |
| 9 | ): |
| 10 | thresholds = thresholds[:] |
| 11 | assert thresholds[0] > 0 |
| 12 | thresholds.insert(0, -float("inf")) |
| 13 | thresholds.append(float("inf")) |
| 14 | assert all(low <= high for (low, high) in zip(thresholds[:-1], thresholds[1:])) |
| 15 | assert all(l in [-1, 0, 1] for l in labels) |
| 16 | assert len(labels) == len(thresholds) - 1 |
| 17 | |
| 18 | self.low_quality_thrshold = 0.3 |
| 19 | self.thresholds = thresholds |
| 20 | self.labels = labels |
| 21 | self.allow_low_quality_matches = allow_low_quality_matches |
| 22 | |
| 23 | |
| 24 | def _iou(self, boxes, box): |
| 25 | iw = torch.clamp(boxes[:, 2], max=box[2]) - torch.clamp(boxes[:, 0], min=box[0]) |
| 26 | ih = torch.clamp(boxes[:, 3], max=box[3]) - torch.clamp(boxes[:, 1], min=box[1]) |
| 27 | |
| 28 | inter = torch.clamp(iw, min=0) * torch.clamp(ih, min=0) |
| 29 | |
| 30 | areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) |
| 31 | area = (box[2] - box[0]) * (box[3] - box[1]) |
| 32 | |
| 33 | iou = inter / (areas + area - inter) |
| 34 | return iou |
| 35 | |
| 36 | def __call__(self, gt_boxes, anchors): |
| 37 | if len(gt_boxes) == 0: |
| 38 | default_matches = torch.zeros((len(anchors)), dtype=torch.int64).to(anchors.tensor.device) |
| 39 | default_match_labels = torch.zeros((len(anchors)), dtype=torch.int8).to(anchors.tensor.device) + self.labels[0] |
| 40 | return default_matches, default_match_labels |
| 41 | |
| 42 | gt_boxes_tensor = gt_boxes.tensor |
| 43 | anchors_tensor = anchors.tensor |
| 44 | |
| 45 | max_ious = torch.zeros((len(anchors))).to(anchors_tensor.device) |
| 46 | matched_inds = torch.zeros((len(anchors)), dtype=torch.long).to(anchors_tensor.device) |
| 47 | gt_ious = torch.zeros((len(gt_boxes))).to(anchors_tensor.device) |
| 48 | |
| 49 | for i in range(len(gt_boxes)): |
| 50 | ious = self._iou(anchors_tensor, gt_boxes_tensor[i]) |
| 51 | gt_ious[i] = ious.max() |
| 52 | matched_inds = torch.where(ious > max_ious, torch.zeros(1, dtype=torch.long, device=matched_inds.device)+i, matched_inds) |
| 53 | max_ious = torch.max(ious, max_ious) |
| 54 | del(ious) |
| 55 | |
| 56 | matched_vals = max_ious |
| 57 | matches = matched_inds |
| 58 | |
| 59 | match_labels = matches.new_full(matches.size(), 1, dtype=torch.int8) |
| 60 | |
| 61 | for (l, low, high) in zip(self.labels, self.thresholds[:-1], self.thresholds[1:]): |
| 62 | low_high = (matched_vals >= low) & (matched_vals < high) |
| 63 | match_labels[low_high] = l |