指定一个类以在预训练模型上使用 YOLOv8 进行检测[关闭]

Question

我是 YOLOv8 的新手，我只想让模型仅检测某些类，而不是模型训练的所有 80 个类。如何指定 YOLOv8 模型仅检测一类？例如仅人。

from ultralytics import YOLO
model = YOLO('YOLOv8m.pt')

我记得我们可以用 YOLOv5 做到这一点，但我不能用 YOLOv8 做到同样的事情：

model = torch.hub.load("ultralytics/yolov5", 'custom', path='yolov5s.pt')
model.classes = [0]  # Only person
model.conf = 0.6

Answer 1

只需在

classes

中指定

predict

以及您想要预测的类 ID

from ultralytics.yolo.engine.model import YOLO

model = YOLO("yolov8n.pt")
model.predict(source="0", show=True, stream=True, classes=0)  # [0, 3, 5] for multiple classes
for i, (result) in enumerate(results):
    print('Do something with class 0')