利用pytorch训练好的模型测试单张图片

pytorch网络输入图片的格式是[B,C,H,W],分别为批大小(batchsize),图片通道数(channel),图片高(height),图片宽(width)。

图片读取方式主要有两种:(1)通过PIL进行读取;(2)通过opencv进行读取。分别进行介绍。

 

(1)通过PIL进行读取:

通过PIL的Image读取的图片是一个图片对象,可以进行裁剪翻转等torchvision.transforms变换。

torchvision.transforms可以对图像对象进行一系列裁剪、翻转等转换操作,其中也包括转换为tensor张量。(transforms.ToTensor())

(2)通过opencv进行读取

opencv读取的是ndarray格式,不能进行torchvision.transforms变换。

 

两种读取图像的测试代码如下,其中Net为自己随便写的一个网络,此处输出的类别为7类,模型替换为自己的,同时修改class_names和权重文件。注意,读取的图片路径也要修改。

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'

#(1)此处为使用PIL进行测试的代码
transform_valid = transforms.Compose([
    transforms.Resize((56, 56), interpolation=2),
    transforms.ToTensor()
  ]
)
img = Image.open(img_path)
img_ = transform_valid(img).unsqueeze(0) #拓展维度

##(2)此处为使用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]])

 

  • 33
    点赞
  • 163
    收藏
    觉得还不错? 一键收藏
  • 15
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 15
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值