|
|
|
|
| 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
|
|
|
