|
|
|
|
|
|
|
|
|
from typing import Any, Dict, List, Tuple
|
|
import torch
|
|
from torch.nn import functional as F
|
|
|
|
from detectron2.config import CfgNode
|
|
from detectron2.structures import Instances
|
|
|
|
from densepose.converters.base import IntTupleBox
|
|
from densepose.data.utils import get_class_to_mesh_name_mapping
|
|
from densepose.modeling.cse.utils import squared_euclidean_distance_matrix
|
|
from densepose.structures import DensePoseDataRelative
|
|
|
|
from .densepose_base import DensePoseBaseSampler
|
|
|
|
|
|
class DensePoseCSEBaseSampler(DensePoseBaseSampler):
|
|
"""
|
|
Base DensePose sampler to produce DensePose data from DensePose predictions.
|
|
Samples for each class are drawn according to some distribution over all pixels estimated
|
|
to belong to that class.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
cfg: CfgNode,
|
|
use_gt_categories: bool,
|
|
embedder: torch.nn.Module,
|
|
count_per_class: int = 8,
|
|
):
|
|
"""
|
|
Constructor
|
|
|
|
Args:
|
|
cfg (CfgNode): the config of the model
|
|
embedder (torch.nn.Module): necessary to compute mesh vertex embeddings
|
|
count_per_class (int): the sampler produces at most `count_per_class`
|
|
samples for each category
|
|
"""
|
|
super().__init__(count_per_class)
|
|
self.embedder = embedder
|
|
self.class_to_mesh_name = get_class_to_mesh_name_mapping(cfg)
|
|
self.use_gt_categories = use_gt_categories
|
|
|
|
def _sample(self, instance: Instances, bbox_xywh: IntTupleBox) -> Dict[str, List[Any]]:
|
|
"""
|
|
Sample DensPoseDataRelative from estimation results
|
|
"""
|
|
if self.use_gt_categories:
|
|
instance_class = instance.dataset_classes.tolist()[0]
|
|
else:
|
|
instance_class = instance.pred_classes.tolist()[0]
|
|
mesh_name = self.class_to_mesh_name[instance_class]
|
|
|
|
annotation = {
|
|
DensePoseDataRelative.X_KEY: [],
|
|
DensePoseDataRelative.Y_KEY: [],
|
|
DensePoseDataRelative.VERTEX_IDS_KEY: [],
|
|
DensePoseDataRelative.MESH_NAME_KEY: mesh_name,
|
|
}
|
|
|
|
mask, embeddings, other_values = self._produce_mask_and_results(instance, bbox_xywh)
|
|
indices = torch.nonzero(mask, as_tuple=True)
|
|
selected_embeddings = embeddings.permute(1, 2, 0)[indices].cpu()
|
|
values = other_values[:, indices[0], indices[1]]
|
|
k = values.shape[1]
|
|
|
|
count = min(self.count_per_class, k)
|
|
if count <= 0:
|
|
return annotation
|
|
|
|
index_sample = self._produce_index_sample(values, count)
|
|
closest_vertices = squared_euclidean_distance_matrix(
|
|
selected_embeddings[index_sample], self.embedder(mesh_name)
|
|
)
|
|
closest_vertices = torch.argmin(closest_vertices, dim=1)
|
|
|
|
sampled_y = indices[0][index_sample] + 0.5
|
|
sampled_x = indices[1][index_sample] + 0.5
|
|
|
|
_, _, w, h = bbox_xywh
|
|
x = (sampled_x / w * 256.0).cpu().tolist()
|
|
y = (sampled_y / h * 256.0).cpu().tolist()
|
|
|
|
annotation[DensePoseDataRelative.X_KEY].extend(x)
|
|
annotation[DensePoseDataRelative.Y_KEY].extend(y)
|
|
annotation[DensePoseDataRelative.VERTEX_IDS_KEY].extend(closest_vertices.cpu().tolist())
|
|
return annotation
|
|
|
|
def _produce_mask_and_results(
|
|
self, instance: Instances, bbox_xywh: IntTupleBox
|
|
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Method to get labels and DensePose results from an instance
|
|
|
|
Args:
|
|
instance (Instances): an instance of `DensePoseEmbeddingPredictorOutput`
|
|
bbox_xywh (IntTupleBox): the corresponding bounding box
|
|
|
|
Return:
|
|
mask (torch.Tensor): shape [H, W], DensePose segmentation mask
|
|
embeddings (Tuple[torch.Tensor]): a tensor of shape [D, H, W],
|
|
DensePose CSE Embeddings
|
|
other_values (Tuple[torch.Tensor]): a tensor of shape [0, H, W],
|
|
for potential other values
|
|
"""
|
|
densepose_output = instance.pred_densepose
|
|
S = densepose_output.coarse_segm
|
|
E = densepose_output.embedding
|
|
_, _, w, h = bbox_xywh
|
|
embeddings = F.interpolate(E, size=(h, w), mode="bilinear")[0]
|
|
coarse_segm_resized = F.interpolate(S, size=(h, w), mode="bilinear")[0]
|
|
mask = coarse_segm_resized.argmax(0) > 0
|
|
other_values = torch.empty((0, h, w), device=E.device)
|
|
return mask, embeddings, other_values
|
|
|
|
def _resample_mask(self, output: Any) -> torch.Tensor:
|
|
"""
|
|
Convert DensePose predictor output to segmentation annotation - tensors of size
|
|
(256, 256) and type `int64`.
|
|
|
|
Args:
|
|
output: DensePose predictor output with the following attributes:
|
|
- coarse_segm: tensor of size [N, D, H, W] with unnormalized coarse
|
|
segmentation scores
|
|
Return:
|
|
Tensor of size (S, S) and type `int64` with coarse segmentation annotations,
|
|
where S = DensePoseDataRelative.MASK_SIZE
|
|
"""
|
|
sz = DensePoseDataRelative.MASK_SIZE
|
|
mask = (
|
|
F.interpolate(output.coarse_segm, (sz, sz), mode="bilinear", align_corners=False)
|
|
.argmax(dim=1)
|
|
.long()
|
|
.squeeze()
|
|
.cpu()
|
|
)
|
|
return mask
|
|
|
