基于paddle自定义卷积神经网络进行垃圾分类,本竞赛所用训练和测试图片均来自生活场景。总共四十个类别,类别和标签对应关系在训练集中的dict文件里。图片中垃圾的类别,格式是“一级类别/二级类别”,二级类别是具体的垃圾物体类别,也就是训练数据中标注的类别,比如一次性快餐盒、果皮果肉、旧衣服等。一级类别有四种类别:可回收物、厨余垃圾、有害垃圾和其他垃圾。 数据文件包括训练集(有标注)和测试集(无标注),训练集的所有图片分别保存在train文件夹下面的0-39个文件夹中,文件名即类别标签,测试集共有400张待分类的垃圾图片在test文件夹下,testpath.txt保存了所有测试集文件的名称,格式为:name+\n。提交结果的格式如下: 每一行为:图像名 标签 test1.jpg 29写出相关代码
时间: 2023-06-15 08:03:08 浏览: 119
由于数据集比较大,建议在使用时将数据集下载到本地进行操作。下面是基于PaddlePaddle框架的卷积神经网络代码:
```python
import os
import random
import numpy as np
import paddle
import paddle.nn as nn
import paddle.optimizer as opt
from paddle.io import Dataset
from PIL import Image
# 定义数据集
class TrashDataset(Dataset):
def __init__(self, mode):
self.imgs = []
self.labels = []
self.mode = mode
self.class_dict = {}
# 获取类别和标签对应关系
with open('train/class_dict.csv', 'r', encoding='utf-8') as f:
for line in f.readlines():
class_name, class_id = line.strip().split(',')
self.class_dict[class_name] = int(class_id)
if self.mode == 'train':
# 读取训练集图片
for i in range(40):
class_name = str(i)
class_path = 'train/' + class_name
for img_name in os.listdir(class_path):
img_path = class_path + '/' + img_name
self.imgs.append(img_path)
self.labels.append(self.class_dict[class_name])
else:
# 读取测试集图片
with open('test/testpath.txt', 'r', encoding='utf-8') as f:
for line in f.readlines():
img_name = line.strip()
self.imgs.append('test/' + img_name)
def __getitem__(self, idx):
img_path = self.imgs[idx]
label = self.labels[idx]
# 读取图片并进行预处理
img = Image.open(img_path).convert('RGB')
img = img.resize((224, 224), Image.ANTIALIAS)
img = np.array(img).astype('float32')
img = img.transpose((2, 0, 1))
img = img / 255.0
if self.mode == 'train':
img = img + np.random.normal(loc=0.0, scale=0.01, size=img.shape)
return img, label
def __len__(self):
return len(self.imgs)
# 定义卷积神经网络模型
class CNN(nn.Layer):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2D(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1)
self.bn1 = nn.BatchNorm2D(num_features=64)
self.relu1 = nn.ReLU()
self.pool1 = nn.Pool2D(pool_size=2, pool_stride=2, pool_type='max')
self.conv2 = nn.Conv2D(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1)
self.bn2 = nn.BatchNorm2D(num_features=128)
self.relu2 = nn.ReLU()
self.pool2 = nn.Pool2D(pool_size=2, pool_stride=2, pool_type='max')
self.conv3 = nn.Conv2D(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1)
self.bn3 = nn.BatchNorm2D(num_features=256)
self.relu3 = nn.ReLU()
self.pool3 = nn.Pool2D(pool_size=2, pool_stride=2, pool_type='max')
self.fc1 = nn.Linear(in_features=256*28*28, out_features=1024)
self.relu4 = nn.ReLU()
self.fc2 = nn.Linear(in_features=1024, out_features=40)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu1(x)
x = self.pool1(x)
x = self.conv2(x)
x = self.bn2(x)
x = self.relu2(x)
x = self.pool2(x)
x = self.conv3(x)
x = self.bn3(x)
x = self.relu3(x)
x = self.pool3(x)
x = paddle.flatten(x, start_axis=1, stop_axis=-1)
x = self.fc1(x)
x = self.relu4(x)
x = self.fc2(x)
return x
# 训练函数
def train(model, train_loader, epoch, optimizer):
model.train()
for batch_id, data in enumerate(train_loader()):
x_data = data[0]
y_data = data[1]
predicts = model(x_data)
loss = nn.functional.cross_entropy(predicts, y_data)
acc = paddle.metric.accuracy(predicts, y_data)
loss.backward()
optimizer.step()
optimizer.clear_grad()
if batch_id % 10 == 0:
print('Epoch [{}/{}], Batch [{}/{}], Loss: {:.4f}, Acc: {:.4f}'.format(
epoch+1, num_epochs, batch_id+1, len(train_loader), loss.numpy(), acc.numpy()))
# 预测函数
def predict(model, test_loader):
model.eval()
predicts = []
for batch_id, data in enumerate(test_loader()):
x_data = data[0]
y_data = data[1]
output = model(x_data)
output = nn.functional.softmax(output, axis=1)
output = np.argmax(output.numpy(), axis=1).tolist()
predicts.extend(output)
return predicts
# 训练函数
def train_model():
# 定义超参数
batch_size = 32
num_epochs = 10
learning_rate = 0.001
# 加载数据集
train_dataset = TrashDataset(mode='train')
train_loader = paddle.io.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dataset = TrashDataset(mode='test')
test_loader = paddle.io.DataLoader(test_dataset, batch_size=batch_size)
# 初始化模型和优化器
model = CNN()
optimizer = opt.Adam(learning_rate=learning_rate, parameters=model.parameters())
# 开始训练
for epoch in range(num_epochs):
train(model, train_loader, epoch, optimizer)
# 预测测试集并保存结果
predicts = predict(model, test_loader)
with open('result.txt', 'w', encoding='utf-8') as f:
for i, p in enumerate(predicts):
f.write('test_{}.jpg {}\n'.format(i+1, p))
# 主函数
if __name__ == '__main__':
train_model()
```
上面的代码定义了一个CNN模型,使用Adam优化器进行训练,并在测试集上进行预测并将结果保存到文件中。在训练过程中,通过打印损失和精度可以观察模型的训练情况。