| 57 | |
| 58 | |
| 59 | def reduce_straight_alpha( |
| 60 | rgba: Float[Tensor, "batch 4 height width"], |
| 61 | ) -> Float[Tensor, "batch 4"]: |
| 62 | color, alpha = rgba.split((3, 1), dim=1) |
| 63 | |
| 64 | # Color becomes a weighted average of color (weighted by alpha). |
| 65 | weighted_color = reduce(color * alpha, "b c h w -> b c", "sum") |
| 66 | alpha_sum = reduce(alpha, "b c h w -> b c", "sum") |
| 67 | color = weighted_color / (alpha_sum + 1e-10) |
| 68 | |
| 69 | # Alpha becomes mean alpha. |
| 70 | alpha = reduce(alpha, "b c h w -> b c", "mean") |
| 71 | |
| 72 | return torch.cat((color, alpha), dim=-1) |
| 73 | |
| 74 | |
| 75 | @torch.no_grad() |