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]
