import torch
from torch import nn
from PIL import Image
from torchvision import transforms, datasets
import cv2
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, padding=1, kernel_size=3)
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, padding=1, kernel_size=3)
self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, padding=1, kernel_size=3)
self.conv4 = nn.Conv2d(in_channels=128, out_channels=128, padding=1, kernel_size=3)
self.relu = nn.ReLU()
self.flatten = nn.Flatten()
self.linear1 = nn.Linear(128 * 7 * 7, 512)
self.linear2 = nn.Linear(512, 6)
self.softmax = nn.Softmax()
def forward(self, x):
x = self.relu(self.conv1(x))
x = self.pool1(x)
x = self.relu(self.conv2(x))
x = self.pool1(x)
x = self.relu(self.conv3(x))
x = self.pool1(x)
x = self.relu(self.conv4(x))
x = self.flatten(x)
x = self.relu(self.linear1(x))
y = self.linear2(x)
return y
class_names = ['GC', 'GL', 'NL', 'RC', 'RL', 'UK'] # 这个顺序很重要,要和训练时候的类名顺序一致
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
##载入模型并读取权重
model = Net()
model.load_state_dict(torch.load("./data/detect_light.pt"))
model.to(device)
model.eval()
img_path = '/home/jwd/dataset/roi455.jpg'
#==========使用PIL进行测试的代码=====================================
transform_valid = transforms.Compose([transforms.Resize((56, 56), interpolation=2),
transforms.ToTensor()])
img = Image.open(img_path)
img_ = transform_valid(img).unsqueeze(0) # 拓展维度
#==========使用opencv读取图像的测试代码,若使用opencv进行读取,将上面(1)注释掉即可==========
# img = cv2.imread(img_path)
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# img = cv2.resize(img, (56, 56))
# img_ = torch.from_numpy(img).float().permute(2, 0, 1).unsqueeze(0)/255
img_ = img_.to(device)
outputs = model(img_)
# 输出概率最大的类别
_, indices = torch.max(outputs, 1)
percentage = torch.nn.functional.softmax(outputs, dim=1)[0] * 100
perc = percentage[int(indices)].item()
result = class_names[indices]
print('predicted:', result)
# 得到预测结果,并且从大到小排序
# _, indices = torch.sort(outputs, descending=True)
# 返回每个预测值的百分数
# percentage = torch.nn.functional.softmax(outputs, dim=1)[0] * 100
# print([(class_names[idx], percentage[idx].item()) for idx in indices[0][:5]])
pytorch模型推理单张图片读取方式
最新推荐文章于 2022-11-09 20:44:42 发布