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()