| 1 | import math |
| 2 | |
| 3 | import torch |
| 4 | import torch.nn.functional as F |
| 5 | from einops import rearrange |
| 6 | |
| 7 | |
| 8 | class BlurDownsample(torch.nn.Module): |
| 9 | """ |
| 10 | Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel. |
| 11 | Applies only on H,W. Works for dims=2 or dims=3 (per-frame). |
| 12 | """ |
| 13 | |
| 14 | def __init__(self, dims: int, stride: int, kernel_size: int = 5) -> None: |
| 15 | super().__init__() |
| 16 | assert dims in (2, 3) |
| 17 | assert isinstance(stride, int) |
| 18 | assert stride >= 1 |
| 19 | assert kernel_size >= 3 |
| 20 | assert kernel_size % 2 == 1 |
| 21 | self.dims = dims |
| 22 | self.stride = stride |
| 23 | self.kernel_size = kernel_size |
| 24 | |
| 25 | # 5x5 separable binomial kernel using binomial coefficients [1, 4, 6, 4, 1] from |
| 26 | # the 4th row of Pascal's triangle. This kernel is used for anti-aliasing and |
| 27 | # provides a smooth approximation of a Gaussian filter (often called a "binomial filter"). |
| 28 | # The 2D kernel is constructed as the outer product and normalized. |
| 29 | k = torch.tensor([math.comb(kernel_size - 1, k) for k in range(kernel_size)]) |
| 30 | k2d = k[:, None] @ k[None, :] |
| 31 | k2d = (k2d / k2d.sum()).float() # shape (kernel_size, kernel_size) |
| 32 | self.register_buffer("kernel", k2d[None, None, :, :]) # (1, 1, kernel_size, kernel_size) |
| 33 | |
| 34 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 35 | if self.stride == 1: |
| 36 | return x |
| 37 | |
| 38 | if self.dims == 2: |
| 39 | return self._apply_2d(x) |
| 40 | else: |
| 41 | # dims == 3: apply per-frame on H,W |
| 42 | b, _, f, _, _ = x.shape |
| 43 | x = rearrange(x, "b c f h w -> (b f) c h w") |
| 44 | x = self._apply_2d(x) |
| 45 | h2, w2 = x.shape[-2:] |
| 46 | x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f, h=h2, w=w2) |
| 47 | return x |
| 48 | |
| 49 | def _apply_2d(self, x2d: torch.Tensor) -> torch.Tensor: |
| 50 | c = x2d.shape[1] |
| 51 | weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) # depthwise |
| 52 | x2d = F.conv2d(x2d, weight=weight, bias=None, stride=self.stride, padding=self.kernel_size // 2, groups=c) |
| 53 | return x2d |
| 54 |