john
/
asteroid-automator
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Mostly functional again after imagepipeline module refactor

This commit is contained in:
John McCardle 2021-12-28 19:46:52 -05:00
parent 1435752a30
commit f02313da39
3 changed files with 118 additions and 105 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

58
gamemodel.py
View File

@ -12,13 +12,13 @@ class GameModel:
def __init__(self, io:gameio.AbstractGameIO): def __init__(self, io:gameio.AbstractGameIO):
self.gameio = io self.gameio = io
self.asteroids = [ self.asteroids = [
CVImage("big", color = False, filename = "images/game_assets/rock-big.png"), CVImage("big", filename = "images/game_assets/rock-big.png"),
CVImage("normal", color = False, filename = "images/game_assets/rock-normal.png"), CVImage("normal", filename = "images/game_assets/rock-normal.png"),
CVImage("small", color = False, filename = "images/game_assets/rock-small.png") CVImage("small", filename = "images/game_assets/rock-small.png")
] ]
self.ships = [ self.ships = [
("ship_off", cv2.imread("images/game_assets/spaceship-off.png", 0)), CVImage("ship_off", filename = "images/game_assets/spaceship-off.png"),
("ship_on", cv2.imread("images/game_assets/spaceship-on.png", 0)) CVImage("ship_on", filename = "images/game_assets/spaceship-on.png")
] ]
#self.missile = ("missile", cv2.imread("images/game_assets/missile.png", 0)) #self.missile = ("missile", cv2.imread("images/game_assets/missile.png", 0))
self.frame = None self.frame = None
@ -37,19 +37,16 @@ class GameModel:
def inner(self, *args, **kwargs): def inner(self, *args, **kwargs):
if self.frame is None: if self.frame is None:
## #print("Fetching frame.") ## #print("Fetching frame.")
## sshot = self.gameio.fetch_sshot() sshot = self.gameio.fetch_sshot()
## open_cv_image = np.array(sshot) open_cv_image = np.array(sshot)
## # Convert RGB to BGR # Convert RGB to BGR
## self.frame = open_cv_image[:, :, ::-1].copy() array = open_cv_image[:, :, ::-1].copy()
## self.color_frame = np.copy(self.frame) self.color_frame = CVImage("gameio frame", np.copy(array))
## self.frame = cv2.cvtColor(self.frame, cv2.COLOR_BGR2GRAY) ## 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 = CVImage("BW frame", self.color_frame.copy())
self.frame.image = self.frame.convert_color(False) self.frame.image = self.frame.convert_color(False)
print(self.frame)
self.mask_frame() self.mask_frame()
print(self.frame)
return fn(self, *args, **kwargs) return fn(self, *args, **kwargs)
return inner return inner
@ -73,8 +70,8 @@ class GameModel:
self.lives_img = CVImage("lives", self.frame.snip(self.lives_rect)) self.lives_img = CVImage("lives", self.frame.snip(self.lives_rect))
self.frame.image = self.frame.mask(self.lives_rect) self.frame.image = self.frame.mask(self.lives_rect)
self.score_img = CVImage("score", self.frame.snip(self.score_mask)) self.score_img = CVImage("score", self.frame.snip(self.score_rect))
self.frame.image = self.frame.mask(self.score_mask) self.frame.image = self.frame.mask(self.score_rect)
def clear_frame(self): def clear_frame(self):
self.prev_frame = frame 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: ## 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)) ## asteroid_rects.append((pt, (pt[0] + w, pt[1] + h), label))
## return asteroid_rects ## return asteroid_rects
results = [self.frame.template_detect(i, results = []
for a in self.asteroids:
r = self.frame.template_detect(a,
self.cv_template_thresh, self.cv_template_thresh,
self.duplicate_dist_thresh) self.duplicate_dist_thresh)
for i in self.asteroids] results.extend(r)
return results
@with_frame @with_frame
def display_results(self, rects = [], pointsets = [], circles = []): def display_results(self, rects = [], pointsets = [], circles = []):
@ -131,8 +131,7 @@ class GameModel:
displayable.draw_circle(center, radius) displayable.draw_circle(center, radius)
displayable.draw_text(label, center, (255, 255, 0)) displayable.draw_text(label, center, (255, 255, 0))
cv2.imshow("Results", displayable) displayable.show()
cv2.waitKey(0)
@with_frame @with_frame
def frame_sift(self): def frame_sift(self):
@ -151,7 +150,7 @@ class GameModel:
## #return { "matchsets": matchsets, ## #return { "matchsets": matchsets,
## # "kp_desc": kp_desc ## # "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}") #print(f"max radius^2: {ship_rsq}")
#clusters = pointcluster.cluster_set([k.pt for k in frame_kp], sqrt(ship_rsq)) #clusters = pointcluster.cluster_set([k.pt for k in frame_kp], sqrt(ship_rsq))
#return clusters #return clusters
@ -168,7 +167,13 @@ class GameModel:
## if not ship_rects or squared_distance(ship_rects[-1][0], pt) > self.duplicate_dist_thresh: ## 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)) ## ship_rects.append((pt, (pt[0] + w, pt[1] + h), label))
## return ship_rects ## 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 @with_frame
def find_missiles(self): def find_missiles(self):
@ -202,7 +207,7 @@ class GameModel:
p = CVImage.blob_params(minThreshold = 10, maxThreshold = 200, p = CVImage.blob_params(minThreshold = 10, maxThreshold = 200,
maxArea = 100, maxArea = 100,
minConvexity = 0.95, minConvexity = 0.95,
minIntertiaRatio = 0.4) minInertiaRatio = 0.4)
return self.frame.blob_detect(size=9, params=p) return self.frame.blob_detect(size=9, params=p)
#im_with_keypoints = cv2.drawKeypoints(self.frame, keypoints, np.array([]), #im_with_keypoints = cv2.drawKeypoints(self.frame, keypoints, np.array([]),
# (0,0,255), cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS) # (0,0,255), cv2.DRAW_MATCHES_FLAGS_DRAW_RICH_KEYPOINTS)
@ -278,9 +283,12 @@ if __name__ == '__main__':
ship_results = gm.find_ships() ship_results = gm.find_ships()
polygons = [c.points for c in s_results] polygons = [c.points for c in s_results]
##circles = [(c.center, c.max_distance, f"cluster_{i}") for i, c in enumerate(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() missile_results = gm.find_missiles()
##m_circles = [(pt, 10, f"missile_{i}") for i, pt in enumerate(missiles)] ##m_circles = [(pt, 10, f"missile_{i}") for i, pt in enumerate(missiles)]
##pprint(a_results+ship_results+missile_results) ##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() gm.analyse_frame()

