|
|
|
|
|
import random |
|
from typing import Optional, Tuple |
|
import torch |
|
|
|
from densepose.converters import ToChartResultConverterWithConfidences |
|
|
|
from .densepose_base import DensePoseBaseSampler |
|
|
|
|
|
class DensePoseConfidenceBasedSampler(DensePoseBaseSampler): |
|
""" |
|
Samples DensePose data from DensePose predictions. |
|
Samples for each class are drawn using confidence value estimates. |
|
""" |
|
|
|
def __init__( |
|
self, |
|
confidence_channel: str, |
|
count_per_class: int = 8, |
|
search_count_multiplier: Optional[float] = None, |
|
search_proportion: Optional[float] = None, |
|
): |
|
""" |
|
Constructor |
|
|
|
Args: |
|
confidence_channel (str): confidence channel to use for sampling; |
|
possible values: |
|
"sigma_2": confidences for UV values |
|
"fine_segm_confidence": confidences for fine segmentation |
|
"coarse_segm_confidence": confidences for coarse segmentation |
|
(default: "sigma_2") |
|
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__(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): an array of size [n, k] that contains |
|
estimated values (U, V, confidences); |
|
n: number of channels (U, V, 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[2]) |
|
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_labels_and_results(self, instance) -> Tuple[torch.Tensor, torch.Tensor]: |
|
""" |
|
Method to get labels and DensePose results from an instance, with confidences |
|
|
|
Args: |
|
instance (Instances): an instance of `DensePoseChartPredictorOutputWithConfidences` |
|
|
|
Return: |
|
labels (torch.Tensor): shape [H, W], DensePose segmentation labels |
|
dp_result (torch.Tensor): shape [3, H, W], DensePose results u and v |
|
stacked with the confidence channel |
|
""" |
|
converter = ToChartResultConverterWithConfidences |
|
chart_result = converter.convert(instance.pred_densepose, instance.pred_boxes) |
|
labels, dp_result = chart_result.labels.cpu(), chart_result.uv.cpu() |
|
dp_result = torch.cat( |
|
(dp_result, getattr(chart_result, self.confidence_channel)[None].cpu()) |
|
) |
|
|
|
return labels, dp_result |
|
|
