Spaces:
Sleeping
Sleeping
File size: 1,996 Bytes
538cbfd | 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import cv2
import torch
import numpy as np
import requests
from ultralytics import YOLO
import time
from pathlib import Path
def object_detection_image(path: str = None, url: str = None, model_id: str = "yolov8l.pt", output_dir: str = "."):
# GPU Detection
if torch.cuda.is_available():
device = "cuda"
print("✅ GPU:", torch.cuda.get_device_name(0))
imgsz=1280
elif torch.backends.mps.is_available():
device = "mps"
print("✅ Using Apple Silicon GPU")
imgsz=1280
else:
device = "cpu"
print("Using CPU as GPU is not available")
imgsz=640
model_id = "yolov8s.pt"
# Load model ON GPU
model = YOLO(model_id)
model.to(device)
if device == "cuda":
model.fuse() # speedup
# Load image from URL or path
if url is not None:
resp = requests.get(url, timeout=10)
image_np = np.frombuffer(resp.content, np.uint8)
image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
elif path is not None:
image = cv2.imread(path)
else:
url = "https://ultralytics.com/images/bus.jpg"
resp = requests.get(url, timeout=10)
image_np = np.frombuffer(resp.content, np.uint8)
image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
assert image is not None, "❌ Image decode failed"
# Run model
results = model.predict(
source=image,
imgsz=imgsz,
conf=0.25,
iou=0.5,
device=device,
verbose=True
)
# Draw & save
annotated = results[0].plot()
timestamp = int(time.time())
if not output_dir.endswith("/"):
output_dir += "/"
# Creating output directory if it doesn't exist
Path(output_dir).mkdir(parents=True, exist_ok=True)
output_path = f"{output_dir}/yolo_image_{timestamp}.png"
cv2.imwrite(output_path, annotated)
print("✅ Detection complete. Image saved at:", output_path)
return annotated
|