def convert_raw_prediction(raw_pred: Sequence[np.ndarray], record: BaseRecord,
detection_threshold: float):
raw_bboxes, raw_masks = raw_pred
scores, labels, bboxes = _unpack_raw_bboxes(raw_bboxes)
keep_mask = scores > detection_threshold
keep_scores = scores[keep_mask]
keep_labels = labels[keep_mask]
keep_bboxes = [BBox.from_xyxy(*o) for o in bboxes[keep_mask]]
keep_masks = MaskArray(np.vstack(raw_masks)[keep_mask])
pred = BaseRecord((
ScoresRecordComponent(),
ImageRecordComponent(),
InstancesLabelsRecordComponent(),
BBoxesRecordComponent(),
MasksRecordComponent(),
))
pred.detect.set_class_map(record.detect.class_map)
pred.detect.set_scores(keep_scores)
pred.detect.set_labels_by_id(keep_labels)
pred.detect.set_bboxes(keep_bboxes)
pred.detect.set_masks(keep_masks)
pred.above_threshold = keep_mask
return Prediction(pred=pred, ground_truth=record)
def convert_raw_prediction(raw_pred: Sequence[np.ndarray],
detection_threshold: float):
raw_bboxes, raw_masks = raw_pred
scores, labels, bboxes = _unpack_raw_bboxes(raw_bboxes)
keep_mask = scores > detection_threshold
keep_scores = scores[keep_mask]
keep_labels = labels[keep_mask]
keep_bboxes = [BBox.from_xyxy(*o) for o in bboxes[keep_mask]]
keep_masks = MaskArray(np.vstack(raw_masks)[keep_mask])
return {
"scores": keep_scores,
"labels": keep_labels,
"bboxes": keep_bboxes,
"masks": keep_masks,
}
def convert_raw_prediction(
sample,
raw_pred: dict,
record: BaseRecord,
detection_threshold: float,
keep_image: bool = False,
):
# convert predictions
raw_bboxes, raw_masks = raw_pred
scores, labels, bboxes = _unpack_raw_bboxes(raw_bboxes)
keep_mask = scores > detection_threshold
keep_scores = scores[keep_mask]
keep_labels = labels[keep_mask]
keep_bboxes = [BBox.from_xyxy(*o) for o in bboxes[keep_mask]]
keep_masks = MaskArray(np.vstack(raw_masks)[keep_mask])
keep_labels = convert_background_from_last_to_zero(
label_ids=keep_labels, class_map=record.detection.class_map)
pred = BaseRecord((
ScoresRecordComponent(),
ImageRecordComponent(),
InstancesLabelsRecordComponent(),
BBoxesRecordComponent(),
MasksRecordComponent(),
))
pred.detection.set_class_map(record.detection.class_map)
pred.detection.set_scores(keep_scores)
pred.detection.set_labels_by_id(keep_labels)
pred.detection.set_bboxes(keep_bboxes)
pred.detection.set_masks(keep_masks)
pred.above_threshold = keep_mask
if keep_image:
image = sample["img"]
image = image.detach().cpu().numpy().transpose(1, 2, 0)
pred.set_img(image)
record.set_img(image)
return Prediction(pred=pred, ground_truth=record)