我写了一个脚本,实现了图片分类问题的全自动化训练

众所周知,图片分类问题属于计算机视觉中比较容易解决的问题之一

这几天被数据集的问题搞得焦头烂额,
照理说分类问题的数据集应该比较好制作

如果之前没有现成的数据集
也会变得比较麻烦
直到我偶然发现了一个HuggingFace的图片搜索API
无限次调用
而且不需要身份验证
真香
以后再也不怕组建数据集麻烦了
我在Colab上写了一个脚本,可以自动从API查询图片下载图片,构建数据集.最后使用YOLOv5的图片分类模型进行训练

特色

  • 只需写下想要识别的物体,后台会自动组建数据集
  • 自动化Yolov5模型训练

使用方法

在search_terms中填入我们想构建数据集的包含的类别

由于我想做一个花卉分类数据集,这里我填入的都是花卉的名称

在这里插入图片描述
然后依次运行我们的图片下载模块,数据集构建模块,YOLOv5训练模块即可完成对模型的训练

以下是通过脚本自动组建的数据集
我写了一个脚本,实现了图片分类问题的全自动化训练_第1张图片

以下是使用训练好的模型进行预测
我写了一个脚本,实现了图片分类问题的全自动化训练_第2张图片

脚本仓库地址

https://github.com/ZYM66/AutoTrain_ImageClassify_model_ByYOLOv5

你可能感兴趣的:(笔记,实用工具分享,分类,自动化,计算机视觉,yolov5)