def fast_rcnn_inference_single_image(boxes, scores, image_shape, score_thresh, nms_thresh, topk_per_image, class_logits=None, estimate_uncertainty=False, variance=torch.Tensor([])): """ Single-image inference. Return bounding-box detection results by thresholding on scores and applying non-maximum suppression (NMS). Args: Same as `fast_rcnn_inference`, but with boxes, scores, and image shapes per image. Returns: Same as `fast_rcnn_inference`, but for only one image. """ valid_mask = torch.isfinite(boxes).all(dim=1) & torch.isfinite(scores).all( dim=1) if not valid_mask.all(): boxes = boxes[valid_mask] scores = scores[valid_mask] scores = scores[:, :-1] num_bbox_reg_classes = boxes.shape[1] // 4 # Convert to Boxes to use the `clip` function ... boxes = Boxes(boxes.reshape(-1, 4)) boxes.clip(image_shape) boxes = boxes.tensor.view(-1, num_bbox_reg_classes, 4) # R x C x 4 # Filter results based on detection scores filter_mask = scores > score_thresh # R x K # R' x 2. First column contains indices of the R predictions; # Second column contains indices of classes. # Get box ID with predicted class label: [box id, class label] filter_inds = filter_mask.nonzero() import numpy as np class_id = np.argmax(scores.cpu().numpy(), axis=1) class_id = np.array([np.arange(1000), class_id]) class_id = np.swapaxes(class_id, 1, 0) boxes_one_class = boxes[class_id[:, 0], class_id[:, 1], :].cpu().numpy() scores_one_class = np.max(scores.cpu().numpy(), axis=1) if not class_logits == None: class_logits = class_logits[filter_inds[:, 0]] predicted_probs = scores[filter_inds[:, 0]] if num_bbox_reg_classes == 1: boxes = boxes[filter_inds[:, 0], 0] else: boxes = boxes[filter_mask] scores_filtered = scores[filter_mask] # Apply per-class NMS keep = batched_nms(boxes, scores_filtered, filter_inds[:, 1], nms_thresh) if topk_per_image >= 0: keep = keep[:topk_per_image] boxes_final, scores_final, filter_inds_final = boxes[ keep], scores_filtered[keep], filter_inds[keep] result = Instances(image_shape) result.pred_boxes = Boxes(boxes_final) result.scores = scores_final result.pred_classes = filter_inds_final[:, 1] # Jamie # Save out logits if not class_logits == None: #result.class_logits = class_logits[filter_inds_final[:,0]] result.class_logits = class_logits[keep] result.prob_score = predicted_probs[keep] #class_logits = class_logits[filter_inds_final[:,0]] #result.class_logits = class_logits[keep] if estimate_uncertainty: # std from 1000 proposals #stds = nms_calc_uncertainty(boxes_final.cpu().numpy(), scores_final.cpu().numpy(), boxes_one_class, scores_one_class, 0.75) # std from bbox with class confidence score higher than threshold stds = nms_calc_uncertainty(boxes_final.cpu().numpy(), scores_final.cpu().numpy(), boxes.cpu().numpy(), scores_filtered.cpu().numpy(), 0.9) result.stds = torch.Tensor(stds).cuda() if len(variance) > 0: result.vars = variance[keep] return result, filter_inds_final[:, 0]
def fast_rcnn_inference_single_image_with_overlap( boxes, scores, overlap_boxes, overlap_probs, image_shape, score_thresh, nms_thresh, topk_per_image, allow_oob=False, ): """ Single-image inference. Return bounding-box detection results by thresholding on scores and applying non-maximum suppression (NMS). Args: Same as `fast_rcnn_inference`, but with boxes, scores, and image shapes per image. Returns: Same as `fast_rcnn_inference`, but for only one image. """ valid_mask = torch.isfinite(boxes).all(dim=1) & torch.isfinite(scores).all( dim=1) if not valid_mask.all(): boxes = boxes[valid_mask] scores = scores[valid_mask] overlap_boxes = overlap_boxes[valid_mask] overlap_probs = overlap_probs[valid_mask] scores = scores[:, :-1] num_bbox_reg_classes = boxes.shape[1] // 4 # Convert to Boxes to use the `clip` function ... if not allow_oob: boxes = Boxes(boxes.reshape(-1, 4)) boxes.clip(image_shape) boxes = boxes.tensor.view(-1, num_bbox_reg_classes, 4) # R x C x 4 assert (overlap_boxes.size(1) == 4 ), "overlap boxes prediction has no category, but: {}".format( overlap_boxes.size()) overlap_boxes = Boxes(overlap_boxes) overlap_boxes.clip(image_shape) overlap_boxes = overlap_boxes.tensor else: boxes = boxes.view(-1, num_bbox_reg_classes, 4) # Filter results based on detection scores filter_mask = scores > score_thresh # R x K # R' x 2. First column contains indices of the R predictions; # Second column contains indices of classes. filter_inds = filter_mask.nonzero() if num_bbox_reg_classes == 1: boxes = boxes[filter_inds[:, 0], 0] overlap_boxes = overlap_boxes[filter_inds[:, 0]] else: boxes = boxes[filter_mask] overlap_boxes = overlap_boxes[filter_inds[:, 0]] scores = scores[filter_mask] overlap_probs = overlap_probs[filter_mask] # Apply per-class NMS self_defined_nms_on = True # False if self_defined_nms_on: boxes = np.ascontiguousarray(boxes.cpu()) scores = np.ascontiguousarray(scores.cpu()) overlap_probs = np.ascontiguousarray(overlap_probs.cpu()) overlap_boxes = np.ascontiguousarray(overlap_boxes.cpu()) keep = batched_noh_nms(boxes, scores, overlap_probs, overlap_boxes, Nt=nms_thresh, thresh=0.01, method=3) boxes = torch.from_numpy(boxes).cuda() scores = torch.from_numpy(scores).cuda() overlap_probs = torch.from_numpy(overlap_probs).cuda() overlap_boxes = torch.from_numpy(overlap_boxes).cuda() keep = keep[scores[keep].argsort(descending=True)] else: from torchvision.ops import nms keep = nms(boxes, scores, nms_thresh) if topk_per_image >= 0: keep = keep[:topk_per_image] boxes, scores, overlap_boxes, overlap_probs, filter_inds = ( boxes[keep], scores[keep], overlap_boxes[keep], overlap_probs[keep], filter_inds[keep], ) result = Instances(image_shape) result.pred_boxes = Boxes(boxes) result.scores = scores result.overlap_boxes = Boxes(overlap_boxes) result.overlap_probs = overlap_probs result.pred_classes = filter_inds[:, 1] return result, filter_inds[:, 0]