pytorch加载一张图片在训练好的模型上测试(test single image with pytorch)

pytorch加载一张图片在训练好的模型上测试(test single image with pytorch)

import argparse
import os
import shutil
import time
import sys
import numpy as np

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import model as models
from torchvision import transforms,datasets 
from PIL import Image

from print_log import make_print_to_file

parser = argparse.ArgumentParser(description='Pedestrian Attribute Framework')
parser.add_argument('--experiment', default='peta', type=str, required=True, help='(default=%(default)s)')
parser.add_argument('--approach', default='Resnet101', type=str, required=True, help='(default=%(default)s)')
parser.add_argument('--epochs', default=60, type=int, required=False, help='(default=%(default)d)')
parser.add_argument('--batch_size', default=64, type=int, required=False, help='(default=%(default)d)')
parser.add_argument('--lr', '--learning-rate', default=0.0001, type=float, required=False, help='(default=%(default)f)')
parser.add_argument('--optimizer', default='adam', type=str, required=False, help='(default=%(default)s)')
parser.add_argument('--momentum', default=0.9, type=float, required=False, help='(default=%(default)f)')
parser.add_argument('--weight_decay', default=0.0005, type=float, required=False, help='(default=%(default)f)')
parser.add_argument('--start-epoch', default=0, type=int, required=False, help='(default=%(default)d)')
parser.add_argument('--print_freq', default=100, type=int, required=False, help='(default=%(default)d)')
parser.add_argument('--save_freq', default=10, type=int, required=False, help='(default=%(default)d)')
parser.add_argument('--resume', default='', type=str, required=False, help='(default=%(default)s)')
parser.add_argument('--decay_epoch', default=(20,40), type=eval, required=False, help='(default=%(default)d)')
parser.add_argument('--prefix', default='', type=str, required=False, help='(default=%(default)s)')
parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true', required=False, help='evaluate model on validation set')

# Seed
np.random.seed(1)
torch.manual_seed(1)
if torch.cuda.is_available(): torch.cuda.manual_seed(1)
else: print('[CUDA unavailable]'); sys.exit()

#####################################################################################################


def main():
    make_print_to_file(path='./log/')
    global args, best_accu
    args = parser.parse_args()

    print('=' * 100)
    print('Arguments = ')
    for arg in vars(args):
        print('\t' + arg + ':', getattr(args, arg))
    print('=' * 100)
    
    attr_num = 35
    attr_names = ['0:Age16-30','1:Age31-45','2:Age46-60','3:AgeAbove61','4:Backpack','5:CarryingOther',
                  '6:Casual lower','7:Casual upper','8:Formal lower','9:Formal upper','10:Hat',
                  '11:Jacket','12:Jeans','13:Leather Shoes','14:Logo','15:Longhair',
                  '16:Male','17:Messenger Bag','18:Muffler','19:No accessory','20:No carrying',
                  '21:Plaid','22:PlasticBags','23:Sandals','24:Shoes','25:Shorts',
                  '26:Short sleeve','27:skirt','28:Sneaker','29:stripes','30:Sunglasses',
                  '31:Trousers','32:Tshirt','33:UpperOther','34:V-Neck']
    
    # create model
    #model = models.__dict__[args.approach](pretrained=True, num_classes=attr_num)
    model = models.__dict__[args.approach](num_labels=attr_num)

    # get the number of model parameters
    print('Number of model parameters: {}'.format(
        sum([p.data.nelement() for p in model.parameters()])))
    print('')

    # for training on multiple GPUs.
    # Use CUDA_VISIBLE_DEVICES=0,1 to specify which GPUs to use
    model = torch.nn.DataParallel(model).cuda()

    # optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']
            best_accu = checkpoint['best_accu']
            model.load_state_dict(checkpoint['state_dict'])
            print("=> loaded checkpoint '{}' (epoch {})"
                  .format(args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))

    cudnn.benchmark = False
    cudnn.deterministic = True
    
    ## test single image with pytorch
    img_path = '/home/lyp/Datasets/PETA/PETA_dataset/GRID/archive/0002_2_25027_160_75_118_274.jpeg'
    #img_label = [0 0 1 0 0 0 0 0 1 1 0 0 0 1 0 0 1 1 0 1 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0]
    
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    transform_test = transforms.Compose([
        transforms.Resize(size=(224, 224)),
        transforms.ToTensor(),
        normalize
        ])

    img = Image.open(img_path)
    img_tensor = transform_test(img).unsqueeze(0) #拓展维度

    if args.evaluate:
        test(img_tensor, model, attr_num, attr_names)
        return 
    

def test(img_tensor, model, attr_num, attr_names):
    model.eval()
    
    input = img_tensor.cuda(non_blocking=True)
    output = model(input)
    output = torch.sigmoid(output.data).cpu().numpy()
    print('output1 = ', output[0])
    output = np.where(output > 0.5, 1, 0)
    print('output2 = ', output[0])
    
    # print attribute name
    class_name = []
    for i in range(0,attr_num):
        if output[0][i] == 1:
            class_name.append(i)
    print(class_name)
    labels = []
    for i in range(0, len(class_name)):
        labels.append(attr_names[class_name[i]])
    print('Predicted results:', labels)
        
    

if __name__ == '__main__':
    main()

因为我是基于PETA数据集训练的模型，很多部分的内容需要大家参考自己的模型进行更改