毕业设计要做一个垃圾分类系统,需要识别可回收垃圾、厨余垃圾、有害垃圾和其他垃圾等四个大类,在网上找到了很多开源的数据集,但是质量参差不齐,而且有坏图的存在,所以我就将这些数据集还有自己爬取的数据一起清洗了一遍,全部保存为了jpg的格式,一共有245个小类和4个大类。模型训练使用的是tensorflow2.3,其中mobilenet的准确率有82%,并使用pyqt5构建了图形化界面。
代码结构
images 目录主要是放置一些图片,包括测试的图片和ui界面使用的图片
models 目录下放置训练好的两组模型,分别是cnn模型和mobilenet的模型
results 目录下放置的是训练的训练过程的一些可视化的图,两个txt文件是训练过程中的输出,两个图是两个模型训练过程中训练集和验证集准确率和loss变化曲线
utils 是主要是我测试的时候写的一些文件,对这个项目没有实际的用途
mainwindow.py 是界面文件,主要是利用pyqt5完成的界面,通过上传图片可以对图片种类进行预测
testmodel.py 是测试文件,主要是用于测试两组模型在验证集上的准确率,这个信息你从results的txt的输出中也能获取
train_cnn.py 是训练cnn模型的代码
train_mobilenet.py 是训练mobilenet模型的代码
效果
数据集整理
放置到相应的子文件夹
数据集收集完成之后,我们还需要对数据集进行整理,如果是爬虫爬取的图片可能会有一些质量比较差的图片,那么整理之前还需要进行数据的清洗,删除质量不好的图片,数据集整理其实很简单,我们只需要将数据集进行归类即可,即相同类别的图片放在一个文件夹下,比如下面的这个数据集,白菜的文件夹下放的全是白菜的图片,土豆的文件夹下则放的全是土豆的图片。
划分训练集和测试集
注:如果是使用的开源数据集,开源数据集可能已经进行了数据集的划分,直接使用即可,不需要再次进行划分,比如这里是我下载到的农作物病虫害的数据集,已经分别提供了训练集、测试集和验证集,就不需要再次进行数据集的划分。
为了方便我们进行数据集的加载,我们还需要将图片划分为训练集和测试集,如果需要的话你还需要划分出验证集,验证集在一般的任务中是可选的,因为是自己收集的数据集的话,数据量比较少,如果再划分验证集的话可能会导致训练量不够,这里我写了一段数据集划分的代码逻辑,大家输入原始的数据集位置和划分之后的数据集位置,指定数据集划分的比例,即可完成数据集的划分。
import os
import random
from shutil import copy2
def data_set_split(src_data_folder, target_data_folder, train_scale=0.8, val_scale=0.0, test_scale=0.2):
'''
读取源数据文件夹,生成划分好的文件夹,分为trian、val、test三个文件夹进行
:param src_data_folder: 源文件夹 E:/biye/gogogo/note_book/torch_note/data/utils_test/data_split/src_data
:param target_data_folder: 目标文件夹 E:/biye/gogogo/note_book/torch_note/data/utils_test/data_split/target_data
:param train_scale: 训练集比例
:param val_scale: 验证集比例
:param test_scale: 测试集比例
:return:
'''
print("开始数据集划分")
class_names = os.listdir(src_data_folder)
# 在目标目录下创建文件夹
split_names = ['train', 'val', 'test']
for split_name in split_names:
split_path = os.path.join(target_data_folder, split_name)
if os.path.isdir(split_path):
pass
else:
os.mkdir(split_path)
# 然后在split_path的目录下创建类别文件夹
for class_name in class_names:
class_split_path = os.path.join(split_path, class_name)
if os.path.isdir(class_split_path):
pass
else:
os.mkdir(class_split_path)
# 按照比例划分数据集,并进行数据图片的复制
# 首先进行分类遍历
for class_name in class_names:
current_class_data_path = os.path.join(src_data_folder, class_name)
current_all_data = os.listdir(current_class_data_path)
current_data_length = len(current_all_data)
current_data_index_list = list(range(current_data_length))
random.shuffle(current_data_index_list)
train_folder = os.path.join(os.path.join(target_data_folder, 'train'), class_name)
val_folder = os.path.join(os.path.join(target_data_folder, 'val'), class_name)
test_folder = os.path.join(os.path.join(target_data_folder, 'test'), class_name)
train_stop_flag = current_data_length * train_scale
val_stop_flag = current_data_length * (train_scale + val_scale)
current_idx = 0
train_num = 0
val_num = 0
test_num = 0
for i in current_data_index_list:
src_img_path = os.path.join(current_class_data_path, current_all_data[i])
if current_idx <= train_stop_flag:
copy2(src_img_path, train_folder)
# print("{}复制到了{}".format(src_img_path, train_folder))
train_num = train_num + 1
elif (current_idx > train_stop_flag) and (current_idx <= val_stop_flag):
copy2(src_img_path, val_folder)
# print("{}复制到了{}".format(src_img_path, val_folder))
val_num = val_num + 1
else:
copy2(src_img_path, test_folder)
# print("{}复制到了{}".format(src_img_path, test_folder))
test_num = test_num + 1
current_idx = current_idx + 1
print("*********************************{}*************************************".format(class_name))
print(
"{}类按照{}:{}:{}的比例划分完成,一共{}张图片".format(class_name, train_scale, val_scale, test_scale, current_data_length))
print("训练集{}:{}张".format(train_folder, train_num))
print("验证集{}:{}张".format(val_folder, val_num))
print("测试集{}:{}张".format(test_folder, test_num))
if __name__ == '__main__':
src_data_folder = "C:/Users/Scm97/Desktop/dejahu/data" # todo 原始数据集目录
target_data_folder = "C:/Users/Scm97/Desktop/dejahu/split_data" # todo 数据集分割之后存放的目录
data_set_split(src_data_folder, target_data_folder)
完整源码下载: