Commit b634945d authored by limm's avatar limm
Browse files

support v0.6

parent 5b3792fc
# Copyright (c) Facebook, Inc. and its affiliates.
from typing import Any, Tuple, Type
import torch
class BaseConverter:
"""
Converter base class to be reused by various converters.
Converter allows one to convert data from various source types to a particular
destination type. Each source type needs to register its converter. The
registration for each source type is valid for all descendants of that type.
"""
@classmethod
def register(cls, from_type: Type, converter: Any = None):
"""
Registers a converter for the specified type.
Can be used as a decorator (if converter is None), or called as a method.
Args:
from_type (type): type to register the converter for;
all instances of this type will use the same converter
converter (callable): converter to be registered for the given
type; if None, this method is assumed to be a decorator for the converter
"""
if converter is not None:
cls._do_register(from_type, converter)
def wrapper(converter: Any) -> Any:
cls._do_register(from_type, converter)
return converter
return wrapper
@classmethod
def _do_register(cls, from_type: Type, converter: Any):
cls.registry[from_type] = converter # pyre-ignore[16]
@classmethod
def _lookup_converter(cls, from_type: Type) -> Any:
"""
Perform recursive lookup for the given type
to find registered converter. If a converter was found for some base
class, it gets registered for this class to save on further lookups.
Args:
from_type: type for which to find a converter
Return:
callable or None - registered converter or None
if no suitable entry was found in the registry
"""
if from_type in cls.registry: # pyre-ignore[16]
return cls.registry[from_type]
for base in from_type.__bases__:
converter = cls._lookup_converter(base)
if converter is not None:
cls._do_register(from_type, converter)
return converter
return None
@classmethod
def convert(cls, instance: Any, *args, **kwargs):
"""
Convert an instance to the destination type using some registered
converter. Does recursive lookup for base classes, so there's no need
for explicit registration for derived classes.
Args:
instance: source instance to convert to the destination type
Return:
An instance of the destination type obtained from the source instance
Raises KeyError, if no suitable converter found
"""
instance_type = type(instance)
converter = cls._lookup_converter(instance_type)
if converter is None:
if cls.dst_type is None: # pyre-ignore[16]
output_type_str = "itself"
else:
output_type_str = cls.dst_type
raise KeyError(f"Could not find converter from {instance_type} to {output_type_str}")
return converter(instance, *args, **kwargs)
IntTupleBox = Tuple[int, int, int, int]
def make_int_box(box: torch.Tensor) -> IntTupleBox:
int_box = [0, 0, 0, 0]
int_box[0], int_box[1], int_box[2], int_box[3] = tuple(box.long().tolist())
return int_box[0], int_box[1], int_box[2], int_box[3]
# Copyright (c) Facebook, Inc. and its affiliates.
from ..structures import DensePoseChartPredictorOutput, DensePoseEmbeddingPredictorOutput
from . import (
HFlipConverter,
ToChartResultConverter,
ToChartResultConverterWithConfidences,
ToMaskConverter,
densepose_chart_predictor_output_hflip,
densepose_chart_predictor_output_to_result,
densepose_chart_predictor_output_to_result_with_confidences,
predictor_output_with_coarse_segm_to_mask,
predictor_output_with_fine_and_coarse_segm_to_mask,
)
ToMaskConverter.register(
DensePoseChartPredictorOutput, predictor_output_with_fine_and_coarse_segm_to_mask
)
ToMaskConverter.register(
DensePoseEmbeddingPredictorOutput, predictor_output_with_coarse_segm_to_mask
)
ToChartResultConverter.register(
DensePoseChartPredictorOutput, densepose_chart_predictor_output_to_result
)
ToChartResultConverterWithConfidences.register(
DensePoseChartPredictorOutput, densepose_chart_predictor_output_to_result_with_confidences
)
HFlipConverter.register(DensePoseChartPredictorOutput, densepose_chart_predictor_output_hflip)
# Copyright (c) Facebook, Inc. and its affiliates.
from dataclasses import fields
import torch
from densepose.structures import DensePoseChartPredictorOutput, DensePoseTransformData
def densepose_chart_predictor_output_hflip(
densepose_predictor_output: DensePoseChartPredictorOutput,
transform_data: DensePoseTransformData,
) -> DensePoseChartPredictorOutput:
"""
Change to take into account a Horizontal flip.
"""
if len(densepose_predictor_output) > 0:
PredictorOutput = type(densepose_predictor_output)
output_dict = {}
for field in fields(densepose_predictor_output):
field_value = getattr(densepose_predictor_output, field.name)
# flip tensors
if isinstance(field_value, torch.Tensor):
setattr(densepose_predictor_output, field.name, torch.flip(field_value, [3]))
densepose_predictor_output = _flip_iuv_semantics_tensor(
densepose_predictor_output, transform_data
)
densepose_predictor_output = _flip_segm_semantics_tensor(
densepose_predictor_output, transform_data
)
for field in fields(densepose_predictor_output):
output_dict[field.name] = getattr(densepose_predictor_output, field.name)
return PredictorOutput(**output_dict)
else:
return densepose_predictor_output
def _flip_iuv_semantics_tensor(
densepose_predictor_output: DensePoseChartPredictorOutput,
dp_transform_data: DensePoseTransformData,
) -> DensePoseChartPredictorOutput:
point_label_symmetries = dp_transform_data.point_label_symmetries
uv_symmetries = dp_transform_data.uv_symmetries
N, C, H, W = densepose_predictor_output.u.shape
u_loc = (densepose_predictor_output.u[:, 1:, :, :].clamp(0, 1) * 255).long()
v_loc = (densepose_predictor_output.v[:, 1:, :, :].clamp(0, 1) * 255).long()
Iindex = torch.arange(C - 1, device=densepose_predictor_output.u.device)[
None, :, None, None
].expand(N, C - 1, H, W)
densepose_predictor_output.u[:, 1:, :, :] = uv_symmetries["U_transforms"][Iindex, v_loc, u_loc]
densepose_predictor_output.v[:, 1:, :, :] = uv_symmetries["V_transforms"][Iindex, v_loc, u_loc]
for el in ["fine_segm", "u", "v"]:
densepose_predictor_output.__dict__[el] = densepose_predictor_output.__dict__[el][
:, point_label_symmetries, :, :
]
return densepose_predictor_output
def _flip_segm_semantics_tensor(
densepose_predictor_output: DensePoseChartPredictorOutput, dp_transform_data
):
if densepose_predictor_output.coarse_segm.shape[1] > 2:
densepose_predictor_output.coarse_segm = densepose_predictor_output.coarse_segm[
:, dp_transform_data.mask_label_symmetries, :, :
]
return densepose_predictor_output
# Copyright (c) Facebook, Inc. and its affiliates.
from typing import Dict
import torch
from torch.nn import functional as F
from detectron2.structures.boxes import Boxes, BoxMode
from ..structures import (
DensePoseChartPredictorOutput,
DensePoseChartResult,
DensePoseChartResultWithConfidences,
)
from . import resample_fine_and_coarse_segm_to_bbox
from .base import IntTupleBox, make_int_box
def resample_uv_tensors_to_bbox(
u: torch.Tensor,
v: torch.Tensor,
labels: torch.Tensor,
box_xywh_abs: IntTupleBox,
) -> torch.Tensor:
"""
Resamples U and V coordinate estimates for the given bounding box
Args:
u (tensor [1, C, H, W] of float): U coordinates
v (tensor [1, C, H, W] of float): V coordinates
labels (tensor [H, W] of long): labels obtained by resampling segmentation
outputs for the given bounding box
box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
Return:
Resampled U and V coordinates - a tensor [2, H, W] of float
"""
x, y, w, h = box_xywh_abs
w = max(int(w), 1)
h = max(int(h), 1)
u_bbox = F.interpolate(u, (h, w), mode="bilinear", align_corners=False)
v_bbox = F.interpolate(v, (h, w), mode="bilinear", align_corners=False)
uv = torch.zeros([2, h, w], dtype=torch.float32, device=u.device)
for part_id in range(1, u_bbox.size(1)):
uv[0][labels == part_id] = u_bbox[0, part_id][labels == part_id]
uv[1][labels == part_id] = v_bbox[0, part_id][labels == part_id]
return uv
def resample_uv_to_bbox(
predictor_output: DensePoseChartPredictorOutput,
labels: torch.Tensor,
box_xywh_abs: IntTupleBox,
) -> torch.Tensor:
"""
Resamples U and V coordinate estimates for the given bounding box
Args:
predictor_output (DensePoseChartPredictorOutput): DensePose predictor
output to be resampled
labels (tensor [H, W] of long): labels obtained by resampling segmentation
outputs for the given bounding box
box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
Return:
Resampled U and V coordinates - a tensor [2, H, W] of float
"""
return resample_uv_tensors_to_bbox(
predictor_output.u,
predictor_output.v,
labels,
box_xywh_abs,
)
def densepose_chart_predictor_output_to_result(
predictor_output: DensePoseChartPredictorOutput, boxes: Boxes
) -> DensePoseChartResult:
"""
Convert densepose chart predictor outputs to results
Args:
predictor_output (DensePoseChartPredictorOutput): DensePose predictor
output to be converted to results, must contain only 1 output
boxes (Boxes): bounding box that corresponds to the predictor output,
must contain only 1 bounding box
Return:
DensePose chart-based result (DensePoseChartResult)
"""
assert len(predictor_output) == 1 and len(boxes) == 1, (
f"Predictor output to result conversion can operate only single outputs"
f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes"
)
boxes_xyxy_abs = boxes.tensor.clone()
boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
box_xywh = make_int_box(boxes_xywh_abs[0])
labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0)
uv = resample_uv_to_bbox(predictor_output, labels, box_xywh)
return DensePoseChartResult(labels=labels, uv=uv)
def resample_confidences_to_bbox(
predictor_output: DensePoseChartPredictorOutput,
labels: torch.Tensor,
box_xywh_abs: IntTupleBox,
) -> Dict[str, torch.Tensor]:
"""
Resamples confidences for the given bounding box
Args:
predictor_output (DensePoseChartPredictorOutput): DensePose predictor
output to be resampled
labels (tensor [H, W] of long): labels obtained by resampling segmentation
outputs for the given bounding box
box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
Return:
Resampled confidences - a dict of [H, W] tensors of float
"""
x, y, w, h = box_xywh_abs
w = max(int(w), 1)
h = max(int(h), 1)
confidence_names = [
"sigma_1",
"sigma_2",
"kappa_u",
"kappa_v",
"fine_segm_confidence",
"coarse_segm_confidence",
]
confidence_results = {key: None for key in confidence_names}
confidence_names = [
key for key in confidence_names if getattr(predictor_output, key) is not None
]
confidence_base = torch.zeros([h, w], dtype=torch.float32, device=predictor_output.u.device)
# assign data from channels that correspond to the labels
for key in confidence_names:
resampled_confidence = F.interpolate(
getattr(predictor_output, key), (h, w), mode="bilinear", align_corners=False
)
result = confidence_base.clone()
for part_id in range(1, predictor_output.u.size(1)):
if resampled_confidence.size(1) != predictor_output.u.size(1):
# confidence is not part-based, don't try to fill it part by part
continue
result[labels == part_id] = resampled_confidence[0, part_id][labels == part_id]
if resampled_confidence.size(1) != predictor_output.u.size(1):
# confidence is not part-based, fill the data with the first channel
# (targeted for segmentation confidences that have only 1 channel)
result = resampled_confidence[0, 0]
confidence_results[key] = result
return confidence_results # pyre-ignore[7]
def densepose_chart_predictor_output_to_result_with_confidences(
predictor_output: DensePoseChartPredictorOutput, boxes: Boxes
) -> DensePoseChartResultWithConfidences:
"""
Convert densepose chart predictor outputs to results
Args:
predictor_output (DensePoseChartPredictorOutput): DensePose predictor
output with confidences to be converted to results, must contain only 1 output
boxes (Boxes): bounding box that corresponds to the predictor output,
must contain only 1 bounding box
Return:
DensePose chart-based result with confidences (DensePoseChartResultWithConfidences)
"""
assert len(predictor_output) == 1 and len(boxes) == 1, (
f"Predictor output to result conversion can operate only single outputs"
f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes"
)
boxes_xyxy_abs = boxes.tensor.clone()
boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
box_xywh = make_int_box(boxes_xywh_abs[0])
labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0)
uv = resample_uv_to_bbox(predictor_output, labels, box_xywh)
confidences = resample_confidences_to_bbox(predictor_output, labels, box_xywh)
return DensePoseChartResultWithConfidences(labels=labels, uv=uv, **confidences)
# Copyright (c) Facebook, Inc. and its affiliates.
from typing import Any
from .base import BaseConverter
class HFlipConverter(BaseConverter):
"""
Converts various DensePose predictor outputs to DensePose results.
Each DensePose predictor output type has to register its convertion strategy.
"""
registry = {}
dst_type = None
@classmethod
def convert(cls, predictor_outputs: Any, transform_data: Any, *args, **kwargs):
"""
Performs an horizontal flip on DensePose predictor outputs.
Does recursive lookup for base classes, so there's no need
for explicit registration for derived classes.
Args:
predictor_outputs: DensePose predictor output to be converted to BitMasks
transform_data: Anything useful for the flip
Return:
An instance of the same type as predictor_outputs
"""
return super(HFlipConverter, cls).convert(
predictor_outputs, transform_data, *args, **kwargs
)
# Copyright (c) Facebook, Inc. and its affiliates.
from typing import Any
import torch
from torch.nn import functional as F
from detectron2.structures import BitMasks, Boxes, BoxMode
from .base import IntTupleBox, make_int_box
from .to_mask import ImageSizeType
def resample_coarse_segm_tensor_to_bbox(coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox):
"""
Resample coarse segmentation tensor to the given
bounding box and derive labels for each pixel of the bounding box
Args:
coarse_segm: float tensor of shape [1, K, Hout, Wout]
box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
corner coordinates, width (W) and height (H)
Return:
Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
"""
x, y, w, h = box_xywh_abs
w = max(int(w), 1)
h = max(int(h), 1)
labels = F.interpolate(coarse_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
return labels
def resample_fine_and_coarse_segm_tensors_to_bbox(
fine_segm: torch.Tensor, coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox
):
"""
Resample fine and coarse segmentation tensors to the given
bounding box and derive labels for each pixel of the bounding box
Args:
fine_segm: float tensor of shape [1, C, Hout, Wout]
coarse_segm: float tensor of shape [1, K, Hout, Wout]
box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
corner coordinates, width (W) and height (H)
Return:
Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
"""
x, y, w, h = box_xywh_abs
w = max(int(w), 1)
h = max(int(h), 1)
# coarse segmentation
coarse_segm_bbox = F.interpolate(
coarse_segm, (h, w), mode="bilinear", align_corners=False
).argmax(dim=1)
# combined coarse and fine segmentation
labels = (
F.interpolate(fine_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
* (coarse_segm_bbox > 0).long()
)
return labels
def resample_fine_and_coarse_segm_to_bbox(predictor_output: Any, box_xywh_abs: IntTupleBox):
"""
Resample fine and coarse segmentation outputs from a predictor to the given
bounding box and derive labels for each pixel of the bounding box
Args:
predictor_output: DensePose predictor output that contains segmentation
results to be resampled
box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
corner coordinates, width (W) and height (H)
Return:
Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
"""
return resample_fine_and_coarse_segm_tensors_to_bbox(
predictor_output.fine_segm,
predictor_output.coarse_segm,
box_xywh_abs,
)
def predictor_output_with_coarse_segm_to_mask(
predictor_output: Any, boxes: Boxes, image_size_hw: ImageSizeType
) -> BitMasks:
"""
Convert predictor output with coarse and fine segmentation to a mask.
Assumes that predictor output has the following attributes:
- coarse_segm (tensor of size [N, D, H, W]): coarse segmentation
unnormalized scores for N instances; D is the number of coarse
segmentation labels, H and W is the resolution of the estimate
Args:
predictor_output: DensePose predictor output to be converted to mask
boxes (Boxes): bounding boxes that correspond to the DensePose
predictor outputs
image_size_hw (tuple [int, int]): image height Himg and width Wimg
Return:
BitMasks that contain a bool tensor of size [N, Himg, Wimg] with
a mask of the size of the image for each instance
"""
H, W = image_size_hw
boxes_xyxy_abs = boxes.tensor.clone()
boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
N = len(boxes_xywh_abs)
masks = torch.zeros((N, H, W), dtype=torch.bool, device=boxes.tensor.device)
for i in range(len(boxes_xywh_abs)):
box_xywh = make_int_box(boxes_xywh_abs[i])
box_mask = resample_coarse_segm_tensor_to_bbox(predictor_output[i].coarse_segm, box_xywh)
x, y, w, h = box_xywh
masks[i, y : y + h, x : x + w] = box_mask
return BitMasks(masks)
def predictor_output_with_fine_and_coarse_segm_to_mask(
predictor_output: Any, boxes: Boxes, image_size_hw: ImageSizeType
) -> BitMasks:
"""
Convert predictor output with coarse and fine segmentation to a mask.
Assumes that predictor output has the following attributes:
- coarse_segm (tensor of size [N, D, H, W]): coarse segmentation
unnormalized scores for N instances; D is the number of coarse
segmentation labels, H and W is the resolution of the estimate
- fine_segm (tensor of size [N, C, H, W]): fine segmentation
unnormalized scores for N instances; C is the number of fine
segmentation labels, H and W is the resolution of the estimate
Args:
predictor_output: DensePose predictor output to be converted to mask
boxes (Boxes): bounding boxes that correspond to the DensePose
predictor outputs
image_size_hw (tuple [int, int]): image height Himg and width Wimg
Return:
BitMasks that contain a bool tensor of size [N, Himg, Wimg] with
a mask of the size of the image for each instance
"""
H, W = image_size_hw
boxes_xyxy_abs = boxes.tensor.clone()
boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
N = len(boxes_xywh_abs)
masks = torch.zeros((N, H, W), dtype=torch.bool, device=boxes.tensor.device)
for i in range(len(boxes_xywh_abs)):
box_xywh = make_int_box(boxes_xywh_abs[i])
labels_i = resample_fine_and_coarse_segm_to_bbox(predictor_output[i], box_xywh)
x, y, w, h = box_xywh
masks[i, y : y + h, x : x + w] = labels_i > 0
return BitMasks(masks)
# Copyright (c) Facebook, Inc. and its affiliates.
from typing import Any
from detectron2.structures import Boxes
from ..structures import DensePoseChartResult, DensePoseChartResultWithConfidences
from .base import BaseConverter
class ToChartResultConverter(BaseConverter):
"""
Converts various DensePose predictor outputs to DensePose results.
Each DensePose predictor output type has to register its convertion strategy.
"""
registry = {}
dst_type = DensePoseChartResult
@classmethod
def convert(cls, predictor_outputs: Any, boxes: Boxes, *args, **kwargs) -> DensePoseChartResult:
"""
Convert DensePose predictor outputs to DensePoseResult 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
Return:
An instance of DensePoseResult. If no suitable converter was found, raises KeyError
"""
return super(ToChartResultConverter, cls).convert(predictor_outputs, boxes, *args, **kwargs)
class ToChartResultConverterWithConfidences(BaseConverter):
"""
Converts various DensePose predictor outputs to DensePose results.
Each DensePose predictor output type has to register its convertion strategy.
"""
registry = {}
dst_type = DensePoseChartResultWithConfidences
@classmethod
def convert(
cls, predictor_outputs: Any, boxes: Boxes, *args, **kwargs
) -> DensePoseChartResultWithConfidences:
"""
Convert DensePose predictor outputs to DensePoseResult with confidences
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 with confidences
to be converted to BitMasks
boxes (Boxes): bounding boxes that correspond to the DensePose
predictor outputs
Return:
An instance of DensePoseResult. If no suitable converter was found, raises KeyError
"""
return super(ToChartResultConverterWithConfidences, cls).convert(
predictor_outputs, boxes, *args, **kwargs
)
# Copyright (c) Facebook, Inc. and its affiliates.
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
)
# Copyright (c) Facebook, Inc. and its affiliates.
from .meshes import builtin
from .build import (
build_detection_test_loader,
build_detection_train_loader,
build_combined_loader,
build_frame_selector,
build_inference_based_loaders,
has_inference_based_loaders,
BootstrapDatasetFactoryCatalog,
)
from .combined_loader import CombinedDataLoader
from .dataset_mapper import DatasetMapper
from .inference_based_loader import InferenceBasedLoader, ScoreBasedFilter
from .image_list_dataset import ImageListDataset
from .utils import is_relative_local_path, maybe_prepend_base_path
# ensure the builtin datasets are registered
from . import datasets
# ensure the bootstrap datasets builders are registered
from . import build
__all__ = [k for k in globals().keys() if not k.startswith("_")]
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment