2020-03-28 17:22:46 +00:00
|
|
|
import cv2
|
|
|
|
import numpy as np
|
|
|
|
from PIL import Image
|
|
|
|
|
|
|
|
|
|
|
|
class Template:
|
2020-04-09 08:30:40 +00:00
|
|
|
def __init__(self, file):
|
2020-03-28 17:22:46 +00:00
|
|
|
"""
|
|
|
|
Args:
|
2020-04-09 08:30:40 +00:00
|
|
|
file (str): Filepath of template file.
|
2020-03-28 17:22:46 +00:00
|
|
|
"""
|
2020-04-09 08:30:40 +00:00
|
|
|
self.file = file
|
|
|
|
self.image = np.array(Image.open(file))
|
2020-03-28 17:22:46 +00:00
|
|
|
|
2020-04-09 08:30:40 +00:00
|
|
|
def match(self, image, similarity=0.85):
|
2020-03-28 17:22:46 +00:00
|
|
|
"""
|
|
|
|
Args:
|
2020-04-09 08:30:40 +00:00
|
|
|
image:
|
|
|
|
similarity (float): 0 to 1.
|
2020-03-28 17:22:46 +00:00
|
|
|
|
2020-04-09 08:30:40 +00:00
|
|
|
Returns:
|
|
|
|
bool: If matches.
|
|
|
|
"""
|
|
|
|
res = cv2.matchTemplate(np.array(image), self.image, cv2.TM_CCOEFF_NORMED)
|
|
|
|
_, sim, _, _ = cv2.minMaxLoc(res)
|
2020-04-12 05:54:06 +00:00
|
|
|
# print(self.file, sim)
|
2020-04-09 08:30:40 +00:00
|
|
|
return sim > similarity
|
2020-03-28 17:22:46 +00:00
|
|
|
|
2020-04-09 08:30:40 +00:00
|
|
|
def match_multi(self, image, similarity=0.85):
|
2020-03-28 17:22:46 +00:00
|
|
|
"""
|
|
|
|
Args:
|
2020-04-09 08:30:40 +00:00
|
|
|
image:
|
|
|
|
similarity (float): 0 to 1.
|
2020-03-28 17:22:46 +00:00
|
|
|
|
|
|
|
Returns:
|
2020-04-09 08:30:40 +00:00
|
|
|
np.ndarray: np.array([[x0, y0], [x1, y1])
|
2020-03-28 17:22:46 +00:00
|
|
|
"""
|
2020-04-09 08:30:40 +00:00
|
|
|
result = cv2.matchTemplate(np.array(image), self.image, cv2.TM_CCOEFF_NORMED)
|
|
|
|
result = np.array(np.where(result > similarity)).T
|
2020-03-28 17:22:46 +00:00
|
|
|
|
2020-04-09 08:30:40 +00:00
|
|
|
return result
|