Spaces:
Running
on
Zero
Running
on
Zero
File size: 8,373 Bytes
61c2d32 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
# Copyright (c) Facebook, Inc. and its affiliates.
from typing import Any
import torch
from torch.nn import functional as F
from detectron2.config import CfgNode
from detectron2.layers import ConvTranspose2d
from ...structures import decorate_predictor_output_class_with_confidences
from ..confidence import DensePoseConfidenceModelConfig, DensePoseUVConfidenceType
from ..utils import initialize_module_params
class DensePoseChartConfidencePredictorMixin:
"""
Predictor contains the last layers of a DensePose model that take DensePose head
outputs as an input and produce model outputs. Confidence predictor mixin is used
to generate confidences for segmentation and UV tensors estimated by some
base predictor. Several assumptions need to hold for the base predictor:
1) the `forward` method must return SIUV tuple as the first result (
S = coarse segmentation, I = fine segmentation, U and V are intrinsic
chart coordinates)
2) `interp2d` method must be defined to perform bilinear interpolation;
the same method is typically used for SIUV and confidences
Confidence predictor mixin provides confidence estimates, as described in:
N. Neverova et al., Correlated Uncertainty for Learning Dense Correspondences
from Noisy Labels, NeurIPS 2019
A. Sanakoyeu et al., Transferring Dense Pose to Proximal Animal Classes, CVPR 2020
"""
def __init__(self, cfg: CfgNode, input_channels: int):
"""
Initialize confidence predictor using configuration options.
Args:
cfg (CfgNode): configuration options
input_channels (int): number of input channels
"""
# we rely on base predictor to call nn.Module.__init__
super().__init__(cfg, input_channels) # pyre-ignore[19]
self.confidence_model_cfg = DensePoseConfidenceModelConfig.from_cfg(cfg)
self._initialize_confidence_estimation_layers(cfg, input_channels)
self._registry = {}
initialize_module_params(self) # pyre-ignore[6]
def _initialize_confidence_estimation_layers(self, cfg: CfgNode, dim_in: int):
"""
Initialize confidence estimation layers based on configuration options
Args:
cfg (CfgNode): configuration options
dim_in (int): number of input channels
"""
dim_out_patches = cfg.MODEL.ROI_DENSEPOSE_HEAD.NUM_PATCHES + 1
kernel_size = cfg.MODEL.ROI_DENSEPOSE_HEAD.DECONV_KERNEL
if self.confidence_model_cfg.uv_confidence.enabled:
if self.confidence_model_cfg.uv_confidence.type == DensePoseUVConfidenceType.IID_ISO:
self.sigma_2_lowres = ConvTranspose2d( # pyre-ignore[16]
dim_in, dim_out_patches, kernel_size, stride=2, padding=int(kernel_size / 2 - 1)
)
elif (
self.confidence_model_cfg.uv_confidence.type
== DensePoseUVConfidenceType.INDEP_ANISO
):
self.sigma_2_lowres = ConvTranspose2d(
dim_in, dim_out_patches, kernel_size, stride=2, padding=int(kernel_size / 2 - 1)
)
self.kappa_u_lowres = ConvTranspose2d( # pyre-ignore[16]
dim_in, dim_out_patches, kernel_size, stride=2, padding=int(kernel_size / 2 - 1)
)
self.kappa_v_lowres = ConvTranspose2d( # pyre-ignore[16]
dim_in, dim_out_patches, kernel_size, stride=2, padding=int(kernel_size / 2 - 1)
)
else:
raise ValueError(
f"Unknown confidence model type: "
f"{self.confidence_model_cfg.confidence_model_type}"
)
if self.confidence_model_cfg.segm_confidence.enabled:
self.fine_segm_confidence_lowres = ConvTranspose2d( # pyre-ignore[16]
dim_in, 1, kernel_size, stride=2, padding=int(kernel_size / 2 - 1)
)
self.coarse_segm_confidence_lowres = ConvTranspose2d( # pyre-ignore[16]
dim_in, 1, kernel_size, stride=2, padding=int(kernel_size / 2 - 1)
)
def forward(self, head_outputs: torch.Tensor):
"""
Perform forward operation on head outputs used as inputs for the predictor.
Calls forward method from the base predictor and uses its outputs to compute
confidences.
Args:
head_outputs (Tensor): head outputs used as predictor inputs
Return:
An instance of outputs with confidences,
see `decorate_predictor_output_class_with_confidences`
"""
# assuming base class returns SIUV estimates in its first result
base_predictor_outputs = super().forward(head_outputs) # pyre-ignore[16]
# create output instance by extending base predictor outputs:
output = self._create_output_instance(base_predictor_outputs)
if self.confidence_model_cfg.uv_confidence.enabled:
if self.confidence_model_cfg.uv_confidence.type == DensePoseUVConfidenceType.IID_ISO:
# assuming base class defines interp2d method for bilinear interpolation
output.sigma_2 = self.interp2d(self.sigma_2_lowres(head_outputs)) # pyre-ignore[16]
elif (
self.confidence_model_cfg.uv_confidence.type
== DensePoseUVConfidenceType.INDEP_ANISO
):
# assuming base class defines interp2d method for bilinear interpolation
output.sigma_2 = self.interp2d(self.sigma_2_lowres(head_outputs))
output.kappa_u = self.interp2d(self.kappa_u_lowres(head_outputs)) # pyre-ignore[16]
output.kappa_v = self.interp2d(self.kappa_v_lowres(head_outputs)) # pyre-ignore[16]
else:
raise ValueError(
f"Unknown confidence model type: "
f"{self.confidence_model_cfg.confidence_model_type}"
)
if self.confidence_model_cfg.segm_confidence.enabled:
# base predictor outputs are assumed to have `fine_segm` and `coarse_segm` attributes
# base predictor is assumed to define `interp2d` method for bilinear interpolation
output.fine_segm_confidence = (
F.softplus(
self.interp2d(self.fine_segm_confidence_lowres(head_outputs)) # pyre-ignore[16]
)
+ self.confidence_model_cfg.segm_confidence.epsilon
)
output.fine_segm = base_predictor_outputs.fine_segm * torch.repeat_interleave(
output.fine_segm_confidence, base_predictor_outputs.fine_segm.shape[1], dim=1
)
output.coarse_segm_confidence = (
F.softplus(
self.interp2d(
self.coarse_segm_confidence_lowres(head_outputs) # pyre-ignore[16]
)
)
+ self.confidence_model_cfg.segm_confidence.epsilon
)
output.coarse_segm = base_predictor_outputs.coarse_segm * torch.repeat_interleave(
output.coarse_segm_confidence, base_predictor_outputs.coarse_segm.shape[1], dim=1
)
return output
def _create_output_instance(self, base_predictor_outputs: Any):
"""
Create an instance of predictor outputs by copying the outputs from the
base predictor and initializing confidence
Args:
base_predictor_outputs: an instance of base predictor outputs
(the outputs type is assumed to be a dataclass)
Return:
An instance of outputs with confidences
"""
PredictorOutput = decorate_predictor_output_class_with_confidences(
type(base_predictor_outputs) # pyre-ignore[6]
)
# base_predictor_outputs is assumed to be a dataclass
# reassign all the fields from base_predictor_outputs (no deep copy!), add new fields
output = PredictorOutput(
**base_predictor_outputs.__dict__,
coarse_segm_confidence=None,
fine_segm_confidence=None,
sigma_1=None,
sigma_2=None,
kappa_u=None,
kappa_v=None,
)
return output
|