| | |
| |
|
| | import random |
| | from typing import Optional, 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_cse_base import DensePoseCSEBaseSampler |
| |
|
| |
|
| | class DensePoseCSEConfidenceBasedSampler(DensePoseCSEBaseSampler): |
| | """ |
| | Samples DensePose data from DensePose predictions. |
| | Samples for each class are drawn using confidence value estimates. |
| | """ |
| |
|
| | def __init__( |
| | self, |
| | cfg: CfgNode, |
| | use_gt_categories: bool, |
| | embedder: torch.nn.Module, |
| | confidence_channel: str, |
| | count_per_class: int = 8, |
| | search_count_multiplier: Optional[float] = None, |
| | search_proportion: Optional[float] = None, |
| | ): |
| | """ |
| | Constructor |
| | |
| | Args: |
| | cfg (CfgNode): the config of the model |
| | embedder (torch.nn.Module): necessary to compute mesh vertex embeddings |
| | confidence_channel (str): confidence channel to use for sampling; |
| | possible values: |
| | "coarse_segm_confidence": confidences for coarse segmentation |
| | (default: "coarse_segm_confidence") |
| | count_per_class (int): the sampler produces at most `count_per_class` |
| | samples for each category (default: 8) |
| | search_count_multiplier (float or None): if not None, the total number |
| | of the most confident estimates of a given class to consider is |
| | defined as `min(search_count_multiplier * count_per_class, N)`, |
| | where `N` is the total number of estimates of the class; cannot be |
| | specified together with `search_proportion` (default: None) |
| | search_proportion (float or None): if not None, the total number of the |
| | of the most confident estimates of a given class to consider is |
| | defined as `min(max(search_proportion * N, count_per_class), N)`, |
| | where `N` is the total number of estimates of the class; cannot be |
| | specified together with `search_count_multiplier` (default: None) |
| | """ |
| | super().__init__(cfg, use_gt_categories, embedder, count_per_class) |
| | self.confidence_channel = confidence_channel |
| | self.search_count_multiplier = search_count_multiplier |
| | self.search_proportion = search_proportion |
| | assert (search_count_multiplier is None) or (search_proportion is None), ( |
| | f"Cannot specify both search_count_multiplier (={search_count_multiplier})" |
| | f"and search_proportion (={search_proportion})" |
| | ) |
| |
|
| | def _produce_index_sample(self, values: torch.Tensor, count: int): |
| | """ |
| | Produce a sample of indices to select data based on confidences |
| | |
| | Args: |
| | values (torch.Tensor): a tensor of length k that contains confidences |
| | k: number of points labeled with part_id |
| | count (int): number of samples to produce, should be positive and <= k |
| | |
| | Return: |
| | list(int): indices of values (along axis 1) selected as a sample |
| | """ |
| | k = values.shape[1] |
| | if k == count: |
| | index_sample = list(range(k)) |
| | else: |
| | |
| | |
| | |
| | _, sorted_confidence_indices = torch.sort(values[0]) |
| | if self.search_count_multiplier is not None: |
| | search_count = min(int(count * self.search_count_multiplier), k) |
| | elif self.search_proportion is not None: |
| | search_count = min(max(int(k * self.search_proportion), count), k) |
| | else: |
| | search_count = min(count, k) |
| | sample_from_top = random.sample(range(search_count), count) |
| | index_sample = sorted_confidence_indices[-search_count:][sample_from_top] |
| | return index_sample |
| |
|
| | 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 |
| | `DensePoseEmbeddingPredictorOutputWithConfidences` |
| | 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: a tensor of shape [1, H, W], DensePose CSE confidence |
| | """ |
| | _, _, w, h = bbox_xywh |
| | densepose_output = instance.pred_densepose |
| | mask, embeddings, _ = super()._produce_mask_and_results(instance, bbox_xywh) |
| | other_values = F.interpolate( |
| | getattr(densepose_output, self.confidence_channel), |
| | size=(h, w), |
| | mode="bilinear", |
| | )[0].cpu() |
| | return mask, embeddings, other_values |
| |
|
