File size: 2,023 Bytes
2915166 54fb288 2915166 54fb288 2915166 54fb288 |
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 |
from ultralytics import YOLO
from PIL import Image
import numpy as np
import cv2
import gradio as gr
import pandas as pd
trained_model_path = 'best.pt'
model = YOLO(trained_model_path)
def inference_function(image):
"""
Performs object detection inference on an input image using the loaded YOLOv8 model.
Returns:
tuple: (PIL Image with bounding boxes, Pandas DataFrame of detections)
OR ("No objects detected", empty DataFrame) if no detections
"""
pil_image = Image.fromarray(np.uint8(image)).convert('RGB')
results_list = model(pil_image, verbose=False)
results = results_list[0]
detections = []
image_np = np.array(pil_image)
image_with_boxes_np = image_np.copy()
if results.boxes:
for *xyxy, conf, cls in results.boxes.data:
x1, y1, x2, y2 = map(int, xyxy)
class_id = int(cls)
confidence = float(conf)
class_name = model.names[class_id]
detections.append({
'class_name': class_name,
'confidence': confidence,
'bbox': [x1, y1, x2, y2]
})
label = f"{class_name} {confidence:.2f}"
cv2.rectangle(image_with_boxes_np, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(image_with_boxes_np, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
detections_df = pd.DataFrame(detections)
image_with_boxes_pil = Image.fromarray(image_with_boxes_np)
return image_with_boxes_pil, detections_df
else:
return "No objects detected", pd.DataFrame()
# Gradio Interface
iface = gr.Interface(
fn=inference_function,
inputs=gr.Image(type="pil"),
outputs=[
gr.Image(label="Detected Objects"),
gr.DataFrame(label="Detections Table")
],
examples=['3_jpg.rf.c46999f92d679a134503e2c481ca4d9d.jpg', '10_jpg.rf.438680c2fe8e5e96e41932b3276761e2.jpg', '100_jpg.rf.37be664a78a2a0e331cef5cb77186ca7.jpg']
)
iface.launch() |