import gradio as gr
import kornia as K
from kornia.core import Tensor
import torch
import numpy as np

# Define Functions
def process_image(file):
    if isinstance(file, np.ndarray):
        # If the input is already a numpy array, convert it to a tensor
        img = K.image_to_tensor(file).float() / 255.0
    else:
        # If it's a file path, load it using kornia
        img = K.io.load_image(file, K.io.ImageLoadType.RGB32)
    return img.unsqueeze(0)  # Add batch dimension: 1xCxHxW

def box_blur_fn(file, box_blur):
    img = process_image(file)
    x_out: Tensor = K.filters.box_blur(img, (int(box_blur), int(box_blur)))
    return K.utils.tensor_to_image(x_out.squeeze())

def blur_pool2d_fn(file, blur_pool2d):
    img = process_image(file)
    x_out: Tensor = K.filters.blur_pool2d(img, int(blur_pool2d))
    return K.utils.tensor_to_image(x_out.squeeze())

def gaussian_blur_fn(file, gaussian_blur2d):
    img = process_image(file)
    x_out: Tensor = K.filters.gaussian_blur2d(img,
                                              (int(gaussian_blur2d), int(gaussian_blur2d)),
                                              (float(gaussian_blur2d)/2, float(gaussian_blur2d)/2))
    return K.utils.tensor_to_image(x_out.squeeze())

def max_blur_pool2d_fn(file, max_blur_pool2d):
    img = process_image(file)
    x_out: Tensor = K.filters.max_blur_pool2d(img, int(max_blur_pool2d))
    return K.utils.tensor_to_image(x_out.squeeze())

def median_blur_fn(file, median_blur):
    img = process_image(file)
    x_out: Tensor = K.filters.median_blur(img, (int(median_blur), int(median_blur)))
    return K.utils.tensor_to_image(x_out.squeeze())

# Define Examples
examples = [
    ["examples/monkey.jpg", 1],
    ["examples/pikachu.jpg", 1]
]

# Define Demos
box_blur_demo = gr.Interface(
    box_blur_fn,
    [
        gr.Image(type="numpy"),
        gr.Slider(minimum=1, maximum=20, step=1, value=10, label="Box Blur")
    ],
    "image",
    examples=examples,
)

blur_pool2d_demo = gr.Interface(
    blur_pool2d_fn,
    [
        gr.Image(type="numpy"),
        gr.Slider(minimum=1, maximum=40, step=1, value=20, label="Blur Pool")
    ],
    "image",
    examples=examples,
)

gaussian_blur_demo = gr.Interface(
    gaussian_blur_fn,
    [
        gr.Image(type="numpy"),
        gr.Slider(minimum=1, maximum=30, step=2, value=15, label="Gaussian Blur")
    ],
    "image",
    examples=examples,
)

max_blur_pool2d_demo = gr.Interface(
    max_blur_pool2d_fn,
    [
        gr.Image(type="numpy"),
        gr.Slider(minimum=1, maximum=40, step=1, value=20, label="Max Pool")
    ],
    "image",
    examples=examples,
)

median_blur_demo = gr.Interface(
    median_blur_fn,
    [
        gr.Image(type="numpy"),
        gr.Slider(minimum=1, maximum=30, step=2, value=15, label="Median Blur")
    ],
    "image",
    examples=examples,
)

# Create Interface
demo = gr.TabbedInterface(
    [
        box_blur_demo,
        blur_pool2d_demo,
        gaussian_blur_demo,
        max_blur_pool2d_demo,
        median_blur_demo
    ],
    [
        "Box Blur",
        "Blur Pool",
        "Gaussian Blur",
        "Max Pool",
        "Median Blur"
    ]
)

demo.launch()