diff --git a/gamemodel.py b/gamemodel.py index 05de4f7..e05ee11 100644 --- a/gamemodel.py +++ b/gamemodel.py @@ -5,15 +5,16 @@ 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", color = False, filename = "images/game_assets/rock-big.png"), + CVImage("normal", color = False, filename = "images/game_assets/rock-normal.png"), + CVImage("small", color = False, filename = "images/game_assets/rock-small.png") ] self.ships = [ ("ship_off", cv2.imread("images/game_assets/spaceship-off.png", 0)), @@ -35,105 +36,111 @@ 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) +## #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) +## self.mask_frame() + self.color_frame = CVImage("gameio frame") + self.color_frame.from_pil(self.gameio.fetch_sshot()) + self.frame = CVImage("BW frame", self.color_frame.copy()) + self.frame.image = self.frame.convert_color(False) + print(self.frame) self.mask_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.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.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_mask)) + self.frame.image = self.frame.mask(self.score_mask) -## 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 +## 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 = [self.frame.template_detect(i, + self.cv_template_thresh, + self.duplicate_dist_thresh) + for i in self.asteroids] @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) +## 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) + displayable = CVImage("GameModel results", self.color_frame.copy()) #else: # displayable = np.copy(self.color_frame) - for pt, wh, label in rects: - color = { "big": (255, 0, 0), + 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: +## cv2.rectangle(displayable, pt, wh, color, 1) +## cv2.putText(displayable, label, pt, +## cv2.FONT_HERSHEY_PLAIN, +## 1.0, color) + 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) +## 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) +## 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) + displayable.draw_circle(center, radius) + displayable.draw_text(label, center, (255, 255, 0)) cv2.imshow("Results", displayable) cv2.waitKey(0) @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"] +## 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) @@ -144,61 +151,68 @@ class GameModel: ## #return { "matchsets": matchsets, ## # "kp_desc": kp_desc ## # } - ship_rsq = rect_radius_squared(*self.ships[0][1].shape) * 0.85 + ship_r = sqrt(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 + #clusters = pointcluster.cluster_set([k.pt for k in frame_kp], sqrt(ship_rsq)) + #return clusters + 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 +## 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 + return [self.frame.template_detect(a, self.cv_template_thresh, self.duplicate_dist_thresh) for a in self.ships] @with_frame def find_missiles(self): - # Setup SimpleBlobDetector parameters. - params = cv2.SimpleBlobDetector_Params() +## # 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 - # 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 + p = CVImage.blob_params(minThreshold = 10, maxThreshold = 200, + maxArea = 100, + minConvexity = 0.95, + minIntertiaRatio = 0.4) + return self.frame.blob_detect(size=9, params=p) #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] +## 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] def analyse_frame(self): rocks = self.find_asteroids() diff --git a/imagepipeline.py b/imagepipeline.py index bc0011c..1c103ee 100644 --- a/imagepipeline.py +++ b/imagepipeline.py @@ -1,27 +1,108 @@ import cv2 import numpy as np +import typing +import pointcluster + +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 __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) class CVImage: - def __init__(self, label="", img=None, color=False, **kwargs): + """Dummy definition to allow recursive type hints""" + pass + +class CVImage: + def __init__(self, label="", img:np.ndarray=None, color:bool=False, **kwargs): + """You can provide a 'filename' keyword arg to automatically load a file.""" self.label = label self.image = img self.iscolor = color + self._init_kwargs = kwargs if kwargs: - kwargs["color"] = color - self.load(**kwargs) + load_kwargs = dict(kwargs) # copy + load_kwargs["color"] = color # share arg between both functions + self.load(**load_kwargs) - def load(self, filename, color=False, label=None): + def load(self, filename:str, color:bool=False, label:str=None): + """Load an image from file. You can optionally set the 'label' keyword.""" self.image = cv2.imread(filename, int(color)) if label: self.label = label return self + def from_pil(self, pil_img, color=False): + self.image = np.array(pil_img) + self.image = self.image[:, :, ::-1].copy() + self.color = None # force check in cv2.cvtColor + self.image = self.convert_color(color) + + def convert_color(self, color:bool): + if color == self.iscolor: return self.image + 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][0]:rect[1][0], - rect[0][1]:rect[1][1]] + 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)) @@ -33,12 +114,13 @@ class CVImage: cv2.rectangle(mask, *rect, color=mask_color, thickness=cv2.FILLED) return cv2.bitwise_and(self.image, mask) - def sift_clusters(self, cluster_radius): + 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) - def blob_params(self, minThreshold = 10, maxThreshold = 200, + @staticmethod + def blob_params(cls, *, minThreshold = 10, maxThreshold = 200, minArea = None, maxArea = None, minCircularity = None, maxCircularity = None, minConvexity = None, maxConvexity = None, @@ -64,18 +146,69 @@ class CVImage: if maxCircularity: p.maxCircularity = maxCircularity return p - def blob_detect(self, params=None, invert=False): - if params is None: params = self.blob_params() + 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) - return detector.detect(cv2.bitwise_not(self.image) if invert else self.image) + 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]: + h, w = template.image.shape + res = cv2.matchTemplate(self.image, template.image, cv2.TM_CCOEFF_NORMED) + loc = np.where(rec >= 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(*pt, w, h, label=template.label)) 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.iscolor 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.iscolor else 255 + cv2.polylines(self.image, np.int32([points]), closed, color) + + def draw_circle(self, center, radius, thickness = 1): + if color is None: + color = (255, 255, 255) if self.iscolor else 255 + cv2.circle(self.image, np.int32(center), 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() +