| 17 | |
| 18 | |
| 19 | class BackboneExplicitDepth(Backbone[BackboneExplicitDepthCfg]): |
| 20 | def __init__( |
| 21 | self, |
| 22 | cfg: BackboneExplicitDepthCfg, |
| 23 | num_frames: int | None, |
| 24 | image_shape: tuple[int, int] | None, |
| 25 | ) -> None: |
| 26 | super().__init__(cfg, num_frames=num_frames, image_shape=image_shape) |
| 27 | depth = torch.full( |
| 28 | (num_frames, *image_shape), cfg.initial_depth, dtype=torch.float32 |
| 29 | ) |
| 30 | self.depth = nn.Parameter(depth) |
| 31 | weights = torch.full((num_frames - 1, *image_shape), 0, dtype=torch.float32) |
| 32 | self.weights = nn.Parameter(weights) |
| 33 | |
| 34 | def forward(self, batch: Batch, flows: Flows) -> BackboneOutput: |
| 35 | b, _, _, _, _ = batch.videos.shape |
| 36 | assert b == 1 |
| 37 | |
| 38 | return BackboneOutput( |
| 39 | self.depth[None], |
| 40 | (self.cfg.weight_sensitivity * self.weights).sigmoid()[None], |
| 41 | ) |
nothing calls this directly
no outgoing calls
no test coverage detected