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, "Please upload both images."
        cv_images = [cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) for img in [image1, image2]]
    else:  # Multiple Images
        if len(images) < 2:
            return None, "Please upload at least 2 images."
        cv_images = [cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) for img in images]
    
    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, "Stitching successful!"
    except Exception as e:
        return None, f"Stitching failed: {str(e)}"

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)

with gr.Blocks() as iface:
    gr.Markdown("# Advanced Image Stitcher")
    gr.Markdown("Upload images to create a panorama. Adjust parameters for better results.")
    
    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")
        image2 = gr.Image(label="Image 2")
        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")
    
    with gr.Row():
        output_image = gr.Image(type="numpy", label="Stitched Panorama")
        status = gr.Textbox(label="Status")
    
    stitch_button.click(
        stitch_images,
        inputs=[image_mode, image1, image2, images, detector, matcher, estimator, blender_type, crop, wave_correct],
        outputs=[output_image, status]
    )
    
    gr.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],
    )

iface.launch()