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
import numpy as np

def inference(img_1, img_2):
    # Convert numpy arrays to Tensors and ensure correct shape
    img_1: Tensor = K.image_to_tensor(img_1, keepdim=False).float() / 255.0
    img_2: Tensor = K.image_to_tensor(img_2, keepdim=False).float() / 255.0
    
    # Ensure 3D tensors (C, H, W)
    if img_1.ndim == 2:
        img_1 = img_1.unsqueeze(0)
    if img_2.ndim == 2:
        img_2 = img_2.unsqueeze(0)
    
    # Ensure 3 channel images
    if img_1.shape[0] == 1:
        img_1 = img_1.repeat(3, 1, 1)
    if img_2.shape[0] == 1:
        img_2 = img_2.repeat(3, 1, 1)
    
    # Add batch dimension
    img_1 = img_1.unsqueeze(0)
    img_2 = img_2.unsqueeze(0)
    
    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'],
]

with gr.Blocks(theme='huggingface') as demo_app:
    gr.Markdown("# Image Stitching using Kornia and LoFTR")
    with gr.Row():
        input_image1 = gr.Image(label="Input Image 1")
        input_image2 = gr.Image(label="Input Image 2")
    output_image = gr.Image(label="Output Image")
    stitch_button = gr.Button("Stitch Images")
    stitch_button.click(fn=inference, inputs=[input_image1, input_image2], outputs=output_image)
    gr.Examples(examples=examples, inputs=[input_image1, input_image2])

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