import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np

# Load YOLO model
model = YOLO("best.pt")

# Define the prediction function
def detect_species(files):
    """
    Detect species in uploaded images using the YOLO model.
    
    Args:
        files (list): List of uploaded image files.
    
    Returns:
        list: Annotated images or error messages for each file.
    """
    annotated_images = []
    for file in files:  # Loop through the list of uploaded files
        try:
            # Open and process the image
            image = Image.open(file)
            image = np.array(image)  # Convert to numpy array
            results = model(image)  # Run the YOLO model
            annotated_image = results[0].plot()  # Generate annotated image
            annotated_images.append(annotated_image)  # Add to results
        except Exception as e:
            # Handle errors (e.g., unsupported or corrupt files)
            error_message = f"Error processing file {file.name}: {str(e)}"
            print(error_message)  # Log the error
            annotated_images.append(np.zeros((100, 100, 3)))  # Placeholder image
    return annotated_images

# Create the Gradio app
app = gr.Interface(
    fn=detect_species,
    inputs=gr.Files(file_types=["image"], label="Upload Images"),  # Allow multiple image uploads
    outputs=gr.Gallery(label="Detection Results"),  # Display results in a gallery
    title="Finnish Meadow Plants Detection",
    description="Upload one or more leaf images, and the model will detect the species.",
    examples=["example1.jpg", "example2.jpg"],  # Optional: Add example images
)

# Launch the app
if __name__ == "__main__":
    app.launch(share=True)  # Set share=True to get a public URL