Initial commit

module/base/template.py

@@ -0,0 +1,65 @@
+import cv2
+import numpy as np
+from PIL import Image
+
+
+class Template:
+    def __init__(self, area, color, button, file):
+        """
+        Args:
+            file(str): Relative path to file.
+        """
+
+        # self.area = image.getbbox()
+        self.area = area
+        self.color = color
+        self.button = button
+        image = Image.open(file)
+        self.image = np.array(image.crop(self.area))
+        self.similarity = 0.85
+        self.preprocess_func = None
+
+    def set_preprocess_func(self, func):
+        """
+        Args:
+            func: Image preprocess function
+        """
+        self.preprocess_func = func
+        self.image = self._preprocess(self.image)
+
+    def _preprocess(self, image):
+        image = image.astype(float)
+        image = self.preprocess_func(image)
+        image[image > 255] = 255
+        image[image < 0] = 0
+        image = image.astype('uint8')
+        return image
+
+    def match(self, image):
+        """
+        Args:
+            image: Full-screen screenshot.
+
+        Returns:
+            bool: True if template matched.
+        """
+        image = np.array(image.crop(self.area))
+        if self.preprocess_func is not None:
+            image = self._preprocess(image)
+
+        res = cv2.matchTemplate(np.array(image), self.image, cv2.TM_CCOEFF_NORMED)
+        _, similarity, _, _ = cv2.minMaxLoc(res)
+
+        return similarity > self.similarity
+
+
+def preprocess_func_example(image):
+    """
+     Args:
+         image (np.ndarray):
+
+     Returns:
+         np.ndarray
+     """
+    image = (image - 64) / 0.75
+    return image