Spaces:
Sleeping
Sleeping
File size: 1,004 Bytes
b570453 | 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 | import torch
from ultralytics import YOLO
from PIL import Image
import numpy as np
def load_yolov8_model():
# Load the YOLOv8 model
model = YOLO('yolov8n.pt') # Using the smallest version for speed; adjust as needed
return model
def run_object_detection(model, image_path):
# Load the image
# image = Image.open(image_path).convert('RGB')
image_np = np.array(image_path)
# Run inference
results = model(image_np)
# Process results
detections = []
for result in results:
boxes = result.boxes
for box in boxes:
x1, y1, x2, y2 = box.xyxy[0].tolist() # Bounding box coordinates
conf = box.conf[0].item() # Confidence score
cls = int(box.cls[0].item()) # Class ID
label = model.names[cls] # Class name
detections.append({
'box': [x1, y1, x2, y2],
'confidence': conf,
'label': label
})
return detections |