I wrote an adaptive color thresholding function in Python (because OpenCV's cv2.adaptiveThreshold didn't fit my needs) and it is way too slow. I've made it as efficient as I can, but it still takes almost 500 ms on a 1280x720 image. I would greatly appreciate any suggestions that will make this function more efficient!
Here's what the function does:
It uses a cross-shape of one-pixel thickness as the structuring element. For each pixel in the image, it computes the average values of ksize
adjacent pixels in four directions independently (i.e. the average of ksize
pixels in the same row to the left, in the same column above, in the same row to the right, and in the same column below). I end with four average values, one for each direction. A pixel only meets the threshold criterion if it is brighter than either both the left AND right averages or both the top AND bottom averages (plus some constant C
).
I compute those averages incrementally for all pixels at the same time using numpy.roll()
, but I still need to do this ksize
times. The ksize
will usually be 20-50.
Here is the code, the relevant part is really just what happens inside the for
-loop:
def bilateral_adaptive_threshold(img, ksize=20, C=0, mode='floor', true_value=255, false_value=0):
mask = np.full(img.shape, false_value, dtype=np.int16)
left_thresh = np.zeros_like(img, dtype=np.float32) #Store the right-side average of each pixel here
right_thresh = np.zeros_like(img, dtype=np.float32) #Store the left-side average of each pixel here
up_thresh = np.zeros_like(img, dtype=np.float32) #Store the top-side average of each pixel here
down_thresh = np.zeros_like(img, dtype=np.float32) #Store the bottom-side average of each pixel here
for i in range(1, ksize+1):
roll_left = np.roll(img, -i, axis=1)
roll_right = np.roll(img, i, axis=1)
roll_up = np.roll(img, -i, axis=0)
roll_down = np.roll(img, i, axis=0)
roll_left[:,-i:] = 0
roll_right[:,:i] = 0
roll_up[-i:,:] = 0
roll_down[:i,:] = 0
left_thresh += roll_right
right_thresh += roll_left
up_thresh += roll_down
down_thresh += roll_up
left_thresh /= ksize
right_thresh /= ksize
up_thresh /= ksize
down_thresh /= ksize
if mode == 'floor':
mask[((img > left_thresh+C) & (img > right_thresh+C)) | ((img > up_thresh+C) & (img > down_thresh+C))] = true_value
elif mode == 'ceil':
mask[((img < left_thresh-C) & (img < right_thresh-C)) | ((img < up_thresh-C) & (img < down_thresh-C))] = true_value
else: raise ValueError("Unexpected mode value. Expected value is 'floor' or 'ceil'.")
return mask