import gradio as gr
import kornia as K
from kornia.core import Tensor
from kornia.contrib import ImageStitcher
import kornia.feature as KF
import torch

def inference(file_1, file_2): 
    img_1: Tensor = K.io.load_image(file_1.name, K.io.ImageLoadType.RGB32)
    img_1 = img_1[None]  # 1xCxHxW / fp32 / [0, 1]
    img_2: Tensor = K.io.load_image(file_2.name, K.io.ImageLoadType.RGB32)
    img_2 = img_2[None]  # 1xCxHxW / fp32 / [0, 1]
    
    IS = ImageStitcher(KF.LoFTR(pretrained='outdoor'), estimator='ransac')
    with torch.no_grad():
        result = IS(img_1, img_2)
        
    return K.tensor_to_image(result[0])
    

examples = [
    ['examples/foto1B.jpg',
    'examples/foto1A.jpg'],
]


inputs = [
    gr.inputs.Image(type='file', label='Input Image'),
    gr.inputs.Image(type='file', label='Input Image'),
]

outputs = [
    gr.outputs.Image(type='file', label='Output Image'),

]

title = "Image Stitching using Kornia and LoFTR"

demo_app = gr.Interface(
    fn=inference,
    inputs=inputs,
    outputs=outputs,
    title=title,
    examples=examples,
    cache_examples=True,
    live=True,
    theme='huggingface',
)
demo_app.launch()