(
flow: Float[Tensor, "frame height width 2"],
)
| 19 | |
| 20 | |
| 21 | def flow_with_key( |
| 22 | flow: Float[Tensor, "frame height width 2"], |
| 23 | ) -> Float[Tensor, "3 height vis_width"]: |
| 24 | _, h, w, _ = flow.shape |
| 25 | length = min(h, w) |
| 26 | x = torch.linspace(-1, 1, length, device=flow.device) |
| 27 | y = torch.linspace(-1, 1, length, device=flow.device) |
| 28 | key = torch.stack(torch.meshgrid((x, y), indexing="xy"), dim=0) |
| 29 | flow = rearrange(flow, "f h w xy -> f xy h w") |
| 30 | return hcat( |
| 31 | *(flow_to_color(flow) / 255), |
| 32 | flow_to_color(key) / 255, |
| 33 | ) |
| 34 | |
| 35 | |
| 36 | @dataclass |