基于CNN＋PyTorch实现视觉检测分类 原创

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.sampler import WeightedRandomSampler
from torch.utils.data import random_split
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
import os
import warnings
warnings.filterwarnings("ignore")

# 设备配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using {device} device")

# 指定超参数
minibatch_size = 10
learning_rate = 0.01
epochs = 60

#定义加载器函数
def custom_loader(path):
with open(path, 'rb') as f:
img = Image.open(f)
img = img.crop((50, 60, 750, 460))  #Size: 700x400 px
img.load()
return img

# 图像路径（本地路径以加速加载）
path = "data/Coil_Vision/01_train_val_test"

#用于加载的转换函数
transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5), (0.5))])

# 在文件夹结构中创建数据集
dataset = datasets.ImageFolder(path, transform=transform, loader=custom_loader)
train_set, val_set, test_set = random_split(dataset, [round(0.5*len(dataset)), 
round(0.3*len(dataset)), 
round(0.2*len(dataset))])

# 定义一个采样器来平衡这些类
# training dataset
lbls = [dataset[idx][1] for idx in train_set.indices]
bc = np.bincount(lbls)
p_nOK = bc.sum()/bc[0]
p_OK = bc.sum()/bc[1]
lst_train = [p_nOK if lbl==0 else p_OK for lbl in lbls]
train_sampler = WeightedRandomSampler(weights=lst_train, num_samples=len(lbls))

# 用批尺寸定义加载器
train_loader = DataLoader(dataset=train_set, batch_size=minibatch_size, sampler=train_sampler)
val_loader = DataLoader(dataset=val_set, batch_size=minibatch_size, shuffle=True)
test_loader = DataLoader(dataset=test_set, shuffle=True)

# 图形和轴对象
fig, axs = plt.subplots(nrows=2, ncols=5, figsize=(20,7), sharey=True, sharex=True)

count_OK = 0
count_nOK = 0

# 循环遍历加载器批次
for (batch_data, batch_lbls) in train_loader:

#循环遍历batch_lbls
for i, lbl in enumerate(batch_lbls):

# 如果标签为0 (nOK)，在第1行绘制图形
if (lbl.item() == 0) and (count_nOK < 5):
axs[1, count_nOK].imshow(batch_data[i][0], cmap='gray')
axs[1, count_nOK].set_title(f"nOK Part#: {str(count_nOK)}", fontsize=14)
count_nOK += 1

#如果标签为1 (OK)，在第0行绘制图形
elif (lbl.item() == 1) and (count_OK < 5):
axs[0, count_OK].imshow(batch_data[i][0], cmap='gray')
axs[0, count_OK].set_title(f"OK Part#: {str(count_OK)}", fontsize=14)
count_OK += 1

#如果两个计数器都是>=5停止循环
if (count_OK >=5) and (count_nOK >=5):
break

# 配置绘图画布
fig.suptitle("Sample plot of OK and nonOK Parts", fontsize=24)
plt.setp(axs, xticks=[], yticks=[]) 
plt.show()

class CNN(nn.Module):

def __init__(self):
super().__init__()

# Define model layers
self.model_layers = nn.Sequential(

nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),

nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),

nn.Flatten(),
nn.Linear(16*97*172, 120),
nn.ReLU(),
nn.Linear(120, 2)
)

def forward(self, x):
out = self.model_layers(x)
return out

# 在cpu或gpu上定义模型
model = CNN().to(device)

#损失函数和优化器
loss = nn.CrossEntropyLoss()

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

# Count number of parameters / thereof trainable
num_param = sum([p.numel() for p in model.parameters()])
num_param_trainable = sum([p.numel() for p in model.parameters() if p.requires_grad == True])

print(f"Our model has {num_param:,} parameters. Thereof trainable are {num_param_trainable:,}!")

def val_test(dataloader, model):
# 获取数据集大小
dataset_size = len(dataloader.dataset)

# 关闭梯度计算以进行验证
with torch.no_grad():
# 循环数据集
correct = 0
wrong_preds = []
for (images, labels) in dataloader:
images, labels = images.to(device), labels.to(device)

#从模型中获取原始值
output = model(images)

# 推导预测
y_pred = output.argmax(1)

# 对所有批次进行正确的分类计数
correct += (y_pred == labels).type(torch.float32).sum().item()

# Save wrong predictions (image, pred_lbl, true_lbl)
for i, _ in enumerate(labels):
if y_pred[i] != labels[i]:
wrong_preds.append((images[i], y_pred[i], labels[i]))

# Calculate accuracy
acc = correct / dataset_size

return acc, wrong_preds

acc_train = {}
acc_val = {}
# 对世代进行迭代处理
for epoch in range(epochs):

