# PyScreeze
# by Al Sweigart
# https://github.com/asweigart/pyscreeze
# BSD license
"""
So , apparently Pillow support on Ubuntu 64 - bit has several additional steps since it doesn ' t have JPEG/PNG support out of the box. Description here:
https : / / stackoverflow . com / questions / 7648200 / pip - install - pil - e - tickets - 1 - no - jpeg - png - support
http : / / ubuntuforums . org / showthread . php ? t = 1751455
"""
__version__ = ' 0.1.21 '
import collections
import datetime
import os
import subprocess
import sys
import time
import errno
try :
from PIL import Image
from PIL import ImageOps
except ImportError :
pass
from contextlib import contextmanager
try :
import cv2 , numpy
useOpenCV = True
RUNNING_CV_2 = cv2 . __version__ [ 0 ] < ' 3 '
except ImportError :
useOpenCV = False
RUNNING_PYTHON_2 = sys . version_info [ 0 ] == 2
if useOpenCV :
if RUNNING_CV_2 :
LOAD_COLOR = cv2 . CV_LOAD_IMAGE_COLOR
LOAD_GRAYSCALE = cv2 . CV_LOAD_IMAGE_GRAYSCALE
else :
LOAD_COLOR = cv2 . IMREAD_COLOR
LOAD_GRAYSCALE = cv2 . IMREAD_GRAYSCALE
GRAYSCALE_DEFAULT = False
# For version 0.1.19 I changed it so that ImageNotFoundException was raised
# instead of returning None. In hindsight, this change came too late, so I'm
# changing it back to returning None. But I'm also including this option for
# folks who would rather have it raise an exception.
USE_IMAGE_NOT_FOUND_EXCEPTION = False
scrotExists = False
try :
if sys . platform not in ( ' java ' , ' darwin ' , ' win32 ' ) :
whichProc = subprocess . Popen (
[ ' which ' , ' scrot ' ] , stdout = subprocess . PIPE , stderr = subprocess . PIPE )
scrotExists = whichProc . wait ( ) == 0
except OSError as ex :
if ex . errno == errno . ENOENT :
# if there is no "which" program to find scrot, then assume there
# is no scrot.
pass
else :
raise
if sys . platform == ' win32 ' :
from ctypes import windll
# win32 DC(DeviceContext) Manager
@contextmanager
def __win32_openDC ( hWnd ) :
hDC = windll . user32 . GetDC ( hWnd )
if hDC == 0 : #NULL
raise WindowsError ( " windll.user32.GetDC failed : return NULL " )
try :
yield hDC
finally :
if windll . user32 . ReleaseDC ( hWnd , hDC ) == 0 :
raise WindowsError ( " windll.user32.ReleaseDC failed : return 0 " )
Box = collections . namedtuple ( ' Box ' , ' left top width height ' )
Point = collections . namedtuple ( ' Point ' , ' x y ' )
RGB = collections . namedtuple ( ' RGB ' , ' red green blue ' )
class ImageNotFoundException ( Exception ) :
pass # This is an exception class raised when the locate functions fail.
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 ) :
# The function imread loads an image from the specified file and
# returns it. If the image cannot be read (because of missing
# file, improper permissions, unsupported or invalid format),
# the function returns an empty matrix
# http://docs.opencv.org/3.0-beta/modules/imgcodecs/doc/reading_and_writing_images.html
if grayscale :
img_cv = cv2 . imread ( img , LOAD_GRAYSCALE )
else :
img_cv = cv2 . imread ( img , LOAD_COLOR )
if img_cv is None :
raise IOError ( " Failed to read %s because file is missing, "
" has improper permissions, or is an "
" unsupported or invalid format " % img )
elif isinstance ( img , numpy . ndarray ) :
# don't try to convert an already-gray image to gray
if grayscale and len ( img . shape ) == 3 : # and img.shape[2] == 3:
img_cv = cv2 . cvtColor ( img , cv2 . COLOR_BGR2GRAY )
else :
img_cv = img
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 ,
confidence = 0.999 ) :
""" faster but more memory-intensive than pure python
step 2 skips every other row and column = ~ 3 x faster but prone to miss ;
to compensate , the algorithm automatically reduces the confidence
threshold by 5 % ( which helps but will not avoid all misses ) .
limitations :
- OpenCV 3. x & python 3. x not tested
- RGBA images are treated as RBG ( ignores alpha channel )
"""
if grayscale is None :
grayscale = GRAYSCALE_DEFAULT
confidence = float ( confidence )
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 dimension(s) exceed the haystack image or region dimensions ' )
if step == 2 :
confidence * = 0.95
needleImage = needleImage [ : : step , : : step ]
haystackImage = haystackImage [ : : step , : : step ]
else :
step = 1
# 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 :
if USE_IMAGE_NOT_FOUND_EXCEPTION :
raise ImageNotFoundException ( ' Could not locate the image (highest confidence = %.3f ) ' % result . max ( ) )
else :
return
# use a generator for API consistency:
matchx = matches [ 1 ] * step + region [ 0 ] # vectorized
matchy = matches [ 0 ] * step + region [ 1 ]
for x , y in zip ( matchx , matchy ) :
yield Box ( x , y , needleWidth , needleHeight )
def _locateAll_python ( needleImage , haystackImage , grayscale = None , limit = None , region = None , step = 1 ) :
# setup all the arguments
if grayscale is None :
grayscale = GRAYSCALE_DEFAULT
needleFileObj = None
if isinstance ( needleImage , str ) :
# 'image' is a filename, load the Image object
needleFileObj = open ( needleImage , ' rb ' )
needleImage = Image . open ( needleFileObj )
haystackFileObj = None
if isinstance ( haystackImage , str ) :
# 'image' is a filename, load the Image object
haystackFileObj = open ( haystackImage , ' rb ' )
haystackImage = Image . open ( haystackFileObj )
if region is not None :
haystackImage = haystackImage . crop ( ( region [ 0 ] , region [ 1 ] , region [ 0 ] + region [ 2 ] , region [ 1 ] + region [ 3 ] ) )
else :
region = ( 0 , 0 ) # set to 0 because the code always accounts for a region
if grayscale : # if grayscale mode is on, convert the needle and haystack images to grayscale
needleImage = ImageOps . grayscale ( needleImage )
haystackImage = ImageOps . grayscale ( haystackImage )
else :
# if not using grayscale, make sure we are comparing RGB images, not RGBA images.
if needleImage . mode == ' RGBA ' :
needleImage = needleImage . convert ( ' RGB ' )
if haystackImage . mode == ' RGBA ' :
haystackImage = haystackImage . convert ( ' RGB ' )
# setup some constants we'll be using in this function
needleWidth , needleHeight = needleImage . size
haystackWidth , haystackHeight = haystackImage . size
needleImageData = tuple ( needleImage . getdata ( ) )
haystackImageData = tuple ( haystackImage . getdata ( ) )
needleImageRows = [ needleImageData [ y * needleWidth : ( y + 1 ) * needleWidth ] for y in range ( needleHeight ) ] # LEFT OFF - check this
needleImageFirstRow = needleImageRows [ 0 ]
assert len ( needleImageFirstRow ) == needleWidth , ' For some reason, the calculated width of first row of the needle image is not the same as the width of the image. '
assert [ len ( row ) for row in needleImageRows ] == [ needleWidth ] * needleHeight , ' For some reason, the needleImageRows aren \' t the same size as the original image. '
numMatchesFound = 0
# NOTE: After running tests/benchmarks.py on the following code, it seem that having a step
# value greater than 1 does not give *any* significant performance improvements.
# Since using a step higher than 1 makes for less accurate matches, it will be
# set to 1.
step = 1 # hard-code step as 1 until a way to improve it can be figured out.
if step == 1 :
firstFindFunc = _kmp
else :
firstFindFunc = _steppingFind
for y in range ( haystackHeight ) : # start at the leftmost column
for matchx in firstFindFunc ( needleImageFirstRow , haystackImageData [ y * haystackWidth : ( y + 1 ) * haystackWidth ] , step ) :
foundMatch = True
for searchy in range ( 1 , needleHeight , step ) :
haystackStart = ( searchy + y ) * haystackWidth + matchx
if needleImageData [ searchy * needleWidth : ( searchy + 1 ) * needleWidth ] != haystackImageData [ haystackStart : haystackStart + needleWidth ] :
foundMatch = False
break
if foundMatch :
# Match found, report the x, y, width, height of where the matching region is in haystack.
numMatchesFound + = 1
yield Box ( matchx + region [ 0 ] , y + region [ 1 ] , needleWidth , needleHeight )
if limit is not None and numMatchesFound > = limit :
# Limit has been reached. Close file handles.
if needleFileObj is not None :
needleFileObj . close ( )
if haystackFileObj is not None :
haystackFileObj . close ( )
return
# There was no limit or the limit wasn't reached, but close the file handles anyway.
if needleFileObj is not None :
needleFileObj . close ( )
if haystackFileObj is not None :
haystackFileObj . close ( )
if numMatchesFound == 0 :
if USE_IMAGE_NOT_FOUND_EXCEPTION :
raise ImageNotFoundException ( ' Could not locate the image. ' )
else :
return
def locate ( needleImage , haystackImage , * * kwargs ) :
# Note: The gymnastics in this function is because we want to make sure to exhaust the iterator so that the needle and haystack files are closed in locateAll.
kwargs [ ' limit ' ] = 1
points = tuple ( locateAll ( needleImage , haystackImage , * * kwargs ) )
if len ( points ) > 0 :
return points [ 0 ]
else :
if USE_IMAGE_NOT_FOUND_EXCEPTION :
raise ImageNotFoundException ( ' Could not locate the image. ' )
else :
return None
def locateOnScreen ( image , minSearchTime = 0 , * * kwargs ) :
""" minSearchTime - amount of time in seconds to repeat taking
screenshots and trying to locate a match . The default of 0 performs
a single search .
"""
start = time . time ( )
while True :
try :
screenshotIm = screenshot ( region = None ) # the locateAll() function must handle cropping to return accurate coordinates, so don't pass a region here.
retVal = locate ( image , screenshotIm , * * kwargs )
try :
screenshotIm . fp . close ( )
except AttributeError :
# Screenshots on Windows won't have an fp since they came from
# ImageGrab, not a file. Screenshots on Linux will have fp set
# to None since the file has been unlinked
pass
if retVal or time . time ( ) - start > minSearchTime :
return retVal
except ImageNotFoundException :
if time . time ( ) - start > minSearchTime :
if USE_IMAGE_NOT_FOUND_EXCEPTION :
raise
else :
return None
def locateAllOnScreen ( image , * * kwargs ) :
screenshotIm = screenshot ( region = None ) # the locateAll() function must handle cropping to return accurate coordinates, so don't pass a region here.
retVal = locateAll ( image , screenshotIm , * * kwargs )
try :
screenshotIm . fp . close ( )
except AttributeError :
# Screenshots on Windows won't have an fp since they came from
# ImageGrab, not a file. Screenshots on Linux will have fp set
# to None since the file has been unlinked
pass
return retVal
def locateCenterOnScreen ( image , * * kwargs ) :
coords = locateOnScreen ( image , * * kwargs )
return center ( coords )
def showRegionOnScreen ( region , outlineColor = ' red ' , filename = ' _showRegionOnScreen.png ' ) :
from PIL import ImageDraw # this is the only function that needs this, and it's rarely called
screenshotIm = screenshot ( )
draw = ImageDraw . Draw ( screenshotIm )
region = ( region [ 0 ] , region [ 1 ] , region [ 2 ] + region [ 0 ] , region [ 3 ] + region [ 1 ] ) # convert from (left, top, right, bottom) to (left, top, width, height)
draw . rectangle ( region , outline = outlineColor )
screenshotIm . save ( filename )
def _screenshot_win32 ( imageFilename = None , region = None ) :
try :
im = ImageGrab . grab ( )
except NameError :
raise ImportError ( ' Pillow module must be installed to use screenshot functions on Windows. ' )
if region is not None :
assert len ( region ) == 4 , ' region argument must be a tuple of four ints '
region = [ int ( x ) for x in region ]
im = im . crop ( ( region [ 0 ] , region [ 1 ] , region [ 2 ] + region [ 0 ] , region [ 3 ] + region [ 1 ] ) )
if imageFilename is not None :
im . save ( imageFilename )
return im
def _screenshot_osx ( imageFilename = None , region = None ) :
if imageFilename is None :
tmpFilename = ' screenshot %s .png ' % ( datetime . datetime . now ( ) . strftime ( ' % Y- % m %d _ % H- % M- % S- %f ' ) )
else :
tmpFilename = imageFilename
subprocess . call ( [ ' screencapture ' , ' -x ' , tmpFilename ] )
im = Image . open ( tmpFilename )
if region is not None :
assert len ( region ) == 4 , ' region argument must be a tuple of four ints '
region = [ int ( x ) for x in region ]
im = im . crop ( ( region [ 0 ] , region [ 1 ] , region [ 2 ] + region [ 0 ] , region [ 3 ] + region [ 1 ] ) )
os . unlink ( tmpFilename ) # delete image of entire screen to save cropped version
im . save ( tmpFilename )
else :
# force loading before unlinking, Image.open() is lazy
im . load ( )
if imageFilename is None :
os . unlink ( tmpFilename )
return im
def _screenshot_linux ( imageFilename = None , region = None ) :
if not scrotExists :
raise NotImplementedError ( ' " scrot " must be installed to use screenshot functions in Linux. Run: sudo apt-get install scrot ' )
if imageFilename is None :
tmpFilename = ' .screenshot %s .png ' % ( datetime . datetime . now ( ) . strftime ( ' % Y- % m %d _ % H- % M- % S- %f ' ) )
else :
tmpFilename = imageFilename
if scrotExists :
subprocess . call ( [ ' scrot ' , tmpFilename ] )
im = Image . open ( tmpFilename )
if region is not None :
assert len ( region ) == 4 , ' region argument must be a tuple of four ints '
region = [ int ( x ) for x in region ]
im = im . crop ( ( region [ 0 ] , region [ 1 ] , region [ 2 ] + region [ 0 ] , region [ 3 ] + region [ 1 ] ) )
os . unlink ( tmpFilename ) # delete image of entire screen to save cropped version
im . save ( tmpFilename )
else :
# force loading before unlinking, Image.open() is lazy
im . load ( )
if imageFilename is None :
os . unlink ( tmpFilename )
return im
else :
raise Exception ( ' The scrot program must be installed to take a screenshot with PyScreeze on Linux. Run: sudo apt-get install scrot ' )
def _kmp ( needle , haystack , _dummy ) : # Knuth-Morris-Pratt search algorithm implementation (to be used by screen capture)
# build table of shift amounts
shifts = [ 1 ] * ( len ( needle ) + 1 )
shift = 1
for pos in range ( len ( needle ) ) :
while shift < = pos and needle [ pos ] != needle [ pos - shift ] :
shift + = shifts [ pos - shift ]
shifts [ pos + 1 ] = shift
# do the actual search
startPos = 0
matchLen = 0
for c in haystack :
while matchLen == len ( needle ) or \
matchLen > = 0 and needle [ matchLen ] != c :
startPos + = shifts [ matchLen ]
matchLen - = shifts [ matchLen ]
matchLen + = 1
if matchLen == len ( needle ) :
yield startPos
def _steppingFind ( needle , haystack , step ) :
for startPos in range ( 0 , len ( haystack ) - len ( needle ) + 1 ) :
foundMatch = True
for pos in range ( 0 , len ( needle ) , step ) :
if haystack [ startPos + pos ] != needle [ pos ] :
foundMatch = False
break
if foundMatch :
yield startPos
def center ( coords ) :
return Point ( coords [ 0 ] + int ( coords [ 2 ] / 2 ) , coords [ 1 ] + int ( coords [ 3 ] / 2 ) )
def pixelMatchesColor ( x , y , expectedRGBColor , tolerance = 0 ) :
pix = pixel ( x , y )
if len ( pix ) == 3 or len ( expectedRGBColor ) == 3 : #RGB mode
r , g , b = pix [ : 3 ]
exR , exG , exB = expectedRGBColor [ : 3 ]
return ( abs ( r - exR ) < = tolerance ) and ( abs ( g - exG ) < = tolerance ) and ( abs ( b - exB ) < = tolerance )
elif len ( pix ) == 4 and len ( expectedRGBColor ) == 4 : #RGBA mode
r , g , b , a = pix
exR , exG , exB , exA = expectedRGBColor
return ( abs ( r - exR ) < = tolerance ) and ( abs ( g - exG ) < = tolerance ) and ( abs ( b - exB ) < = tolerance ) and ( abs ( a - exA ) < = tolerance )
else :
assert False , ' Color mode was expected to be length 3 (RGB) or 4 (RGBA), but pixel is length %s and expectedRGBColor is length %s ' % ( len ( pix ) , len ( expectedRGBColor ) )
def pixel ( x , y ) :
if sys . platform == ' win32 ' :
# On Windows, calling GetDC() and GetPixel() is twice as fast as using our screenshot() function.
with __win32_openDC ( 0 ) as hdc : # handle will be released automatically
color = windll . gdi32 . GetPixel ( hdc , x , y )
if color < 0 :
raise WindowsError ( " windll.gdi32.GetPixel faild : return {} " . format ( color ) )
# color is in the format 0xbbggrr https://msdn.microsoft.com/en-us/library/windows/desktop/dd183449(v=vs.85).aspx
bbggrr = " {:0>6x} " . format ( color ) # bbggrr => 'bbggrr' (hex)
b , g , r = ( int ( bbggrr [ i : i + 2 ] , 16 ) for i in range ( 0 , 6 , 2 ) )
return ( r , g , b )
else :
# Need to select only the first three values of the color in
# case the returned pixel has an alpha channel
return RGB ( * ( screenshot ( ) . getpixel ( ( x , y ) ) [ : 3 ] ) )
# set the screenshot() function based on the platform running this module
if sys . platform . startswith ( ' java ' ) :
raise NotImplementedError ( ' Jython is not yet supported by PyScreeze. ' )
elif sys . platform == ' darwin ' :
screenshot = _screenshot_osx
elif sys . platform == ' win32 ' :
screenshot = _screenshot_win32
try :
from PIL import ImageGrab
except ImportError :
pass
else :
screenshot = _screenshot_linux
grab = screenshot # for compatibility with Pillow/PIL's ImageGrab module.
# 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