| 25 | |
| 26 | |
| 27 | class FlowPredictor(nn.Module, ABC, Generic[T]): |
| 28 | def __init__(self, cfg: T) -> None: |
| 29 | super().__init__() |
| 30 | self.cfg = cfg |
| 31 | |
| 32 | @abstractmethod |
| 33 | def forward( |
| 34 | self, |
| 35 | videos: Float[Tensor, "batch frame 3 height width"], |
| 36 | ) -> Float[Tensor, "batch frame-1 height width 2"]: |
| 37 | pass |
| 38 | |
| 39 | @staticmethod |
| 40 | def rescale_flow( |
| 41 | flow: Float[Tensor, "batch frame height width 2"], |
| 42 | shape: tuple[int, int], |
| 43 | ) -> Float[Tensor, "batch frame height_scaled width_scaled 2"]: |
| 44 | b, f, _, _, _ = flow.shape |
| 45 | flow = rearrange(flow, "b f h w xy -> (b f) xy h w") |
| 46 | flow = F.interpolate(flow, shape, mode="bilinear", align_corners=False) |
| 47 | return rearrange(flow, "(b f) xy h w -> b f h w xy", b=b, f=f) |
| 48 | |
| 49 | @staticmethod |
| 50 | def rescale_mask( |
| 51 | mask: Float[Tensor, "batch frame height width"], |
| 52 | shape: tuple[int, int], |
| 53 | ) -> Float[Tensor, "batch frame height_scaled width_scaled"]: |
| 54 | b, f, _, _ = mask.shape |
| 55 | flow = rearrange(mask, "b f h w -> (b f) () h w") |
| 56 | flow = F.interpolate(flow, shape, mode="bilinear", align_corners=False) |
| 57 | return rearrange(flow, "(b f) () h w -> b f h w", b=b, f=f) |
| 58 | |
| 59 | @staticmethod |
| 60 | def compute_consistency_mask( |
| 61 | videos: Float[Tensor, "batch frame 3 height width"], |
| 62 | flow: Float[Tensor, "batch frame-1 height width 2"], |
| 63 | ) -> Float[Tensor, "batch frame-1 height width"]: |
| 64 | source, target, b, f = split_videos(videos) |
| 65 | |
| 66 | # Sample a target pixel for each source pixel. |
| 67 | _, _, h, w = source.shape |
| 68 | source_xy, _ = sample_image_grid((h, w), source.device) |
| 69 | target_xy = source_xy + rearrange(flow, "b f h w xy -> (b f) h w xy") |
| 70 | target_pixels = F.grid_sample( |
| 71 | target, |
| 72 | target_xy * 2 - 1, |
| 73 | mode="bilinear", |
| 74 | padding_mode="zeros", |
| 75 | align_corners=False, |
| 76 | ) |
| 77 | |
| 78 | # Map pixel color differences to mask weights. |
| 79 | deltas = (source - target_pixels).abs().max(dim=1).values |
| 80 | return rearrange((1 - deltas) ** 8, "(b f) h w -> b f h w", b=b, f=f - 1) |
| 81 | |
| 82 | def compute_bidirectional_flow( |
| 83 | self, |
| 84 | batch: Batch, |
nothing calls this directly
no outgoing calls
no test coverage detected