def generate_detections( self, proposal_bboxes: Tensor, proposal_classes: Tensor, proposal_transformers: Tensor, image_width: int, image_height: int) -> Tuple[Tensor, Tensor, Tensor]: proposal_transformers = proposal_transformers.view( -1, self.num_classes, 4) mean = self._transformer_normalize_mean.repeat( 1, self.num_classes, 1) std = self._transformer_normalize_std.repeat( 1, self.num_classes, 1) proposal_transformers = proposal_transformers * std - mean proposal_bboxes = proposal_bboxes.view(-1, 1, 4).repeat( 1, self.num_classes, 1) detection_bboxes = BBox.apply_transformer( proposal_bboxes.view(-1, 4), proposal_transformers.view(-1, 4)) detection_bboxes = detection_bboxes.view(-1, self.num_classes, 4) detection_bboxes[:, :, [0, 2]] = detection_bboxes[:, :, [0, 2]].clamp( min=0, max=image_width) detection_bboxes[:, :, [1, 3]] = detection_bboxes[:, :, [1, 3]].clamp( min=0, max=image_height) proposal_probs = F.softmax(proposal_classes, dim=1) detection_bboxes = detection_bboxes.cpu() proposal_probs = proposal_probs.cpu() generated_bboxes = [] generated_classes = [] generated_probs = [] for c in range(1, self.num_classes): detection_class_bboxes = detection_bboxes[:, c, :] proposal_class_probs = proposal_probs[:, c] _, sorted_indices = proposal_class_probs.sort(descending=True) detection_class_bboxes = detection_class_bboxes[sorted_indices] proposal_class_probs = proposal_class_probs[sorted_indices] kept_indices = NMS.suppress(detection_class_bboxes.cuda(), threshold=0.3) detection_class_bboxes = detection_class_bboxes[kept_indices] proposal_class_probs = proposal_class_probs[kept_indices] generated_bboxes.append(detection_class_bboxes) generated_classes.append( torch.ones(len(kept_indices), dtype=torch.int) * c) generated_probs.append(proposal_class_probs) generated_bboxes = torch.cat(generated_bboxes, dim=0) generated_classes = torch.cat(generated_classes, dim=0) generated_probs = torch.cat(generated_probs, dim=0) return generated_bboxes, generated_classes, generated_probs
def generate_proposals(self, anchor_bboxes: Tensor, objectnesses: Tensor, transformers: Tensor, image_width: int, image_height: int) -> Tensor: proposal_score = objectnesses[:, 1] _, sorted_indices = torch.sort(proposal_score, dim=0, descending=True) sorted_transformers = transformers[sorted_indices] sorted_anchor_bboxes = anchor_bboxes[sorted_indices] proposal_bboxes = BBox.apply_transformer(sorted_anchor_bboxes, sorted_transformers.detach()) proposal_bboxes = BBox.clip(proposal_bboxes, 0, 0, image_width, image_height) proposal_bboxes = proposal_bboxes[:self._pre_nms_top_n] kept_indices = NMS.suppress(proposal_bboxes, threshold=0.7) proposal_bboxes = proposal_bboxes[kept_indices] proposal_bboxes = proposal_bboxes[:self._post_nms_top_n] return proposal_bboxes
def forward(self, features: Tensor, image_width: int, image_height: int) -> Tuple[Tensor, Tensor, Tensor, Tensor]: anchor_bboxes = RegionProposalNetwork._generate_anchors(image_width, image_height, num_x_anchors=features.shape[3], num_y_anchors=features.shape[2]).cuda() features = self._features(features) objectnesses = self._objectness(features) transformers = self._transformer(features) objectnesses = objectnesses.permute(0, 2, 3, 1).contiguous().view(-1, 2) transformers = transformers.permute(0, 2, 3, 1).contiguous().view(-1, 4) proposal_bboxes = RegionProposalNetwork._generate_proposals(anchor_bboxes, objectnesses, transformers, image_width, image_height) proposal_bboxes = proposal_bboxes[:12000 if self.training else 6000] keep_indices = NMS.suppress(proposal_bboxes, threshold=0.7) proposal_bboxes = proposal_bboxes[keep_indices] proposal_bboxes = proposal_bboxes[:2000 if self.training else 300] return anchor_bboxes, objectnesses, transformers, proposal_bboxes
def _run_nms(self, bboxes): start = time.time() keep_indices = NMS.suppress(bboxes.contiguous(), threshold=0.7) print('%s in %.3fs, %d -> %d' % (self.id(), time.time() - start, len(bboxes), len(keep_indices))) return keep_indices