Caffe用訓練好的模型測試圖片
阿新 • • 發佈:2019-02-03
這是一個python指令碼,用訓練好的caffemodel來測試圖片,接下來直接上程式碼,裡面有詳細解釋,大部分你要修改的只是路徑,另外在這個指令碼的基礎上你可以根據自己的需要進行改動。
需要的東西:訓練好的caffemodel,deploy.prototxt(可以從你的train.prototxt修改得到),可以用的caffe,待測試的影象(比如jpg)
import sys
caffe_root='/your/caffe/root/' #修改成你的Caffe專案路徑
sys.path.append(caffe_root+'python')
import caffe
caffe.set_mode_gpu() #設定為GPU執行
from pylab import *
# 修改成你的deploy.prototxt檔案路徑
model_def = 'your/deploy.prototxt/path'
model_weights = '/your/caffemodel/path' # 修改成你的caffemodel檔案的路徑
net = caffe.Net(model_def, # defines the structure of the model
model_weights, # contains the trained weights
caffe.TEST) # use test mode (e.g., don't perform dropout)
#這是一個由mean.binaryproto檔案生成mean.npy檔案的函式
def convert_mean(binMean,npyMean):
blob = caffe.proto.caffe_pb2.BlobProto()
bin_mean = open(binMean, 'rb' ).read()
blob.ParseFromString(bin_mean)
arr = np.array( caffe.io.blobproto_to_array(blob) )
npy_mean = arr[0]
np.save(npyMean, npy_mean )
binMean='your/mean.binaryproto/path' #修改成你的mean.binaryproto檔案的路徑
npyMean='which/path/you/want/to/save/mean.npy' #你想把生成的mean.npy檔案放在哪個路徑下
convert_mean(binMean,npyMean)
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
transformer.set_transpose('data', (2,0,1)) # 通道變換,例如從(530,800,3) 變成 (3,530,800)
transformer.set_mean('data', np.load(npyMean).mean(1).mean(1)) #如果你在訓練模型的時候沒有對輸入做mean操作,那麼這邊也不需要
transformer.set_raw_scale('data', 255) # rescale from [0, 1] to [0, 255]
transformer.set_channel_swap('data', (2, 1, 0)) # swap channels from RGB to BGR
with open('your/txt file/path') as image_list: # 修改成你要測試的txt檔案的路徑,這個txt檔案的內容一般是:每行表示影象的路徑,然後空格,然後是標籤,也就是說每行都是兩列
with open('predict/result/you/want/to/save','w') as result: # 如果你想把預測的結果寫到一個txt檔案中,那麼把這個路徑修改成你想儲存這個txt檔案的路徑
count_right=0
count_all=0
while 1:
list_name=image_list.readline()
if list_name == '\n' or list_name == '': #如果txt檔案都讀完了則跳出迴圈
break
image_type=list_name[0:-3].split('.')[-1]
if image_type == 'gif': #這裡我對gif個數的影象直接跳過
continue
image = caffe.io.load_image('/your/image/path/'+list_name)
# 這裡要新增你的影象所在的路徑,根據你的list_name靈活調整,總之就是影象路徑
#imshow(image)
transformed_image = transformer.preprocess('data', image)
# 用轉換後的影象代替net.blob中的data
net.blobs['data'].data[...] = transformed_image
net.blobs['data'].reshape(1, 3, 224, 224)
### perform classification
output = net.forward()
# 讀取預測結果和真實label
output_prob = net.blobs['prob'].data[0]
true_label = int(list_name[-2:-1])
# 如果預測結果和真實label一樣,則count_right+1
if(output_prob.argmax()==true_label):
count_right=count_right+1
count_all=count_all+1
# 儲存預測結果,這個可選
result.writelines(list_name[0:-1]+' '+str(output_prob.argmax())+'\n')
#可以每預測完100個樣本就列印一些,這樣好知道預測的進度,尤其是要預測幾萬或更多樣本的時候,否則你還以為程式碼卡死了
if(count_all%100==0):
print count_all
# 列印總的預測結果
print 'Accuracy: '+ str(float(count_right)/float(count_all))
print 'count_all: ' + str(count_all)
print 'count_right: ' + str(count_right)
print 'count_wrong: ' + str(count_all-count_right)