103
imagepipeline.py
View File

@ -3,97 +3,42 @@ import numpy as np
import typing import typing
import pointcluster import pointcluster
class Rect: from shapes import Rect
def __init__(self, *args, label=None, **kwargs): from utility import *
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"<Rect label={repr(self.label)}, (({self.x}, {self.y}), ({self.w}, {self.h}))>"
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: class CVImage:
"""Dummy definition to allow recursive type hints""" """Dummy definition to allow recursive type hints"""
pass pass
class CVImage: 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.""" """You can provide a 'filename' keyword arg to automatically load a file."""
self.label = label self.label = label
self.image = img self.image = img
self.iscolor = color
self._init_kwargs = kwargs self._init_kwargs = kwargs
if kwargs: if kwargs:
load_kwargs = dict(kwargs) # copy load_kwargs = dict(kwargs) # copy
load_kwargs["color"] = color # share arg between both functions
self.load(**load_kwargs) 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.""" """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 if label: self.label = label
return self 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): 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) return cv2.cvtColor(self.image, cv2.COLOR_GRAY2BGR if color else cv2.COLOR_BGR2GRAY)
def __repr__(self): def __repr__(self):
if self._init_kwargs: if self._init_kwargs:
kwargstr = ", " + ", ".join([f"{k}={repr(self._init_kwargs[k])}" for k in self._init_kwargs]) kwargstr = ", " + ", ".join([f"{k}={repr(self._init_kwargs[k])}" for k in self._init_kwargs])
else: else:
kwargstr = '' kwargstr = ''
return f"<CVImage label={repr(self.label)}, image={self.image.shape} px, iscolor={self.iscolor}{kwargstr}>" return f"<CVImage label={repr(self.label)}, image={self.image.shape} px, is_color={self.is_color}{kwargstr}>"
def copy(self): def copy(self):
return np.copy(self.image) return np.copy(self.image)
@ -107,9 +52,9 @@ class CVImage:
def mask(self, rect, mask_color=None, nonmask_color=None): 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)) assert all((len(rect)==2, len(rect[0])==2, len(rect[1])==2)) #((x,y),(w,h))
if mask_color is None: 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: 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) mask = np.full(self.image.shape, nonmask_color, dtype=np.uint8)
cv2.rectangle(mask, *rect, color=mask_color, thickness=cv2.FILLED) cv2.rectangle(mask, *rect, color=mask_color, thickness=cv2.FILLED)
return cv2.bitwise_and(self.image, mask) 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) return pointcluster.cluster_set([k.pt for k in keypoints], cluster_radius)
@staticmethod @staticmethod
def blob_params(cls, *, minThreshold = 10, maxThreshold = 200, def blob_params(*, minThreshold = 10, maxThreshold = 200,
minArea = None, maxArea = None, minArea = None, maxArea = None,
minCircularity = None, maxCircularity = None, minCircularity = None, maxCircularity = None,
minConvexity = None, maxConvexity = None, minConvexity = None, maxConvexity = None,
@ -137,7 +82,7 @@ class CVImage:
if minConvexity: p.minConvexity = minConvexity if minConvexity: p.minConvexity = minConvexity
if maxConvexity: p.maxConvexity = maxConvexity if maxConvexity: p.maxConvexity = maxConvexity
if minInertiaRatio or maxInertiaRatio: if minInertiaRatio or maxInertiaRatio:
p.filterByInertiaRatio = True p.filterByInertia = True
if minInertiaRatio: p.minInertiaRatio = minInertiaRatio if minInertiaRatio: p.minInertiaRatio = minInertiaRatio
if maxInertiaRatio: p.maxInertiaRatio = maxInertiaRatio if maxInertiaRatio: p.maxInertiaRatio = maxInertiaRatio
if minCircularity or maxCircularity: if minCircularity or maxCircularity:
@ -159,14 +104,20 @@ class CVImage:
return rects return rects
def template_detect(self, template:CVImage, threshold:int, dupe_spacing:int) -> typing.List[Rect]: 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 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) res = cv2.matchTemplate(self.image, template.image, cv2.TM_CCOEFF_NORMED)
loc = np.where(rec >= threshold) loc = np.where(res >= threshold)
rects = [] rects = []
for pt in zip(*loc[::-1]): for pt in zip(*loc[::-1]):
if len(rects) > 0: if len(rects) > 0:
if squared_distance(rects[-1][0], pt) < dupe_spacing: continue 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): def show(self, delay=0):
cv2.imshow(self.label, self.image) 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): def draw_rect(self, rect:Rect, color=None, text_color=None, text:bool=True, thickness=1):
if color is 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.rectangle(self.image, rect.point, rect.point2, color, thickness) cv2.rectangle(self.image, rect.point, rect.point2, color, thickness)
if text: if text:
self.draw_text(rect.label, rect.point, text_color if text_color else color) 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): def draw_poly(self, points:typing.List[typing.Tuple], closed=True, color=None):
if color is 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) 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: if color is None:
color = (255, 255, 255) if self.iscolor else 255 color = (255, 255, 255) if self.is_color else 255
cv2.circle(self.image, np.int32(center), radius, color, thickness) cv2.circle(self.image, np.int32(center), np.int32(radius), color, thickness)
def draw_text(self, text, point, color): def draw_text(self, text, point, color):
cv2.putText(self.image, text, np.int32(point), cv2.FONT_HERSHEY_PLAIN, 1.0, color) cv2.putText(self.image, text, np.int32(point), cv2.FONT_HERSHEY_PLAIN, 1.0, color)

54
shapes.py Normal file
View File

@ -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"<Rect label={repr(self.label)}, (({self.x}, {self.y}), ({self.w}, {self.h}))>"
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)