diff --git a/水稻叶片病害分类/大作业/exe.py b/水稻叶片病害分类/大作业/exe.py new file mode 100644 index 0000000..83ac6b5 --- /dev/null +++ b/水稻叶片病害分类/大作业/exe.py @@ -0,0 +1,201 @@ +import numpy as np +import os +import cv2 +import matplotlib.pyplot as plt +import pandas as pd +import torch +import torchvision.transforms as transforms +import torch.nn as nn +import time +from torch.utils.data import DataLoader, Dataset +from sklearn.model_selection import train_test_split +from torch.utils.tensorboard import SummaryWriter + +# 导入数据 +path = [] +data_set = [] # 全部数据 +label = [] # 全部标签 +path.append("Rice Disease/Bacterial leaf blight") +path.append("Rice Disease/Brown spot") +path.append("Rice Disease/Leaf smut") +for i in range(3): + img_dir = sorted(os.listdir(path[i])) + for file in img_dir: + img = cv2.imread(os.path.join(path[i], file)) + if img.shape[0] > img.shape[1]: + img = np.rot90(img).copy() + data_set.append(cv2.resize(img, (1024, 512))) + label.append(i) + + +# 划分训练集和验证集,训练集与验证集比例为7:3 +data_train, data_val, label_train, label_val = train_test_split( + data_set, label, test_size=0.3 +) + +# 训练集数据增强操作 +train_set_transform = transforms.Compose( + [ + transforms.ToPILImage(), + transforms.RandomRotation(15), # 随机旋转 + transforms.ToTensor() + # Tensor格式为[3,512,1024] + ] +) +# 验证集无需旋转 +val_set_transorm = transforms.Compose([transforms.ToPILImage(), transforms.ToTensor()]) + + +class ImgDataset(Dataset): + def __init__(self, x, y, transform=None): + self.x = x + self.y = torch.LongTensor(y) + self.transform = transform + + def __len__(self): + return len(self.x) + + def __getitem__(self, index): + X = self.transform(self.x[index]) + Y = self.y[index] + return X, Y + + +batch_size = 5 + +TrainSet = ImgDataset(data_train, label_train, train_set_transform) +ValidationSet = ImgDataset(data_val, label_val, val_set_transorm) + +TrainLoader = DataLoader(TrainSet, batch_size, shuffle=True) +ValidationLoader = DataLoader(ValidationSet, batch_size, shuffle=False) + +print("DataSet complete") + + +# 卷积神经网络结构 +class Net(nn.Module): + def __init__(self): + super(Net, self).__init__() + # input维度[3,512,1024] + # torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding) + # torch.nn.MaxPool2d(kernel_size, stride, padding) + + # 卷积层 + self.cnn = nn.Sequential( + # 第一层卷积与池化 + nn.Conv2d(3, 32, 7, 1, 3), # [32,512,1024] + nn.BatchNorm2d(32), # 归一化操作 + nn.ReLU(), + nn.MaxPool2d(2, 2, 0), # [32,256,512] + # 第二次 + nn.Conv2d(32, 64, 3, 1, 1), # [64,256,512] + nn.BatchNorm2d(64), + nn.ReLU(), + nn.MaxPool2d(2, 2, 0), # [64,128,256] + # 第三次 + nn.Conv2d(64, 128, 3, 1, 1), # [128,128,256] + nn.BatchNorm2d(128), + nn.ReLU(), + nn.MaxPool2d(2, 2, 0), # [128,64,128] + # 第四次 + nn.Conv2d(128, 128, 3, 1, 1), # [128,64,128] + nn.BatchNorm2d(128), + nn.ReLU(), + nn.MaxPool2d(2, 2, 0), # [128,32,64] + # 第五次 + nn.Conv2d(128, 256, 3, 1, 1), # [256,32,64] + nn.BatchNorm2d(256), + nn.ReLU(), + nn.MaxPool2d(2, 2, 0), # [256,16,32] + # 第六次 + nn.Conv2d(256, 256, 3, 1, 1), # [256,16,32] + nn.BatchNorm2d(256), + nn.ReLU(), + nn.MaxPool2d(2, 2, 0), # [256,8,16] + # 第七次 + nn.Conv2d(256, 512, 3, 1, 1), # [512,8,16] + nn.BatchNorm2d(512), + nn.ReLU(), + nn.MaxPool2d(2, 2, 0), # [512,4,8] + ) + # 全连接层 + self.fc = nn.Sequential( + nn.Linear(512 * 4 * 8, 1024), + nn.ReLU(), + nn.Linear(1024, 512), + nn.ReLU(), + nn.Linear(512, 64), + nn.ReLU(), + nn.Linear(64, 3), + ) + + # 前向传播 + def forward(self, x): + out = self.cnn(x) + out = out.view(out.size()[0], -1) # 将多维数据改成一维 + out = self.fc(out) + return out + + +print("Net complete") + + +model = Net().cuda() +loss = nn.CrossEntropyLoss() # 交叉熵损失函数 +optimizer = torch.optim.Adam(model.parameters(), lr=0.000001) # Adam优化器 +num_epoch = 80 # 迭代次数 +writer = SummaryWriter(log_dir="runs/result") # 实例化SummaryWriter,用于记录训练过程中的数据并做可视化 +# init_img用于带入网络从而便于记录网络结构 +init_img = torch.zeros(1, 3, 512, 1024).cuda() +writer.add_graph(model, init_img) +for epoch in range(num_epoch): + epoch_start_time = time.time() + + train_acc = 0.0 + val_acc = 0.0 + train_loss = 0.0 + val_loss = 0.0 + + model.train() + for i, data in enumerate(TrainLoader): + optimizer.zero_grad() + train_pred = model(data[0].cuda()) + batch_loss = loss(train_pred, data[1].cuda()) + batch_loss.backward() + optimizer.step() + + train_acc += np.sum( + np.argmax(train_pred.cpu().data.numpy(), axis=1) == data[1].numpy() + ) + train_loss += batch_loss.item() + + model.eval() + with torch.no_grad(): + for i, data in enumerate(ValidationLoader): + val_pred = model(data[0].cuda()) + batch_loss = loss(val_pred, data[1].cuda()) + + val_acc += np.sum( + np.argmax(val_pred.cpu().data.numpy(), axis=1) == data[1].numpy() + ) + val_loss += batch_loss.item() + + # 將結果 print 出來 + print( + "[%03d/%03d] %2.2f sec(s) Train Acc: %3.6f Loss: %3.6f | Validation Acc: %3.6f Loss: %3.6f" + % ( + epoch + 1, + num_epoch, + time.time() - epoch_start_time, + train_acc / TrainSet.__len__(), + train_loss / TrainSet.__len__(), + val_acc / ValidationSet.__len__(), + val_loss / ValidationSet.__len__(), + ) + ) + + # 记录训练过程中的数据 + writer.add_scalar("Train Acc", train_acc / TrainSet.__len__(), epoch) + writer.add_scalar("Train Loss", train_loss / TrainSet.__len__(), epoch) + writer.add_scalar("Validation Acc", val_acc / ValidationSet.__len__(), epoch) + writer.add_scalar("Validation Loss", val_loss / ValidationSet.__len__(), epoch)