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)
