PyTorch——调用ResNet的预训练模型做图片分类

import torch
from torchvision import models
from PIL import Image
from torchvision import transforms

dir(models)
img=Image.open("C:/Users/Dell/Desktop/446598877159358796.jpg")
#img.show()
preprocess=transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(
    mean=[0.485,0.456,0.406],
    std=[0.229,0.224,0.225]
    )
])
img_t=preprocess(img)
#print(img_t)
batch_t=torch.unsqueeze(img_t,0)
#unsqueeze() 增维
resnet=models.resnet101(pretrained=True)
resnet.eval()
out=resnet(batch_t)
with open('E:/imagenet_classes.txt') as f:
    labels=[line.strip() for line in f.readlines()]
#strip() 方法删除任何前导(开头的空格)和尾随(结尾的空格)字符(空格是要删除的默认前导字符)。
#readlines() 方法返回一个列表,其中包含文件中的每一行作为列表项。
#请使用 hint 参数来限制返回的行数。如果返回的字节总数超过了指定的数目,则不会再返回任何行。
_, index=torch.max(out,1)
# 当函数有两个及两个以上的返回值时可以用_,来选择返回其中某个值
# 
# def calculate(a,b):
#     c = a+b
#     d = a-b
#     e = a*b
#     return c,d,e
# f,_,_ = calculate(1,2)
# _,g,_ = calculate(1,2)
# _,_,h = calculate(1,2)
# 
# print(f) #此时输出的是c的值
# print(g) #此时输出的是d的值
# print(h) #此时输出的是e的值
#print(index)
#output:tensor([207])
print(labels[index[0]])
#output:golden retriever
percentage=torch.nn.functional.softmax(out,dim=1)[0]*100
print(percentage[index[0]].item())
#output:97.98258972167969
_,indices=torch.sort(out,descending=True)
[(labels[idx],percentage[idx].item()) for idx in indices[0][:5]]
#output:[('golden retriever', 97.98258972167969), ('flat-coated retriever', 1.1356502771377563), ('tennis ball', 0.28936538100242615), ('Irish setter, red setter', 0.1782001107931137), ('Sussex spaniel', 0.05873418599367142)]

你可能感兴趣的:(PyTorch,python)