diff --git a/gamemodel.py b/gamemodel.py index 1094588..af91961 100644 --- a/gamemodel.py +++ b/gamemodel.py @@ -2,20 +2,26 @@ import gameio import cv2 import numpy as np +def squared_distance(vec1, vec2): + """returns distance-squared between two x, y point tuples""" + return (vec1[0] - vec2[0])**2 + (vec1[1] - vec2[1])**2 + class GameModel: """Platform-independent representation of the game's state.""" def __init__(self, io:gameio.AbstractGameIO): self.gameio = io - self.asteroids = [cv2.imread("images/game_assets/rock-big.png", 0), - cv2.imread("images/game_assets/rock-normal.png", 0), - cv2.imread("images/game_assets/rock-small.png", 0) - ] + self.asteroids = [ + ("big", cv2.imread("images/game_assets/rock-big.png", 0)), + ("normal", cv2.imread("images/game_assets/rock-normal.png", 0)), + ("small", cv2.imread("images/game_assets/rock-small.png", 0)) + ] self.frame = None - self.thresh = 0.6 # reconfigurable at runtime + self.cv_template_thresh = 0.6 # reconfigurable at runtime + self.duplicate_dist_thresh = 10 def with_frame(fn): """Decorator to process screenshot to cv2 format once upon first requirement, then reuse.""" - def inner(self): + def inner(self, *args, **kwargs): if self.frame is None: print("Fetching frame.") sshot = self.gameio.fetch_sshot() @@ -23,7 +29,7 @@ class GameModel: # Convert RGB to BGR self.frame = open_cv_image[:, :, ::-1].copy() self.frame = cv2.cvtColor(self.frame, cv2.COLOR_BGR2GRAY) - return fn(self) + return fn(self, *args, **kwargs) return inner def clear_frame(self): @@ -32,16 +38,27 @@ class GameModel: @with_frame def find_asteroids(self): asteroid_rects = [] - displayable = np.copy(self.frame) - for a in self.asteroids: + for label, a in self.asteroids: h, w = a.shape res = cv2.matchTemplate(self.frame, a, cv2.TM_CCOEFF_NORMED) - loc = np.where( res >= self.thresh) - ## Example code for displaying detected asteroid locations - #for pt in zip(*loc[::-1]): - # cv2.rectangle(displayable, pt, (pt[0] + w, pt[1] + h), 255, 1) - #cv2.imshow("Found asteroids", displayable) - #cv2.waitKey(0) + loc = np.where( res >= self.cv_template_thresh) + for pt in zip(*loc[::-1]): + if not asteroid_rects or squared_distance(asteroid_rects[-1][0], pt) > self.duplicate_dist_thresh: + asteroid_rects.append((pt, (pt[0] + w, pt[1] + h), label)) + return asteroid_rects + + @with_frame + def display_results(self, results): + """Draws results on the current frame for test purposes.""" + displayable = np.copy(self.frame) + for pt, wh, label in results: + cv2.rectangle(displayable, pt, wh, 255, 1) + cv2.putText(displayable, label, pt, + cv2.FONT_HERSHEY_PLAIN, + 1.0, 255) + cv2.imshow("Results", displayable) + cv2.waitKey(0) + if __name__ == '__main__': import platform @@ -53,8 +70,18 @@ if __name__ == '__main__': else: io = gameio.LinuxGameIO() - input("Press to locate the game at the start screen.") + #input("Press to locate the game at the start screen.") gm = GameModel(io) - input("Press to detect asteroids on screen.") - gm.find_asteroids() + # for testing purposes, populating window location at top-left of my screen + # io.loc is None when the title screen isn't found. + # manually setting io.loc crops all screenshots as if the title was found. + import pyscreeze + io.loc = pyscreeze.Box(0, 25, 800, 599) + + #input("Press to detect asteroids on screen.") + results = gm.find_asteroids() + print(f"Found {len(results)} asteroids") + for a in results: + print(a[0]) # position tuple + gm.display_results(results)