def convert_raw_prediction(
raw_pred: dict, record: BaseRecord, detection_threshold: float
):
pred = faster_convert_raw_prediction(
raw_pred=raw_pred, record=record, detection_threshold=detection_threshold
)
above_threshold = pred.detect.above_threshold
kps = raw_pred["keypoints"][above_threshold]
keypoints = []
for k in kps:
k = k.cpu().numpy()
k = list(chain.from_iterable(k))
# `if sum(k) > 0` prevents empty `KeyPoints` objects to be instantiated.
# E.g. `k = [0, 0, 0, 0, 0, 0]` is a flattened list of 2 points `(0, 0, 0)` and `(0, 0, 0)`. We don't want a `KeyPoints` object to be created on top of this list.
if sum(k) > 0:
keypoints.append(KeyPoints.from_xyv(k, None))
pred.pred.add_component(KeyPointsRecordComponent())
pred.pred.detect.add_keypoints(keypoints)
pred.pred.detect.keypoints_scores = (
raw_pred["keypoints_scores"][above_threshold].detach().cpu().numpy()
)
return pred
def convert_raw_prediction(
sample,
raw_pred: dict,
record: Sequence[BaseRecord],
detection_threshold: float,
mask_threshold: float,
keep_image: bool = False,
):
pred = faster_convert_raw_prediction(
sample=sample,
raw_pred=raw_pred,
record=record,
detection_threshold=detection_threshold,
keep_image=keep_image,
)
above_threshold = pred.detection.above_threshold
masks_probs = raw_pred["masks"][above_threshold]
masks_probs = masks_probs.detach().cpu().numpy()
# convert probabilities to 0 or 1 based on mask_threshold
masks = masks_probs > mask_threshold
masks = MaskArray(masks.squeeze(1))
pred.pred.add_component(InstanceMasksRecordComponent())
pred.detection.set_mask_array(masks)
if keep_image:
# HACK: quick fix for when we have to add masks back
if len(sample) > 1:
tensor_image, label = sample
mask = MaskArray(np.array(label["masks"].cpu().numpy()))
pred.ground_truth.detection.set_mask_array(mask)
return pred
def convert_raw_prediction(
sample,
raw_pred: dict,
record: Sequence[BaseRecord],
detection_threshold: float,
mask_threshold: float,
keep_image: bool = False,
):
pred = faster_convert_raw_prediction(
sample=sample,
raw_pred=raw_pred,
record=record,
detection_threshold=detection_threshold,
keep_image=keep_image,
)
above_threshold = pred.detection.above_threshold
masks_probs = raw_pred["masks"][above_threshold]
masks_probs = masks_probs.detach().cpu().numpy()
# convert probabilities to 0 or 1 based on mask_threshold
masks = masks_probs > mask_threshold
masks = MaskArray(masks.squeeze(1))
pred.pred.add_component(MasksRecordComponent())
pred.detection.set_masks(masks)
return pred
def convert_raw_prediction(raw_pred: dict, detection_threshold: float,
mask_threshold: float):
preds = faster_convert_raw_prediction(
raw_pred=raw_pred, detection_threshold=detection_threshold)
above_threshold = preds["above_threshold"]
masks_probs = raw_pred["masks"][above_threshold]
masks_probs = masks_probs.detach().cpu().numpy()
# convert probabilities to 0 or 1 based on mask_threshold
masks = masks_probs > mask_threshold
masks = MaskArray(masks.squeeze(1))
return {"masks": masks, **preds}
