(x: torch.Tensor, radius: int)
| 53 | |
| 54 | |
| 55 | def _wavelet_blur(x: torch.Tensor, radius: int) -> torch.Tensor: |
| 56 | assert x.dim() == 4, 'x 必须是 (N, C, H, W)' |
| 57 | N, C, H, W = x.shape |
| 58 | base = _make_gaussian3x3_kernel(x.dtype, x.device) |
| 59 | weight = base.view(1, 1, 3, 3).repeat(C, 1, 1, 1) |
| 60 | pad = radius |
| 61 | x_pad = F.pad(x, (pad, pad, pad, pad), mode='replicate') |
| 62 | out = F.conv2d(x_pad, weight, bias=None, stride=1, padding=0, dilation=radius, groups=C) |
| 63 | return out |
| 64 | |
| 65 | |
| 66 | def _wavelet_decompose(x: torch.Tensor, levels: int = 5) -> Tuple[torch.Tensor, torch.Tensor]: |
no test coverage detected