(
exports: ModelExports,
z_value: float,
)
| 39 | |
| 40 | |
| 41 | def get_frustums( |
| 42 | exports: ModelExports, |
| 43 | z_value: float, |
| 44 | ) -> Float[Tensor, "frame endpoint=2 line=8 xyz=3"]: |
| 45 | # Generate xy points at the corners. |
| 46 | xy, _ = sample_image_grid((2, 2), device=device) |
| 47 | xy = xy * 2 - 0.5 |
| 48 | |
| 49 | # Un-project the corners to the specified Z value. |
| 50 | rays = unproject( |
| 51 | xy, |
| 52 | torch.ones_like(xy[..., 0]), |
| 53 | rearrange(exports.intrinsics, "b f i j -> b f () () i j"), |
| 54 | ) |
| 55 | rays = rays / rays[..., -1:] * z_value |
| 56 | |
| 57 | # Add the camera origin. |
| 58 | rays = rearrange(rays[0, 0], "h w xyz -> (h w) xyz") |
| 59 | rays = torch.cat((torch.zeros_like(rays[:1]), rays), dim=0) |
| 60 | |
| 61 | # Convert the points into world space. |
| 62 | rays = einsum( |
| 63 | exports.extrinsics[0], |
| 64 | homogenize_points(rays), |
| 65 | "f i j, p j -> f p i", |
| 66 | )[..., :3] |
| 67 | |
| 68 | # Aggregate the lines needed to make frustums. |
| 69 | o, a, b, c, d = rays.unbind(dim=-2) |
| 70 | lines = [ |
| 71 | [a, b], |
| 72 | [b, d], |
| 73 | [d, c], |
| 74 | [c, a], |
| 75 | [o, c], |
| 76 | [o, a], |
| 77 | [o, b], |
| 78 | [o, d], |
| 79 | ] |
| 80 | return torch.stack([torch.stack(line, dim=1) for line in lines], dim=-2) |
| 81 | |
| 82 | |
| 83 | def project( |
no test coverage detected