YoLo/runing.py

33 lines
1.0 KiB
Python
Executable File

import cv2
from ultralytics import YOLO
# Load model đã huấn luyện
model = YOLO("src/server/train5/weights/best.pt")
# Đọc ảnh
image_path = "dataset/images/s-l1600.png"
img = cv2.imread(image_path)
# Chạy mô hình để dự đoán
results = model(image_path, conf=0.5) # Hạ conf xuống để giữ nhiều phát hiện hơn
# Xử lý kết quả dự đoán
for r in results:
for box in r.boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0]) # Lấy tọa độ bbox
# print(x1, y1, x2, y2)
conf = float(box.conf[0]) # Độ tin cậy
cls = int(box.cls[0]) # Nhãn class ID
label = f"{model.names[cls]} {conf:.2f}" # Tạo nhãn hiển thị
# Vẽ bounding box (màu xanh lá cây)
cv2.rectangle(img, (x1, y1), (x2, y2), (199, 80, 105), 2)
# Vẽ nhãn (tên object)
cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
# Hiển thị ảnh với bounding boxes
cv2.imshow("Detected Logos", img)
cv2.waitKey(0)
cv2.destroyAllWindows()