本篇将代码放在最后
目录
1.数据集制作
可以参考我的另一篇文章采用爬虫制作所需数据集
YOLO和MobileNetV3实现四种天气分类(雨,雪,雾,晴)
此次分类虽说是铺装道路与非铺装道路两类,但从整体项目上看,数据集的搜集要考虑以下几方面:
- 光照条件,白天和夜间
- 天气情况,如雨天,雪天,将路面覆盖后均按照未铺装道路分类
考虑到这两点,在搜集数据集时就可以爬取白天和夜间的公路,以及白天,夜间,恶劣天气条件下的土路,这样构成了两类数据集。再对其进行裁剪,修改成同一大小备用:
划分数据集:split_data.py

将图片放在images中的两类文件夹下,paved和unpaved,运行split_data.py将会划分出训练集和验证集。
2.开始训练
打开classify中的train.py

修改这三个地方,改为自己存放分类模型的位置,存放数据集的位置,训练轮数

训练之后会在runs里生成权重文件
3.测试
在predict.py中做如下修改,就可以运行测试了,同样也是修改路径

因为后续想要在同一张图上输出多个标签,可以修改一下标签的大小和颜色
可以根据RGB常用颜色表 将标签修改为自己想要的颜色:



结果展示:

修改此处可修改标签位置和颜色:

修改后输出:

如图所示,标签变到另一边了
将模型转为onnx模型
python export.py --weights F:/yolov5/runs/train-cls/exp/weights/best.pt --include onnx
4.代码附录
(1)爬虫代码
# -*- coding:utf8 -*-
import requests
import json
from urllib import parse
import os
import time
class BaiduImageSpider(object):
def __init__(self):
self.json_count = 0 # 请求到的json文件数量(一个json文件包含30个图像文件)
self.url = 'https://image.baidu.com/search/acjson?tn=resultjson_com&logid=5179920884740494226&ipn=rj&ct' \
'=201326592&is=&fp=result&queryWord={' \
'}&cl=2&lm=-1&ie=utf-8&oe=utf-8&adpicid=&st=-1&z=&ic=0&hd=&latest=©right=&word={' \
'}&s=&se=&tab=&width=&height=&face=0&istype=2&qc=&nc=1&fr=&expermode=&nojc=&pn={' \
'}&rn=30&gsm=1e&1635054081427= '
self.directory = r"D:/pythonproject/reptile/image/{}" # 存储目录 这里需要修改为自己希望保存的目录 {}不要丢
self.header = {
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) '
'Chrome/95.0.4638.54 Safari/537.36 Edg/95.0.1020.30 '
}
# 创建存储文件夹
def create_directory(self, name):
self.directory = self.directory.format(name)
# 如果目录不存在则创建
if not os.path.exists(self.directory):
os.makedirs(self.directory)
self.directory += r'\{}'
# 获取图像链接
def get_image_link(self, url):
list_image_link = []
strhtml = requests.get(url, headers=self.header) # Get方式获取网页数据
jsonInfo = json.loads(strhtml.text)
for index in range(30):
list_image_link.append(jsonInfo['data'][index]['thumbURL'])
return list_image_link
# 下载图片
def save_image(self, img_link, filename):
res = requests.get(img_link, headers=self.header)
if res.status_code == 404:
print(f"图片{img_link}下载出错------->")
with open(filename, "wb") as f:
f.write(res.content)
print("存储路径:" + filename)
# 入口函数
def run(self):
searchName = input("查询内容:")
searchName_parse = parse.quote(searchName) # 编码
self.create_directory(searchName)
pic_number = 0 # 图像数量
for index in range(self.json_count):
pn = (index+1)*30
request_url = self.url.format(searchName_parse, searchName_parse, str(pn))
list_image_link = self.get_image_link(request_url)
for link in list_image_link:
pic_number += 1
self.save_image(link, self.directory.format(str(pic_number)+'.jpg'))
time.sleep(0.2) # 休眠0.2秒,防止封ip
print(searchName+"----图像下载完成--------->")
if __name__ == '__main__':
spider = BaiduImageSpider()
spider.json_count = 10 # 定义下载10组图像,也就是三百张
spider.run()
(2)划分数据集 split_data.py
import os
from shutil import copy, rmtree
import random
def mk_file(file_path: str):
if os.path.exists(file_path):
# 如果文件夹存在,则先删除原文件夹再重新创建
rmtree(file_path)
os.makedirs(file_path)
def main():
# 保证随机可复现
random.seed(0)
# 将数据集中10%的数据划分到验证集中
split_rate = 0.1
# 指向解压后的flower_photos文件夹
# getcwd():该函数不需要传递参数,获得当前所运行脚本的路径
cwd = os.getcwd()
# join():用于拼接文件路径,可以传入多个路径
data_root = os.path.join(cwd, "F:/yolov5/")
origin_flower_path = os.path.join(data_root, "roadDATA")
# 确定路径存在,否则反馈错误
assert os.path.exists(origin_flower_path), "path '{}' does not exist.".format(origin_flower_path)
# isdir():判断某一路径是否为目录
# listdir():返回指定的文件夹包含的文件或文件夹的名字的列表
flower_class = [cla for cla in os.listdir(origin_flower_path)
if os.path.isdir(os.path.join(origin_flower_path, cla))]
# 创建训练集train文件夹,并由类名在其目录下创建子目录
train_root = os.path.join("F:/yolov5/roadDATA/train")
mk_file(train_root)
for cla in flower_class:
# 建立每个类别对应的文件夹
mk_file(os.path.join(train_root, cla))
# 创建验证集val文件夹,并由类名在其目录下创建子目录
val_root = os.path.join("F:/yolov5/roadDATA/val")
mk_file(val_root)
for cla in flower_class:
# 建立每个类别对应的文件夹
mk_file(os.path.join(val_root, cla))
# 遍历所有类别的图像并按比例分成训练集和验证集
for cla in flower_class:
cla_path = os.path.join(origin_flower_path, cla)
# iamges列表存储了该目录下所有图像的名称
images = os.listdir(cla_path)
num = len(images)
# 随机采样验证集的索引
# 从images列表中随机抽取k个图像名称
# random.sample:用于截取列表的指定长度的随机数,返回列表
# eval_index保存验证集val的图像名称
eval_index = random.sample(images, k=int(num * split_rate))
for index, image in enumerate(images):
if image in eval_index:
# 将分配至验证集中的文件复制到相应目录
image_path = os.path.join(cla_path, image)
new_path = os.path.join(val_root, cla)
copy(image_path, new_path)
else:
# 将分配至训练集中的文件复制到相应目录
image_path = os.path.join(cla_path, image)
new_path = os.path.join(train_root, cla)
copy(image_path, new_path)
# '\r'回车,回到当前行的行首,而不会换到下一行,如果接着输出,本行以前的内容会被逐一覆盖
# end="":将print自带的换行用end中指定的str代替
print("\r[{}] processing [{}/{}]".format(cla, index + 1, num), end="")
print()
print("processing done!")
if __name__ == '__main__':
main()
本文详细介绍了如何通过爬虫获取数据、制作数据集、训练YOLO和MobileNetV3模型进行天气和道路类型分类,并提供了划分数据集和测试模型的代码示例。
1300

被折叠的 条评论
为什么被折叠?



