218 lines
8.0 KiB
Python
218 lines
8.0 KiB
Python
|
# Screenshot-related features of PyAutoGUI
|
||
|
|
|
||
|
|
"""
|
||
|
|
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
|
||
|
|
"""
|
||
|
|
|
||
|
|
import datetime
|
||
|
|
import os
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
from PIL import Image
|
||
|
|
from PIL import ImageOps
|
||
|
|
|
||
|
|
RUNNING_PYTHON_2 = sys.version_info[0] == 2
|
||
|
|
|
||
|
|
scrotExists = False
|
||
|
|
maimExists = 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:
|
||
|
|
# if there is no "which" program to find scrot, then assume there is no scrot.
|
||
|
|
pass
|
||
|
|
|
||
|
|
try:
|
||
|
|
if sys.platform not in ('java', 'darwin', 'win32'):
|
||
|
|
whichProc = subprocess.Popen(['which', 'maim'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||
|
|
maimExists = whichProc.wait() == 0
|
||
|
|
except:
|
||
|
|
# if there is no "which" program to find maim, then assume there is no maim.
|
||
|
|
pass
|
||
|
|
|
||
|
|
def locateAll(needleImage, haystackImage, grayscale=False, limit=None):
|
||
|
|
needleFileObj = None
|
||
|
|
haystackFileObj = None
|
||
|
|
if isinstance(needleImage, str):
|
||
|
|
# 'image' is a filename, load the Image object
|
||
|
|
needleFileObj = open(needleImage, 'rb')
|
||
|
|
needleImage = Image.open(needleFileObj)
|
||
|
|
if isinstance(haystackImage, str):
|
||
|
|
# 'image' is a filename, load the Image object
|
||
|
|
haystackFileObj = open(haystackImage, 'rb')
|
||
|
|
haystackImage = Image.open(haystackFileObj)
|
||
|
|
|
||
|
|
|
||
|
|
if grayscale:
|
||
|
|
needleImage = ImageOps.grayscale(needleImage)
|
||
|
|
haystackImage = ImageOps.grayscale(haystackImage)
|
||
|
|
|
||
|
|
needleWidth, needleHeight = needleImage.size
|
||
|
|
haystackWidth, haystackHeight = haystackImage.size
|
||
|
|
|
||
|
|
needleImageData = tuple(needleImage.getdata()) # TODO - rename to needleImageData??
|
||
|
|
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
|
||
|
|
assert [len(row) for row in needleImageRows] == [needleWidth] * needleHeight
|
||
|
|
|
||
|
|
numMatchesFound = 0
|
||
|
|
|
||
|
|
for y in range(haystackHeight):
|
||
|
|
for matchx in _kmp(needleImageFirstRow, haystackImageData[y * haystackWidth:(y+1) * haystackWidth]):
|
||
|
|
foundMatch = True
|
||
|
|
for searchy in range(1, needleHeight):
|
||
|
|
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 (matchx, y, 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()
|
||
|
|
|
||
|
|
|
||
|
|
# 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()
|
||
|
|
|
||
|
|
|
||
|
|
def locate(needleImage, haystackImage, grayscale=False):
|
||
|
|
# 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.
|
||
|
|
points = tuple(locateAll(needleImage, haystackImage, grayscale, 1))
|
||
|
|
if len(points) > 0:
|
||
|
|
return points[0]
|
||
|
|
else:
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def locateOnScreen(image, grayscale=False,region=None):
|
||
|
|
screenshotIm = screenshot(region=region)
|
||
|
|
retVal = locate(image, screenshotIm, grayscale)
|
||
|
|
if 'fp' in dir(screenshotIm) and screenshotIm.fp is not None:
|
||
|
|
screenshotIm.fp.close() # Screenshots on Windows won't have an fp since they came from ImageGrab, not a file.
|
||
|
|
return retVal
|
||
|
|
|
||
|
|
|
||
|
|
def locateAllOnScreen(image, grayscale=False, limit=None, region=None):
|
||
|
|
screenshotIm = screenshot(region=region)
|
||
|
|
retVal = locateAll(image, screenshotIm, grayscale, limit)
|
||
|
|
if 'fp' in dir(screenshotIm) and screenshotIm.fp is not None:
|
||
|
|
screenshotIm.fp.close() # Screenshots on Windows won't have an fp since they came from ImageGrab, not a file.
|
||
|
|
return retVal
|
||
|
|
|
||
|
|
|
||
|
|
def locateCenterOnScreen(image, grayscale=False, region=None):
|
||
|
|
return center(locateOnScreen(image, grayscale, region))
|
||
|
|
|
||
|
|
|
||
|
|
def _screenshot_win32(imageFilename=None):
|
||
|
|
im = ImageGrab.grab()
|
||
|
|
if imageFilename is not None:
|
||
|
|
im.save(imageFilename)
|
||
|
|
return im
|
||
|
|
|
||
|
|
|
||
|
|
def _screenshot_osx(imageFilename=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 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:
|
||
|
|
if not region:
|
||
|
|
subprocess.call(['scrot', tmpFilename])
|
||
|
|
else:
|
||
|
|
if not maimExists:
|
||
|
|
raise NotImplementedError('"maim" must be installed to use screenshot functions with region in Linux. Run: sudo apt-get install maim')
|
||
|
|
left,top,width,height = [str(x) for x in region]
|
||
|
|
subprocess.call(['maim','-x',left,'-y',top,'-w',width,'-h',height, tmpFilename])
|
||
|
|
im = Image.open(tmpFilename)
|
||
|
|
if imageFilename is None:
|
||
|
|
os.unlink(tmpFilename)
|
||
|
|
return im
|
||
|
|
|
||
|
|
else:
|
||
|
|
raise Exception('The scrot program must be installed to take a screenshot with PyAutoGUI on Linux. Run: sudo apt-get install scrot')
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
def _kmp(needle, haystack): # 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 center(coords):
|
||
|
|
return (coords[0] + int(coords[2] / 2), coords[1] + int(coords[3] / 2))
|
||
|
|
|
||
|
|
|
||
|
|
def pixelMatchesColor(x, y, expectedRGBColor, tolerance=0):
|
||
|
|
r, g, b = screenshot().getpixel((x, y))
|
||
|
|
exR, exG, exB = expectedRGBColor
|
||
|
|
|
||
|
|
return (abs(r - exR) <= tolerance) and (abs(g - exG) <= tolerance) and (abs(b - exB) <= tolerance)
|
||
|
|
|
||
|
|
|
||
|
|
def pixel(x, y):
|
||
|
|
return screenshot().getpixel((x, y))
|
||
|
|
|
||
|
|
|
||
|
|
# set the screenshot() function based on the platform running this module
|
||
|
|
if sys.platform.startswith('java'):
|
||
|
|
raise NotImplementedError('Jython is not yet supported by PyAutoGUI.')
|
||
|
|
elif sys.platform == 'darwin':
|
||
|
|
screenshot = _screenshot_osx
|
||
|
|
elif sys.platform == 'win32':
|
||
|
|
screenshot = _screenshot_win32
|
||
|
|
from PIL import ImageGrab
|
||
|
|
else:
|
||
|
|
screenshot = _screenshot_linux
|
||
|
|
|
||
|
|
|
||
|
|
grab = screenshot # for compatibility with Pillow/PIL's ImageGrab module.
|