Spaces:
Sleeping
Sleeping
File size: 1,153 Bytes
75f48fa |
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 |
from ultralytics import YOLO # type: ignore
import cv2
import numpy as np
# Load a pre-trained YOLOv8 model (nano version = small & fast)
model = YOLO("yolov8n.pt")
def detect_objects(image):
"""
Run YOLO on the input image.
Returns:
- annotated image with bounding boxes
- list of detected object names
"""
# Handle different image formats and channel counts
if isinstance(image, np.ndarray):
# If image has 4 channels (RGBA), convert to RGB
if image.shape[-1] == 4:
image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
# If image has 1 channel (grayscale), convert to RGB
elif len(image.shape) == 2 or image.shape[-1] == 1:
image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
results = model(image)
annotated_img = results[0].plot()
# Extract detected object names
detected_objects = []
for box in results[0].boxes:
cls_id = int(box.cls[0].item()) # class ID
label = results[0].names[cls_id] # class name
detected_objects.append(label)
return annotated_img, detected_objects
|