import gradio as gr
import kornia as K
import numpy as np
import kornia.geometry as KG
import torch


def geometry_transform(images: list,
                       translation: float,
                       scale: float,
                       angle: float) -> np.ndarray:

    file_names: list = [f.name for f in images]
    image_list: list = [K.io.load_image(f, K.io.ImageLoadType(0)).float().unsqueeze(0)/255 for f in file_names]
    if len(image_list) > 1:
        image_list = [K.geometry.resize(x, x.shape[-2:], antialias=True) for x in image_list]
        image_batch: torch.Tensor = torch.cat(image_list, 0)
    center: torch.Tensor = torch.tensor([x.shape[1:] for x in image_batch])/2
    translation = torch.tensor(translation).repeat(len(image_list), 2)
    scale = torch.tensor(scale).repeat(len(image_list), 2)
    angle = torch.tensor(angle).repeat(len(image_list))
    affine_matrix: torch.Tensor = KG.get_affine_matrix2d(translation, center, scale, angle)
    with torch.inference_mode():
        transformed: torch.Tensor = KG.transform.warp_affine(image_batch, affine_matrix[:, :2], dsize=image_batch.shape[2:])
    concat_images: list = torch.cat([x for x in transformed], dim=-1)
    final_images: np.ndarray = K.tensor_to_image(concat_images*255).astype(np.uint8) 

    return final_images

def main():

    title = """
    <h1 align="center">
    Geometry Image Transforms with Kornia!
    </h1>
    """

    with gr.Blocks() as demo:
        gr.Markdown(title)

        with gr.Row():
            images_input = gr.Files()
            with gr.Column():
                translation = gr.Number(label= "Translation")
                scale = gr.Number(label = "Scale", value= 1.0)
                angle = gr.Number(label = "Angle")
        
        button = gr.Button('Transform')
        result = gr.Image()

        button.click(
            geometry_transform,
            inputs=[images_input,translation, scale, angle],
            outputs=result
        )
    demo.launch()


if __name__ == '__main__':
    main()