Yolov5训练路况分类模型(公路,土路)

本文详细介绍了如何通过爬虫获取数据、制作数据集、训练YOLO和MobileNetV3模型进行天气和道路类型分类,并提供了划分数据集和测试模型的代码示例。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

本篇将代码放在最后

目录

1.数据集制作

2.开始训练

3.测试

4.代码附录

(1)爬虫代码

(2)划分数据集 split_data.py


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=&copyright=&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()

评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值