[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:
parent
fd1767cfe1
commit
1435752a30
240
gamemodel.py
240
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,49 +36,45 @@ 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)
|
||||
## 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)
|
||||
|
||||
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)
|
||||
|
||||
## print("Displaying images for testing purposes")
|
||||
## cv2.imshow("Original", self.color_frame)
|
||||
## cv2.waitKey(0)
|
||||
## cv2.imshow("Masked", self.frame)
|
||||
## cv2.waitKey(0)
|
||||
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)
|
||||
|
||||
def clear_frame(self):
|
||||
self.prev_frame = frame
|
||||
|
|
@ -85,55 +82,65 @@ class GameModel:
|
|||
|
||||
@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()
|
||||
|
|
|
|||
155
imagepipeline.py
155
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"<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:
|
||||
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"<CVImage label={repr(self.label)}, image={self.image.shape} px, iscolor={self.iscolor}{kwargstr}>"
|
||||
|
||||
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()
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue