create_motion_blur_kernel.py

from typing import Tuple, Union

import numpy as np
from skimage.draw import line_aa
import tensorflow as tf


def create_motion_blur_kernel(dim: int, angle: float) -> tf.Tensor:
    # Define a disk which contains the dim x dim square
    radius = np.sqrt(0.5 * dim**2)
    x = int(radius + radius * np.cos(angle))
    y = int(radius + radius * np.sin(angle))
    
    # Draw the line from the center of the disk
    rr, cc, val = line_aa(int(radius), int(radius), x, y)
    
    # Make the outer square which contains the disk
    size = int(np.ceil(radius * 2))
    kernel = np.zeros((size, size), dtype=np.float32)
    
    # Burn the elements according to the drawn line
    kernel[rr, cc] = val
    
    # Switch to the inner square
    crop1 = int(radius) - dim // 2
    crop2 = int(radius) + dim // 2 + 1
    kernel = kernel[crop1:crop2, crop1:crop2]
    
    # Normalize to 1
    kernel /= np.sum(kernel)
    
    # Tiling and reshaping to fit tf.nn.depthwise_conv2d()
    return tf.tile(tf.constant(kernel)[:, :, None, None], [1, 1, 3, 1])