ict-app-1 / app.py
xgenatiik's picture
Update app.py
bdc6c71 verified
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