87 lines
2.8 KiB
Python
87 lines
2.8 KiB
Python
import torch
|
||
import torchvision
|
||
from torch.utils.data import Dataset, DataLoader
|
||
from torchvision import transforms,datasets
|
||
from torchvision.transforms import ToTensor
|
||
import numpy as np
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
import torch.optim as optim
|
||
import matplotlib.pyplot as plt
|
||
from PIL import Image
|
||
|
||
|
||
device=torch.device("cpu")
|
||
|
||
cropbox=(0,0,316,316)
|
||
batch_size = 1
|
||
|
||
class CNN(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.features=nn.Sequential(
|
||
nn.Conv2d(3,6, kernel_size=5),
|
||
nn.ReLU(inplace=True),
|
||
nn.MaxPool2d(kernel_size=2,stride=2),
|
||
nn.Conv2d(6,12, kernel_size=5),
|
||
nn.ReLU(inplace=True),
|
||
nn.MaxPool2d(kernel_size=2,stride=2),
|
||
nn.Conv2d(12,24,kernel_size=5,stride=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(24,24,kernel_size=5,stride=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.Conv2d(24,48,kernel_size=5,stride=1),
|
||
nn.ReLU(inplace=True),
|
||
nn.MaxPool2d(kernel_size=2,stride=2),
|
||
)
|
||
|
||
# 请把之前 Network1 的 self.fc1 复制到这里：
|
||
# self.fc1 = nn.Linear(in_features=XXX, out_features=120)
|
||
|
||
self.classifier=nn.Sequential(
|
||
nn.Linear(32*32*48,2048),
|
||
nn.ReLU(inplace=True),
|
||
#nn.Dropout(0.5),
|
||
nn.Linear(2048,1024),
|
||
nn.ReLU(inplace=True),
|
||
#nn.Dropout(0.5),
|
||
nn.Linear(1024,1),
|
||
)
|
||
|
||
|
||
def forward(self, t):
|
||
t=self.features(t)
|
||
# 特别注意：最后一个卷积层的输出在流入全连接层之前，要对激活值的维度进行变换，变换后的张量
|
||
# 是二维张量，其中 dim0 对应该批次中的不同样本，dim1 对应每个样本流向全连接层的输入激活向量。
|
||
# 请完成如下代码。
|
||
|
||
# dim0：一个批次的样本数
|
||
# dim1：每一个样本的输入激活向量长度
|
||
|
||
# t = t.reshape(XXX,XXX)
|
||
t = t.reshape(batch_size,32*32*48)
|
||
|
||
t = self.classifier(t)
|
||
|
||
return t
|
||
|
||
const_avg=3.6156537532806396
|
||
confidence_margin = 5.0
|
||
|
||
if __name__ == '__main__':
|
||
network=CNN()
|
||
network = torch.load("Dev_netv1.0_CNN.pth") # PATH same as above
|
||
network.eval()
|
||
print("模型已导入完成。")
|
||
path=input("请输入待识别图片地址：")
|
||
img = Image.open(path)
|
||
img_2=img.crop(cropbox)
|
||
img_ten=transforms.ToTensor()(np.array(img_2))
|
||
pred=network(img_ten).item()
|
||
if pred<=const_avg:
|
||
print("该图片正常。预测输出：{:3f}".format(pred))
|
||
elif pred>=confidence_margin:
|
||
print("该图片异常。预测输出：{:3f}".format(pred))
|
||
else:
|
||
output=(pred-const_avg)/(confidence_margin-const_avg)
|
||
print("无法完全确定该图片是否正常。参考值：{:.3f}。参考值靠近0说明正常可能性更大，靠近1说明异常可能性更大。\n预测输出：{:3f}".format(output,pred)) |