OpenCV version of locateAll

pyscreeze/__init__.py

@@ -20,6 +20,14 @@
 from PIL import Image
 from PIL import ImageOps
 
+try:
+    import cv2, numpy
+    useOpenCV = True
+    CONFIDENCE = 0.999  # for thresholding matches in _locateAll_opencv
+except ImportError:
+    useOpenCV = False
+
+
 RUNNING_PYTHON_2 = sys.version_info[0] == 2
 
 RAISE_IF_NOT_FOUND = False
@@ -40,7 +48,81 @@
     else:
         raise
 
-def locateAll(needleImage, haystackImage, grayscale=None, limit=None, region=None, step=1):
+
+class ImageNotFoundException(Exception):
+    pass
+
+
+def _load_cv2(img, grayscale=None):
+    # load images if given filename, or convert as needed to opencv
+    # Alpha layer just causes failures at this point, so flatten to RGB.
+    # RGBA: load with -1 * cv2.CV_LOAD_IMAGE_COLOR to preserve alpha
+    # to matchTemplate, need template and image to be the same wrt having alpha
+
+    if grayscale is None:
+        grayscale = GRAYSCALE_DEFAULT
+    if isinstance(img, str):
+        if grayscale:
+            img_cv = cv2.imread(img, cv2.CV_LOAD_IMAGE_GRAYSCALE)
+        else:
+            img_cv = cv2.imread(img, cv2.CV_LOAD_IMAGE_COLOR)  # -1 gets RGBA
+    elif isinstance(img, numpy.ndarray):
+        # don't try to convert an already-gray image
+        if grayscale and len(img.shape) == 3:  # and img.shape[2] == 3:
+            img_cv = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    elif hasattr(img, 'convert'):
+        # assume its a PIL.Image, convert to cv format
+        img_array = numpy.array(img.convert('RGB'))
+        img_cv = img_array[:, :, ::-1].copy()  # -1 does RGB -> BGR
+        if grayscale:
+            img_cv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
+    else:
+        raise TypeError('expected an image filename, OpenCV numpy array, or PIL image')
+    return img_cv
+
+
+def _locateAll_opencv(needleImage, haystackImage, grayscale=None, limit=10000, region=None, step=1):
+    """ faster but more memory-intensive than pure python
+        step=2 skips every other row and column, ~3x faster but less accurate
+        limitations:
+          - OpenCV 3.x & python 3.x not tested
+          - RGBA images are not handled / untested
+          - step=2 not fully tested
+    """
+    if grayscale is None:
+        grayscale = GRAYSCALE_DEFAULT
+
+    needleImage = _load_cv2(needleImage, grayscale)
+    needleHeight, needleWidth = needleImage.shape[:2]
+    haystackImage = _load_cv2(haystackImage, grayscale)
+
+    if region:
+        haystackImage = haystackImage[region[1]:region[1]+region[3],
+                                      region[0]:region[0]+region[2]]
+    else:
+        region = (0, 0)  # full image; these values used in the yield statement
+    if (haystackImage.shape[0] < needleImage.shape[0] or
+        haystackImage.shape[1] < needleImage.shape[1]):
+        # avoid semi-cryptic OpenCV error below if bad size
+        raise ValueError('needle dimensions exceed the haystack image (or region)')
+    if step == 2:
+        needleImage = needleImage[::step, ::step]
+        haystackImage = haystackImage[::step, ::step]
+
+    # get all matches at once, credit: https://stackoverflow.com/questions/7670112/finding-a-subimage-inside-a-numpy-image/9253805#9253805
+    result = cv2.matchTemplate(haystackImage, needleImage, cv2.TM_CCOEFF_NORMED)
+    match_indices = numpy.arange(result.size)[(result > CONFIDENCE).flatten()]
+    matches = numpy.unravel_index(match_indices[:limit], result.shape)
+
+    if len(matches[0]) == 0 and RAISE_IF_NOT_FOUND:
+        raise ImageNotFoundException('Could not locate the image.')
+
+    # use a generator for API consistency (everything is computed already):
+    for y, matchx in zip(matches[0], matches[1]):
+        yield (step * matchx + region[0], step * y + region[1], needleWidth, needleHeight)
+
+
+def _locateAll_python(needleImage, haystackImage, grayscale=None, limit=None, region=None, step=1):
     if grayscale is None:
         grayscale = GRAYSCALE_DEFAULT
     needleFileObj = None
@@ -303,6 +385,10 @@ def pixel(x, y):
 
 grab = screenshot # for compatibility with Pillow/PIL's ImageGrab module.
 
-
-class ImageNotFoundException(Exception):
-    pass
+# set the locateAll function to use opencv if possible; python 3 needs opencv 3.0+
+if useOpenCV:
+    locateAll = _locateAll_opencv
+    if not RUNNING_PYTHON_2 and cv2.__version__ < '3':
+        locateAll = _locateAll_python
+else:
+    locateAll = _locateAll_python