File size: 953 Bytes
938e515 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
# Copyright (c) Facebook, Inc. and its affiliates.
from detectron2.structures import BitMasks, Instances
from densepose.converters import ToMaskConverter
class MaskFromDensePoseSampler:
"""
Produce mask GT from DensePose predictions
This sampler simply converts DensePose predictions to BitMasks
that a contain 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 ToMaskConverter.convert(
instances.pred_densepose, instances.pred_boxes, instances.image_size
)
|