#! /usr/bin/env python3
# -*- coding: utf-8 -*-
# File : smooth.py
# Author : Jiayuan Mao
# Email : maojiayuan@gmail.com
# Date : 01/04/2019
#
# This file is part of Jacinle.
# Distributed under terms of the MIT license.
import math
import torch
from jacinle.utils.argument import get_2dshape
from .conv import CustomKernel, MaxPoolingKernelDef
__all__ = ['NormalizedBoxSmooth', 'normalized_box_smooth', 'GaussianSmooth', 'GaussianSmoothTruncated', 'gaussian_smooth', 'gaussian_smooth_truncated', 'MaximumSmooth', 'maximum_smooth']
[docs]
class NormalizedBoxSmooth(CustomKernel):
[docs]
def __init__(self, kernel_size, border_mode='reflect'):
self.kernel_size = get_2dshape(kernel_size)
super().__init__(self._gen_kernel(), border_mode=border_mode)
def _gen_kernel(self):
kernel = torch.ones(self.kernel_size, dtype=torch.float32)
kernel /= kernel.sum()
return kernel
[docs]
def normalized_box_smooth(image, kernel_size, border_mode='reflect'):
return NormalizedBoxSmooth(kernel_size, border_mode=border_mode).to(image.device)(image)
[docs]
class GaussianSmooth(CustomKernel):
[docs]
def __init__(self, kernel_size, sigma, border_mode='reflect'):
assert type(kernel_size) is int, 'GaussianSmooth supports only square kernel.'
self.kernel_size = kernel_size
self.sigma = float(sigma)
super().__init__(self._gen_kernel(), border_mode=border_mode)
def _gen_kernel(self):
# Create a x, y coordinate grid of shape (kernel_size, kernel_size, 2)
x_cord = torch.arange(self.kernel_size).float()
x_grid = x_cord.repeat(self.kernel_size).view(self.kernel_size, self.kernel_size)
y_grid = x_grid.t()
xy_grid = torch.stack([x_grid, y_grid], dim=-1)
mean = (self.kernel_size - 1) / 2.
variance = self.sigma ** 2.
# Calculate the 2-dimensional gaussian kernel which is
# the product of two gaussian distributions for two different
# variables (in this case called x and y)
Z = (1. / (2. * math.pi * variance))
gaussian_kernel = Z * torch.exp(
-torch.sum((xy_grid - mean) ** 2., dim=-1) /
(2 * variance)
)
# Make sure sum of values in gaussian kernel equals 1.
gaussian_kernel /= gaussian_kernel.sum()
return gaussian_kernel
[docs]
class GaussianSmoothTruncated(GaussianSmooth):
[docs]
def __init__(self, sigma, truncate=4, border_mode='reflect'):
sigma = float(sigma)
kernel_size = int(sigma * truncate + 0.5)
kernel_size = 2 * kernel_size + 1
super().__init__(kernel_size, sigma, border_mode=border_mode)
[docs]
def gaussian_smooth(image, kernel_size, sigma, border_mode='reflect'):
return GaussianSmooth(kernel_size, sigma, border_mode=border_mode).to(image.device)(image)
[docs]
def gaussian_smooth_truncated(image, sigma, truncate=4, border_mode='reflect'):
return GaussianSmoothTruncated(sigma, truncate=truncate, border_mode=border_mode).to(image.device)(image)
[docs]
class MaximumSmooth(CustomKernel):
[docs]
def __init__(self, kernel_size, border_mode='reflect'):
self.kernel_size = kernel_size
super().__init__(MaxPoolingKernelDef(self.kernel_size), border_mode=border_mode)
[docs]
def maximum_smooth(image, kernel_size, border_mode='reflect'):
return MaximumSmooth(kernel_size, border_mode=border_mode).to(image.device)(image)