【ひとりで】画像分類できるもん【yolov8】

画像分類するなら、ラクで精度が良いこれで

画像分類っていざするとなるとどのモデルでやればいいかわかんなくないですか？
そんなあなたにこれ。

インストール

pip install ultralytics

トレーニング

以下の形式でフォルダにデータを分けます。

cifar-10-/
|
|-- train/
|   |-- cat/
|   |   |-- 10008_cat.png
|   |   |-- 10009_cat.png
|   |   |-- ...
|   |
|   |-- person/
|   |   |-- 1000_person.png
|   |   |-- 1001_person.png
|   |   |-- ...
|
|-- val/
|   |-- cat/
|   |   |-- 20008_cat.png
|   |   |-- 20009_cat.png
|   |   |-- ...
|   |
|   |-- person/
|   |   |-- 2000_person.png
|   |   |-- 2001_person.png
|   |   |-- ...

あとは、これでok.

from ultralytics import YOLO

# Load a model
model = YOLO("yolov8m-cls.pt")  # load a pretrained model (recommended for training)

# Train the model
results = model.train(data="my_dataset", epochs=100, imgsz=640)

データセットのパスを指定して実行するだけ。
runs/classify/train/weights/にbest.ptファイルが保存されます。

推論

model = YOLO(best.pt)
results = model("image.jpg")
cls = results[0].probs.top1
label = results[0].names[cls]

0
"cat"

1
"person"