1. 程式人生 > >Caffe用訓練好的模型測試圖片

Caffe用訓練好的模型測試圖片

這是一個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)