yolov8快速实现场景分类

场景识别/场景分类（scene classification/recognition）目的是预测单帧2D图像的场景类别，比如足球场，滑雪场，沙滩等。基于yolov8，大概5-10分钟就可以实现。

框架

yolov8 第八代了，目前是ultralytics在持续开发，代码抽象程度越来越高了，好处就是高级封装后使用更简单，相对也更难魔改。如果有优化需求的建议结合NNI，具体可以看我过往发布的关于NNI的文章。

# Install the ultralytics package from PyPI
pip install ultralytics

数据集

SUN397，一共397各列别，10万左右的数据量。下载链接如下，请遵守数据原作者的开源协议，本次分享若侵权，请联系删除。

点击下载

代码

from ultralytics import YOLO
import argparse

argparser = argparse.ArgumentParser("SceneCls based on yolov8")
argparser.add_argument("--pred", action="store_true")
argparser.add_argument("--images", type=str)

# parse args
args = argparser.parse_args()

# Train the model
if not args.pred:
    model = YOLO('yolov8n-cls.pt')  # load a pretrained model (recommended for training)
    results = model.train(data='../Datasets/SUN397', 
                          epochs=50, 
                          imgsz=224, 
                          batch=512, 
                          project='SceneCls', 
                          name='20231221_v1.0.1',
                          save=True,
                          save_period=10,
                          workers=8,
                          device='3',
                          plots=True)
# Prediciton
else:
    assert args.images, "images must be non-None"
    model = YOLO('./SceneCls/20231221_v1.0.1/weights/best.pt')
    predicts = model([args.images])

实验结果