(
points: Float[Tensor, "*batch xyz=3"],
midpoint: Float[Tensor, "xyz=3"],
scale: float,
canvas_size: float | int,
)
| 81 | |
| 82 | |
| 83 | def project( |
| 84 | points: Float[Tensor, "*batch xyz=3"], |
| 85 | midpoint: Float[Tensor, "xyz=3"], |
| 86 | scale: float, |
| 87 | canvas_size: float | int, |
| 88 | ) -> tuple[ |
| 89 | Float[Tensor, "*batch xy=2"], # projected (x, y) |
| 90 | Int64[Tensor, " *batch"], # index (in depth ordering) |
| 91 | ]: |
| 92 | # Project the points (isometric). |
| 93 | projection = [ |
| 94 | [2, 1], |
| 95 | [0, 2.25], |
| 96 | [2, -1], |
| 97 | ] |
| 98 | projection = torch.tensor(projection, dtype=torch.float32, device=device) |
| 99 | xy = einsum(projection, points - midpoint, "i j, ... i -> ... j") |
| 100 | xy = xy * scale + 0.5 * canvas_size |
| 101 | |
| 102 | # Figure out the correct depth ordering. |
| 103 | look = torch.tensor((1, -1, -1), dtype=torch.float32, device=device) |
| 104 | depth = einsum(points, look, "... xyz, xyz -> ...") |
| 105 | ordering = depth.view(-1).argsort().view(depth.shape) |
| 106 | |
| 107 | return xy, ordering |
| 108 | |
| 109 | |
| 110 | def render_frustums( |
nothing calls this directly
no outgoing calls
no test coverage detected