motion_blur_with_tf.function.py

from PIL import Image

def create_motion_blur_func(dim: int, angle: float):
    kernel = create_motion_blur_kernel(dim, angle)

    def motion_blur_func(images):
        return tf.nn.depthwise_conv2d(images, kernel, strides=[1] * 4, padding="SAME")
    
    return motion_blur_func


def load_image(path: str):
    data = np.array(Image.open(path).convert("RGB")).astype(np.float32)
    return tf.constant(data)[None, :, :, :]


def save_image(tensor: Union[np.ndarray, tf.Tensor], path: str):
    if hasattr(tensor, "numpy"):
        tensor = tensor.numpy()
    Image.fromarray(np.squeeze(tensor).astype(np.uint8), "RGB").save(path)


images = load_image("parrot.jpg")
motion_blur = create_motion_blur_func(25, (90 / 180) * np.pi)
save_image(motion_blur(images), "result.jpg")