import numpy as np
def find_image(im, tpl):
im = np.atleast_3d(im)
tpl = np.atleast_3d(tpl)
H, W, D = im.shape[:3]
h, w = tpl.shape[:2]
# Integral image and template sum per channel
sat = im.cumsum(1).cumsum(0)
tplsum = np.array([tpl[:, :, i].sum() for i in range(D)])
# Calculate lookup table for all the possible windows
iA, iB, iC, iD = sat[:-h, :-w], sat[:-h, w:], sat[h:, :-w], sat[h:, w:]
lookup = iD - iB - iC + iA
# Possible matches
possible_match = np.where(np.logical_and.reduce([lookup[..., i] == tplsum[i] for i in range(D)]))
# Find exact match
for y, x in zip(*possible_match):
if np.all(im[y+1:y+h+1, x+1:x+w+1] == tpl):
return (y+1, x+1)
return None
def test_find_image():
while True:
# Generate random 20x20 haystack image
haystack = np.random.randint(0, 255, (20, 20, 3), "uint8")
# Choose a random 5x5 chunk of the haystack image to be used as the needle
y = np.random.randint(0, 20 - 5 + 1)
x = np.random.randint(0, 20 - 5 + 1)
needle = haystack[y:y + 5, x:x + 5]
result = find_image(haystack, needle)
if result == (y, x):
print("SUCCESS:", y, x)
else:
print("FAIL:", y, x)
if __name__ == "__main__":
test_find_image()