import gradio as gr
import cv2
import numpy as np
from typing import Union, List
from pathlib import Path
from PIL import Image
import torch

# Function to resize images

def resize_images(images, scale_percent=50):
    resized_images = []
    for img in images:
        width = int(img.shape[1] * scale_percent / 100)
        height = int(img.shape[0] * scale_percent / 100)
        dim = (width, height)
        resized = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)
        resized_images.append(resized)
    return resized_images

# Function to stitch images

def stitch_images(image_paths, scale_percent=50):
    images = [cv2.imread(path) for path in image_paths]
    resized_images = resize_images(images, scale_percent)
    stitcher = cv2.Stitcher_create()
    status, stitched_image = stitcher.stitch(resized_images)

    if status == cv2.Stitcher_OK:
        print("Stitching successful!")
        return stitched_image
    else:
        print(f"Stitching failed with status code: {status}")
        return None

# Main image processing function

def process_image(image_paths, scale_percent=50):
    stitched_image = stitch_images(image_paths, scale_percent)

    if stitched_image is not None:
        try:
            stitched_image_rgb = cv2.cvtColor(stitched_image, cv2.COLOR_BGR2RGB)
            return stitched_image_rgb
        except Exception as e:
            print(str(e))
            return stitched_image

# Gradio interface function

def gradio_stitch_and_detect(image_files):
    image_paths = [file.name for file in image_files]
    result_image = process_image(image_paths, scale_percent=50)

    if result_image is not None:
        result_image_rgb = result_image
        pil_image = Image.fromarray(result_image_rgb)
        pil_image.save("stitched_image.jpg", "JPEG")
        return pil_image, "stitched_image.jpg"

    return None, None

# Gradio interface
with gr.Blocks() as interface:
    gr.Markdown("<h1 style='color: #2196F3; text-align: center;'>Image Stitcher 🧵</h1>")
    gr.Markdown("<h3 style='color: #2196F3; text-align: center;'>=== Upload the images (.jpg, .png, etc) you want to stitch ===</h3>")

    image_upload = gr.Files(type="filepath", label="Upload Images")
    stitch_button = gr.Button("Stitch", variant="primary")
    stitched_image = gr.Image(type="pil", label="Stitched Image")
    download_button = gr.File(label="Download Stitched Image")

    stitch_button.click(gradio_stitch_and_detect, inputs=image_upload, outputs=[stitched_image, download_button])

interface.launch()