# -*- coding: UTF-8 -*-
import os
import caffe
import numpy as np
root='/home/zf/caffe/'#指定根目錄
deploy=root+'models/bvlc_reference_caffenet/deploy.prototxt'#結構文件
caffe_model=root + 'models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel'
#已經訓練好的model
 
dir =root+'examples/images/'#保存測試圖片的集合
filelist=[]
filenames=os.listdir(dir)
for fn in filenames:
  fullfilename = os.path.join(dir,fn)
  filelist.append(fullfilename)
#filelist.append(fn)
def Test(img):
#加載模型
  net = caffe.Net(deploy,caffe_model,caffe.TEST)
 
# 加載輸入和配置預處理
  transformer = caffe.io.Transformer({'data':net.blobs['data'].data.shape})
  transformer.set_mean('data', np.load('/home/zf/caffe/python/caffe/imagenet/ilsvrc_2012_mean.npy').mean(1).mean(1))
  transformer.set_transpose('data', (2,0,1))
  transformer.set_channel_swap('data', (2,1,0))
  transformer.set_raw_scale('data', 255.0)
 
#注意可以調節預處理批次的大小
#由于是處理一張圖片，所以把原來的10張的批次改為1
  net.blobs['data'].reshape(1,3,227,227)
 
#加載圖片到數據層
  im = caffe.io.load_image(img)
  net.blobs['data'].data[...] = transformer.preprocess('data', im)
 
#前向計算
  out = net.forward()
 
# 其他可能的形式 : out = net.forward_all(data=np.asarray([transformer.preprocess('data', im)]))
 
#預測分類
  print out['prob'].argmax()
 
#打印預測標簽
  labels = np.loadtxt("/home/zf/caffe/data/ilsvrc12/synset_words.txt", str, delimiter='\t')
  top_k = net.blobs['prob'].data[0].flatten().argsort()[-1]
  print 'the class is:',labels[top_k]
  f=file("/home/zhengfeng/caffe/examples/zf/label.txt","a")
  f.writelines(img+' '+labels[top_k]+'\n')
labels_filename=root +'data/ilsvrc12/synset_words.txt'
#循環遍歷文件夾root+'examples/images/'下的所有圖片
for i in range(0,len(filelist)):
  img=filelist[i]
  Test(img)