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