利用卷积神经网络提取图像的特征信息_Pytorch

---

0.前言

特征提取的英文叫做feature extractor，它是将一些原始的输入的数据维度减少或者将原始的特征进行重新组合以便于后续的使用。简单来说有两个作用：减少数据维度，整理已有的数据特征。为后续的图像处理任务提供良好的数据基础．

---

1.引入库

创建一个main.py，代码如下：

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, glob
import h5py
import time
from tqdm import tqdm

2.建立编码器用于提取图像特征信息

代码如下：

class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        VGG = models.vgg16(pretrained=True)#使用与训练权重
        self.feature = VGG.features
    def forward(self, x):
        output = self.feature(x)#输出维度为（512*7*7）
        m = nn.MaxPool2d(7, stride=1)
        output = m(output)#对输出的维度进行平均池化->（512*1*1）
        output = output.view(output.size(0), -1)#512*1
        return output

预训练权重对于不同的数据集是通用的，因为特征是通用的。一般来讲，从0开始训练效果会很差，因为权值太过随机，特征提取效果不明显。简单的理解，预训练权重是别人针对网络用大量的数据进行了长时间的训练，因此使用预训练权重就相当于已经将这个模型优化好了．但是不同网络的预训练权重不同，而且如果改变了网络的结构，预训练权重就可能用不了．

---

3.建立特征提取层

代码如下：

model = Encoder()#定义模型
def extractor(img_path, net, use_gpu, file_name):
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor()]
    )#将图像转化为张量的形式，同时扩充224*224
    img = Image.open(img_path)#读取图像
    img = transform(img)
    x = Variable(torch.unsqueeze(img, dim=0).float(), requires_grad=False)
    if use_gpu:
        x = x.cuda()
        net = net.cuda()
    y = net(x).cpu()
    y = torch.squeeze(y)
    ys = y.data.numpy()
    file_names = file_name
    return file_names, ys

处理数据，并且将数据传入特征提取网络中．返回的是图像的名字和图像的特征信息（每张图像的特征信息为512*1）

---

4.读取图像，存储特征信息

代码如下：

if __name__ == '__main__':
    data_dir = './ImgDB'#图像路径
    files_list = []
    names = []
    features = []
    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))
    use_gpu = torch.cuda.is_available()
    pbar = tqdm(files_list)
    for x_path in files_list and pbar:
        file_name = x_path.split('/')[-1]
        file_names, ys = extractor(x_path, model, use_gpu, file_name)
        features.append(ys)
        names.append(file_names)
        pbar.set_description("Extracting features")
        # 设置进度条右边显示的信息
        pbar.set_postfix(Img_name=file_names)
        time.sleep(0.1)
    features = np.array(features)
    h5f = h5py.File('feature.h5', 'w')#创建h5数据库，用于存储数据
    h5f.create_dataset('tensor', data=features)#创建一个key，里面存储的是所有图像的特征
    h5f.create_dataset('name', data=names)#创建一个key，里面存储的是所有图像的名字
    h5f.close()

5.查看h5文件中的信息

创建h5.py，用于查看之前提取的特征信息．代码如下：

import h5py
f = h5py.File('feature.h5', 'r')
# 读取文件,一定记得加上路径
for key in f.keys():
	print(f[key].name)
# 打印出文件中的关键字
	print(f[key].shape)
# 将key换成某个文件中的关键字,打印出某个数据的大小尺寸
	print(f[key][:])
# 将key换成某个文件中的关键字,打印出该数据(数组)

---

6.结果展示

提取121张图像特征，并将其存到feature.h5文件中:

读取feature.h5文件:

7.总结

如果您在运行此代码时，出现任何问题，欢迎私聊我！我还会尽力为你解答！PS:下次讲针对提取到的特征，进行图像检索．