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

[In Progress] refactoring and debugging imagepipeline 

poorly defined: how to accurately track if an image is color or not? This is causing a bug in the @with_frame decorator in GameModel.

There's also an antipattern of "cvimage.image = cvimage.mod_function()", to perform in-place modifications to images. This is the place for an ImagePipeline object.
This commit is contained in:
John McCardle 2021-12-27 20:30:40 -05:00
parent fd1767cfe1
commit 1435752a30
2 changed files with 271 additions and 124 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

240
gamemodel.py
View File

@ -5,15 +5,16 @@ from functools import wraps
from utility import * from utility import *
import pointcluster import pointcluster
from imagepipeline import CVImage, ImagePipeline
class GameModel: class GameModel:
"""Platform-independent representation of the game's state.""" """Platform-independent representation of the game's state."""
def __init__(self, io:gameio.AbstractGameIO): def __init__(self, io:gameio.AbstractGameIO):
self.gameio = io self.gameio = io
self.asteroids = [ self.asteroids = [
("big", cv2.imread("images/game_assets/rock-big.png", 0)), CVImage("big", color = False, filename = "images/game_assets/rock-big.png"),
("normal", cv2.imread("images/game_assets/rock-normal.png", 0)), CVImage("normal", color = False, filename = "images/game_assets/rock-normal.png"),
("small", cv2.imread("images/game_assets/rock-small.png", 0)) CVImage("small", color = False, filename = "images/game_assets/rock-small.png")
] ]
self.ships = [ self.ships = [
("ship_off", cv2.imread("images/game_assets/spaceship-off.png", 0)), ("ship_off", cv2.imread("images/game_assets/spaceship-off.png", 0)),
@ -35,49 +36,45 @@ class GameModel:
@wraps(fn) @wraps(fn)
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() ## self.frame = open_cv_image[:, :, ::-1].copy()
self.color_frame = np.copy(self.frame) ## self.color_frame = np.copy(self.frame)
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.image = self.frame.convert_color(False)
print(self.frame)
self.mask_frame() self.mask_frame()
return fn(self, *args, **kwargs) return fn(self, *args, **kwargs)
return inner 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): def mask_frame(self):
self.lives_img = self.frame[self.lives_rect[0][0]:self.lives_rect[0][1], ## 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]] ## self.lives_rect[1][0]:self.lives_rect[1][1]]
lives_mask = np.full(self.frame.shape, 255, dtype=np.uint8) ## 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)
cv2.rectangle(lives_mask, self.lives_img = CVImage("lives", self.frame.snip(self.lives_rect))
*self.lives_rect, self.frame.image = self.frame.mask(self.lives_rect)
color=0, thickness=cv2.FILLED) self.score_img = CVImage("score", self.frame.snip(self.score_mask))
self.frame.image = self.frame.mask(self.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)
## 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): def clear_frame(self):
self.prev_frame = frame self.prev_frame = frame
@ -85,55 +82,65 @@ class GameModel:
@with_frame @with_frame
def find_asteroids(self): def find_asteroids(self):
asteroid_rects = [] ## asteroid_rects = []
for label, a in self.asteroids: ## for label, a in self.asteroids:
h, w = a.shape ## h, w = a.shape
res = cv2.matchTemplate(self.frame, a, cv2.TM_CCOEFF_NORMED) ## res = cv2.matchTemplate(self.frame, a, cv2.TM_CCOEFF_NORMED)
loc = np.where( res >= self.cv_template_thresh) ## loc = np.where( res >= self.cv_template_thresh)
for pt in zip(*loc[::-1]): ## for pt in zip(*loc[::-1]):
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,
self.cv_template_thresh,
self.duplicate_dist_thresh)
for i in self.asteroids]
@with_frame @with_frame
def display_results(self, rects = [], pointsets = [], circles = []): def display_results(self, rects = [], pointsets = [], circles = []):
"""Draws results on the current frame for test purposes.""" """Draws results on the current frame for test purposes."""
displayable = np.copy(self.color_frame) ## displayable = np.copy(self.color_frame)
cv2.rectangle(displayable, *self.lives_rect, (255,255,255), 1) ## cv2.rectangle(displayable, *self.lives_rect, (255,255,255), 1)
cv2.rectangle(displayable, *self.score_rect, (255,255,255), 1) ## cv2.rectangle(displayable, *self.score_rect, (255,255,255), 1)
displayable = CVImage("GameModel results", self.color_frame.copy())
#else: #else:
# displayable = np.copy(self.color_frame) # displayable = np.copy(self.color_frame)
for pt, wh, label in rects: label_color = { "big": (255, 0, 0),
color = { "big": (255, 0, 0),
"normal": (0, 255, 0), "normal": (0, 255, 0),
"small": (0, 0, 255), "small": (0, 0, 255),
"missile": (0, 255, 128), "missile": (0, 255, 128),
"ship_on": (0, 0, 128), "ship_on": (0, 0, 128),
"ship_off": (0, 64, 128)}[label] "ship_off": (0, 64, 128)}
cv2.rectangle(displayable, pt, wh, color, 1) for r in rects:
cv2.putText(displayable, label, pt, ## cv2.rectangle(displayable, pt, wh, color, 1)
cv2.FONT_HERSHEY_PLAIN, ## cv2.putText(displayable, label, pt,
1.0, color) ## cv2.FONT_HERSHEY_PLAIN,
## 1.0, color)
displayable.draw_rect(r, color=label_color[r.label])
for ps in pointsets: for ps in pointsets:
color = (0, 255, 255) ## color = (0, 255, 255)
cv2.polylines(displayable, np.int32([ps]), True, color) ## cv2.polylines(displayable, np.int32([ps]), True, color)
displayable.draw_poly(ps, color=(0, 255, 255))
for center, radius, label in circles: for center, radius, label in circles:
color = (255, 255, 0) ## color = (255, 255, 0)
cv2.circle(displayable, np.int32(center), int(radius), color, 1) ## cv2.circle(displayable, np.int32(center), int(radius), color, 1)
cv2.putText(displayable, label, np.int32(center), ## cv2.putText(displayable, label, np.int32(center),
cv2.FONT_HERSHEY_PLAIN, ## cv2.FONT_HERSHEY_PLAIN,
1.0, color) ## 1.0, color)
displayable.draw_circle(center, radius)
displayable.draw_text(label, center, (255, 255, 0))
cv2.imshow("Results", displayable) cv2.imshow("Results", displayable)
cv2.waitKey(0) cv2.waitKey(0)
@with_frame @with_frame
def frame_sift(self): def frame_sift(self):
sift = cv2.SIFT_create() ## sift = cv2.SIFT_create()
kp_desc = {} # dict of (keypoints, descriptions) for all ship sprites ## kp_desc = {} # dict of (keypoints, descriptions) for all ship sprites
kp_desc["frame"] = sift.detectAndCompute(self.frame, None) ## kp_desc["frame"] = sift.detectAndCompute(self.frame, None)
frame_kp, frame_desc = kp_desc["frame"] ## frame_kp, frame_desc = kp_desc["frame"]
## for label, s in self.ships: ## for label, s in self.ships:
## kp_desc[label] = sift.detectAndCompute(s, None) ## kp_desc[label] = sift.detectAndCompute(s, None)
## bf = cv2.BFMatcher(cv2.NORM_L1, crossCheck=True) ## bf = cv2.BFMatcher(cv2.NORM_L1, crossCheck=True)
@ -144,61 +151,68 @@ class GameModel:
## #return { "matchsets": matchsets, ## #return { "matchsets": matchsets,
## # "kp_desc": kp_desc ## # "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}") #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 return self.frame.sift_clusters(cluster_radius = ship_r)
@with_frame @with_frame
def find_ships(self): def find_ships(self):
ship_rects = [] ## ship_rects = []
for label, a in self.ships: ## for label, a in self.ships:
h, w = a.shape ## h, w = a.shape
res = cv2.matchTemplate(self.frame, a, cv2.TM_CCOEFF_NORMED) ## res = cv2.matchTemplate(self.frame, a, cv2.TM_CCOEFF_NORMED)
loc = np.where( res >= self.cv_template_thresh) ## loc = np.where( res >= self.cv_template_thresh)
for pt in zip(*loc[::-1]): ## for pt in zip(*loc[::-1]):
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]
@with_frame @with_frame
def find_missiles(self): def find_missiles(self):
# Setup SimpleBlobDetector parameters. ## # Setup SimpleBlobDetector parameters.
params = cv2.SimpleBlobDetector_Params() ## 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 p = CVImage.blob_params(minThreshold = 10, maxThreshold = 200,
params.minThreshold = 10; maxArea = 100,
params.maxThreshold = 200; minConvexity = 0.95,
minIntertiaRatio = 0.4)
# Filter by Area. return self.frame.blob_detect(size=9, params=p)
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([]), #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)
#cv2.imshow("keypoints", im_with_keypoints) #cv2.imshow("keypoints", im_with_keypoints)
#cv2.waitKey(0) #cv2.waitKey(0)
s = 9 # pixels for the missile ## s = 9 # pixels for the missile
rect_tuple = lambda pt: ((int(pt[0]-s/2),int(pt[1]-s/2)), ## rect_tuple = lambda pt: ((int(pt[0]-s/2),int(pt[1]-s/2)),
(int(pt[0]+s/2), int(pt[1]+s/2)), ## (int(pt[0]+s/2), int(pt[1]+s/2)),
"missile") ## "missile")
return [rect_tuple(k.pt) for k in keypoints] ## return [rect_tuple(k.pt) for k in keypoints]
def analyse_frame(self): def analyse_frame(self):
rocks = self.find_asteroids() rocks = self.find_asteroids()

