| 18 | |
| 19 | |
| 20 | class TrackPredictorCoTracker(TrackPredictor[TrackPredictorCoTrackerCfg]): |
| 21 | def __init__(self, cfg: TrackPredictorCoTrackerCfg) -> None: |
| 22 | super().__init__(cfg) |
| 23 | self.tracker = torch.hub.load( |
| 24 | "facebookresearch/co-tracker:v1.0", "cotracker_w8" |
| 25 | ) |
| 26 | |
| 27 | def forward( |
| 28 | self, |
| 29 | videos: Float[Tensor, "batch frame 3 height width"], |
| 30 | query_frame: int, |
| 31 | ) -> Tracks: |
| 32 | xy, visibility = self.tracker( |
| 33 | videos * 255, |
| 34 | grid_size=self.cfg.grid_size, |
| 35 | grid_query_frame=query_frame, |
| 36 | backward_tracking=True, |
| 37 | ) |
| 38 | |
| 39 | # Normalize the coordinates. |
| 40 | b, f, _, h, w = videos.shape |
| 41 | wh = torch.tensor((w, h), dtype=torch.float32, device=videos.device) |
| 42 | xy = xy / wh |
| 43 | |
| 44 | # Filter visibility based on RGB values. |
| 45 | rgb = F.grid_sample( |
| 46 | rearrange(videos, "b f c h w -> (b f) c h w"), |
| 47 | rearrange(2 * xy - 1, "b f p xy -> (b f) p () xy"), |
| 48 | mode="bilinear", |
| 49 | padding_mode="zeros", |
| 50 | align_corners=False, |
| 51 | ) |
| 52 | rgb = rearrange(rgb, "(b f) c p () -> b f p c", b=b, f=f) |
| 53 | rgb_delta = (rgb[:, [query_frame]] - rgb).abs().norm(dim=-1) |
| 54 | visibility = visibility & (rgb_delta < self.cfg.similarity_threshold) |
| 55 | |
| 56 | return Tracks(xy, visibility, 0) |
nothing calls this directly
no outgoing calls
no test coverage detected