import os
import numpy as np
import torch
import torch.nn
import torchvision.models as models
from torch.autograd import Variable
import torch.cuda
import torchvision.transforms as transforms
from PIL import Image
TARGET_IMG_SIZE = 224
img_to_tensor = transforms.ToTensor()
def make_model():
model = models.vgg16(pretrained=True).features[:28]
model = model.eval()
model.cuda()
return model
def extract_feature(model, imgpath):
model.eval()
img = Image.open(imgpath)
img = img.resize((TARGET_IMG_SIZE, TARGET_IMG_SIZE))
tensor = img_to_tensor(img)
np_tensor=tensor.numpy()
a=[]
a.append(np_tensor)
tensor = torch.Tensor(a)
tensor = tensor.cuda()
result = model(Variable(tensor))
result_npy = result.data.cpu().numpy()
return result_npy[0]
if __name__ == "__main__":
model = make_model()
print(model)
imgpath = './path/to/img.jpg'
tmp = extract_feature(model, imgpath)
print(tmp.shape)
- 得到的tensor shape
-![在这里插入图片描述](https://i-blog.csdnimg.cn/blog_migrate/4f345be7bc1777727b08959232773040.png)