n_correct=0; n_samples=0; n_true_OK=0
for idx, (images, labels) in enumerate(train_loader):
model.train()
# 如果可用，请将数据推送到gpu
images, labels = images.to(device), labels.to(device)

#向前传播
outputs = model(images)
l = loss(outputs, labels)

# 向后传播和优化
optimizer.zero_grad()
l.backward()
optimizer.step()

# 获取预测标签(.max返回(value,index))
_, y_pred = torch.max(outputs.data, 1)

# 计算正确的分类
n_correct += (y_pred == labels).sum().item()
n_true_OK += (labels == 1).sum().item()
n_samples += labels.size(0)

# 在世代结束时：计算准确性和打印信息
if (epoch+1) % 2 == 0:
model.eval()
# 计算准确性
acc_train[epoch+1] = n_correct / n_samples
true_OK = n_true_OK / n_samples
acc_val[epoch+1] = val_test(val_loader, model)[0]

#打印信息
print (f"Epoch [{epoch+1}/{epochs}], Loss: {l.item():.4f}")
print(f"      Training accuracy: {acc_train[epoch+1]*100:.2f}%")
print(f"      True OK: {true_OK*100:.3f}%")
print(f"      Validation accuracy: {acc_val[epoch+1]*100:.2f}%")

# 保存模型和状态词典
torch.save(model, "model.pth")

# 实例化图形和轴对象
fig, ax = plt.subplots(figsize=(10,6))
plt.plot(list(acc_train.keys()), list(acc_train.values()), label="training accuracy")
plt.plot(list(acc_val.keys()), list(acc_val.values()), label="validation accuracy")
plt.title("Accuracies", fontsize=24)
plt.ylabel("%", fontsize=14)
plt.xlabel("Epochs", fontsize=14)
plt.setp(ax.get_xticklabels(), fontsize=14)
plt.legend(loc='best', fontsize=14)
plt.show()

#从文件中读取模型
model = torch.load("model.pth")
model.eval()

print(f"test accuracy: {val_test(test_loader,model)[0]*100:0.1f}%")

%matplotlib inline

# Call test function
tup = val_test(test_loader, model)

#检查是否发生了错误的预测
if len(tup[1])>=1:

# 遍历错误预测的图像
for i, t in enumerate(tup[1]):
plt.figure(figsize=(7,5))
img, y_pred, y_true = t
img = img.to("cpu").reshape(400, 700)
plt.imshow(img, cmap="gray")
plt.title(f"Image {i+1} - Predicted: {y_pred}, True: {y_true}", fontsize=24)
plt.axis("off")
plt.show()
plt.close()
else:
print("No wrong predictions!")

#加载所需的库
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
from PIL import Image
import os

#  设备配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 完全按照第2.8节来定义CNN模型
class CNN(nn.Module):

def __init__(self):
super(CNN, self).__init__()

# 定义模型层
self.model_layers = nn.Sequential(

nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),

nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2),

nn.Flatten(),
nn.Linear(16*97*172, 120),
nn.ReLU(),
nn.Linear(120, 2),
#nn.LogSoftmax(dim=1)
)

def forward(self, x):
out = self.model_layers(x)
return out

# 加载模型的参数
model = torch.load("model.pth")
model.eval()

#定义自定义数据集
class Predict_Set(Dataset):
def __init__(self, img_folder, transform):
self.img_folder = img_folder
self.transform = transform
self.img_lst = os.listdir(self.img_folder)

def __len__(self):
return len(self.img_lst)

def __getitem__(self, idx):
img_path = os.path.join(self.img_folder, self.img_lst[idx])
img = Image.open(img_path)
img = img.crop((50, 60, 750, 460))  #Size: 700x400
img.load()
img_tensor = self.transform(img)
return img_tensor, self.img_lst[idx]

# 指向图像的路径（最好是本地路径，以加速加载）
path = "data/Coil_Vision/02_predict"

# 用于加载的转换函数
transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5), (0.5))])

# 创建数据集作为自定义数据集的实例
predict_set = Predict_Set(path, transform=transform)

# 定义加载程序
predict_loader = DataLoader(dataset=predict_set)

def predict(dataloader, model):

# 关闭梯度计算
with torch.no_grad():

img_lst = []; y_pred_lst = []; name_lst = []
#循环遍历数据加载程序
for image, name in dataloader:
img_lst.append(image)
image = image.to(device)

# 从模型中获取原始值
output = model(image)

#推导预测
y_pred = output.argmax(1)
y_pred_lst.append(y_pred.item())
name_lst.append(name[0])

return img_lst, y_pred_lst, name_lst

# 预测图像的标签
imgs, lbls, names  = predict(predict_loader, model)

