利用YOLOv5进行西红柿成熟度检测

HypnotizeG

已于 2024-06-27 18:55:49 修改

阅读量767

点赞数 29

文章标签： YOLO python 计算机视觉深度学习

于 2024-06-27 18:50:33 首次发布

本文链接：https://blog.csdn.net/HypnotizeG/article/details/140016773

版权

#本文为笔者记录以大一学期大创项目为契机，进行计算机深度学习方面的学习过程的随记文章。发一篇博客的目的不但是为了供大家参考，也想通过它记录笔者自己的成长历程。

#参考文献见文末

一、准备工作

1.基础资源需求

1）Anaconda Prompt

2）Conda 24.5.0

3）Pytorch 2.3.1

4）Python 3.11.7

5）许多西红柿图片

注1：此处为笔者各个资源版本，不代表其他版本不可用。

注2：因为笔者电脑没有GPU，所以并没有用CUDA。

注3：笔者认为有命令行使用经验是非常必要的。

2.YOLOv5

直接通过Github下载压缩包即可，链接：GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

在通过了基础的detect.py公车行人测试后（在runs/detect/exp中找到检测图片），我们可以正式开始准备自己的数据集了！

二、Labelimg

在下面的准备中，我们需要用到.xml（可扩展标记语言）格式的文件，那么我们怎么生成它们呢？——利用Labelimg！

安装

因为我的Base环境下，Python版本是3.11.7，但是Labelimg所需的Python版本为3.8.5，所以利用conda创建一个符合Labelimg使用的Python环境。利用命令行生成Labelimg虚拟环境：

conda create -n labelimg python=3.8.5

再输入如下命令下载labelimg：

pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple

那么现在我们的Labelimg就安装好啦！

检查

使用

在使用Labelimg时，我们输入

labelimg activate

就可以使用Labelimg了。

三、数据集

文件夹

在data文件夹下保存如下几个文件夹：

images、JPEGImages：存放你的原始图片；

注：图片越多，训练效果越好。

Annotations存放接下来生成的.xml文件。

.xml文件生成

利用Labelimg进行.xml文件的生成：

Open Dir打开images文件夹进行框选。

Change Save Dir改变保存路径为Annotations文件夹。

.txt文件生成

step1：

在yolov5的根目录下创建一个trans_txt.py文件，文件代码如下：

import os
import random
trainval_percent = 0.1
train_percent = 0.9
xmlfilepath = 'data/Annotations'
txtsavepath = 'data/ImageSets'
total_xml = os.listdir(xmlfilepath)
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
ftrainval = open('data/ImageSets/trainval.txt', 'w')
ftest = open('data/ImageSets/test.txt', 'w')
ftrain = open('data/ImageSets/train.txt', 'w')
fval = open('data/ImageSets/val.txt', 'w')
for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

step2：

继续创建voc_label.py文件，代码如下：

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
sets = ['train', 'test','val']
classes = ['immature_tomato','mature_tomato']
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)
def convert_annotation(image_id):
    in_file = open('data/Annotations/%s.xml' % (image_id))
    out_file = open('data/labels/%s.txt' % (image_id), 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
wd = getcwd()
print(wd)
for image_set in sets:
    if not os.path.exists('data/labels/'):
        os.makedirs('data/labels/')
    image_ids = open('data/ImageSets/%s.txt' % (image_set)).read().strip().split()
    list_file = open('data/%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write('data/images/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

classes=['immature_tomato','mature_tomato']

注：此处为自己选择识别的种类。

依次运行这两个文件，成功运行后：

1）labels下生成.txt文件；

2）ImageSets下生成四个.txt文件；

3）data下生成三个.txt文件。