|
|
|
|
|
from dataclasses import make_dataclass |
|
from functools import lru_cache |
|
from typing import Any, Optional |
|
import torch |
|
|
|
|
|
@lru_cache(maxsize=None) |
|
def decorate_predictor_output_class_with_confidences(BasePredictorOutput: type) -> type: |
|
""" |
|
Create a new output class from an existing one by adding new attributes |
|
related to confidence estimation: |
|
- sigma_1 (tensor) |
|
- sigma_2 (tensor) |
|
- kappa_u (tensor) |
|
- kappa_v (tensor) |
|
- fine_segm_confidence (tensor) |
|
- coarse_segm_confidence (tensor) |
|
|
|
Details on confidence estimation parameters can be found in: |
|
N. Neverova, D. Novotny, A. Vedaldi "Correlated Uncertainty for Learning |
|
Dense Correspondences from Noisy Labels", p. 918--926, in Proc. NIPS 2019 |
|
A. Sanakoyeu et al., Transferring Dense Pose to Proximal Animal Classes, CVPR 2020 |
|
|
|
The new class inherits the provided `BasePredictorOutput` class, |
|
it's name is composed of the name of the provided class and |
|
"WithConfidences" suffix. |
|
|
|
Args: |
|
BasePredictorOutput (type): output type to which confidence data |
|
is to be added, assumed to be a dataclass |
|
Return: |
|
New dataclass derived from the provided one that has attributes |
|
for confidence estimation |
|
""" |
|
|
|
PredictorOutput = make_dataclass( |
|
BasePredictorOutput.__name__ + "WithConfidences", |
|
fields=[ |
|
("sigma_1", Optional[torch.Tensor], None), |
|
("sigma_2", Optional[torch.Tensor], None), |
|
("kappa_u", Optional[torch.Tensor], None), |
|
("kappa_v", Optional[torch.Tensor], None), |
|
("fine_segm_confidence", Optional[torch.Tensor], None), |
|
("coarse_segm_confidence", Optional[torch.Tensor], None), |
|
], |
|
bases=(BasePredictorOutput,), |
|
) |
|
|
|
|
|
|
|
def slice_if_not_none(data, item): |
|
if data is None: |
|
return None |
|
if isinstance(item, int): |
|
return data[item].unsqueeze(0) |
|
return data[item] |
|
|
|
def PredictorOutput_getitem(self, item): |
|
PredictorOutput = type(self) |
|
base_predictor_output_sliced = super(PredictorOutput, self).__getitem__(item) |
|
return PredictorOutput( |
|
**base_predictor_output_sliced.__dict__, |
|
coarse_segm_confidence=slice_if_not_none(self.coarse_segm_confidence, item), |
|
fine_segm_confidence=slice_if_not_none(self.fine_segm_confidence, item), |
|
sigma_1=slice_if_not_none(self.sigma_1, item), |
|
sigma_2=slice_if_not_none(self.sigma_2, item), |
|
kappa_u=slice_if_not_none(self.kappa_u, item), |
|
kappa_v=slice_if_not_none(self.kappa_v, item), |
|
) |
|
|
|
PredictorOutput.__getitem__ = PredictorOutput_getitem |
|
|
|
def PredictorOutput_to(self, device: torch.device): |
|
""" |
|
Transfers all tensors to the given device |
|
""" |
|
PredictorOutput = type(self) |
|
base_predictor_output_to = super(PredictorOutput, self).to(device) |
|
|
|
def to_device_if_tensor(var: Any): |
|
if isinstance(var, torch.Tensor): |
|
return var.to(device) |
|
return var |
|
|
|
return PredictorOutput( |
|
**base_predictor_output_to.__dict__, |
|
sigma_1=to_device_if_tensor(self.sigma_1), |
|
sigma_2=to_device_if_tensor(self.sigma_2), |
|
kappa_u=to_device_if_tensor(self.kappa_u), |
|
kappa_v=to_device_if_tensor(self.kappa_v), |
|
fine_segm_confidence=to_device_if_tensor(self.fine_segm_confidence), |
|
coarse_segm_confidence=to_device_if_tensor(self.coarse_segm_confidence), |
|
) |
|
|
|
PredictorOutput.to = PredictorOutput_to |
|
return PredictorOutput |
|
|
