概要
YOLOv8で物体検出する際に引数のclassesを調べた.
環境
Windows 11 Pro
Python 3.10.1
torch 2.0.1+cpu
背景
次のようにYOLOv8の既存モデルをCLI上で推論だけすると, デフォルトで様々なクラスラベルにより物体が検出される.
yolo predict model=yolov8x.pt source=source
対象物体に主にラベリングされたクラスだけを抽出して再度推論させたかったためclasses = []にクラスに対応する数値を入力する必要があったが, それの対応表がドキュメントのどこにあるか分からなかった！（知ってる方いたら教えてください）
やったこと&結果
・YOLOv8で物体検出モデルを試してみた
を参考に適当にmodelに画像を読み込ませて results から読み取る.
from ultralytics import YOLO
model = YOLO("yolov8x.pt")
#bus.jpgをカレントディレクトリに置いたがそうでない場合は適当なpathを指定
path = ""
results = model.predict(source=path+"bus.jpg", save=True)
results[0].names
結果は次の通り.
{0: 'person',
1: 'bicycle',
2: 'car',
3: 'motorcycle',
4: 'airplane',
5: 'bus',
6: 'train',
7: 'truck',
8: 'boat',
9: 'traffic light',
10: 'fire hydrant',
11: 'stop sign',
12: 'parking meter',
13: 'bench',
14: 'bird',
15: 'cat',
16: 'dog',
17: 'horse',
18: 'sheep',
19: 'cow',
20: 'elephant',
21: 'bear',
22: 'zebra',
23: 'giraffe',
24: 'backpack',
25: 'umbrella',
26: 'handbag',
27: 'tie',
28: 'suitcase',
29: 'frisbee',
30: 'skis',
31: 'snowboard',
32: 'sports ball',
33: 'kite',
34: 'baseball bat',
35: 'baseball glove',
36: 'skateboard',
37: 'surfboard',
38: 'tennis racket',
39: 'bottle',
40: 'wine glass',
41: 'cup',
42: 'fork',
43: 'knife',
44: 'spoon',
45: 'bowl',
46: 'banana',
47: 'apple',
48: 'sandwich',
49: 'orange',
50: 'broccoli',
51: 'carrot',
52: 'hot dog',
53: 'pizza',
54: 'donut',
55: 'cake',
56: 'chair',
57: 'couch',
58: 'potted plant',
59: 'bed',
60: 'dining table',
61: 'toilet',
62: 'tv',
63: 'laptop',
64: 'mouse',
65: 'remote',
66: 'keyboard',
67: 'cell phone',
68: 'microwave',
69: 'oven',
70: 'toaster',
71: 'sink',
72: 'refrigerator',
73: 'book',
74: 'clock',
75: 'vase',
76: 'scissors',
77: 'teddy bear',
78: 'hair drier',
79: 'toothbrush'}
これにより例えばperson, truck, hot dogだけを検出したい場合はCLI上で次のようにすればいい.
yolo predict model=yolov8x.pt source=source classes=[0, 7, 52]
余談
stream=True の場合はgeneratorが返ってくるよう. 最初こっちでやっててエラーをChatGPTに投げてnext関数というものを知りました.
results = model.predict(source=path+"bus.jpg", stream=True)
result = next(results)
result.names