mirror of https://github.com/RE-OWOD/RE-OWOD
60 lines
2.3 KiB
Python
60 lines
2.3 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
|
|
|
import torch
|
|
|
|
from detectron2.structures import BitMasks, BoxMode, Instances
|
|
|
|
from ..structures import resample_output_to_bbox
|
|
|
|
|
|
def densepose_to_mask(instances: Instances) -> BitMasks:
|
|
"""
|
|
Produce masks from DensePose predictions
|
|
DensePose predictions for a given image, stored in `pred_densepose` field,
|
|
are instances of DensePoseOutput. This sampler takes
|
|
`S` and `I` output tensors (coarse and fine segmentation) and converts
|
|
then to a mask tensor, which is a bool tensor of the size of the input
|
|
image
|
|
|
|
Args:
|
|
instances (Instances): predicted results, expected to have `pred_densepose` field
|
|
that contains `DensePoseOutput` objects
|
|
|
|
Returns:
|
|
`BitMasks` instance with boolean tensors of the size of the input image that have non-zero
|
|
values at pixels that are estimated to belong to the detected objects
|
|
"""
|
|
H, W = instances.image_size
|
|
boxes_xyxy_abs = instances.pred_boxes.tensor.clone().cpu()
|
|
boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
|
|
N = len(boxes_xywh_abs)
|
|
gt_masks = torch.zeros((N, H, W), dtype=torch.bool, device=torch.device("cpu"))
|
|
for i, box_xywh in enumerate(boxes_xywh_abs):
|
|
labels_i, _ = resample_output_to_bbox(instances.pred_densepose[i], box_xywh)
|
|
x, y, w, h = box_xywh.long().tolist()
|
|
gt_masks[i, y : y + h, x : x + w] = labels_i.cpu() > 0
|
|
return BitMasks(gt_masks)
|
|
|
|
|
|
class MaskFromDensePoseSampler:
|
|
"""
|
|
Produce mask GT from DensePose predictions
|
|
DensePose prediction is an instance of DensePoseOutput. This sampler takes
|
|
`S` and `I` output tensors (coarse and fine segmentation) and converts
|
|
then to a mask tensor, which is a bool tensor of the size of the input
|
|
image
|
|
"""
|
|
|
|
def __call__(self, instances: Instances) -> BitMasks:
|
|
"""
|
|
Converts predicted data from `instances` into the GT mask data
|
|
|
|
Args:
|
|
instances (Instances): predicted results, expected to have `pred_densepose` field
|
|
|
|
Returns:
|
|
Boolean Tensor of the size of the input image that has non-zero
|
|
values at pixels that are estimated to belong to the detected object
|
|
"""
|
|
return densepose_to_mask(instances)
|