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