Return the positions of all surface points with forward optical flow applied.
(
surfaces: Float[Tensor, "batch frame *grid xyz=3"],
extrinsics: Float[Tensor, "batch frame 4 4"],
intrinsics: Float[Tensor, "batch frame 3 3"],
)
| 141 | |
| 142 | |
| 143 | def compute_forward_flow( |
| 144 | surfaces: Float[Tensor, "batch frame *grid xyz=3"], |
| 145 | extrinsics: Float[Tensor, "batch frame 4 4"], |
| 146 | intrinsics: Float[Tensor, "batch frame 3 3"], |
| 147 | ) -> Float[Tensor, "batch frame-1 *grid xy=2"]: |
| 148 | """Return the positions of all surface points with forward optical flow applied.""" |
| 149 | |
| 150 | # Since the poses are camera-to-world and transformations are applied right to |
| 151 | # left, this can be understood as follows: First, transform from the earlier |
| 152 | # frame's camera space to world space. Then, transform from world space into the |
| 153 | # later frame's camera space. |
| 154 | forward_transformation = later(extrinsics).inverse() @ earlier(extrinsics) |
| 155 | |
| 156 | singletons = " ".join(["()"] * (surfaces.ndim - 3)) |
| 157 | pattern = f"b f i j -> b f {singletons} i j" |
| 158 | return reproject_points( |
| 159 | earlier(surfaces), |
| 160 | rearrange(forward_transformation, pattern), |
| 161 | rearrange(later(intrinsics), pattern), |
| 162 | ) |
| 163 | |
| 164 | |
| 165 | def compute_backward_flow( |
no test coverage detected