import gradio as gr
import kornia as K
from skimage import io
from kornia.contrib import ImageStitcher
import kornia.feature as KF
import torch

def enhance(file_1, file_2):   
    img_1 = K.image_to_tensor(io.imread(file_1), False).float() / 255.
    img_2 = K.image_to_tensor(io.imread(file_2), False).float() / 255.
    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=enhance,
    inputs=inputs,
    outputs=outputs,
    title=title,
    examples=examples,
    live=True,
    theme='huggingface',
)
demo_app.launch(debug=True)