(window_size: int, sigma)
| 305 | |
| 306 | |
| 307 | def _gaussian(window_size: int, sigma): |
| 308 | if isinstance(sigma, float): |
| 309 | sigma = torch.tensor([[sigma]]) |
| 310 | |
| 311 | batch_size = sigma.shape[0] |
| 312 | |
| 313 | x = (torch.arange(window_size, device=sigma.device, |
| 314 | dtype=sigma.dtype) - window_size // 2).expand(batch_size, -1) |
| 315 | |
| 316 | if window_size % 2 == 0: |
| 317 | x = x + 0.5 |
| 318 | |
| 319 | gauss = torch.exp(-x.pow(2.0) / (2 * sigma.pow(2.0))) |
| 320 | |
| 321 | return gauss / gauss.sum(-1, keepdim=True) |
| 322 | |
| 323 | |
| 324 | def _gaussian_blur2d(input, kernel_size, sigma): |