Depth alignment in a dedicated nn.Module. Wrap scale_map_estimator to perform the conditional logic in a separated torch module outside the forward of RGBGaussianPredictor. This module can be then excluded during symbolic tracing.
| 20 | |
| 21 | |
| 22 | class DepthAlignment(nn.Module): |
| 23 | """Depth alignment in a dedicated nn.Module. |
| 24 | |
| 25 | Wrap scale_map_estimator to perform the conditional logic in a separated torch |
| 26 | module outside the forward of RGBGaussianPredictor. This module can be then |
| 27 | excluded during symbolic tracing. |
| 28 | """ |
| 29 | |
| 30 | def __init__(self, scale_map_estimator: nn.Module | None): |
| 31 | """Initialize DepthAlignmentWrapper. |
| 32 | |
| 33 | Args: |
| 34 | scale_map_estimator: Module to align monodepth to ground truth depth. |
| 35 | """ |
| 36 | super().__init__() |
| 37 | self.scale_map_estimator = scale_map_estimator |
| 38 | |
| 39 | def forward( |
| 40 | self, |
| 41 | monodepth: torch.Tensor, |
| 42 | depth: torch.Tensor, |
| 43 | depth_decoder_features: torch.Tensor | None = None, |
| 44 | ): |
| 45 | """Optionally align monodepth to ground truth with a local scale map. |
| 46 | |
| 47 | Args: |
| 48 | monodepth: The monodepth model with intermediate features to use. |
| 49 | depth: Ground truth depth to align predicted depth to. |
| 50 | depth_decoder_features: The (optional) monodepth decoder features. |
| 51 | """ |
| 52 | if depth is not None and self.scale_map_estimator is not None: |
| 53 | depth_alignment_map = self.scale_map_estimator( |
| 54 | monodepth[:, 0:1], depth, depth_decoder_features |
| 55 | ) |
| 56 | monodepth = depth_alignment_map * monodepth |
| 57 | else: |
| 58 | # Some losses rely on the presence of an alignment map. |
| 59 | # We ensure that they can be computed by creating a fake alignment map. |
| 60 | depth_alignment_map = torch.ones_like(monodepth) |
| 61 | return monodepth, depth_alignment_map |
| 62 | |
| 63 | |
| 64 | class RGBGaussianPredictor(nn.Module): |