20200404316linhao/Day07/pred.py

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

# 加载我们训练的模型
model = YOLO("runs/detect/train5/weights/best.pt")
# 预测结果
result = model("0a6b02a4527e19eef8c6fd19497104e.jpg")  # 利用模型侦测图片中的目标

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

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

boxes = result[0].boxes  # 取一个目标
if len(boxes.cls) > 0:
    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("0a6b02a4527e19eef8c6fd19497104e.jpg")
    # 2. 标注图像
    img = cv2.rectangle(img, (x1, y1), (x2, y2), color=(0, 0, 255), thickness=5)
    # 3. 保存图像
    cv2.imwrite("1.jpg", img)