(
self,
videos: Float[Tensor, "batch frame 3 height width"],
)
| 64 | self.model.load_state_dict(weights, strict=False) |
| 65 | |
| 66 | def forward( |
| 67 | self, |
| 68 | videos: Float[Tensor, "batch frame 3 height width"], |
| 69 | ) -> Float[Tensor, "batch frame-1 height width 2"]: |
| 70 | source, target, b, f = split_videos(videos) |
| 71 | |
| 72 | result = self.model( |
| 73 | source * 255, |
| 74 | target * 255, |
| 75 | attn_splits_list=[2], |
| 76 | corr_radius_list=[-1], |
| 77 | prop_radius_list=[-1], |
| 78 | pred_bidir_flow=False, |
| 79 | ) |
| 80 | flow = result["flow_preds"][-1] |
| 81 | |
| 82 | # Normalize the optical flow. |
| 83 | _, _, h, w = source.shape |
| 84 | wh = torch.tensor((w, h), dtype=torch.float32, device=flow.device) |
| 85 | return rearrange(flow, "(b f) xy h w -> b f h w xy", b=b, f=f - 1) / wh |
nothing calls this directly
no test coverage detected