(
videos: Float[Tensor, "batch frame 3 height width"],
flow: Float[Tensor, "batch frame-1 height width 2"],
)
| 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, |
no test coverage detected