【Pytorch】测试单张图片（调用transforms）

Pytorch测试单张图片（调用transforms）

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
import tqdm
import numpy as np
import argparse
import os.path as osp
import os
# from pytorch_toolbelt.inference import tta
import shutil
import glob
from network import Resnet18, Resnet34, Resnet50, SEnet154, Resnet101, SEResNext50
from PIL import Image
import torchvision.transforms as transforms
from data_process.hp_dataset import HPDataset

net = Resnet50(n_classes=5) 
net.eval()
net.load_state_dict(torch.load(r'/src/model-1021-v5.pth'))
net = net.cuda()

img = Image.open(r"/src/img.tif")
img = img.resize((256, 256), Image.BILINEAR)
img = np.array(img)

to_tensor = transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
img = to_tensor(img)
img = torch.unsqueeze(img, 0) #给最高位添加一个维度，也就是batchsize的大小

img1 = img.cuda()
outputs = net(img1)
outputs = F.softmax(outputs, dim=1)
predicted = torch.max(outputs, dim=1)[1].cpu().item()
outputs = outputs.squeeze(0)
# predicted = torch.max(outputs, dim=1)[1].cpu().numpy()
confidence=outputs[predicted].item()
confidence=round(confidence, 3)
print(confidence)