这里提取灰度图的特征,所以我把它堆叠了三次变成了三通道,vgg-19的预训练模型可以在pytorch提供的官方地址下载
import torch,os
import scipy.io as sio
import numpy as np
import scipy.misc as sc
from torchvision import models
from torch.autograd import Variable
model = models.vgg19()
model.load_state_dict(torch.load('models/vgg19-dcbb9e9d.pth'))
model.cuda()
data_file = os.listdir('img/')
for pic in data_file:
pic_path = 'img/'+pic
print(pic_path)
data = sc.imread(pic_path)
data = np.resize(data,[224,224])
img = np.zeros((1, 3, 224, 224)).astype(np.float32)
for i in range(3):
img[:,i,:,:] = data
img = torch.from_numpy(img)
img = Variable(img).cuda()
feature = model(img)
tmp = []
for key,parm in enumerate(model.classifier.parameters()):
if key == 0 or key == 2 or key == 4:
continue
parm = parm.cpu()
tmp.append(parm.data.numpy())
sio.savemat('feature/'+pic[:-4]+'.mat',{'feature_4096_first':tmp[0],
'feature_4096_second': tmp[1],
'feature_1000': tmp[2]})
# break