一、YOLO V8 Cls在本专栏的前面文章中我们基于YOLO V8 Fine-Tuning训练了自定义的目标检测模型以及15点人脸关键点检测模型从结果中可以看出在模型如此轻量的同时还拥有者如此好的效果。本文基于yolov8n-cls模型实验Fine-Tuning训练花卉图像分类模型。YOLO V8的细节可以参考下面官方的介绍https://docs.ultralytics.com/zh/models/yolov8/#citations-and-acknowledgements本文依旧使用ultralytics框架进行训练和测试其中ultralytics和pytorch的版本如下torch1.13.1cu116ultralytics8.1.37YOLO V8 Cls调用示例如下测试图像这里使用yolov8n-pose模型如果模型不存在会自动下载fromultralyticsimportYOLOfrommatplotlibimportpyplotaspltimportrequests plt.rcParams[font.sans-serif][SimHei]defload_labels():urlhttps://storage.googleapis.com/download.tensorflow.org/data/imagenet_class_index.jsonresponserequests.get(url)ifresponse.status_code200:returnresponse.json()return[]defmain():# 加载模型modelYOLO(yolov8n-cls.pt)# 加载 ImageNet 类别标签labelsload_labels()# 图像预测imageplt.imread(./img/cat.jpg)resultsmodel.predict(image,devicecpu)probsresults[0].probs# top1的分类值classifyprobs.top1# top1的置信度confprobs.top1conf plt.imshow(image)plt.title(预测结果labels[str(classify)][1]概率str(%.2f%(conf)%))plt.show()if__name____main__:main()二、数据集准备和拆分本篇文章使用数千张花卉照片作为数据集共分为5个分类雏菊daisy、蒲公英dandelion、玫瑰roses、向日葵sunflowers、郁金香tulips。数据集下载地址https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz每个分类的图片放在单独的子目录下下载完毕后解压可以看到如下所示图像示例如下如果想要训练自己的图片也可以像这样的方式将每个类别的图片放在相应的子目录下。ultralytics对数据格式的解释如下https://docs.ultralytics.com/zh/datasets/classify/下面拆分80%的图像为训练集20%为验证集importosfromtqdmimporttqdmimportshutil# 数据集位置image_path./data/flower_photos# 训练集的比例training_ratio0.8# 拆分后数据的位置train_dirtrain_datadefsplit_data():forclassifyinos.listdir(image_path):ifnotos.path.isdir(os.path.join(image_path,classify)):continue# 创建目录os.makedirs(os.path.join(train_dir,train,classify),exist_okTrue)os.makedirs(os.path.join(train_dir,val,classify),exist_okTrue)imagesos.listdir(os.path.join(image_path,classify))# 数据拆分train_sizeint(len(images)*training_ratio)train_imagesimages[:train_size]val_imagesimages[train_size:]# 训练数据forimageintqdm(train_images,descDispose Train Classify: classify):shutil.copy(os.path.join(image_path,classify,image),os.path.join(train_dir,train,classify,image))# 验证数据forimageintqdm(val_images,descDispose Val Classify: classify):shutil.copy(os.path.join(image_path,classify,image),os.path.join(train_dir,val,classify,image))if__name____main__:split_data()处理后的结构三、训练使用ultralytics框架训练非常简单仅需三行代码即可完成训练fromultralyticsimportYOLO# Load a modelmodelYOLO(yolov8n-cls.pt)# load a pretrained model (recommended for training)# Train the modelmodel.train(data./train_data,# 数据集地址epochs30,# 训练的周期imgsz640,# 图像的大小device[0],# 设备如果是 cpu 则是 devicecpuworkers0,lr01e-4,# 学习率batch25,# 批次大小ampFalse# 是否启用混合精度训练)运行后可以看到打印的网络结构训练中训练结束后可以在runs目录下面看到训练的结果训练时loss的变化图四、模型预测fromultralyticsimportYOLOimportosimportrandomimportmatplotlib.pyplotasplt plt.rcParams[font.sans-serif][SimHei]defget_data(val_path):images[]forclassifyinos.listdir(val_path):forimageinos.listdir(os.path.join(val_path,classify)):image_pathos.path.join(val_path,classify,image)images.append({path:image_path,label:classify})# 打乱数据random.shuffle(images)# 9个拆分为一组return[images[i:i9]foriinrange(0,len(images),9)]defmain():val_path./train_data/val# 分类标签class_names[daisy,dandelion,roses,sunflowers,tulips]# 加载模型modelYOLO(runs/classify/train/weights/best.pt)# 获取数据集datasget_data(val_path)foritemindatas:plt.figure(figsize(10,10))fori,imageinenumerate(item):plt.subplot(3,3,i1)image_pathimage[path]labelimage[label]imgplt.imread(image_path)# 预测resultsmodel.predict(img,devicecpu)probsresults[0].probs classifyprobs.top1 confprobs.top1conf# 展示plt.imshow(img)plt.title(预测结果class_names[classify]概率str(%.2f%conf)真实结果label)plt.axis(off)plt.show()if__name____main__:main()预测效果