Jeff Boker
Jeff Boker

Reputation: 909

TypeError with Numba when averaging rows in numpy array

I want to take the average of all the rows in a frame and print out the list as shown below.

import numpy as np
import cv2, time
from numba import jit

cap = cv2.VideoCapture(0)

@jit(nopython=True)
def test(frame):
    print(frame.mean(axis=1))     # to take the mean of each row
        
fc = 0
frame_rate = 30
prev = 0
while(cap.isOpened()):
    ret, frame = cap.read()
    time_elapsed = time.time() - prev
    if time_elapsed > 1/frame_rate:
        prev = time.time()
        if ret == True:
            
            cv2.imshow("Video", frame)
            print(frame.shape)   #(480, 640, 3)
            
            test(frame)
            
            if cv2.waitKey(27) & 0xFF == ord('q'):
                break

        # if no frame found
        else:
            break

cap.release()
cv2.destroyAllWindows()

But I keep getting the following error. Numba should work perfectly fine with numerical computations like averaging a numpy array, so I am a little confused about the issue with my function.

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
- Resolution failure for literal arguments:
AssertionError()
- Resolution failure for non-literal arguments:
AssertionError()

During: resolving callee type: BoundFunction(array.mean for array(uint8, 3d, C))
During: typing of call at <ipython-input-2-3d43e63d267e> (11)


File "<ipython-input-2-3d43e63d267e>", line 11:
def test(frame):
    print(frame.mean(axis=1))     # to take the mean of each row

Upvotes: 1

Views: 2294

Answers (1)

max9111
max9111

Reputation: 6482

As already said in the comments np.mean is only supported without any keywords. supported functions. You can quite easily implement it yourself, but if you have some more information about the real problem it is straight forward to optimize it even further.

General solution

import numpy as np
import numba as nb
@nb.njit(fastmath=True)
def mean_gen(frame):
    res=np.zeros((frame.shape[0],3),dtype=np.float64)
    for i in range(frame.shape[0]):
        for j in range(frame.shape[1]):
            for k in range(frame.shape[2]):
                res[i,k]+=frame[i,j,k]
    return res/frame.shape[1]

Solution if the last dimension is always of size 3

In this case the inner loop can be unrolled to get better performance. The assertion `assert frame.shape[2]==3` is not only safety (there is no bounds checking), but also informs the compiler about the exact memory layout, which is often necessary for SIMD-vectorization. Some basic experience in C helps a lot in writing efficient code in Numba.

@nb.njit(fastmath=True,parallel=False)
def mean_3(frame):
    assert frame.shape[2]==3
    res=np.empty((frame.shape[0],3),dtype=np.float64)
    for i in nb.prange(frame.shape[0]):
        acc1=0.
        acc2=0.
        acc3=0.
        for j in range(frame.shape[1]):
            acc1+=frame[i,j,0]
            acc2+=frame[i,j,1]
            acc3+=frame[i,j,2]
        res[i,0]=acc1/frame.shape[1]
        res[i,1]=acc2/frame.shape[1]
        res[i,2]=acc3/frame.shape[1]
    return res

Timings

#generate some date
frame=np.random.rand(1280,720,3)*255
frame=frame.astype(np.uint8)

#Numpy
%timeit np.mean(frame,axis=1)
#17.6 ms ± 557 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit mean_gen(frame)
#2.27 ms ± 22.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit mean_3(frame)
#606 µs ± 11.1 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Upvotes: 1

Related Questions