| 27 | |
| 28 | |
| 29 | def detect_msaa_pixels( |
| 30 | image: Float[Tensor, "batch 4 height width"], |
| 31 | ) -> Bool[Tensor, "batch height width"]: |
| 32 | b, _, h, w = image.shape |
| 33 | |
| 34 | mask = torch.zeros((b, h, w), dtype=torch.bool, device=image.device) |
| 35 | |
| 36 | # Detect horizontal differences. |
| 37 | horizontal = (image[:, :, :, 1:] != image[:, :, :, :-1]).any(dim=1) |
| 38 | mask[:, :, 1:] |= horizontal |
| 39 | mask[:, :, :-1] |= horizontal |
| 40 | |
| 41 | # Detect vertical differences. |
| 42 | vertical = (image[:, :, 1:, :] != image[:, :, :-1, :]).any(dim=1) |
| 43 | mask[:, 1:, :] |= vertical |
| 44 | mask[:, :-1, :] |= vertical |
| 45 | |
| 46 | # Detect diagonal (top left to bottom right) differences. |
| 47 | tlbr = (image[:, :, 1:, 1:] != image[:, :, :-1, :-1]).any(dim=1) |
| 48 | mask[:, 1:, 1:] |= tlbr |
| 49 | mask[:, :-1, :-1] |= tlbr |
| 50 | |
| 51 | # Detect diagonal (top right to bottom left) differences. |
| 52 | trbl = (image[:, :, :-1, 1:] != image[:, :, 1:, :-1]).any(dim=1) |
| 53 | mask[:, :-1, 1:] |= trbl |
| 54 | mask[:, 1:, :-1] |= trbl |
| 55 | |
| 56 | return mask |
| 57 | |
| 58 | |
| 59 | def reduce_straight_alpha( |