import numpy as np
import cv2
from PIL import Image
import torchvision.transforms as transforms
import gradio as gr
from yolov5 import xai_yolov5
from yolov8 import xai_yolov8s

def process_image(image, yolo_versions=["yolov5"]):
    image = np.array(image)
    image = cv2.resize(image, (640, 640))

    result_images = []
    for yolo_version in yolo_versions:
        if yolo_version == "yolov5":
            result_images.append(xai_yolov5(image)) 
        elif yolo_version == "yolov8s":
            result_images.append(xai_yolov8s(image))
        else:
            result_images.append((Image.fromarray(image), f"{yolo_version} not yet implemented."))
    return result_images


interface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil", label="Upload an Image"),
        gr.CheckboxGroup(
            choices=["yolov5", "yolov8s", "yolov10"],
            value=["yolov5"],  # Set default selection to YOLOv5
            label="Select Model(s)",
        )
    ],
    outputs=gr.Gallery(label="Results", elem_id="gallery", rows=2, height=500),
    title="Explainable AI for YOLO Models",
    description="Upload an image to visualize YOLO object detection with Grad-CAM."
)

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