#resnet feature extraction
import os
os.chdir('/root/caffe-master/examples')
import numpy as np
import matplotlib.pyplot as plt
import itertools
import csv
plt.rcParams['figure.figsize']=(10,10)
plt.rcParams['image.interpolation']='nearest'
plt.rcParams['image.cmap']='gray'
import sys
caffe_root='../' #this file should be run from {caffe_root}/examples(otherwise change this line)
import caffe
caffe.set_mode_gpu()
model_def='../atongtest/ResNet-50-deploy.prototxt'
model_weights='../atongtest/ResNet-50-train2_iter_350000.caffemodel'
net=caffe.Net(model_def,model_weights,caffe.TEST)
MEAN_PROTO_PATH=caffe_root+'atongtest/ResNet_mean.binaryproto'
blob=caffe.proto.caffe_pb2.BlobProto()
data=open(MEAN_PROTO_PATH,'rb').read()
blob.ParseFromString(data)
array=np.array(caffe.io.blobproto_to_array(blob))
mu=array[0]
mean=mu.mean(1).mean(1)
net=caffe.Classifier(model_def,model_weights,mean=mean,channel_swap=(2,1,0),raw_scale=255,image_dims=(224,224))
#rootdir 根目录，subdir 根目录下的子目录，若没有 则为[], subfile：根目录下的文件，若没有 则为[]
#os.walk 先遍历根目录，然后将子目录路径添加到原根目录，以rootdir+subdir[i] 做为新的根目录继续遍历 ，由此实现 给定路径下所有文件的遍历
#imgpath:带检测图片路径
#Imglist ：所有图片的路径
imgpath='/root/lfw-gray-224'
imglist=[]
for (rootdir, subdir, subfile) in os.walk(imgpath):
for img_jpg in subfile:
if img_jpg.endswith('.jpg'):
imglist.append(os.path.join(rootdir, img_jpg))
#Pre_feature_ndarray 存放所有图片的特征值，第一行为随机向量，为了可以将后续提取的图片连接，处理完所有图片后删除第一行，
#得到存放所有图片特征值的 ndarray ：Final_feature
pre_feature_ndarray=np.random.randn(1,2048)
for subimg in imglist:
input_image=caffe.io.load_image(subimg)
prediction=net.predict([input_image])
face_feature=net.blobs['pool5'].data[0]
feature_reshape=face_feature.reshape((1,2048))
pre_feature_ndarray=np.r_['0',pre_feature_ndarray,feature_reshape]
final_feature=pre_feature_ndarray[1:]
#特征值 个数
feature_num=len(final_feature)
infile=file('/root/class_in.csv','wb')
outfile=file('/root/class_out.csv','wb')
d1=csv.writer(infile)
d2=csv.writer(outfile)
#计算特征值距离
#itertools.combination('ABCD',2): AB AC AD BC BD CD itertools 迭代器
for img1,img2 in itertools.combinations(range(feature_num),2):
# distance_feature=np.linalg.norm(final_feature[img1]-final_feature[img2])
v1=final_feature[img1].reshape(1,2048)
v2=final_feature[img2].reshape(1,2048)
cos_distance=np.sum((v1*v2),1)/(np.linalg.norm(v1)*np.linalg.norm(v2))
img1_path=imglist[img1]
img2_path=imglist[img2]
if os.path.dirname(img1_path)==os.path.dirname(img2_path):
d1.writerow([cos_distance])
else:
d2.writerow([cos_distance])
infile.close()
outfile.close()