#对分类图像进行迭代
for idx, image in enumerate(imgs):
plt.figure(figsize=(8,6))
plt.imshow(image.squeeze(), cmap="gray")
plt.title(f"\nFile: {names[idx]}, Predicted label: {lbls[idx]}", fontsize=18)
plt.axis("off")
plt.show()
plt.close()

# 设置为空的列表，以存储图层和权重
all_layers = []; conv_weights = []

# 迭代模型的结构
# (First level nn.Sequential)
for _, layer in enumerate(list(model.children())[0]):
if type(layer) == nn.Conv2d:
all_layers.append(layer)
conv_weights.append(layer.weight)
elif type(layer) in [nn.ReLU, nn.MaxPool2d]:
all_layers.append(layer)
conv_weights.append("*")

# 打印层和权重维度信息
for idx, layer in enumerate(all_layers):
print(f"{idx+1}. Layer: {layer}")
if type(layer) == nn.Conv2d:
print(f"          weights: {conv_weights[idx].shape}")
else:
print(f"          weights: {conv_weights[idx]}")
print()

import itertools

#遍历所有层
for idx_out, layer in enumerate(all_layers):

#如果层是一个卷积滤波器
if type(layer) == nn.Conv2d:

# 打印层名称
print(f"\n{idx_out+1}. Layer: {layer} \n")

# 准备绘图并计算出权重
plt.figure(figsize=(25,6))
weights = conv_weights[idx_out][:,0,:,:] # only first input channel
weights = weights.detach().to('cpu')

# 枚举过滤器权重（仅限第一个输入通道）
for idx_in, f in enumerate(weights):
plt.subplot(2,8, idx_in+1)
plt.imshow(f, cmap="gray")
plt.title(f"Filter {idx_in+1}")

# 打印文本
for i, j in itertools.product(range(f.shape[0]), range(f.shape[1])):
if f[i,j] > f.mean():
color = 'black'
else:
color = 'white'
plt.text(j, i, format(f[i, j], '.2f'), horizontalalignment='center', verticalalignment='center', color=color)

plt.axis("off")
plt.show()
plt.close()

# 测试加载程序的批大小为1
img = next(iter(test_loader))[0].to(device)
print(f"\nImage has shape: {img.shape}\n")

# 绘制图像
img_copy = img.to('cpu')
plt.imshow(img_copy.reshape(400,700), cmap="gray")
plt.axis("off")
plt.show()

# 将图像通过第一层
results = [all_layers[0](img)]

# 将上一层的结果传递给下一层
for idx in range(1, len(all_layers)):  # Start at 1, first layer already passed!
results.append(all_layers[idx](results[-1]))  # 将最后一个结果传递给该图层

def gradCAM(x):

# 运行模型并进行预测
logits = model(x)
pred = logits.max(-1)[-1] # 返回最大值（0或1）

# 在最终的conv层上获取激活量
last_conv = model.model_layers[:5]
activations = last_conv(x)

# 计算关于模型预测的梯度
model.zero_grad()
logits[0,pred].backward(retain_graph=True)

# 计算最后一个层每个输出通道的平均梯度
pooled_grads = model.model_layers[3].weight.grad.mean((1,2,3))

#乘以每个输出通道与其对应的平均梯度
for i in range(activations.shape[1]):
activations[:,i,:,:] *= pooled_grads[i]

# 以所有加权输出通道的平均值计算热图
heatmap = torch.mean(activations, dim=1)[0].cpu().detach()

return heatmap

import cv2

def upsampleHeatmap(map, img):
m,M = map.min(), map.max()
i,I = img.min(), img.max()
map = 255 * ((map-m) / (M-m))
img = 255 * ((img-i) / (I-i))
map = np.uint8(map)
img = np.uint8(img)
map = cv2.resize(map, (700,400))
map = cv2.applyColorMap(255-map, cv2.COLORMAP_JET)
map = np.uint8(map)
img = cv2.applyColorMap(255-img, cv2.COLORMAP_JET)
img = np.uint8(img)
map = np.uint8(map*0.7 + img*0.3)
return map

#在数据加载器上迭代
for idx, (image, name) in enumerate(predict_loader):

# 计算热图
image = image.to(device)
heatmap = gradCAM(image)
image = image.cpu().squeeze(0).permute(1,2,0)
heatmap = upsampleHeatmap(heatmap, image)

# 绘制图像和热图
fig = plt.figure(figsize=(14,5))
fig.suptitle(f"\nFile: {names[idx]}, Predicted label: {lbls[idx]}\n", fontsize=24)
plt.subplot(1, 2, 1)
plt.imshow(image, cmap="gray")
plt.title(f"Image", fontsize=14)
plt.axis("off")
plt.subplot(1, 2, 2)
plt.imshow(heatmap)
plt.title(f"Heatmap", fontsize=14)
plt.tight_layout()
plt.axis("off")
plt.show()
plt.close()