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

# Function to stitch uploaded images
def stitch_images(images):
    stitcher = Stitcher()
    panorama = stitcher.stitch(images)  # Stitch the images
    
    # Convert the result to a PIL image
    pil_image = Image.fromarray(cv2.cvtColor(panorama, cv2.COLOR_BGR2RGB))
    
    # Save the image in JPG format
    pil_image.save("stitched_image.jpg", "JPEG")
    
    # Return the stitched image
    return pil_image, "stitched_image.jpg"  # Return the image and the file path

# Create Gradio interface with button
with gr.Blocks() as interface:
    gr.Markdown("<h1 style='color: #2196F3; text-align: center;'>Image Stitcher 🧵</h1>")  # Heading with a new color
    gr.Markdown("<h3 style='color: #2196F3; text-align: center;'>=== Upload the images you want to stitch ===</h3>")  # Subheading
    
    image_upload = gr.Files(type="filepath", label="Upload Images")  # File input for images
    stitch_button = gr.Button("Stitch", variant="primary")  # Button with color
    stitched_image = gr.Image(type="pil", label="Stitched Image")  # Display stitched image
    download_button = gr.File(label="Download Stitched Image")  # Button to download the image

    # Define button interaction
    stitch_button.click(stitch_images, inputs=image_upload, outputs=[stitched_image, download_button])

# Launch the interface
interface.launch()