python:画素の置き換えの計算速度向上

Question

Python初心者です。
画像で黒以外の似たような背景色を全部黒にするプログラムを作っています。
典型的な色として黒以外の画素の最頻値を使っています。
いろいろと苦労してプログラムは動くようになったのですが、
計算速度が遅いです。早くする方法があればご教示ください。
なお、上の画像は閾値20000の場合です。
なお、Kirisakiさんの結果はこちらでした。

該当のソースコード

import cv2, math
import scipy.stats as sstats

bgr=cv2.imread('A.tif') #BGR
imax=bgr.shape[0]
jmax=bgr.shape[1]

rgb = bgr[:, :, [2, 1, 0]]
r,g,b=cv2.split(rgb)
mode_r=sstats.mode(r[r.nonzero()])[0][0]
mode_g=sstats.mode(g[g.nonzero()])[0][0]
mode_b=sstats.mode(b[b.nonzero()])[0][0]
for i in range(imax):
for j in range(jmax):
mr=((int(mode_r)-r[i,j])**2)

     mg=((int(mode_g)-g[i,j])**2)

     mb=((int(mode_b)-b[i,j])**2)

     if(mr+mg+mb < 20000.):
         bgr[i,j]=[0,0,0]
         cv2.imwrite("B.jpg",bgr)

試したこと

for文とif文を使わずに
bgr[mr2　+　mg2　+　mb**2　<　100.]=[0,0,0]
を試したら早く計算が終わったのですが、意図した結果になりませんでした。

Comments

[A_kirisaki]

r と b と g はどこから来ましたか？

[grandchild]

r,g,b=cv2.split(rgb)
申し訳ございません。上の通りです。入力が抜けておりました・・。

Answer 1

置き換えの処理の部分を高速化したいとのことなので、その部分だけ対象として、次のような実装が考えられます。
（cv2.imread(file_path)) の部分は除いています。）

①for-loop と if 文を除く：

import cv2, math
import scipy.stats as sstats


def func1(bgr):
    """for-loop と if 文を除く
    """
    b, g, r = cv2.split(bgr.astype(int))
    mode_r = sstats.mode(r[r.nonzero()])[0][0]
    mode_g = sstats.mode(g[g.nonzero()])[0][0]
    mode_b = sstats.mode(b[b.nonzero()])[0][0]
    bgr[(mode_r  - r)**2 + (mode_g - g)**2 + (mode_b - b)**2 < 100., :] = [0, 0, 0]
    return bgr

②numba.jit を使用する：

import cv2
import scipy.stats as sstats
import numpy
import numba


@numba.jit("void(u1[:,:,:], i4[:,:], i4[:,:], i4[:,:], i4, i4, i4)")
def func2_numba(bgr, b, g, r, mode_b, mode_g, mode_r):
    """numba.jit を適用する処理
    """
    for i in range(b.shape[0]):
        for j in range(b.shape[1]):
            if (mode_r - r[i, j])**2 + (mode_g - g[i, j])**2 + (mode_b - b[i, j])**2 < 100:
                bgr[i, j] = [0, 0, 0]


def func2(bgr):
    """numba.jit を使用する
    """
    b, g, r = cv2.split(bgr.astype(numpy.int32))
    mode_b = sstats.mode(b[b.nonzero()])[0][0]
    mode_g = sstats.mode(g[g.nonzero()])[0][0]
    mode_r = sstats.mode(r[r.nonzero()])[0][0]
    func2_numba(bgr, b, g, r, mode_b, mode_g, mode_r)
    return bgr

オリジナルの実装を func0 としておきます（実装は次の通り）：

import cv2
import scipy.stats as sstats
import numpy


def func0(bgr):
    """オリジナルの実装
    """
    imax = bgr.shape[0]
    jmax = bgr.shape[1]

    rgb = bgr[:, :, [2, 1, 0]]
    r, g, b = cv2.split(rgb)
    mode_r = sstats.mode(r[r.nonzero()])[0][0]
    mode_g = sstats.mode(g[g.nonzero()])[0][0]
    mode_b = sstats.mode(b[b.nonzero()])[0][0]
    for i in range(imax):
        for j in range(jmax):
            mr = (int(mode_r) - r[i, j])**2
            mg = (int(mode_g) - g[i, j])**2
            mb = (int(mode_b) - b[i, j])**2
            if mr + mg + mb < 100.:
                bgr[i, j]=[0, 0, 0]
    return bgr

