1# Copyright 2016-2018 the openage authors. See copying.md for legal info.
2
3# If you wanna boost speed even further:
4# cython: profile=False
5
6
7""" Cython version of the visgrep utility """
8
9import argparse
10cimport cython
11import itertools
12import logging
13import numpy
14cimport numpy
15import sys
16
17from collections import namedtuple
18from libcpp.vector cimport vector
19from PIL import Image
20
21
22TOOL_DESCRIPTION = """Python translation of the visgrep v1.09
23visual grep, greps for images in another image
24Author of the original C version: Steve Slaven - http://hoopajoo.net"""
25
26EPILOG = """The image.png is
27scanned for detect.png starting from X,Y specified above. When detect.png
28is found, then all the match.png files are scanned at an offset of x,y as
29specified above.  If a match is found, then visgrep prints the x,y and
30index of the item.
31
32For example, image.png is a screenshot and match1.png .. to match5.png are
33images of letters a to e.  Each of these letters is enclosed in a blue box,
34so detect.png is an image of the upper left corner of the box.  This box is
35not included in the match*.png files, so they are actually offset 5 pixels
36down and 4 pixels to the left.  You might run it like this then:
37
38  visgrep -b -t50 -x-4 -y5 image.png match_corner.png match_a.png match_b.png ...
39
40Etc, with all matches listed.  Now suppose the screen showed 'ace' so
41visgrep might output:
42
430 10,10 0
4412 50,10 2
457 90,10 4
46
47Showing that match_a.png (index 0) is at 10,10 on the screen.  If no match
48is found even though the detection image is found, the index will be -1.
49
50The first match was 100%% accurate, while the second and third were very slightly
51inaccurate, probably due to anti-aliasing on the fonts.
52
53Exit status is 0 for successful match, 1 for no match, and 2 for error.
54
55See the examples page for use cases for different flags"""
56
57
58ctypedef numpy.uint8_t pixel_t
59
60cdef struct pixel:
61    pixel_t r
62    pixel_t g
63    pixel_t b
64    pixel_t a
65
66ctypedef pixel_t[:, :, :] image_t
67
68Point = namedtuple('Point', ['x', 'y'])
69Size = namedtuple('Size', ['width', 'height'])
70FoundResult = namedtuple('FoundResult', ['badness', 'point'])
71
72GeomParams = namedtuple('GeomParams', ['off', 'src', 'sub', 'start', 'pattern_off', 'pattern'])
73GeomParams.__new__.__defaults__ = (Point(0, 0),
74                                   Size(None, None),
75                                   Size(1, 1),
76                                   Point(0, 0),
77                                   None,
78                                   None)
79
80MetricParams = namedtuple('MetricParams', ['tolerance'])
81ImgParams = namedtuple('ImgParams', ['img', 'detect', 'match', 'scan_all'])
82
83
84cdef numpy.ndarray img_to_array(img):
85    """
86    Convert a PIL image to a numpy array.
87    """
88
89    if not isinstance(img, Image.Image):
90        raise ValueError("PIL image required, not '%s'" % type(img))
91
92    return numpy.array(img)
93
94
95cpdef numpy.ndarray crop_array(array, corners):
96    """
97    crop a numpy array by the absolute corner coordinates:
98
99    (x0, y0, x1, y1) = (left, upper, right, lower)
100
101    Those corners are all INCLUSIVE.
102
103    Array must have shape (h, w, 4), i.e. an RGBA image.
104    """
105
106    x0, y0, x1, y1 = corners
107    return array[y0:y1, x0:x1]
108
109
110@cython.boundscheck(False)
111@cython.wraparound(False)
112cdef inline pixel img_pixel_get(image_t img, Py_ssize_t x, Py_ssize_t y):
113    """
114    Get pixel color or zero if outside bounds.
115    Totally speed boosted.
116    """
117
118    if x < img.shape[1] and y < img.shape[0]:
119        return pixel(img[y, x, 0],
120                     img[y, x, 1],
121                     img[y, x, 2],
122                     img[y, x, 3])
123
124    return pixel(0, 0, 0, 0)
125
126
127@cython.boundscheck(False)
128@cython.wraparound(False)
129cdef img_subimage_find(image_t master, image_t find,
130                       start_from, float tolerance, find_next):
131    """
132    Fuzzily find a part of an image that matches the pattern.
133    """
134
135    cdef Py_ssize_t x_end = master.shape[1] - find.shape[1]
136    cdef Py_ssize_t y_end = master.shape[0] - find.shape[0]
137
138    if find_next:
139        start_from = Point(start_from.x + 1, start_from.y)
140
141    cdef Py_ssize_t ymax = start_from.y
142    cdef Py_ssize_t xmax = start_from.x
143    cdef Py_ssize_t y_it
144    cdef Py_ssize_t x_it
145    cdef float badness
146    # Loop the whole freakin image looking for this sub image, but not past edges
147    for y_it in range(ymax, y_end + 1):
148        for x_it in range(xmax, x_end + 1):
149            badness = img_subimage_cmp(master, find, x_it, y_it, tolerance)
150            if badness <= tolerance:
151                return FoundResult(badness, Point(x_it, y_it))
152
153    # No match
154    return FoundResult(-1, Point(-1, -1))
155
156
157#@cython.profile(False)
158cdef inline float img_pixel_cmp(const pixel &pix, const pixel &other_pix):
159    """
160    o is the compare from pixel, assumed to be from a pattern. It's transparency
161    is the transparency used to modify the tolerance value
162      return( memcmp( &p, &o, sizeof( PIXEL ) ) );
163    make tolerance mean something
164    """
165    cdef int difference = abs(pix.r - other_pix.r) + abs(pix.g - other_pix.g) + abs(pix.b - other_pix.b)
166    cdef int transparentness = other_pix.a
167    return difference * (transparentness / 255)
168
169
170@cython.boundscheck(False)
171@cython.wraparound(False)
172@cython.nonecheck(False)
173cdef float img_subimage_cmp(image_t master, image_t subimage,
174                            Py_ssize_t where_x, Py_ssize_t where_y,
175                            float tolerance):
176    """
177    Returns 0 if subimage is inside master at where, like *cmp usually does for other stuff
178    otherwise returns an integer of how different the match is, for each color component
179    value off.  tolerance is how high to go before bailing.  set lower to avoid processing
180    lots of extra pixels, it will just ret when tolerance is met
181    """
182
183    logging.debug("Comparing subimage where=%d,%d", where_x, where_y)
184
185    # Check if subimage even fits in masterimage at POINT
186    if ((where_x + subimage.shape[1]) > master.shape[1] or
187        (where_y + subimage.shape[0]) > master.shape[0]):
188        # Superbad
189        logging.debug("Subimage would not fit here")
190        return 1000
191
192    cdef float badness = 0
193
194    cdef Py_ssize_t sptx
195    cdef Py_ssize_t spty
196
197    for sptx in range(subimage.shape[1]):
198        for spty in range(subimage.shape[0]):
199            # Map U/V to X/Y.
200            # Grab pels and see if they match
201            mpx = img_pixel_get(master, sptx + where_x, spty + where_y)
202            spx = img_pixel_get(subimage, sptx, spty)
203
204            badness += abs(img_pixel_cmp(mpx, spx))
205
206            if badness > tolerance:
207                logging.debug("Bail out early, badness > tolerance %d > %d", badness, tolerance)
208                # No match here, bail early
209                return badness
210
211    # Matched all of subimage
212    logging.debug("Image match ok, badness = %d", badness)
213    return badness
214
215
216cdef do_output(show_badness, badness, point, idx):
217    """
218    Print match coordinates and score.
219    """
220    if show_badness:
221        print("%d %d,%d %d" % (badness, point.x, point.y, idx))
222    else:
223        print("%d,%d %d" % (point.x, point.y, idx))
224
225
226cdef advance_point(point, img, start_x):
227    """
228    Move across the image.
229    """
230    next_point = point(point.x + 1, point.y)
231
232    if next_point.x > img.shape[1]:
233        next_point = point(start_x, next_point.y + 1)
234        if next_point.y > img.shape[0]:
235            # done, bail
236            next_point = point(-1, -1)
237
238    return next_point
239
240
241cdef img_match_any(image_t img, patterns, off, tolerance, pt_match):
242    """
243    Move across the image.
244    """
245    gotmatch = None
246
247    tmp_pt_x = pt_match.point.x + off.x
248    tmp_pt_y = pt_match.point.y + off.y
249
250    for cnt, pattern in enumerate(patterns):
251        if gotmatch is None:
252            logging.info(" Testing for pattern %d", cnt)
253            logging.info(" %d,%d ", tmp_pt_x, tmp_pt_y)
254            badness = img_subimage_cmp(img,
255                                       pattern,
256                                       tmp_pt_x,
257                                       tmp_pt_y,
258                                       tolerance)
259
260            if badness <= tolerance:
261                logging.info("  YES")
262
263                gotmatch = (FoundResult(badness, Point(tmp_pt_x, tmp_pt_y)), cnt)
264
265                # Fall out
266                break
267            else:
268                logging.info("  NO")
269
270    return gotmatch
271
272
273cpdef visgrep(image, pattern, tolerance, height=None):
274    """
275    Return points where pattern is found in the image.
276
277    image and pattern are numpy arrays of RGBA images.
278    """
279
280    if not (isinstance(image, numpy.ndarray) and
281            isinstance(pattern, numpy.ndarray)):
282        raise ValueError("image and pattern must be a numpy array")
283
284    geom_params = GeomParams(src=Size(None, height))
285    metric_params = MetricParams(tolerance)
286
287    img_params = ImgParams(image, pattern, [pattern], False)
288
289    ret = list()
290    for find in visgrep_cli(geom_params, metric_params, img_params):
291        ret.append(find[0])
292
293    return ret
294
295
296cdef visgrep_cli(geom_params, metric_params, img_params):
297    """
298    Perform search, return list of results.
299    """
300
301    cdef image_t img, find
302    cdef vector[image_t] matches
303
304    pt_match = FoundResult(0, Point(0, 0))
305    results = []
306    find_next = False
307
308    # we'll search in this image.
309    img = crop_array(img_params.img,
310                     (0, 0,
311                      geom_params.src.width or img_params.img.shape[1],
312                      geom_params.src.height or img_params.img.shape[0]))
313
314    if geom_params.pattern_off is not None:
315        patterns_crop = (geom_params.pattern_off.x,
316                         geom_params.pattern_off.y,
317                         geom_params.pattern.width + geom_params.pattern_off.x,
318                         geom_params.pattern.height + geom_params.pattern_off.y)
319
320        find = crop_array(img_params.detect, patterns_crop)
321
322        matches = [crop_array(match, patterns_crop)
323                   for match in img_params.match]
324    else:
325        find = img_params.detect
326        matches = img_params.match
327
328    logging.info("Detecting offsets...")
329    pt_match = FoundResult(pt_match.badness, Point(geom_params.start.x, geom_params.start.y))
330    find_next = False
331    while pt_match.point.x != -1:
332        if img_params.scan_all:
333            # fake match here
334            if find_next:
335                next_point = advance_point(pt_match.point, img, geom_params.start.x)
336
337                # increment counters
338                pt_match = FoundResult(pt_match.badness, next_point)
339        else:
340            pt_match = img_subimage_find(img, find, pt_match.point,
341                                         metric_params.tolerance,
342                                         find_next)
343
344        # Not first time anymore
345        find_next = True
346
347        if pt_match.point.x != -1:
348            logging.info("  Found match at %d,%d", pt_match.point.x, pt_match.point.y)
349
350            if len(img_params.match) == 1 and\
351               numpy.array_equal(img_params.match[0], img_params.detect):
352                # Detection pattern is the single match pattern
353                results.append((pt_match, 0))
354                continue
355
356            # Try and identify what thing it is
357            gotmatch = None
358            for tmp_off_x, tmp_off_y in itertools.product(range(geom_params.sub.width),
359                                                          range(geom_params.sub.height)):
360                gotmatch = img_match_any(img,
361                                         matches,
362                                         Point(geom_params.off.x + tmp_off_x,
363                                               geom_params.off.y + tmp_off_y),
364                                         metric_params.tolerance,
365                                         pt_match)
366
367                if gotmatch is not None:
368                    results.append(gotmatch)
369                    break
370
371            # Notify of no match
372            if gotmatch is None:
373                logging.info(" NO ITEMS MATCHED!")
374                if not img_params.scan_all:
375                    results.append((pt_match, -1))
376
377    return results
378
379
380def main():
381    """
382    Visual grep, greps for images in another image.
383    """
384    parser = argparse.ArgumentParser(description=TOOL_DESCRIPTION, epilog=EPILOG,
385                                     formatter_class=argparse.RawDescriptionHelpFormatter)
386
387    parser.add_argument("-x", dest="off_x", default=0, metavar="X_OFF", type=int,
388                        help="Set x offset for detection matching")
389    parser.add_argument("-y", dest="off_y", default=0, metavar="Y_OFF", type=int,
390                        help="Set y offset for detection matching")
391    parser.add_argument("-I", dest="src_width", default=0, metavar="WIDTH", type=int,
392                        help="Set width for cropping the source image")
393    parser.add_argument("-J", dest="src_height", default=0, metavar="HEIGHT", type=int,
394                        help="Set height for cropping the source image")
395    parser.add_argument("-W", dest="sub_width", default=0, metavar="X_OFF_WIDTH", type=int,
396                        help="Set x offset width for detection matching")
397    parser.add_argument("-H", dest="sub_height", default=0, metavar="Y_OFF_HEIGHT", type=int,
398                        help="Set y offset height for detection matching")
399    parser.add_argument("-X", dest="start_x", default=0, metavar="START_X_OFF", type=int,
400                        help="Start scanning at X")
401    parser.add_argument("-Y", dest="start_y", default=0, metavar="START_Y_OFF", type=int,
402                        help="Start scanning at Y")
403    parser.add_argument("-u", dest="off_u", default=0, metavar="U_OFF", type=int,
404                        help="Set u offset for cropping patterns before use")
405    parser.add_argument("-v", dest="off_v", default=0, metavar="V_OFF", type=int,
406                        help="Set y offset for cropping patterns before use")
407    parser.add_argument("-M", dest="pattern_width", metavar="U_WIDTH", type=int,
408                        help="Set width of the cropped patterns")
409    parser.add_argument("-N", dest="pattern_height", metavar="V_HEIGHT", type=int,
410                        help="Set height of the cropped patterns")
411    parser.add_argument("-a", dest="scan_all", help="""Scan all patterns, not just after matching
412                        the detection pattern note: this method is much slower because we scan for
413                        all images at every pixel instead of just at detection points. Also, in
414                        this mode the detection image is ignored, there will be no matches
415                        for tile -1""", action="store_true")
416    parser.add_argument("-t", dest="tolerance", default=0, metavar="TOLERANCE", type=int,
417                        help="Set tolerance for 'fuzzy' matches, higher numbers are more tolerant")
418    parser.add_argument("-b", dest="show_badness",
419                        help="""Display 'badness' value, higher numbers mean match is less accurate,
420                        a badness value of 0 means the match is pixel-perfect""",
421                        action="store_true")
422    parser.add_argument("-d", metavar="DEBUGLEVEL", type=int, help="Print debug messages")
423    parser.add_argument("image", metavar='image.png', help="Image to search in")
424    parser.add_argument("detect", metavar='detect.png', help="Pattern to search")
425    parser.add_argument("match", metavar='match.png', nargs="*",
426                        help="Images to compare if detected")
427
428    args = parser.parse_args()
429
430    logging.basicConfig(level=args.d)
431
432    pattern_crop_params = [v is not None for v in
433                           (args.off_u, args.off_v, args.pattern_width, args.pattern_height)]
434
435    if any(pattern_crop_params) and not all(pattern_crop_params):
436        parser.error('-u, -v, -M, -N must be given together')
437
438    image = img_to_array(Image.open(args.image).convert('RGBA'))
439    find = img_to_array(Image.open(args.detect).convert('RGBA'))
440    match = [img_to_array(Image.open(fname).convert('RGBA'))
441             for fname in args.match]
442
443    geom_params = GeomParams(Point(args.off_x, args.off_y),
444                             Size(args.src_width, args.src_height),
445                             Size(args.sub_width, args.sub_height),
446                             Point(args.start_x, args.start_y),
447                             Point(args.off_u, args.off_v),
448                             Size(args.pattern_width, args.pattern_height))
449    metric_params = MetricParams(args.tolerance)
450    img_params = ImgParams(image, find, match, args.scan_all)
451
452    results = visgrep_cli(geom_params, metric_params, img_params)
453
454    exit_status = 1
455
456    for result in results:
457        do_output(args.show_badness,
458                  result[0].badness,
459                  result[0].point,
460                  result[1])
461
462        if result[1] != -1:
463            exit_status = 0
464
465    sys.exit(exit_status)
466
467
468if __name__ == "__main__":
469    main()
470