numpy.where()を使った配列処理

画像処理などで、Numpyの配列に対して条件により処理を行うとき、for文をまわして要素にアクセスする場合とnumpy.where()を使う場合と速さを比較しました。 

 

 numpy.where(条件, 満たした場合の値, 満たさない場合の値)

という使い方です。

 

まず操作するデータを作成します。

import numpy as np
import time

# randintで画素数256x256の100枚のデータをランダムで作成

np.random.seed(seed=1)

array1 = np.random.randint(0, 256, size=6553600).reshape(100, 256, 256)

（これはnumpy.randint()を使って作ったものなので画像としては無茶苦茶です）

 

このデータ配列に対して、画素値が50より大でかつ200以下のものを抽出した画像を作る操作を行います（つまり50～200以外は画素値0となります）。 

 

for文でこれを処理します。

# 画素値が50より大で200以下はそのまま,それ以外は画素値を0の処理を行い新しい配列として格納
# その処理時間を測定

start = time.time()

array2 = np.zeros_like(array1)

for i in range(0, 256):
    for j in range(0, 256):
        for k in range(0, 100):
            if (array1[k, j, i] > 50) and (array1[k, j, i] <= 200):
                array2[k, j, i] = array1[k, j, i]


t = time.time() - start
print(t)

結果は 10.067001104354858secでした。

 

 

次に同じ処理をnumpy.where()を使って行います。

# 同じ操作をnumpy.where()を使って処理
# その処理時間を測定

start= time.time()

array3 = np.where((array1 > 50) & (array1 <= 200), array1, 0)

t = time.time() - start
print(t)

結果は 0.08173227310180664secと、numpy.where()を使うほうが圧倒的に速く、この処理自体のコードは1行のみでした。

 

 上のコードのように条件が複数ある場合はかっこでくくって、&（かつ）や|（または）でつなげます。andやorではないので注意を。