diff --git a/gamemodel.py b/gamemodel.py index e05ee11..7e4a78b 100644 --- a/gamemodel.py +++ b/gamemodel.py @@ -12,13 +12,13 @@ class GameModel: def __init__(self, io:gameio.AbstractGameIO): self.gameio = io self.asteroids = [ - 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") + 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 @@ -37,19 +37,16 @@ class GameModel: 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) + sshot = self.gameio.fetch_sshot() + open_cv_image = np.array(sshot) + # Convert RGB to BGR + array = open_cv_image[:, :, ::-1].copy() + self.color_frame = CVImage("gameio frame", np.copy(array)) ## 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() + print(self.frame) return fn(self, *args, **kwargs) return inner @@ -73,8 +70,8 @@ class GameModel: 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) + self.score_img = CVImage("score", self.frame.snip(self.score_rect)) + self.frame.image = self.frame.mask(self.score_rect) def clear_frame(self): self.prev_frame = frame @@ -91,10 +88,13 @@ class GameModel: ## 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] + 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 = []): @@ -130,9 +130,8 @@ class GameModel: ## 1.0, color) displayable.draw_circle(center, radius) displayable.draw_text(label, center, (255, 255, 0)) - - cv2.imshow("Results", displayable) - cv2.waitKey(0) + + displayable.show() @with_frame def frame_sift(self): @@ -151,7 +150,7 @@ class GameModel: ## #return { "matchsets": matchsets, ## # "kp_desc": kp_desc ## # } - ship_r = sqrt(rect_radius_squared(*self.ships[0][1].shape) * 0.85) + ship_r = sqrt(rect_radius_squared(*self.ships[0].image.shape[:2]) * 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 @@ -168,7 +167,13 @@ class GameModel: ## 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] + 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): @@ -202,7 +207,7 @@ class GameModel: p = CVImage.blob_params(minThreshold = 10, maxThreshold = 200, maxArea = 100, minConvexity = 0.95, - minIntertiaRatio = 0.4) + minInertiaRatio = 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) @@ -278,9 +283,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 index 1c103ee..c494078 100644 --- a/imagepipeline.py +++ b/imagepipeline.py @@ -3,97 +3,42 @@ 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) +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, color:bool=False, **kwargs): + 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.iscolor = color self._init_kwargs = kwargs if kwargs: load_kwargs = dict(kwargs) # copy - load_kwargs["color"] = color # share arg between both functions self.load(**load_kwargs) - def load(self, filename:str, color:bool=False, label:str=None): + @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, int(color)) + self.image = cv2.imread(filename) 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"" + return f"" def copy(self): return np.copy(self.image) @@ -107,9 +52,9 @@ class CVImage: 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.iscolor else 0 + mask_color = (0, 0, 0) if self.is_color else 0 if nonmask_color is None: - nonnmask_color = (255, 255, 255) if self.iscolor else 255 + 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) @@ -120,7 +65,7 @@ class CVImage: return pointcluster.cluster_set([k.pt for k in keypoints], cluster_radius) @staticmethod - def blob_params(cls, *, minThreshold = 10, maxThreshold = 200, + def blob_params(*, minThreshold = 10, maxThreshold = 200, minArea = None, maxArea = None, minCircularity = None, maxCircularity = None, minConvexity = None, maxConvexity = None, @@ -137,7 +82,7 @@ class CVImage: if minConvexity: p.minConvexity = minConvexity if maxConvexity: p.maxConvexity = maxConvexity if minInertiaRatio or maxInertiaRatio: - p.filterByInertiaRatio = True + p.filterByInertia = True if minInertiaRatio: p.minInertiaRatio = minInertiaRatio if maxInertiaRatio: p.maxInertiaRatio = maxInertiaRatio if minCircularity or maxCircularity: @@ -159,14 +104,20 @@ class CVImage: return rects def template_detect(self, template:CVImage, threshold:int, dupe_spacing:int) -> typing.List[Rect]: - h, w = template.image.shape + 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(rec >= threshold) + 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(*pt, w, h, label=template.label)) + 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) @@ -174,20 +125,20 @@ class CVImage: 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 + 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.iscolor else 255 + 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, thickness = 1): + def draw_circle(self, center, radius, color=None, 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) + 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) 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)