-
Notifications
You must be signed in to change notification settings - Fork 1.2k
/
screenshotUtil.py
217 lines (170 loc) · 7.96 KB
/
screenshotUtil.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
# 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.