def create_instances(predictions, image_size, conf_threshold=0.5):
ret = Instances(image_size)
score = np.asarray([x["score"] for x in predictions])
chosen = (score > conf_threshold).nonzero()[0]
score = score[chosen]
bbox = np.asarray([predictions[i]["bbox"] for i in chosen]).reshape(-1, 4)
labels = np.asarray(
[dataset_id_map(predictions[i]["category_id"]) for i in chosen])
ret.scores = score
ret.pred_boxes = Boxes(
BoxMode.convert(bbox, BoxMode.XYWH_ABS, BoxMode.XYXY_ABS))
ret.pred_classes = labels
try:
ret.pred_masks = [predictions[i]["segmentation"] for i in chosen]
except KeyError:
pass
return ret
def detector_postprocess(results, output_height, output_width, mask_threshold=0.5):
if isinstance(output_width, torch.Tensor):
output_width_tmp = output_width.float()
else:
output_width_tmp = output_width
if isinstance(output_height, torch.Tensor):
output_height_tmp = output_height.float()
else:
output_height_tmp = output_height
scale_x, scale_y = (
output_width_tmp / results.image_size[1],
output_height_tmp / results.image_size[0]
)
results = Instances((output_height, output_width), **results.get_fields())
if results.has("pred_boxes"):
output_boxes = results.pred_boxes
elif results.has("proposal_boxes"):
output_boxes = results.proposal_boxes
output_boxes.scale(scale_x, scale_y)
output_boxes.clip(results.image_size)
results = results[output_boxes.nonempty()]
if results.has("pred_masks"):
results.pred_masks = retry_if_cuda_oom(paste_masks_in_image)(
results.pred_masks[:, 0, :, :],
results.pred_boxes,
results.image_size,
threshold=mask_threshold,
)
if results.has("pred_keypoints"):
results.pred_keypoints[:, :, 0] *= scale_x
results.pred_keypoints[:, :, 1] *= scale_y
return results