PythonのOpenCVで2重forループを避ける

Question

PythonにおけるOpenCVの実装方法に関する質問です。

以下のような画素(x,y)を網羅するような2重forループを使用すると非常に遅いです。
（処理そのものは深い意味はないですが、BGRの特定の値に対応して条件分岐を伴って色を変換する想定）

import cv2

img = cv2.imread('./hoge.png')

height, width, _ = img.shape

for x in range(height):
    for y in range(width):
        if 0 <= img[x][y][0] < 100:
            img[x][y][2] = img[x][y][0]
        elif 100 <= img[x][y][0] < 200:
            img[x][y][2] = img[x][y][1]
        else:
            img[x][y][2] = img[x][y][2]

cv2.imwrite('output1.png', img)

以下のようなforループを使用しない処理を考えましたが、
もう少し良い（直感的な）書き方は出来ないものでしょうか？

layer1~3を上記のif~elif~elseに対応させて作っておき、
cv2.bitwise_andのmask機能を使ってifの条件に一致する部分だけを合体しています。

import cv2
import numpy as np

img = cv2.imread('./hoge.png')
output = np.zeros_like(img)

b,g,r = cv2.split(img)

layer1 = output.copy()
layer1[:,:,2] = b
mask1 = cv2.inRange(b, 0, 99)

layer2 = output.copy()
layer2[:,:,2] = g
mask2 = cv2.inRange(b, 100, 199)

layer3 = output.copy()
layer3[:,:,2] = r
mask3 = cv2.inRange(b, 200, 255)

output[:,:,0] = img[:,:,0]
output[:,:,1] = img[:,:,1]

output += cv2.bitwise_and(layer1, layer1, mask=mask1)
output += cv2.bitwise_and(layer2, layer2, mask=mask2)
output += cv2.bitwise_and(layer3, layer3, mask=mask3)

cv2.imwrite('output2.png', output)

質問は、forを使わずに高速で、出来るだけ可読性が高い書き方は何か、です。
よろしくお願い致します。

Answer 1

numpy の bool indexing を利用して、代入を行えばよいです。
質問の1個目のコードは numpy で書くと以下のように書き換えられます。

for x in range(height):
    for y in range(width):
        if 0 <= img[x][y][0] < 100:
            img[x][y][2] = img[x][y][0]
        elif 100 <= img[x][y][0] < 200:
            img[x][y][2] = img[x][y][1]
        else:
            img[x][y][2] = img[x][y][2]

↓

b, g, r = cv2.split(img)  # チャンネルごとに分割

img[(0 <= b) & (b < 100), 2] = b[(0 <= b) & (b < 100)]
img[(100 <= b) & (b < 200), 2] = g[(100 <= b) & (b < 200)]

ループを消すことで300倍程度高速化できました。(ループバージョンと値が一致することは numpy.array_equal() で確認)

58.8 ms ± 103 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
181 µs ± 140 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

今回のようなケースでは以下を検討するといいと思います。

1. numpy の機能を使えないか検討する
2. numpy で書くことができなくて、ループが発生してしまうのであれば numba で高速化する

NumPy配列のブールインデックス参照 | hydroculのメモ
Pythonで高速化処理！numbaとCythonの実行速度を比較してみた。 - Np-Urのデータ分析教室

Comments

[sin_250]

私の環境でも期待通りの書き方、動作を確認しました。
まさに求めていたご回答です。ありがとうございます！