Pycaffe提取图片特征

代码:

import sys
import numpy as np
import caffe
import os
import time
import shutil
from PIL import Image

def get_image_files():
#获取path路径下所有图片的名字
    path = '/home/fwei/fdata/zhangye/img'
    files = os.listdir(path)
    s=[]
    for file in files:
        s.append(file)
    return s

def remove_bad_images():
#对于不能正常读取的图片,从文件夹下移除
    s = get_image_files()
    for image in s:
        try:
            image_file = '/home/fwei/fdata/zhangye/img/' +image
            img = Image.open(image_file)
            img.verify()
            #img = caffe.io.load_image(image_file)
        except IOError:
            print(image_file)
            shutil.move(image_file,'/home/fwei/fdata/errimg')

def extra_image_feature():
    model_file = '/home/fwei/caffe/models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel'
    deploy_prototxt = '/home/fwei/caffe/models/bvlc_reference_caffenet/deploy.prototxt'
    net = caffe.Net(deploy_prototxt,model_file,caffe.TEST)
    layer = 'fc7' #提取fc7层的特征
    if layer not in net.blobs:
        raise TypeError("Invalid layer name: " + layer)
    imagemean_file = '/home/fwei/caffe/python/caffe/imagenet/ilsvrc_2012_mean.npy'
    transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
    transformer.set_mean('data', np.load(imagemean_file).mean(1).mean(1))
    transformer.set_transpose('data', (2,0,1))
    transformer.set_raw_scale('data', 255.0)
    net.blobs['data'].reshape(1,3,227,227)   
    image_files = get_image_files()
    i=0
    caffe.set_device(0)
    caffe.set_mode_gpu()
    for input_image_file in image_files:
            img_file = '/home/fwei/fdata/zhangye/img/'+input_image_file
            print ('No '+str(i))
            print img_file
            print '_______________________________'
            img = caffe.io.load_image(img_file)
            net.blobs['data'].data[...] = transformer.preprocess('data', img)
            net.forward()
            feature_data = net.blobs[layer].data[0]
            feature_data = (feature_data - feature_data.min())/(feature_data.max() - feature_data.min())
            output_file = input_image_file+'.txt'
            with open(output_file, 'w') as f:
                np.savetxt(f,feature_data,fmt='%.10f',delimiter='\n')
            f.close()
            i+=1
     print i

remove_bad_images()
extra_image_feature()


你可能感兴趣的:(Pycaffe提取图片特征)