Рассматривали ли вы использовать Numpy? Часто это хорошее промежуточное звено между Python списками и Numba, что-то вроде:
h1 = (cropped.sum(axis=-1) == 765).sum()
или
h1 = (cropped == 255).all(axis=-1).sum()
Пример кода, который вы предоставляете, не является действительным Numba. Ваша подпись также неверна, поскольку входные данные являются трехмерным массивом, а выходные - целым числом, вероятно, оно должно быть следующим:
@njit(int64(uint8[:,:,:]))
Циклирование массива, как и вы, не является допустимым кодом. Точный перевод вашего кода будет выглядеть примерно так:
@njit(int64(uint8[:,:,:]))
def findWhite(cropped):
h1 = int64(0)
ys, xs, n_bands = cropped.shape
for i in range(ys):
for j in range(xs):
if cropped[i, j, :].sum() == 765:
h1 += 1
return h1
Но это не очень быстро и не побеждает Numpy на моей машине. С Numba можно явно l oop для каждого элемента в массиве, это уже намного быстрее:
@njit(int64(uint8[:,:,:]))
def findWhite_numba(cropped):
h1 = int64(0)
ys, xs, zs = cropped.shape
for i in range(ys):
for j in range(xs):
incr = 1
for k in range(zs):
if cropped[i, j, k] != 255:
incr = 0
break
h1 += incr
return h1
Для массива 5000x5000x3 это результат для меня:
Numpy (h1 = (cropped == 255).all(axis=-1).sum()
):
427 ms ± 6.37 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
findWhite:
612 ms ± 6.16 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
findWhite_numba:
31 ms ± 1.51 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
Преимущество метода Numpy заключается в том, что оно обобщает на любое количество измерений.