pytorch学习一：AlexNet颜色分类
作者：互联网
数据集
百度AI下载的颜色块数据集，一共7种颜色，不过都是合在一起的
通过将图片转化到HSV空间，利用阈值提取不同颜色的图片，存放到不同文件夹，制作分类数据集。
根据表格数据，提取颜色代码
import numpy as np
import collections
import os
import cv2
# 定义字典存放颜色分量上下限
# 例如：{颜色: [min分量, max分量]}
# {'red': [array([160, 43, 46]), array([179, 255, 255])]}
def getColorList():
dict = collections.defaultdict(list)
# 黑色
lower_black = np.array([0, 0, 0])
upper_black = np.array([180, 255, 46])
color_list = []
color_list.append(lower_black)
color_list.append(upper_black)
dict['black'] = color_list
#灰色
lower_gray = np.array([0, 0, 46])
upper_gray = np.array([180, 43, 220])
color_list = []
color_list.append(lower_gray)
color_list.append(upper_gray)
dict['gray']=color_list
# 白色
lower_white = np.array([0, 0, 221])
upper_white = np.array([180, 30, 255])
color_list = []
color_list.append(lower_white)
color_list.append(upper_white)
dict['white'] = color_list
# 红色
lower_red = np.array([156, 43, 46])
upper_red = np.array([180, 255, 255])
color_list = []
color_list.append(lower_red)
color_list.append(upper_red)
dict['red'] = color_list
# 红色2
lower_red = np.array([0, 43, 46])
upper_red = np.array([10, 255, 255])
color_list = []
color_list.append(lower_red)
color_list.append(upper_red)
dict['red2'] = color_list
# 橙色
lower_orange = np.array([11, 43, 46])
upper_orange = np.array([25, 255, 255])
color_list = []
color_list.append(lower_orange)
color_list.append(upper_orange)
dict['orange'] = color_list
# 黄色
lower_yellow = np.array([26, 43, 46])
upper_yellow = np.array([34, 255, 255])
color_list = []
color_list.append(lower_yellow)
color_list.append(upper_yellow)
dict['yellow'] = color_list
# 绿色
lower_green = np.array([35, 43, 46])
upper_green = np.array([77, 255, 255])
color_list = []
color_list.append(lower_green)
color_list.append(upper_green)
dict['green'] = color_list
# 青色
lower_cyan = np.array([78, 43, 46])
upper_cyan = np.array([99, 255, 255])
color_list = []
color_list.append(lower_cyan)
color_list.append(upper_cyan)
dict['cyan'] = color_list
# 蓝色
lower_blue = np.array([100, 43, 46])
upper_blue = np.array([124, 255, 255])
color_list = []
color_list.append(lower_blue)
color_list.append(upper_blue)
dict['blue'] = color_list
# 紫色
lower_purple = np.array([125, 43, 46])
upper_purple = np.array([155, 255, 255])
color_list = []
color_list.append(lower_purple)
color_list.append(upper_purple)
dict['purple'] = color_list
return dict
if __name__ == '__main__':
color_dict = getColorList()
pic_dir = 'color_train_all/'
for images in os.listdir(pic_dir):
img = cv2.imread(pic_dir+images)
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
mask_white = cv2.inRange(hsv, color_dict['white'][0], color_dict['white'][1])
# 背景颜色为黑色，加入图片有像素，则保存
if np.max(mask_white)>0:
res = cv2.bitwise_and(img, img, mask=mask_white)
cv2.imwrite('train/white/'+images,res)
mask_green = cv2.inRange(hsv, color_dict['green'][0], color_dict['green'][1])
if np.max(mask_green)>0:
res = cv2.bitwise_and(img, img, mask=mask_green)
cv2.imwrite('train/green/'+images,res)
mask_purple = cv2.inRange(hsv, color_dict['purple'][0], color_dict['purple'][1])
if np.max(mask_purple)>0:
res = cv2.bitwise_and(img, img, mask=mask_purple)
cv2.imwrite('train/purple/'+images,res)
mask_blue = cv2.inRange(hsv, color_dict['blue'][0], color_dict['blue'][1])
if np.max(mask_blue)>0:
res = cv2.bitwise_and(img, img, mask=mask_blue)
cv2.imwrite('train/blue/'+images,res)
mask_cyan = cv2.inRange(hsv, color_dict['cyan'][0], color_dict['cyan'][1])
if np.max(mask_cyan)>0:
res = cv2.bitwise_and(img, img, mask=mask_cyan)
cv2.imwrite('train/cyan/'+images,res)
mask_red = cv2.inRange(hsv, color_dict['red2'][0], color_dict['red2'][1])
if np.max(mask_red)>0:
res = cv2.bitwise_and(img, img, mask=mask_red)
cv2.imwrite('train/red/'+images,res)
mask_yellow = cv2.inRange(hsv, color_dict['yellow'][0], color_dict['yellow'][1])
if np.max(mask_yellow)>0:
res = cv2.bitwise_and(img, img, mask=mask_yellow)
cv2.imwrite('train/yellow/'+images,res)
提取后文件，颜色放在单独的颜色文件下,并统一放在train文件夹下
train
- blue
–001.jpg
–002.jpg
–…- cyan
- green
- purple
- red
- white
- yellow
同理制作test文件夹
pytorch建立AlexNet模型
model.py
import torch.nn as nn
import torch
class AlexNet(nn.Module):
def __init__(self, num_classes=1000, init_weights=False):
super(AlexNet, self).__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 48, kernel_size=11, stride=4, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(48, 128, kernel_size=5, padding=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.Conv2d(128, 192, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(192, 192, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(192, 128, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2),
self.classifier = nn.Sequential(
nn.Dropout(p=0.5),
nn.Linear(128 * 6 * 6, 2048),
nn.ReLU(inplace=True),
nn.Dropout(p=0.5),
nn.Linear(2048, 2048),
nn.ReLU(inplace=True),
nn.Linear(2048, num_classes),
)
if init_weights:
self._initialize_weights()
def forward(self, x):
x = self.features(x)
x = torch.flatten(x, start_dim=1)
x = self.classifier(x)
return x
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
if m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight, 0, 0.01)
nn.init.constant_(m.bias, 0)
模型训练
train.py
import json
import torch
import torch.nn as nn
from torchvision import transforms, datasets, utils
import matplotlib.pyplot as plt
import numpy as np
import torch.optim as optim
from tqdm import tqdm
from model import AlexNet
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("using {} device.".format(device))
batch_size = 32
# torchvision自带的图片预处理
# transforms.Resize()将图片调整为指定大小
# transforms.RandomHorizontalFlip()图片随机水平翻转
# ......参考torchvision官网教学
data_transform = {
"train": transforms.Compose([transforms.Resize((224, 224)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),
"val": transforms.Compose([transforms.Resize((224, 224)), # cannot 224, must (224, 224)
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])}
# torchvison中的datasets加载训练集
train_dataset = datasets.ImageFolder(root="../train/",transform=data_transform["train"])
train_num = len(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
# {'blue': 0, 'cyan': 1, 'green': 2, 'purple': 3, 'red': 4, 'white': 5, 'yellow': 6}
flower_list = train_dataset.class_to_idx
# 将字典进行编码，最终生成class_indices.json文件
cla_dict = dict((val, key) for key, val in flower_list.items())
json_str = json.dumps(cla_dict, indent=4)
with open('class_indices.json', 'w') as json_file:
json_file.write(json_str)
# torchvison中的datasets加载验证集
validate_dataset = datasets.ImageFolder(root='../val',transform=data_transform["val"])
val_num = len(validate_dataset)
validate_loader = torch.utils.data.DataLoader(validate_dataset,batch_size=batch_size, shuffle=False)
print("using {} images for training, {} images for validation.".format(train_num,val_num))
# # 训练集图像可视化，可注释
# test_data_iter = iter(validate_loader)
# test_image, test_label = test_data_iter.__next__()
# def imshow(img):
# img = img / 2 + 0.5 # unnormalize
# print(img.shape)
# npimg = img.numpy()
# plt.imshow(np.transpose(npimg, (1, 2, 0)))
# plt.show()
# print(' '.join('%5s' % cla_dict[test_label[j].item()] for j in range(4)))
# # make_grid的作用是将若干幅图像拼成一幅图像,在需要展示一批数据时很有用
# imshow(utils.make_grid(test_image))
# num_classes=分类个数 init_weights=初始化权重
net = AlexNet(num_classes=7, init_weights=True)
net.to(device)
loss_function = nn.CrossEntropyLoss() # 多分类常用的损失函数
optimizer = optim.Adam(net.parameters(), lr=0.0002) #优化器
epochs = 5
best_acc = 0.0 # 更新准确率最高的数值
best_loss = 1.0 # 更新损失最低的数值
train_steps = len(train_loader)
for epoch in range(epochs):
# 通过net.train()可以保证dropout/BatchNormal只在训练时候起作用
net.train()
running_loss = 0.0 # 统计训练过程中的损失
train_bar = tqdm(train_loader) # 119
for step, data in enumerate(train_bar):
# images: (batchsize,3,224,224)
# labels: batchsize
images, labels = data
optimizer.zero_grad()
# outputs: (batchsize,classes)
outputs = net(images.to(device))
loss = loss_function(outputs, labels.to(device)) # 计算预测值与真实值
loss.backward() # 损失反向传播
optimizer.step() # 更新参数
# 计算一共多少损失
running_loss += loss.item()
train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch+1,epochs,loss)
# 验证
net.eval()
acc = 0.0 # 计算精度公式 number / epoch
with torch.no_grad(): # 进制pytorch对参数跟踪
val_bar = tqdm(validate_loader)
for val_data in val_bar:
val_images, val_labels = val_data
outputs = net(val_images.to(device))
predict_y = torch.max(outputs, dim=1)[1]
acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
val_accurate = acc / val_num
print('[epoch %d] train_loss: %.3f val_accuracy: %.3f' %(epoch + 1, running_loss/train_steps, val_accurate))
val_loss = running_loss/train_steps
if val_loss<best_loss:
best_loss = val_loss
if val_accurate > best_acc:
best_acc = val_accurate
torch.save(net.state_dict(), './epoch%d_train_loss_%.2fval_accuracy_%.2f.pth'
%(epoch + 1, running_loss/train_steps, val_accurate))
print('Finished Training')
if __name__ == '__main__':
main()
模型预测
import os
import cv2
import json
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
from model import AlexNet
img_path = 'purple.jpg'
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 和训练时图像处理一样：不过只需要resize、ToTensor、Normalize
data_transform = transforms.Compose(
[transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# 读取json文件并解码
json_path = 'class_indices.json'
assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
json_file = open(json_path, "r")
class_indict = json.load(json_file)
# 创建模型
model = AlexNet(num_classes=7).to(device)
# 加载训练好的权重
weights_path = "epoch4train_loss_0.00val_accuracy_1.00.pth"
assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(weights_path)
model.load_state_dict(torch.load(weights_path))
model.eval() # 验证模式
''''''''''
OpenCV转换成PIL.Image格式
image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
Image转换成OpenCV格式
img = cv2.cvtColor(numpy.asarray(image), cv2.COLOR_RGB2BGR)
'''''''''''
image = cv2.imread(img_path)
img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
img = data_transform(img)
img = torch.unsqueeze(img, dim=0) # 维度扩张，增加一个batchsize维度
with torch.no_grad():
output = torch.squeeze(model(img.to(device))).cpu() # 降维
predict = torch.softmax(output, dim=0)
predict_cla = torch.argmax(predict).numpy()
print_res = "class:{} confidence:{:.2}".format(class_indict[str(predict_cla)],predict[predict_cla].numpy())
cv2.putText(image, print_res, (0, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
cv2.imshow('image',image)
cv2.waitKey(0)
if __name__ == '__main__':
main()
标签：nn,color,list,分类,cv2,pytorch,train,dict,AlexNet 来源： https://blog.csdn.net/weixin_45930948/article/details/120104737