python用YOLOv8对图片进行分类

用yolov8的模型进行分类

先上效果图
图片资源
模型下载地址
 https://github.com/ultralytics/ultralytics
代码

import matplotlib.pyplot as plt
from ultralytics import YOLO
from PIL import Image
import cv2

model = YOLO('../ultralytics/yolov8n.pt')

# print(model.names)

img_path = '../ultralytics/ultralytics/assets/bus.jpg'

img = cv2.imread(img_path)
results = model(img_path)

preds = results[0].boxes.xyxy.cpu().numpy().astype('uint32')

colors = plt.cm.get_cmap('hsv', len(model.model.names))

for index, pred in enumerate(preds):
    i = int(results[0].boxes.cls[index].item())
    color = colors(i)
    color = (color[0] * 200, color[1] * 255, color[2] * 200)
    img = cv2.rectangle(img, (int(pred[0]), int(pred[1])), (int(pred[2]), int(pred[3])), color, 2)
    img = cv2.putText(img, results[0].names[i], (int(pred[0]), int(pred[1])), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)

cv2.imshow("1", img)
cv2.waitKey()
cv2.destroyAllWindows()

'''
print('..........')
print(len(results[0].names))                        # 所有的分类名
print(len(results[0].boxes.cls))                    # 结果分类
print(results[0].boxes.conf)                        # 置信度
print(results[0].boxes.cls)                         # 分类类别
print(type(int(results[0].boxes.cls[0].item())))    # 类别索引
print(results[0].names[0])                          # 分类的名称（可以根据类别索引进行获取）
print(results[0].boxes.xyxy)                        # 左上角和右下角坐标
'''