本篇博客的所有代码已上传至GitHub仓库,后续会更新各个文件夹及文件的详细说明,用者自取
由于卷积神经网络训练花卉识别分类器博客已将模型的训练、测试代码写好,且可以通过这篇博客获取到大神训练好的模型,这里我只简单记录一下copy并运行时遇到的问题。
一、VGG16_test.py
遇到的问题及解决方法
1. 不知道电脑是否有CUDA
解决:在Terminal的python中输入下列代码
if torch.cuda.is_available():
print("Has cuda")
else:
print("No cuda")
2. 需要将加载模型的设备改为CPU
解决:将原代码中的代码片A对应地更改为代码片B
# 代码片A
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.load_state_dict(torch.load("VGG16_flower_200.pkl"))
net.to(device)
# 代码片B
device = torch.device("cpu")
net.load_state_dict(torch.load("VGG16_flower_200.pkl", map_location=device))
3. 提示各种文件路径不正确
解决:将所有的路径改为本机适用的,并将下载好的后缀名为.pkl
的模型copy到./模型训练/VGG16
文件夹中
至此测试数据集路径、模型路径的程序就可以运行起来了,且可以在CLI中看到识别花卉的正确性
二、自写分类.py
1. 大致思路
for循环遍历图片数据集
net网络得到score
softmax得到各种花的概率
将最大概率的花的index映射成花名
将图片内容和花名作为一对集合写进文件夹
2. 分类.py
大部分代码仍copy原图片测试.py
# 原代码
from PIL import Image
import torchvision.transforms as transforms
from torchvision import models #人家的模型
from torch.autograd import Variable
import torch
#from torchvision.datasets import ImageFolder
from torch import nn
#import VGG16_model
#数据预处理
data_transform = transforms.Compose([
transforms.Resize((224,224), 2), #对图像大小统一
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[ #图像归一化
0.229, 0.224, 0.225])
])
#类别
#这个类别是我在训练的过程输出的训练集的类别,是按照训练的顺序排列的
data_classes = ['Cerasus', 'Dianthus', 'Digitalis_purpurea', 'Eschscholtzia',
'Gazania', 'Jasminum', 'Matthiola', 'Narcissus', 'Nymphaea',
'Pharbitis', 'Rhododendron', 'Rosa', 'Tithonia', 'Tropaeolum_majus',
'daisy', 'dandelion', 'peach_blossom', 'roses', 'sunflowers', 'tulips']
#读取数据
img = Image.open('./图片/向日葵.jpg')
img=data_transform(img)#这里经过转换后输出的input格式是[C,H,W],网络输入还需要增加一维批量大小B
img = img.unsqueeze(0)#增加一维,输出的img格式为[1,C,H,W]
#类别
#train_dataset = ImageFolder(root='work/data/train/',transform=data_transform)
#data_classes = train_dataset.classes
#选择CPU还是GPU的操作
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#选择模型
net = models.vgg16()
net.classifier = nn.Sequential(nn.Linear(25088, 4096), #vgg16
nn.ReLU(),
nn.Dropout(p=0.5),
nn.Linear(4096, 4096),
nn.ReLU(),
nn.Dropout(p=0.5),
nn.Linear(4096, 20))
#读取参数
net.load_state_dict(torch.load("VGG16_flower_200.pkl",map_location=torch.device('cpu')))
net.eval()
net.to(device)
img = Variable(img)
score = net(img)#将图片输入网络得到输出
probability = nn.functional.softmax(score,dim=1)#计算softmax,即该图片属于各类的概率
max_value,index = torch.max(probability,1)#找到最大概率对应的索引号,该图片即为该索引号对应的类别
print()
print("识别为'{}'的概率为{}".format(data_classes[index.item()],max_value.item()))
3. 选择加载设备部分与读取参数部分的更改
解决:将原代码中的代码片A对应地更改为代码片B
# 代码片A
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.load_state_dict(torch.load("VGG16_flower_200.pkl"))
net.to(device)
# 代码片B
device = torch.device("cpu")
net.load_state_dict(torch.load("VGG16_flower_200.pkl", map_location=device))
4. 批量读取识别图片
1. 设置照片根目录并加载
root_path = "/Users/stone/PycharmProjects/pythonProject/flower_det/flower_recognition/data/myTest"
files = os.listdir(root_path)
2. for循环遍历识别图片,得到最大概率的花名
for file in files:
if not os.path.isdir(file):
img = Image.open(root_path+"/"+file)
img = data_transform(img) # 图片预处理
img = img.unsqueeze(0) # 增加一维,输出的img格式为[1,C,H,W]
img = Variable(img) # 转换格式
score = net(img) # 将图片输入网络得到输出
probability = nn.functional.softmax(score, dim=1) # 计算softmax,即该图片属于各类的概率
max_value, index = torch.max(probability, 1) # 找到最大概率对应的索引号,该图片即为该索引号对应的类别
3. 在for循环中利用CV2保存图片
cv_img = cv2.imread(root_path+"/"+file)
dir_name = "/Users/stone/PycharmProjects/pythonProject/flower_det/flower_recognition/data" \
+ "/" \
+ data_classes[index.item()]
if not os.path.exists(dir_name):
os.mkdir(dir_name)
images = os.listdir(dir_name)
num_images = len(images)
filename = dir_name + "/" + str(num_images+1) + ".jpg"
cv2.imwrite(filename, cv_img)
至此,分类.py
算是完成了基础功能,运行后,在dir_name
文件夹中会保存识别分类好的花卉