155
imagepipeline.py
View File

@ -1,27 +1,108 @@
import cv2 import cv2
import numpy as np 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"<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:
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.label = label
self.image = img self.image = img
self.iscolor = color self.iscolor = color
self._init_kwargs = kwargs
if kwargs: if kwargs:
kwargs["color"] = color load_kwargs = dict(kwargs) # copy
self.load(**kwargs) 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)) self.image = cv2.imread(filename, int(color))
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):
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"<CVImage label={repr(self.label)}, image={self.image.shape} px, iscolor={self.iscolor}{kwargstr}>"
def copy(self): def copy(self):
return np.copy(self.image) return np.copy(self.image)
def snip(self, rect): def snip(self, rect):
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))
return self.image[rect[0][0]:rect[1][0], return self.image[rect[0][1]:rect[0][1]+rect[1][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): 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))
@ -33,12 +114,13 @@ class CVImage:
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)
def sift_clusters(self, cluster_radius): def sift_clusters(self, cluster_radius) -> pointcluster.PointCluster:
sift = cv2.SIFT_create() sift = cv2.SIFT_create()
keypoints, descriptions = sift.detectAndCompute(self.image, None) keypoints, descriptions = sift.detectAndCompute(self.image, None)
return pointcluster.cluster_set([k.pt for k in keypoints], cluster_radius) 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, minArea = None, maxArea = None,
minCircularity = None, maxCircularity = None, minCircularity = None, maxCircularity = None,
minConvexity = None, maxConvexity = None, minConvexity = None, maxConvexity = None,
@ -64,18 +146,69 @@ class CVImage:
if maxCircularity: p.maxCircularity = maxCircularity if maxCircularity: p.maxCircularity = maxCircularity
return p return p
def blob_detect(self, params=None, invert=False): def blob_detect(self, size:int, params=None, invert:bool=False, label:str=None) -> typing.List[Rect]:
if params is None: params = self.blob_params() if params is None: params = CVImage.blob_params()
detector = cv2.SimpleBlobDetector_create(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): def show(self, delay=0):
cv2.imshow(self.label, self.image) cv2.imshow(self.label, self.image)
cv2.waitKey(delay) 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: class ImagePipeline:
def __init__(self): def __init__(self):
pass 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()