图片分类作为cv基础,可以当作基础案例进行练习,这个是我自己做的一个猫狗猴的一个多分类,作为初学者可以试试,玩一下
import cv2
import warnings
import torch
import torchvision
import torchvision.transforms as transforms
import os
import random
import numpy as np
warnings.filterwarnings("ignore")
def transform_img(img):
# 将图片尺寸缩放道 224x224
img = cv2.resize(img, (32, 32))
# 读入的图像数据格式是[H, W, C]
# 使用转置操作将其变成[C, H, W]
img = np.transpose(img, (2, 0, 1))
img = img.astype('float32')
# 将数据范围调整到[-1.0, 1.0]之间
img = img / 255.
img = img * 2.0 - 1.0
return img
# 定义训练集数据读取器
def data_loader(datadir, batch_size=10, mode='train'):
# 将datadir目录下的文件列出来,每条文件都要读入
filenames = os.listdir(datadir)
def reader():
if mode == 'train':
# 训练时随机打乱数据顺序
random.shuffle(filenames)
batch_imgs = []
batch_labels = []
for name in filenames:
filepath = os.path.join(datadir, name)
img = cv2.imread(filepath)
img = transform_img(img)
if name.split('_')[0] == 'D':
label = 0
elif name.split('_')[0] == 'H':
label = 1
elif name.split('_')[0] == 'T':
# P开头的是病理性近视,属于正样本,标签为1
label = 2
else:
raise ('Not excepted file name')
# 每读取一个样本的数据,就将其放入数据列表中
batch_imgs.append(img)
batch_labels.append(label)
if len(batch_imgs) == batch_size:
# 当数据列表的长度等于batch_size的时候,
# 把这些数据当作一个mini-batch,并作为数据生成器的一个输出
imgs_array = np.array(batch_imgs).astype('float32')
labels_array = np.array(batch_labels).astype('int64').reshape(-1, 1)
yield imgs_array, labels_array
batch_imgs = []
batch_labels = []
if len(batch_imgs) > 0:
# 剩余样本数目不足一个batch_size的数据,一起打包成一个mini-batch
imgs_array = np.array(batch_imgs).astype('float32')
labels_array = np.array(batch_labels).astype('int64').reshape(-1, 1)
yield imgs_array, labels_array
return reader
DATADIR = 'data/dog_cat_pig'
train_loader = data_loader(DATADIR, batch_size=2, mode='train')
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 3)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
net = Net()
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
for epoch in range(100): # 多批次循环
running_loss = 0.0
for batch_id, data in enumerate(train_loader(),0):
x_data, y_data = data
x_data=torch.from_numpy(x_data)
y_data=torch.from_numpy(y_data)
inputs=x_data.to(device)
labels=y_data.to(device).view(2)
# 梯度置0
optimizer.zero_grad()
# 正向传播,反向传播,优化
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# 打印状态信息
running_loss += loss.item()
if epoch / 20 == 0:
print(running_loss)
torch.save(net.state_dict(), 'palm.pdparams')
print('Finished Training')
def predict(model, filepath, params_file_path):
# 开启0号GPU预估
use_gpu = False
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('start evaluation .......')
# 加载模型参数
model_state_dict = torch.load(params_file_path)
model.eval()
batch_imgs = []
img = cv2.imread(filepath)
img = transform_img(img)
imgs_array = np.array(img).astype('float32')
batch_imgs.append(imgs_array)
img = torch.tensor(batch_imgs)
# 计算预测和精度
result = model(img)
pred = F.sigmoid(result)
print(pred)
dict_angle = {0: '狗', 1: '猴', 2: '猫'}
for i, values in enumerate(pred[0]):
if pred[0].max() == values:
print(dict_angle[i])
return result
predict(net, 'data/dog_cat_pig/T_4.jpg', params_file_path="palm.pdparams")
start evaluation …
猫
最后预测出来时一个猫,说明模型效果还可以,大家可以按照自己的想法,进行不同的分类,如有需要数据集和代码加群753035545