Rewrite screenshotmatcher in python

.gitignore

@@ -0,0 +1,8 @@
+www/results/*
+!www/results/.gitkeep
+
+__pycache__/
+*.py[cod]
+*$py.class
+
+Thumbs.db

screenshotmatcher/common/config.py

@@ -0,0 +1,8 @@
+import common.utils
+
+APP_NAME = 'Screenshot Matcher'
+ICON_PATH = 'gui/icon.png'
+HOST = common.utils.getCurrentIPAddress()
+PORT = 49049
+SERVICE_URL = 'http://{}:{}'.format(HOST, PORT)
+ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}

screenshotmatcher/common/utils.py

@@ -0,0 +1,27 @@
+import socket
+import os
+import platform
+import subprocess
+
+def getCurrentIPAddress():
+  s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+  s.connect(("8.8.8.8", 80))
+  ipAddr = s.getsockname()[0]
+  s.close()
+  return ipAddr
+
+def getScriptDir(filename):
+  return os.path.dirname(os.path.realpath(filename))
+
+
+def allowed_file(filename):
+  return '.' in filename and \
+          filename.rsplit('.', 1)[1].lower() in {'png', 'jpg', 'jpeg'}
+
+def open_file_or_dir(path):
+    if platform.system() == "Windows":
+        os.startfile(path)
+    elif platform.system() == "Darwin":
+        subprocess.Popen(["open", path])
+    else:
+        subprocess.Popen(["xdg-open", path])

screenshotmatcher/gui/qrcode.py

@@ -0,0 +1,45 @@
+from wx import (  # pylint: disable=no-name-in-module
+  BoxSizer, 
+  App, 
+  Icon, 
+  Menu, 
+  Frame, 
+  CallAfter, 
+  Image,
+  EmptyImage,
+  Bitmap,
+  StaticBitmap,
+  NullBitmap,
+  ID_ANY,
+  EVT_MENU, 
+  BITMAP_TYPE_PNG, 
+  HORIZONTAL, 
+)
+
+from wx.adv import TaskBarIcon # pylint: disable=no-name-in-module
+import qrcode
+
+import common.config
+
+def static_bitmap_from_pil_image(caller, pil_image):
+  wx_image = Image(pil_image.size[0], pil_image.size[1])
+  wx_image.SetData(pil_image.convert("RGB").tobytes())
+  bitmap = Bitmap(wx_image)
+  static_bitmap = StaticBitmap(caller, ID_ANY, NullBitmap)
+  static_bitmap.SetBitmap(bitmap)
+  return static_bitmap
+
+class QRCodeFrame(Frame):
+  def __init__(self, app, parent, id, title):
+    Frame.__init__(self, parent, id, title, (-1, -1), (300, 300))
+
+    sizer = BoxSizer(HORIZONTAL)
+
+    img = qrcode.make(common.config.SERVICE_URL)
+
+    static_bitmap = static_bitmap_from_pil_image(self, img)
+    sizer.Add(static_bitmap)
+    self.SetSizer(sizer)
+    self.Centre()
+    self.Layout()
+

screenshotmatcher/gui/tray.py

@@ -0,0 +1,37 @@
+from wx import App, Icon, Menu, Frame, BITMAP_TYPE_PNG, CallAfter, EVT_MENU # pylint: disable=no-name-in-module
+from wx.adv import TaskBarIcon # pylint: disable=no-name-in-module
+
+import gui.qrcode
+import common.utils
+
+class Tray(TaskBarIcon):
+
+    def __init__(self, app, app_name, icon_path):
+        TaskBarIcon.__init__(self)
+        self.app = app
+        self.app_name = app_name
+        self.SetIcon(Icon(icon_path, BITMAP_TYPE_PNG), app_name)
+        self.Bind(EVT_MENU, self.OnTaskbarQR, id=1)
+        self.Bind(EVT_MENU, self.OnTaskbarResults, id=2)
+        self.Bind(EVT_MENU, self.OnTaskbarClose, id=3)
+
+
+    def CreatePopupMenu(self):
+        menu = Menu()
+        menu.Append(1, 'Show QR Code')
+        menu.Append(2, 'Show Results')
+        menu.Append(3, 'Close')
+        return menu
+
+    def OnTaskbarQR(self, event):
+        # Initialize QR code window
+        self.qr_frame = gui.qrcode.QRCodeFrame(self, None, -1, self.app_name)
+        self.qr_frame.Show(True)
+
+    def OnTaskbarResults(self, event):
+        common.utils.open_file_or_dir(common.utils.getScriptDir(__file__) + '/../../www/results')
+
+    def OnTaskbarClose(self, event):
+        if hasattr(self, 'qr_frame') and self.qr_frame:
+            CallAfter(self.qr_frame.Destroy)
+        self.Destroy()

screenshotmatcher/main.py

@@ -0,0 +1,38 @@
+# Import third party modules
+from wx import App
+import threading
+import logging
+import time
+
+# Import app config
+from common.config import (
+    APP_NAME, 
+    ICON_PATH,
+)
+
+# Import app parts
+import gui.tray
+import common.utils
+from matching.matcher import Matcher
+from server.server import Server
+
+# Main App
+class MainApp(App):
+
+    def OnInit(self):
+        # Initialize tray icon/menu
+        self.tray = gui.tray.Tray(self, APP_NAME, ICON_PATH)
+        return True
+
+if __name__ == "__main__":
+
+    # Init Server and GUI
+    server = Server()
+    app = MainApp()
+
+    # Start server in different thread
+    x = threading.Thread(target=server.start, args=(), daemon=True)
+    x.start()
+
+    # Start GUI event loop
+    app.MainLoop()

screenshotmatcher/matching/matcher.py

@@ -0,0 +1,182 @@
+import pyscreenshot as ImageGrab
+import time
+import numpy as np
+from matplotlib import pyplot as plt
+
+from cv2 import ( # pylint: disable=no-name-in-module
+  perspectiveTransform,
+  findHomography,
+  RANSAC,
+  FlannBasedMatcher,
+  imread,
+  imwrite,
+  IMREAD_COLOR,
+  IMREAD_GRAYSCALE,
+  DescriptorMatcher_create,
+  ORB_create,
+)
+
+from cv2.xfeatures2d import ( # pylint: disable=no-name-in-module,import-error
+  SURF_create,
+)
+
+
+class Matcher():
+
+  def __init__(self, match_uid, photo):
+
+    self.match_uid = match_uid
+    self.match_dir = '../www/results/' + match_uid
+
+    self.screenshot_file = 'screenshot.png'
+    self.screenshot = ImageGrab.grab()
+    self.screenshot.save(self.match_dir + '/' + self.screenshot_file)
+
+    self.photo_file = photo
+
+  def match(self, algorithm='SURF'):
+
+    start_time = time.perf_counter()
+
+    # Load pictures
+
+    photo = imread( '{}/{}'.format(self.match_dir, self.photo_file), IMREAD_GRAYSCALE )
+    screen = imread( '{}/{}'.format(self.match_dir, self.screenshot_file), IMREAD_GRAYSCALE )
+    screen_colored = imread( '{}/{}'.format(self.match_dir, self.screenshot_file), IMREAD_COLOR )
+
+    # Provisional switch statement
+    if algorithm == 'SURF':
+      match_result = self.algorithm_SURF(photo, screen, screen_colored)
+    if algorithm == 'ORB':
+      match_result = self.algorithm_ORB(photo, screen, screen_colored)
+    else:
+      match_result = self.algorithm_SURF(photo, screen, screen_colored)
+
+
+    print( round( (time.perf_counter() - start_time) * 1000 ), 'ms' )
+
+
+    return match_result
+
+  def algorithm_SURF(self, photo, screen, screen_colored):
+
+    # Init algorithm
+    surf = SURF_create(400)
+    surf.setUpright(True)
+
+    # Detect and compute
+    kp_photo, des_photo = surf.detectAndCompute(photo, None)
+    kp_screen, des_screen = surf.detectAndCompute(screen, None)
+
+    # Descriptor Matcher
+    index_params = dict(algorithm = 0, trees = 5)
+    search_params = dict(checks = 50)
+    flann = FlannBasedMatcher(index_params, search_params)
+
+    # Calc knn Matches
+    matches = flann.knnMatch(des_photo, des_screen, k=2)
+
+    if not matches or len(matches) == 0:
+      return False
+
+    # store all the good matches as per Lowe's ratio test.
+    good = []
+    for m,n in matches:
+        if m.distance < 0.75*n.distance:
+            good.append(m)
+
+    if not good or len(good) < 20:
+      return False
+
+    photo_pts = np.float32([ kp_photo[m.queryIdx].pt for m in good ]).reshape(-1,1,2)
+    screen_pts = np.float32([ kp_screen[m.trainIdx].pt for m in good ]).reshape(-1,1,2)
+
+    M, mask = findHomography(photo_pts, screen_pts, RANSAC, 5.0)
+
+    if len(M) == 0:
+      return False
+
+    h, w = photo.shape
+    pts = np.float32([ [0,0],[0,h-1],[w-1,h-1],[w-1,0] ]).reshape(-1,1,2)
+    dst = perspectiveTransform(pts, M)
+
+
+    minX = dst[0][0][0]
+    minY = dst[0][0][1]
+    maxX = dst[0][0][0]
+    maxY = dst[0][0][1]
+
+    for i in range(4):
+      if dst[i][0][0] < minX:
+        minX = dst[i][0][0]
+      if dst[i][0][0] > maxX:
+        maxX = dst[i][0][0]
+      if dst[i][0][1] < minY:
+        minY = dst[i][0][1]
+      if dst[i][0][1] > maxY:
+        maxY = dst[i][0][1]
+
+    imwrite(self.match_dir + '/result.png', screen_colored[ int(minY):int(maxY), int(minX):int(maxX)])
+
+    return True
+
+
+  def algorithm_ORB(self, photo, screen, screen_colored):
+
+    # Init algorithm
+    orb = ORB_create(800)
+
+    # Detect and compute
+    kp_photo, des_photo = orb.detectAndCompute(photo, None)
+    kp_screen, des_screen = orb.detectAndCompute(screen, None)
+
+    # Descriptor Matcher
+    descriptor_matcher = DescriptorMatcher_create('BruteForce-Hamming')
+
+    # Calc knn Matches
+    matches = descriptor_matcher.knnMatch(des_photo, des_screen, k=2)
+
+    if not matches or len(matches) == 0:
+      return False
+
+    # store all the good matches as per Lowe's ratio test.
+    good = []
+    for m,n in matches:
+        if m.distance < 0.75*n.distance:
+            good.append(m)
+
+    if not good or len(good) < 20:
+      return False
+
+    photo_pts = np.float32([ kp_photo[m.queryIdx].pt for m in good ]).reshape(-1,1,2)
+    screen_pts = np.float32([ kp_screen[m.trainIdx].pt for m in good ]).reshape(-1,1,2)
+
+    M, mask = findHomography(photo_pts, screen_pts, RANSAC, 5.0)
+
+    if len(M) == 0:
+      return False
+
+    h, w = photo.shape
+    pts = np.float32([ [0,0],[0,h-1],[w-1,h-1],[w-1,0] ]).reshape(-1,1,2)
+    dst = perspectiveTransform(pts, M)
+
+
+    minX = dst[0][0][0]
+    minY = dst[0][0][1]
+    maxX = dst[0][0][0]
+    maxY = dst[0][0][1]
+
+    for i in range(4):
+      if dst[i][0][0] < minX:
+        minX = dst[i][0][0]
+      if dst[i][0][0] > maxX:
+        maxX = dst[i][0][0]
+      if dst[i][0][1] < minY:
+        minY = dst[i][0][1]
+      if dst[i][0][1] > maxY:
+        maxY = dst[i][0][1]
+
+    imwrite(self.match_dir + '/result.png', screen_colored[ int(minY):int(maxY), int(minX):int(maxX)])
+
+    return True
+