# 侦测
from ultralytics import YOLO
import cv2
import numpy as np

# 加载我们训练的模型
model = YOLO("best.pt")
# 预测结果
result = model("IMG_3052.JPG")  # 利用模型侦测图片中的目标

# print(result[0])   # [0]表示第一张图像的识别结果

names = result[0].names # 类别对应的名字
print(names)

boxes = result[0].boxes  # 取一个目标

cls = int(boxes.cls[0].cpu().item())   
conf = boxes.conf[0].cpu().item()

x1, y1, x2, y2 = boxes.xyxy[0].cpu().numpy().astype(np.int32)

print(cls)              # 类别ID
print(conf)             # 目标侦测成功的置信度 
print(x1, y1, x2, y2)   # 侦测目标的区域
print(names[cls])   # 获取侦测的目标名字


# 把目标在图像中标注出来
# 1. 打开图像
img = cv2.imread("IMG_3052.JPG")
# 2. 标注图像
img = cv2.rectangle(img, (x1, y1), (x2, y2), color=(0, 0, 255), thickness=5)
# 3. 保存图像
cv2.imwrite("1.jpg", img)