diff --git a/gamemodel.py b/gamemodel.py index 05de4f7..7d27e81 100644 --- a/gamemodel.py +++ b/gamemodel.py @@ -5,19 +5,20 @@ from functools import wraps from utility import * import pointcluster +from imagepipeline import CVImage, ImagePipeline class GameModel: """Platform-independent representation of the game's state.""" def __init__(self, io:gameio.AbstractGameIO): self.gameio = io 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)) + CVImage("big", filename = "images/game_assets/rock-big.png"), + CVImage("normal", filename = "images/game_assets/rock-normal.png"), + CVImage("small", filename = "images/game_assets/rock-small.png") ] self.ships = [ - ("ship_off", cv2.imread("images/game_assets/spaceship-off.png", 0)), - ("ship_on", cv2.imread("images/game_assets/spaceship-on.png", 0)) + CVImage("ship_off", filename = "images/game_assets/spaceship-off.png"), + CVImage("ship_on", filename = "images/game_assets/spaceship-on.png") ] #self.missile = ("missile", cv2.imread("images/game_assets/missile.png", 0)) self.frame = None @@ -35,170 +36,81 @@ class GameModel: @wraps(fn) def inner(self, *args, **kwargs): if self.frame is None: - #print("Fetching frame.") sshot = self.gameio.fetch_sshot() open_cv_image = np.array(sshot) # Convert RGB to BGR - self.frame = open_cv_image[:, :, ::-1].copy() - self.color_frame = np.copy(self.frame) - self.frame = cv2.cvtColor(self.frame, cv2.COLOR_BGR2GRAY) + array = open_cv_image[:, :, ::-1].copy() + self.color_frame = CVImage("gameio frame", np.copy(array)) + self.frame = CVImage("BW frame", self.color_frame.copy()) + self.frame.image = self.frame.convert_color(False) self.mask_frame() + print(self.frame) return fn(self, *args, **kwargs) return inner -## def with_masking(fn): -## """Decorator to cut lives and score into smaller subimages, and mask them out of self.frame.""" -## @wraps(fn) -## def inner(self, *args, **kwargs): -## if self.score_img is None: -## -## return fn(self, *args, **kwargs) -## return inner - def mask_frame(self): - self.lives_img = self.frame[self.lives_rect[0][0]:self.lives_rect[0][1], - self.lives_rect[1][0]:self.lives_rect[1][1]] - lives_mask = np.full(self.frame.shape, 255, dtype=np.uint8) - - cv2.rectangle(lives_mask, - *self.lives_rect, - color=0, thickness=cv2.FILLED) - - self.score_img = self.frame[self.score_rect[0][0]:self.score_rect[0][1], - self.score_rect[1][0]:self.score_rect[1][1]] - score_mask = np.full(self.frame.shape, 255, dtype=np.uint8) - cv2.rectangle(score_mask, - *self.score_rect, - color = 0, thickness=cv2.FILLED) - self.frame = cv2.bitwise_and(self.frame, lives_mask) - self.frame = cv2.bitwise_and(self.frame, score_mask) + self.lives_img = CVImage("lives", self.frame.snip(self.lives_rect)) + self.frame.image = self.frame.mask(self.lives_rect) + self.score_img = CVImage("score", self.frame.snip(self.score_rect)) + self.frame.image = self.frame.mask(self.score_rect) -## print("Displaying images for testing purposes") -## cv2.imshow("Original", self.color_frame) -## cv2.waitKey(0) -## cv2.imshow("Masked", self.frame) -## cv2.waitKey(0) - def clear_frame(self): self.prev_frame = frame self.frame = None @with_frame def find_asteroids(self): - asteroid_rects = [] - 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.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 + results = [] + for a in self.asteroids: + r = self.frame.template_detect(a, + self.cv_template_thresh, + self.duplicate_dist_thresh) + results.extend(r) + return results @with_frame def display_results(self, rects = [], pointsets = [], circles = []): """Draws results on the current frame for test purposes.""" - displayable = np.copy(self.color_frame) - cv2.rectangle(displayable, *self.lives_rect, (255,255,255), 1) - cv2.rectangle(displayable, *self.score_rect, (255,255,255), 1) - #else: - # displayable = np.copy(self.color_frame) - for pt, wh, label in rects: - color = { "big": (255, 0, 0), + displayable = CVImage("GameModel results", self.color_frame.copy()) + label_color = { "big": (255, 0, 0), "normal": (0, 255, 0), "small": (0, 0, 255), "missile": (0, 255, 128), "ship_on": (0, 0, 128), - "ship_off": (0, 64, 128)}[label] - cv2.rectangle(displayable, pt, wh, color, 1) - cv2.putText(displayable, label, pt, - cv2.FONT_HERSHEY_PLAIN, - 1.0, color) + "ship_off": (0, 64, 128)} + for r in rects: + displayable.draw_rect(r, color=label_color[r.label]) for ps in pointsets: - color = (0, 255, 255) - cv2.polylines(displayable, np.int32([ps]), True, color) + displayable.draw_poly(ps, color=(0, 255, 255)) for center, radius, label in circles: - color = (255, 255, 0) - cv2.circle(displayable, np.int32(center), int(radius), color, 1) - cv2.putText(displayable, label, np.int32(center), - cv2.FONT_HERSHEY_PLAIN, - 1.0, color) - - cv2.imshow("Results", displayable) - cv2.waitKey(0) + displayable.draw_circle(center, radius) + displayable.draw_text(label, center, (255, 255, 0)) + + displayable.show() @with_frame def frame_sift(self): - sift = cv2.SIFT_create() - kp_desc = {} # dict of (keypoints, descriptions) for all ship sprites - kp_desc["frame"] = sift.detectAndCompute(self.frame, None) - frame_kp, frame_desc = kp_desc["frame"] -## for label, s in self.ships: -## kp_desc[label] = sift.detectAndCompute(s, None) -## bf = cv2.BFMatcher(cv2.NORM_L1, crossCheck=True) -## matchsets = {} -## for label in kp_desc: -## _, desc = kp_desc[label] -## matchsets[label] = bf.match(frame_desc, desc) -## #return { "matchsets": matchsets, -## # "kp_desc": kp_desc -## # } - ship_rsq = rect_radius_squared(*self.ships[0][1].shape) * 0.85 - #print(f"max radius^2: {ship_rsq}") - clusters = pointcluster.cluster_set([k.pt for k in frame_kp], sqrt(ship_rsq)) - - return clusters + ship_r = sqrt(rect_radius_squared(*self.ships[0].image.shape[:2]) * 0.85) + return self.frame.sift_clusters(cluster_radius = ship_r) @with_frame def find_ships(self): - ship_rects = [] - for label, a in self.ships: - h, w = a.shape - res = cv2.matchTemplate(self.frame, a, cv2.TM_CCOEFF_NORMED) - loc = np.where( res >= self.cv_template_thresh) - for pt in zip(*loc[::-1]): - if not ship_rects or squared_distance(ship_rects[-1][0], pt) > self.duplicate_dist_thresh: - ship_rects.append((pt, (pt[0] + w, pt[1] + h), label)) - return ship_rects + results = [] + for a in self.ships: + r = self.frame.template_detect(a, + self.cv_template_thresh, + self.duplicate_dist_thresh) + results.extend(r) + return results @with_frame def find_missiles(self): - # Setup SimpleBlobDetector parameters. - params = cv2.SimpleBlobDetector_Params() - - # Change thresholds - params.minThreshold = 10; - params.maxThreshold = 200; - - # Filter by Area. - params.filterByArea = True - #params.minArea = 1500 - params.maxArea = 100 - - # Filter by Circularity - #params.filterByCircularity = True - #params.minCircularity = 0.1 - - # Filter by Convexity - params.filterByConvexity = True - params.minConvexity = 0.95 - - # Filter by Inertia - params.filterByInertia = True - params.minInertiaRatio = 0.4 - - detector = cv2.SimpleBlobDetector_create(params) - keypoints = detector.detect(cv2.bitwise_not(self.frame)) # inverted black/white frame - #im_with_keypoints = cv2.drawKeypoints(self.frame, keypoints, np.array([]), - # (0,0,255), cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS) - #cv2.imshow("keypoints", im_with_keypoints) - #cv2.waitKey(0) - s = 9 # pixels for the missile - rect_tuple = lambda pt: ((int(pt[0]-s/2),int(pt[1]-s/2)), - (int(pt[0]+s/2), int(pt[1]+s/2)), - "missile") - return [rect_tuple(k.pt) for k in keypoints] + p = CVImage.blob_params(minThreshold = 10, maxThreshold = 200, + maxArea = 100, + minConvexity = 0.95, + minInertiaRatio = 0.4) + return self.frame.blob_detect(size=9, params=p) def analyse_frame(self): rocks = self.find_asteroids() @@ -264,9 +176,12 @@ if __name__ == '__main__': ship_results = gm.find_ships() polygons = [c.points for c in s_results] ##circles = [(c.center, c.max_distance, f"cluster_{i}") for i, c in enumerate(s_results)] - r_circles = [(c.center, sqrt(rect_radius_squared(*gm.ships[0][1].shape)), f"cluster_{i}") for i, c in enumerate(s_results)] + r_circles = [(c.center, sqrt(rect_radius_squared(*gm.ships[0].image.shape[:2])), f"cluster_{i}") for i, c in enumerate(s_results)] missile_results = gm.find_missiles() ##m_circles = [(pt, 10, f"missile_{i}") for i, pt in enumerate(missiles)] ##pprint(a_results+ship_results+missile_results) - gm.display_results(rects=a_results+ship_results+missile_results, pointsets=polygons, circles=r_circles) + rects = a_results + if ship_results: rects.extend(ship_results) + if missile_results: rects.extend(missile_results) + gm.display_results(rects=rects, pointsets=polygons, circles=r_circles) gm.analyse_frame() diff --git a/imagepipeline.py b/imagepipeline.py new file mode 100644 index 0000000..c494078 --- /dev/null +++ b/imagepipeline.py @@ -0,0 +1,165 @@ +import cv2 +import numpy as np +import typing +import pointcluster + +from shapes import Rect +from utility import * + +class CVImage: + """Dummy definition to allow recursive type hints""" + pass + +class CVImage: + def __init__(self, label="", img:np.ndarray=None, **kwargs): + """You can provide a 'filename' keyword arg to automatically load a file.""" + self.label = label + self.image = img + self._init_kwargs = kwargs + if kwargs: + load_kwargs = dict(kwargs) # copy + self.load(**load_kwargs) + + @property + def is_color(self): + return len(self.image.shape) == 3 + + def load(self, filename:str, label:str=None): + """Load an image from file. You can optionally set the 'label' keyword.""" + self.image = cv2.imread(filename) + if label: self.label = label + return self + + def convert_color(self, color:bool): + return cv2.cvtColor(self.image, cv2.COLOR_GRAY2BGR if color else cv2.COLOR_BGR2GRAY) + + def __repr__(self): + if self._init_kwargs: + kwargstr = ", " + ", ".join([f"{k}={repr(self._init_kwargs[k])}" for k in self._init_kwargs]) + else: + kwargstr = '' + return f"" + + def copy(self): + return np.copy(self.image) + + def snip(self, rect): + assert all((len(rect)==2, len(rect[0])==2, len(rect[1])==2)) #((x,y),(w,h)) + return self.image[rect[0][1]:rect[0][1]+rect[1][1], + rect[0][0]:rect[0][0]+rect[1][0] + ] + + def mask(self, rect, mask_color=None, nonmask_color=None): + assert all((len(rect)==2, len(rect[0])==2, len(rect[1])==2)) #((x,y),(w,h)) + if mask_color is None: + mask_color = (0, 0, 0) if self.is_color else 0 + if nonmask_color is None: + nonmask_color = (255, 255, 255) if self.is_color else 255 + mask = np.full(self.image.shape, nonmask_color, dtype=np.uint8) + cv2.rectangle(mask, *rect, color=mask_color, thickness=cv2.FILLED) + return cv2.bitwise_and(self.image, mask) + + def sift_clusters(self, cluster_radius) -> pointcluster.PointCluster: + sift = cv2.SIFT_create() + keypoints, descriptions = sift.detectAndCompute(self.image, None) + return pointcluster.cluster_set([k.pt for k in keypoints], cluster_radius) + + @staticmethod + def blob_params(*, minThreshold = 10, maxThreshold = 200, + minArea = None, maxArea = None, + minCircularity = None, maxCircularity = None, + minConvexity = None, maxConvexity = None, + minInertiaRatio = None, maxInertiaRatio = None): + p = cv2.SimpleBlobDetector_Params() + p.minThreshold = minThreshold + p.maxThreshold = maxThreshold + if minArea or maxArea: + p.filterByArea = True + if minArea: p.minArea = minArea + if maxArea: p.maxArea = maxArea + if minConvexity or maxConvexity: + p.filterByConvexity = True + if minConvexity: p.minConvexity = minConvexity + if maxConvexity: p.maxConvexity = maxConvexity + if minInertiaRatio or maxInertiaRatio: + p.filterByInertia = True + if minInertiaRatio: p.minInertiaRatio = minInertiaRatio + if maxInertiaRatio: p.maxInertiaRatio = maxInertiaRatio + if minCircularity or maxCircularity: + p.filterByCircularity = True + if minCircularity: p.minCircularity = minCircularity + if maxCircularity: p.maxCircularity = maxCircularity + return p + + def blob_detect(self, size:int, params=None, invert:bool=False, label:str=None) -> typing.List[Rect]: + if params is None: params = CVImage.blob_params() + detector = cv2.SimpleBlobDetector_create(params) + keypoints = detector.detect(cv2.bitwise_not(self.image) if invert else self.image) + rects = [] + s = size / 2.0 + for kp in keypoints: + rects.append(Rect(x=kp.pt[0] - s, y = kp.pt[1] - s, + w = size, h = size, + label = label or "blob")) + return rects + + def template_detect(self, template:CVImage, threshold:int, dupe_spacing:int) -> typing.List[Rect]: + if template.is_color: + h, w, _ = template.image.shape + else: + h, w = template.image.shape + if template.is_color != self.is_color: + template = CVImage(template.label, template.convert_color(not template.is_color)) + res = cv2.matchTemplate(self.image, template.image, cv2.TM_CCOEFF_NORMED) + loc = np.where(res >= threshold) + rects = [] + for pt in zip(*loc[::-1]): + if len(rects) > 0: + if squared_distance(rects[-1][0], pt) < dupe_spacing: continue + rects.append(Rect(x=pt[0], y=pt[1], w=w, h=h, label=template.label)) + return rects + + def show(self, delay=0): + cv2.imshow(self.label, self.image) + cv2.waitKey(delay) + + def draw_rect(self, rect:Rect, color=None, text_color=None, text:bool=True, thickness=1): + if color is None: + color = (255, 255, 255) if self.is_color else 255 + cv2.rectangle(self.image, rect.point, rect.point2, color, thickness) + if text: + self.draw_text(rect.label, rect.point, text_color if text_color else color) + + def draw_poly(self, points:typing.List[typing.Tuple], closed=True, color=None): + if color is None: + color = (255, 255, 255) if self.is_color else 255 + cv2.polylines(self.image, np.int32([points]), closed, color) + + def draw_circle(self, center, radius, color=None, thickness = 1): + if color is None: + color = (255, 255, 255) if self.is_color else 255 + cv2.circle(self.image, np.int32(center), np.int32(radius), color, thickness) + + def draw_text(self, text, point, color): + cv2.putText(self.image, text, np.int32(point), cv2.FONT_HERSHEY_PLAIN, 1.0, color) + + +class ImagePipeline: + def __init__(self): + pass + +# running this module executes tests +if __name__ == '__main__': + # initializer for CVImage can load from file + img = CVImage("test frame", filename="/home/john/Desktop/Screenshot at 2021-12-19 20-55-22.png") + #img.show() + # initializer for CVImage can accept a numpy array + img_no_title = CVImage("test frame", img.snip( ((0,24),(800,600)) )) + #img_no_title.show() + + #standard rectangle format used throughout the class, avoiding ugly splat operator + lives_rect = ((10,10), (190, 65)) + lives = CVImage("lives", img_no_title.snip(lives_rect)) + lives.show() + + diff --git a/shapes.py b/shapes.py new file mode 100644 index 0000000..8e9d14d --- /dev/null +++ b/shapes.py @@ -0,0 +1,54 @@ +class Rect: + def __init__(self, *args, label=None, **kwargs): + if len(args) == 4 and all([type(i) is int or type(i) is float for i in args]): + self.x, self.y, self.w, self.h = args + elif len(args) == 2 and all([type(i) is tuple and len(i) == 2 and all([type(j) is int or type(j) is float for j in i]) for i in args]): + xy, wh = self.args + self.x, self.y = xy + self.w, self.h = wh + elif all([k in kwargs for k in ("x", "y", "w", "h")]): + self.x = kwargs["x"] + self.y = kwargs["y"] + self.w = kwargs["w"] + self.h = kwargs["h"] + elif all([k in kwargs for k in ("x", "y", "x2", "y2")]): + self.x = kwargs["x"] + self.y = kwargs["y"] + self.w = kwargs["x2"] - self.x + self.h = kwargs["y2"] - self.y + elif all([k in kwargs for k in ("x1", "y1", "x2", "y2")]): + self.x = kwargs["x1"] + self.y = kwargs["y1"] + self.w = kwargs["x2"] - self.x + self.h = kwargs["y2"] - self.y + else: + raise RuntimeError("Rect requires 4 values: two coordinates or a coordinate plus width and height.") + self.label = label + + def __repr__(self): + return f"" + + def __iter__(self): + yield (self.x, self.y) + yield (self.w, self.h) + + def __len__(self): + return 2 + + def __getitem__(self, i): + if i == 0: return (self.x, self.y) + elif i == 1: return (self.w, self.h) + else: raise IndexError("Rect only supports index of 0 or 1.") + + def __setitem__(self, i, value): + assert i in (0, 1) and len(value) == 2 + if not i: self.x, self.y = value + else: self.w, self.h = value + + @property + def point(self): + return (self.x, self.y) + + @property + def point2(self): + return (self.x + self.w, self.y + self.h)