手元の PC (Core i7-10510U, 8 GB RAM) で　Jupyter Notebook の %timeit を利用したベンチマークは次の通りです。
処理速度は ② > ① >> original となりました。
※関数のコンパイル時間は除いています。

%timeit func0(bgr) # 1.84 s ± 46 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit func1(bgr) # 34.6 ms ± 1.16 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit func2(bgr) # 30.8 ms ± 158 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

以下、処理速度向上に関するコメントです。
Python のコードの高速化方法はいろいろあります。例えば

1. 対象とする処理コードを関数にする。
2. 上記関数を含むモジュールを実装する。
3. cython を使って処理を定義する。
4. numba.jit を使って関数をジャスト・イン・コンパイルする。
5. 処理部分を他の高速な言語（C/C++, Rust, etc.）で定義し、pybind11 などでPython でも使用できる形に実装する。

などがあります。それぞれ一長一短があります。これらは検索すればすぐに出てくるので、ぜひ調べてみてください。

Comments

[grandchild]

ご回答ありがとうございます。とても勉強になりました。
速度もとても早くなりましたし、結果も同じでしたので、ベストアンサーとさせていただきました。

Kirisakiさんのご回答もとても勉強になりました。ありがとうございました。

Answer 2

こんな感じでどうでしょうか！for ループをなくし行列演算だけにしました。もし結果が異なるようであればサンプル画像を上げてもらえれば幸いです。

import numpy as np
import cv2
import scipy.stats as sstats
import matplotlib.pyplot as plt

bgr = cv2.imread('image/Lenna.png')
rgb = bgr[:, :, [2, 1, 0]]

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

mode_b = sstats.mode(b[b.nonzero()])[0][0]
mode_g = sstats.mode(g[g.nonzero()])[0][0]
mode_r = sstats.mode(r[r.nonzero()])[0][0]

mode = np.array([mode_b, mode_g, mode_r], dtype=int)

a = (mode - bgr)**2
b = a[:, :, 0] + a[:, :, 1] +  a[:, :, 2]
c = (b >= 100).astype(int)
cs = np.stack([c] * 3, axis=2)
result = bgr * cs

plt.imshow(result)
plt.show()

Comments

[grandchild]

ご回答ありがとうございます。とても早くなりました！ただ、青いバラになってしまいました・・・それはそれで、とても綺麗でしたが・・・。

[Surpris]

`matplotlib.pyplot.imshow(src, **kwargs)` は入力データ `src` が３次元の場合、`(height, width, channels)` の shape を持つ RGB 画像での入力が想定されています。
従って BGR 画像を `pyplot.imshow` で正しく表示したい場合は、RGB 画像に変換して入力する必要があります：
```Python
rgb = bgr[:, :, [2, 1, 0]]
plt.imshow(rgb)
```

[grandchild]

青いバラになったのはKirisakiさんの方でした。ありがとうございました。

[A_kirisaki]

まあ保存する時 RGB にするっぽいしええんちゃうん？とか雑に思ってました。

[grandchild]

そうですね。私の知識では、今のところ、どちらの方が内容的にベストアンサーなのか分からないのもあって、申し訳ございません。でも、ありがとうございました。

[Surpris]

（A_kirisaki 様のコメントを受けて）
言葉が足らず誤解を与えたようですみません。
一つ前の私のコメントは、「`pyplot.imshow` で青色に表示されてしまったのはなぜか」という疑問に対する回答にすぎません。
A_kirisaki 様のご回答にある置き換え処理はご質問者様の要望（処理速度の向上）に合致していると私も思います。

[grandchild]

Surprisさんのご回答は自分が試したことの間違いの理由も気づかせてくれましたし、今後処理速度向上について勉強する上での参考になることも書かれており、結果も同じだったので選ばせていただきました。Kirisakiさんの方法は目からうろこのエレガントなご回答でした。お二方に感謝申し上げます。失礼いたします。

[grandchild]

あ、すみません。Kirisakiさんの結果は赤くしても微妙に結果が違いました。何故だかは分かりませんが・・。

[grandchild]

あ、重ね重ねすみません。全く同じになることが確認されました。

[grandchild]

↑間違ってpngで出力してました。