如何限制默认的“yolo”模型识别更少的物体？

Question

我想使用默认的 YOLOv8 模型 (yolov8m.pt) 进行对象检测。我知道默认的 YOLO 模型使用 COCO 数据集，可以检测 80 多个对象。我只想检测其中 5 个，我怎样才能实现这一目标？

Answer 1

要在预测模式下过滤类别，请使用

classes

函数的

predict

参数，您需要在其中指定相关的类别 ID。

类，列表[int]。将预测过滤为一组类 ID。仅返回属于指定类别的检测。对于在多类检测任务中关注相关对象很有用。 https://docs.ultralytics.com/modes/predict/#inference-arguments

COCO 数据集的类列表可以在这里找到：https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/datasets/coco.yaml.

from ultralytics import YOLO

model = YOLO('yolov8n.pt')
results = model.predict(source='img.jpg', classes=[0, 3, 5])