这里主要是使用预训练好的模型进行图片特征的提取，分别使用三个模型进行抽取。
而特征提取是提取神经网络模型的倒数第二层，这里自己对模型的进行了微调，话不多说，直接上代码。
使用resnet
import torch
import torch.nn as nn
from torch.autograd import Variable
from torchvision import models, transforms
from PIL import Image
import numpy as np
import os, globdata_dir = './test' # train
features_dir = './Resnet_features_test' # Resnet_features_train//这里自己修改网络
class net(nn.Module):def __init__(self):super(net, self).__init__()self.net = models.resnet50(pretrained=True)def forward(self, input):output = self.net.conv1(input)output = self.net.bn1(output)output = self.net.relu(output)output = self.net.maxpool(output)output = self.net.layer1(output)output = self.net.layer2(output)output = self.net.layer3(output)output = self.net.layer4(output)output = self.net.avgpool(output)return outputmodel = net()
//加载cuda
model = model.cuda()def extractor(img_path, saved_path, net, use_gpu):transform = transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor()])img = Image.open(img_path)img = transform(img)print(img.shape)x = Variable(torch.unsqueeze(img, dim=0).float(), requires_grad=False)print(x.shape)if use_gpu:x = x.cuda()net = net.cuda()y = net(x).cpu()y = torch.squeeze(y)y = y.data.numpy()print(y.shape)np.savetxt(saved_path, y, delimiter=',')if __name__ == '__main__':extensions = ['jpg', 'jpeg', 'JPG', 'JPEG']files_list = []x = os.walk(data_dir)for path,d,filelist in x:for filename in filelist:file_glob = os.path.join(path, filename)files_list.extend(glob.glob(file_glob))print(files_list)use_gpu = torch.cuda.is_available()for x_path in files_list:print("x_path" + x_path)file_name = x_path.split('/')[-1]fx_path = os.path.join(features_dir, file_name + '.txt')print(fx_path)extractor(x_path, fx_path, model, use_gpu)
使用vgg
import torch
import torch.nn as nn
from torch.autograd import Variable
from torchvision import models, transforms
from PIL import Image
import numpy as np
import os, globdata_dir = './test' # train
features_dir = './Vgg_features_test' # Vgg_features_trainclass Encoder(nn.Module):def __init__(self):super(Encoder, self).__init__()VGG = models.vgg16(pretrained=True)self.feature = VGG.featuresself.classifier = nn.Sequential(*list(VGG.classifier.children())[:-3])pretrained_dict = VGG.state_dict()model_dict = self.classifier.state_dict()pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}model_dict.update(pretrained_dict)self.classifier.load_state_dict(model_dict)def forward(self, x):output = self.feature(x)output = output.view(output.size(0), -1)output = self.classifier(output)return outputmodel = Encoder()
model = model.cuda()def extractor(img_path, saved_path, net, use_gpu):transform = transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor()])img = Image.open(img_path)img = transform(img)print(img.shape)x = Variable(torch.unsqueeze(img, dim=0).float(), requires_grad=False)print(x.shape)if use_gpu:x = x.cuda()net = net.cuda()y = net(x).cpu()y = torch.squeeze(y)y = y.data.numpy()print(y.shape)np.savetxt(saved_path, y, delimiter=',')if __name__ == '__main__':extensions = ['jpg', 'jpeg', 'JPG', 'JPEG']files_list = []x = os.walk(data_dir)for path, d, filelist in x:for filename in filelist:file_glob = os.path.join(path, filename)files_list.extend(glob.glob(file_glob))print(files_list)use_gpu = torch.cuda.is_available()for x_path in files_list:print("x_path" + x_path)file_name = x_path.split('/')[-1]fx_path = os.path.join(features_dir, file_name + '.txt')print(fx_path)extractor(x_path, fx_path, model, use_gpu)
使用Densnet
import torch
import torch.nn as nn
from torch.autograd import Variable
from torchvision import models, transforms
from PIL import Image
import numpy as np
import os, globdata_dir = './test' # train
features_dir = './DenseNet_features_test' # DenseNet_features_trainclass Encoder(nn.Module):def __init__(self):super(Encoder, self).__init__()densnet = models.densenet121(pretrained=True)self.feature = densnet.featuresself.classifier = nn.Sequential(*list(densnet.classifier.children())[:-1])pretrained_dict = densnet.state_dict()model_dict = self.classifier.state_dict()pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}model_dict.update(pretrained_dict)self.classifier.load_state_dict(model_dict)def forward(self, x):output = self.feature(x)avg = nn.AvgPool2d(7, stride=1)output = avg(output)return outputmodel = Encoder()
model = model.cuda()def extractor(img_path, saved_path, net, use_gpu):transform = transforms.Compose([transforms.Resize(256),transforms.CenterCrop(224),transforms.ToTensor()])img = Image.open(img_path)img = transform(img)print(img.shape)x = Variable(torch.unsqueeze(img, dim=0).float(), requires_grad=False)print(x.shape)if use_gpu:x = x.cuda()net = net.cuda()y = net(x).cpu()y = torch.squeeze(y)y = y.data.numpy()print(y.shape)//保存为txtnp.savetxt(saved_path, y, delimiter=',')if __name__ == '__main__':extensions = ['jpg', 'jpeg', 'JPG', 'JPEG']files_list = []x = os.walk(data_dir)for path, d, filelist in x:for filename in filelist:file_glob = os.path.join(path, filename)files_list.extend(glob.glob(file_glob))print(files_list)use_gpu = torch.cuda.is_available()for x_path in files_list:# print("x_path" + x_path)file_name = x_path.split('/')[-1]fx_path = os.path.join(features_dir, file_name + '.txt')# print(fx_path)extractor(x_path, fx_path, model, use_gpu)
以上是使用三个模型的特征提取，可以根据不同的需求进行选择。