import gradio as gr
import kornia as K
from kornia.core import concatenate, Tensor
from PIL import Image
import numpy as np

def rescale_aa(image, height, width):
    # Convert PIL Image to Tensor
    img_np = np.array(image)
    img: Tensor = K.utils.image_to_tensor(img_np).float() / 255.0
    img = img.unsqueeze(0)  # Add batch dimension
    
    img_rescale: Tensor = K.geometry.resize(img, (int(img.shape[2]*height), int(img.shape[3]*width)), antialias=False)
    img_rescale_aa: Tensor = K.geometry.resize(img, (int(img.shape[2]*height), int(img.shape[3]*width)), antialias=True)
    img_out = concatenate([img_rescale, img_rescale_aa], -1)
    
    # Clip the output values from 0 to 1
    return K.utils.tensor_to_image(img_out.clamp_(0, 1))

title = "Kornia Resizing Demo"
description = """
This demo shows the difference between resizing with and without antialiasing using Kornia.
The left half of the output image is resized without antialiasing, and the right half is resized with antialiasing.
"""

with gr.Blocks(title=title) as demo:
    gr.Markdown(f"# {title}")
    gr.Markdown(description)
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Input Image")
            height_slider = gr.Slider(minimum=0.005, maximum=2, step=0.005, value=0.25, label="Height Scale")
            width_slider = gr.Slider(minimum=0.005, maximum=2, step=0.005, value=0.25, label="Width Scale")
        
        image_output = gr.Image(type="pil", label="Resized Image")
    
    rescale_button = gr.Button("Rescale")
    rescale_button.click(
        fn=rescale_aa,
        inputs=[image_input, height_slider, width_slider],
        outputs=image_output
    )
    
    gr.Examples(
        examples=[
            ["examples/a.png", 0.25, 0.25],
            ["examples/iron_man.jpeg", 0.25, 0.25],
        ],
        inputs=[image_input, height_slider, width_slider],
        outputs=image_output,
        fn=rescale_aa,
        cache_examples=True
    )

if __name__ == "__main__":
    demo.launch()