curiouscurrent's picture
Update app.py
fe86db6 verified
raw
history blame
1.87 kB
import gradio as gr
import subprocess
import os
from PIL import Image
# YOLOv8 detection function
def detect_objects(image):
# Save the uploaded image temporarily
input_image_path = "input.jpg"
image.save(input_image_path)
# Define the YOLOv8 command
output_dir = "./runs/detect"
os.makedirs(output_dir, exist_ok=True)
command = [
"yolo", # YOLOv8 CLI command
"task=detect", # Specify the task as detection
"mode=predict", # Set mode to predict
f"model=best.pt", # Path to your YOLOv8 model weights
f"source={input_image_path}", # Input image path (uploaded by user)
f"project={output_dir}", # Output directory
f"name=result", # Subfolder name for results
"exist_ok=True" # Allow overwriting existing results
]
# Run YOLOv8 inference
try:
subprocess.run(command, check=True)
except Exception as e:
return f"Error during YOLO inference: {e}"
# Get the path to the output image
output_image_path = os.path.join(output_dir, "result", os.path.basename(input_image_path))
# Check if the image exists
if os.path.exists(output_image_path):
# Return the image directly to Gradio
return output_image_path # Return file path for Gradio to display
else:
return "Error: Output image not found."
# Create the Gradio interface
interface = gr.Interface(
fn=detect_objects, # The function to call when an image is uploaded
inputs=gr.Image(type="pil"), # Accept images as input
outputs=gr.Image(type="filepath"), # Return image file path for display
title="YOLOv8 Object Detection",
description="Upload an image of floating waste in water, and this app will detect it using YOLOv8."
)
# Launch the Gradio app
interface.launch(share=True) # Use share=True to get a public URL