へっぽこ！リアルタイム画像認識

今回は、caffeのModelZooから学習済みのgoogLeNetをいただいて、それをそのままchainerで使ってリアルタイムの画像認識をします。

拝借するモデルは1000個のラベルについて学習しているので、学習したラベル以外のものについて認識することはできません。
まあ、とりあえず認識できれば嬉しいというへっぽこ感でやっていこうと思います。

モデル読み込み

とりあえず、ここからcaffemodelをダウンロードします。
それから、このスクリプトでラベルファイルなど
をダウンロードします。
chainerでのcaffemodelの読み込みは以下の通り。

import chainer
import chainer.functions as F
from chainer.functions import caffe

googlenet = caffe.CaffeFunction("bvlc_googlenet.caffemodel")

私はこのcaffemodelをpickleで保存して使っています。
こんな感じ。

import cPickle as pickle

# 保存
pickle.dump(func,open('googlenet.pkl','w'))

# 読み込み
func = pickle.load(open('googlenet.pkl', 'rb'))

画像判別

次に画像の判別です。
画像をモデルに通す前に平均画像の差分を取る必要があります。
OpenCVで取り込んだので少し変換する必要があり、変換後差分をとります。
そのあと、モデルを通して1000次元のベクトルを取得し、
最も確率の高い部分のインデックスを持つラベルを返します。

in_size = 224
cropwidth = 256 - in_size
start = cropwidth // 2
stop = start + in_size

# 平均画像の読み込み
meanfile = os.path.join('ilsvrc', 'ilsvrc_2012_mean.npy')
mean_image = np.load(meanfile)
mean_image = mean_image[:, start:stop, start:stop].copy()

def predict(image):
	"""画像を判別"""
	global mean_image, in_size, cropwidth, start, stop
	
	def swap(x):
		x = np.array(x)[:, :, ::-1]
		x = np.swapaxes(x, 0, 2)
		x = np.swapaxes(x, 1, 2)
		return x

	x_data = np.ndarray((1, 3, in_size, in_size), dtype=np.float32)
	
	image = swap(image)
	image = image[:, start:stop, start:stop].copy().astype(np.float32)
	x_data[0] = image-mean_image

	x = chainer.Variable(x_data, volatile=True)
	
	y, = func(inputs={'data': x}, outputs=['loss3/classifier'], train=False)
	synset_i = y.data.argmax(axis=1)
	return synset_words[synset_i]

OpenCVで画像をリアルタイムでキャプチャ

判別と画像出力をスレッドで並列に行うため、リアルタイムでキャプチャするための関数を作ります。

def cap():
    """リアルタイムで画像をキャプチャする"""
    global capture,msg
    fontface=cv2.FONT_ITALIC
    fontscale=1.0
    bold = 4
    color=(255,0,255)
    
    while True:
        ret, image = capture.read()

        if ret == False:
            cv2.destroyAllWindows()
            break

        location=(0,image.shape[0]/2)
        cv2.putText(image,msg,location,fontface,fontscale,color,bold)  
        cv2.imshow("Image Recognition", image)

        if cv2.waitKey(33) >= 0:
    #         cv2.imwrite("image.png", image)
            cv2.destroyAllWindows()
            break
    print 'out'
    cv2.destroyAllWindows()

取り込んだ画像を判別

取り込んだ画像を判別して、グローバル変数にラベルをセットします。
キー入力があるとプログラムを終了します。

def main():
    """取り込まれた画像を判別してラベルをセット"""
    global msg, capture
    while True:
        ret, image = capture.read()

        if ret == False:
            break
        img = Image.fromarray(image[::-1, :, ::-1].copy()).resize((256, 256), Image.ANTIALIAS)

        # ラベルをcapとの共通変数にセット
        msg = predict(img)[0]

        # キー入力で終了
        if cv2.waitKey(33) >= 0:
            break
        if capThread.isAlive() == False:
            break

スレッドの処理

cap()をスレッドで動かして、そのあとmain()を実行します。
これで一応動くはずです。

capture = cv2.VideoCapture(0)
msg = ''
cv2.namedWindow("Image Recognition", cv2.WINDOW_NORMAL)
capture.set(cv2.cv.CV_CAP_PROP_FRAME_WIDTH,324)  #2592
capture.set(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT,228) #1944
capThread = threading.Thread(target=cap, name='cap')
capThread.setDaemon(True)
capThread.start()
main()

まとめ

ということで、畳み込みニューラルネットワークや学習などを説明することなく画像認識できました。Model Zoo 素晴らしいですね。

コードの全体はGithubで公開しています。
GitHub - YasukeXXX/Heppoco

PlayGround