(device: torch.device, birds: LoadedScene)
| 452 | |
| 453 | |
| 454 | def render_flows(device: torch.device, birds: LoadedScene) -> None: |
| 455 | # We just want the flow between images 0 and 4. |
| 456 | raft = FlowPredictorRaft(FlowPredictorRaftCfg("raft", 32, 8)).to(device) |
| 457 | flow = raft.forward( |
| 458 | birds.exports.colors[0, :1], |
| 459 | birds.exports.colors[0, SECOND_FRAME : SECOND_FRAME + 1], |
| 460 | ) |
| 461 | images = flow_to_color(rearrange(flow, "b h w xy -> b xy h w")) / 255 |
| 462 | save_image(images[0], "figures/flow_0_to_4.png") |
| 463 | |
| 464 | # Render a little legend thing |
| 465 | x = torch.linspace(-1, 1, 256, device=flow.device) |
| 466 | y = torch.linspace(-1, 1, 256, device=flow.device) |
| 467 | key = torch.stack(torch.meshgrid((x, y), indexing="xy"), dim=0) |
| 468 | save_image(flow_to_color(key) / 255, "figures/flow_key.png") |
| 469 | |
| 470 | |
| 471 | def render_joint_point_cloud( |
nothing calls this directly
no test coverage detected