import gradio as gr
from stitching import Stitcher
import cv2
import numpy as np

def stitch_images(image_mode, image1, image2, images, detector, matcher, estimator, blender_type, crop, wave_correct):
    if image_mode == "Two Images":
        if image1 is None or image2 is None:
            return None
        cv_images = [cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) for img in [image1, image2]]
    else:  # Multiple Images
        if not images or len(images) < 2:
            return None
        cv_images = []
        for img_path in images:
            img = cv2.imread(img_path.name)
            if img is None:
                return None
            cv_images.append(img)
    
    try:
        stitcher = Stitcher(
            detector=detector,
            matcher_type=matcher,
            estimator=estimator,
            blender_type=blender_type,
            crop=crop,
            wave_correct_kind=wave_correct
        )
        panorama = stitcher.stitch(cv_images)
        
        # Convert back to RGB for display
        panorama_rgb = cv2.cvtColor(panorama, cv2.COLOR_BGR2RGB)
        return panorama_rgb
    except Exception as e:
        return None

def update_image_input(choice):
    if choice == "Two Images":
        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False)
    else:
        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)

title = """<h1 align="center">Stitching</h1>
<p><center>"Upload images to create a panorama. Adjust parameters for better results."</center></p>
"""

with gr.Blocks(theme='bethecloud/storj_theme') as iface:
    
    gr.HTML(title)
    
    with gr.Row():
        image_mode = gr.Radio(["Two Images", "Multiple Images"], label="Image Upload Mode", value="Two Images")
    
    with gr.Row():
        image1 = gr.Image(label="Image 1", height=300)
        image2 = gr.Image(label="Image 2", height=300)
        images = gr.File(file_count="multiple", label="Upload multiple images", visible=False)
    
    image_mode.change(update_image_input, image_mode, [image1, image2, images])
    
    with gr.Row():
        detector = gr.Dropdown(["orb", "sift", "surf", "akaze"], label="Feature Detector", value="orb")
        matcher = gr.Dropdown(["homography", "affine"], label="Matcher Type", value="homography")
        estimator = gr.Dropdown(["homography", "affine"], label="Estimator", value="homography")
    
    with gr.Row():
        blender_type = gr.Dropdown(["multiband", "feather", "no"], label="Blender Type", value="multiband")
        crop = gr.Checkbox(label="Crop Result", value=True)
        wave_correct = gr.Radio(["horiz", "vert", "no"], label="Wave Correction", value="horiz")
    
    stitch_button = gr.Button("Stitch Images")
    
    output_image = gr.Image(type="numpy", label="Stitched Panorama")
    
    stitch_button.click(
        stitch_images,
        inputs=[image_mode, image1, image2, images, detector, matcher, estimator, blender_type, crop, wave_correct],
        outputs=output_image
    )
    
    gr.Examples(
        examples=[
            ["Two Images", "exposure_error_1.jpg", "exposure_error_2.jpg", None, "orb", "homography", "homography", "multiband", True, "horiz"],
            ["Multiple Images", None, None, ["weir_1.jpg", "weir_2.jpg", "weir_3.jpg"], "orb", "homography", "homography", "multiband", True, "horiz"]
        ],
        inputs=[image_mode, image1, image2, images, detector, matcher, estimator, blender_type, crop, wave_correct],
        outputs=output_image,
        fn=stitch_images,
        cache_examples=True,
    )

iface.launch()