Spaces:
Runtime error
Runtime error
File size: 2,201 Bytes
192bcb9 954e12d 192bcb9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | from ultralytics import YOLO
from PIL import Image
import gradio as gr
import yaml
import os
import numpy as np
# Define paths using relative paths for Hugging Face Space deployment
model_weights_path = 'best.pt'
data_yaml_path = 'data.yaml'
# Load the trained YOLOv8 model
model = YOLO(model_weights_path)
# Load the data.yaml file to extract class names
with open(data_yaml_path, 'r') as f:
data_yaml_content = yaml.safe_load(f)
class_names = data_yaml_content['names']
def detect_municipal_problems(image: Image.Image) -> Image.Image:
"""
Performs object detection on an input image using the trained YOLOv8 model
and returns the image with detected bounding boxes and labels.
Args:
image (PIL.Image.Image): The input image to analyze.
Returns:
PIL.Image.Image: The image with detected objects and bounding boxes.
"""
print("Received image for detection.")
# Perform prediction using the loaded model
results = model.predict(source=image, imgsz=640, conf=0.25)
# Assuming only one image is processed at a time
if results:
annotated_image_np = results[0].plot() # plot() returns an RGB numpy array
# Convert the annotated NumPy array (RGB) back to PIL Image
annotated_image = Image.fromarray(annotated_image_np)
print("Detection complete. Image annotated.")
return annotated_image
else:
print("No detections found or an issue occurred during prediction.")
return image # Return original image if no detections or error
# Create the Gradio interface
interface = gr.Interface(
fn=detect_municipal_problems,
inputs=gr.Image(type='pil', label='Upload Image'),
outputs=gr.Image(type='pil', label='Detected Problems'),
title='Municipal Problem Detector using YOLOv8',
description='Upload an image to detect municipal problems like Potholes, Flooding, and Waste Management.',
live=True,
examples=['/content/dataset2/test/images/image_002.jpg'] # Example image from test set
)
# Launch the Gradio application
if __name__ == '__main__':
interface.launch(debug=True, share=True)
|