|
|
|
|
|
from typing import Any, Tuple |
|
|
|
from detectron2.structures import BitMasks, Boxes |
|
|
|
from .base import BaseConverter |
|
|
|
ImageSizeType = Tuple[int, int] |
|
|
|
|
|
class ToMaskConverter(BaseConverter): |
|
""" |
|
Converts various DensePose predictor outputs to masks |
|
in bit mask format (see `BitMasks`). Each DensePose predictor output type |
|
has to register its convertion strategy. |
|
""" |
|
|
|
registry = {} |
|
dst_type = BitMasks |
|
|
|
@classmethod |
|
|
|
|
|
def convert( |
|
cls, |
|
densepose_predictor_outputs: Any, |
|
boxes: Boxes, |
|
image_size_hw: ImageSizeType, |
|
*args, |
|
**kwargs |
|
) -> BitMasks: |
|
""" |
|
Convert DensePose predictor outputs to BitMasks using some registered |
|
converter. Does recursive lookup for base classes, so there's no need |
|
for explicit registration for derived classes. |
|
|
|
Args: |
|
densepose_predictor_outputs: DensePose predictor output to be |
|
converted to BitMasks |
|
boxes (Boxes): bounding boxes that correspond to the DensePose |
|
predictor outputs |
|
image_size_hw (tuple [int, int]): image height and width |
|
Return: |
|
An instance of `BitMasks`. If no suitable converter was found, raises KeyError |
|
""" |
|
return super(ToMaskConverter, cls).convert( |
|
densepose_predictor_outputs, boxes, image_size_hw, *args, **kwargs |
|
) |
|
|
