368 lines
15 KiB
Python
368 lines
15 KiB
Python
import cv2 as cv
|
|
import numpy as np
|
|
from hsvfilter import HsvFilter
|
|
|
|
|
|
class Vision:
|
|
# constants
|
|
TRACKBAR_WINDOW = "Trackbars"
|
|
|
|
# properties
|
|
needle_img = None
|
|
needle_w = 0
|
|
needle_h = 0
|
|
method = None
|
|
|
|
# constructor
|
|
def __init__(self, method=cv.TM_CCOEFF_NORMED):
|
|
# load the image we're trying to match
|
|
# https://docs.opencv.org/4.2.0/d4/da8/group__imgcodecs.html
|
|
# self.needle_img = cv.imread("dig/wtf.jpg", cv.IMREAD_UNCHANGED)
|
|
|
|
# Save the dimensions of the needle image
|
|
#self.needle_w = self.needle_img.shape[1]
|
|
#self.needle_h = self.needle_img.shape[0]
|
|
|
|
# There are 6 methods to choose from:
|
|
# TM_CCOEFF, TM_CCOEFF_NORMED, TM_CCORR, TM_CCORR_NORMED, TM_SQDIFF, TM_SQDIFF_NORMED
|
|
self.method = method
|
|
|
|
|
|
|
|
def find_by_mask_and_validate(self, haystack_img, needle_img, needle_mask, max_results=5):
|
|
# run the OpenCV algorithm
|
|
needle_w = needle_img.shape[1]
|
|
needle_h = needle_img.shape[0]
|
|
|
|
result = cv.matchTemplate(haystack_img, needle_img, cv.TM_CCORR_NORMED, None, needle_mask)
|
|
cv.normalize(result, result, 0, 1, cv.NORM_MINMAX, -1)
|
|
|
|
find_num = 20
|
|
idx_1d = np.argpartition(result.flatten(), -find_num)[-find_num:]
|
|
idx_2d = np.unravel_index(idx_1d, result.shape)
|
|
|
|
tresh_coord = []
|
|
for i in range(0, len(idx_2d[0]), 1):
|
|
y = int(idx_2d[0][i])
|
|
x = int(idx_2d[1][i])
|
|
res1 = result[y][x]
|
|
tresh_coord.append([res1, [y, x]])
|
|
|
|
tresh_coord.sort(reverse=True)
|
|
|
|
rectangles = []
|
|
for tc in tresh_coord:
|
|
y = int(tc[1][0])
|
|
x = int(tc[1][1])
|
|
#for i in range(0, len(idx_2d[0]), 1):
|
|
#y = int(idx_2d[0][i])
|
|
#x = int(idx_2d[1][i])
|
|
rect = [x, y, needle_w, needle_h]
|
|
# Add every box to the list twice in order to retain single (non-overlapping) boxes
|
|
rectangles.append(rect)
|
|
rectangles.append(rect)
|
|
|
|
rectangles, weights = cv.groupRectangles(rectangles, groupThreshold=1, eps=0.5)
|
|
|
|
keep_rects = []
|
|
for rect in rectangles:
|
|
w = rect[0]
|
|
h = rect[1]
|
|
x = rect[2] + w
|
|
y = rect[3] + h
|
|
screenshot_pos = haystack_img[h:y, w:x] # (w, h, x+w, y+h)
|
|
result2 = cv.matchTemplate(screenshot_pos, needle_img, cv.TM_CCOEFF_NORMED)
|
|
_minVal2, _maxVal2, minLoc2, maxLoc2 = cv.minMaxLoc(result2, None)
|
|
#screenshot_pos_img = self.draw_rectangles(screenshot_pos, rectangles)
|
|
#cv.imshow("screenshot_pos", screenshot_pos)
|
|
#cv.waitKey(150)
|
|
if _maxVal2 >= 0.85:
|
|
keep_rects.append(rect)
|
|
if len(keep_rects) >= 5:
|
|
return keep_rects
|
|
|
|
if len(keep_rects) > max_results:
|
|
keep_rects = keep_rects[:max_results]
|
|
|
|
return keep_rects
|
|
|
|
def find_comb(self, haystack_img, needle_img, threshold=0.5, max_results=10, normalize=False, mask=None):
|
|
# run the OpenCV algorithm
|
|
needle_w = needle_img.shape[1]
|
|
needle_h = needle_img.shape[0]
|
|
|
|
if mask is not None:
|
|
result = cv.matchTemplate(haystack_img, needle_img, cv.TM_CCORR_NORMED, None, mask)
|
|
_minVal, _maxVal, minLoc, maxLoc = cv.minMaxLoc(result, None)
|
|
else:
|
|
result = cv.matchTemplate(haystack_img, needle_img, self.method)
|
|
|
|
if normalize:
|
|
cv.normalize(result, result, 0, 1, cv.NORM_MINMAX, -1)
|
|
# Get the all the positions from the match result that exceed our threshold
|
|
_minVal, _maxVal, minLoc, maxLoc = cv.minMaxLoc(result, None)
|
|
if max_results == 1 and _maxVal >= threshold:
|
|
locations = []
|
|
locations.append(maxLoc)
|
|
else:
|
|
locations = np.where(result >= threshold)
|
|
locations = list(zip(*locations[::-1]))
|
|
# print(locations)
|
|
|
|
|
|
#_minVal, _maxVal, minLoc, maxLoc = cv.minMaxLoc(result, None)
|
|
|
|
# if we found no results, return now. this reshape of the empty array allows us to
|
|
# concatenate together results without causing an error
|
|
if not locations:
|
|
return np.array([], dtype=np.int32).reshape(0, 4)
|
|
#while len(locations) > 1000:
|
|
# threshold = threshold + 0.01
|
|
# locations = np.where(result >= threshold)
|
|
# locations = list(zip(*locations[::-1]))
|
|
# print("modified treshhold to:" + str(threshold))
|
|
# print("actual locations:" + str(len(locations)))
|
|
|
|
if len(locations) > 5000:
|
|
return np.array([], dtype=np.int32).reshape(0, 4)
|
|
|
|
# You'll notice a lot of overlapping rectangles get drawn. We can eliminate those redundant
|
|
# locations by using groupRectangles().
|
|
# First we need to create the list of [x, y, w, h] rectangles
|
|
rectangles = []
|
|
for loc in locations:
|
|
rect = [int(loc[0]), int(loc[1]), needle_w, needle_h]
|
|
# Add every box to the list twice in order to retain single (non-overlapping) boxes
|
|
rectangles.append(rect)
|
|
rectangles.append(rect)
|
|
# Apply group rectangles.
|
|
# The groupThreshold parameter should usually be 1. If you put it at 0 then no grouping is
|
|
# done. If you put it at 2 then an object needs at least 3 overlapping rectangles to appear
|
|
# in the result. I've set eps to 0.5, which is:
|
|
# "Relative difference between sides of the rectangles to merge them into a group."
|
|
rectangles, weights = cv.groupRectangles(rectangles, groupThreshold=1, eps=0.5)
|
|
# print(rectangles)
|
|
|
|
# for performance reasons, return a limited number of results.
|
|
# these aren't necessarily the best results.
|
|
if len(rectangles) > max_results:
|
|
#print('Warning: too many results, raise the threshold.')
|
|
rectangles = rectangles[:max_results]
|
|
|
|
return rectangles
|
|
|
|
def find(self, haystack_img, needle_img, threshold=0.5, max_results=10, normalize=False, mask=None):
|
|
# run the OpenCV algorithm
|
|
needle_w = needle_img.shape[1]
|
|
needle_h = needle_img.shape[0]
|
|
|
|
if mask is not None:
|
|
result = cv.matchTemplate(haystack_img, needle_img, cv.TM_CCORR_NORMED, None, mask)
|
|
_minVal, _maxVal, minLoc, maxLoc = cv.minMaxLoc(result, None)
|
|
else:
|
|
result = cv.matchTemplate(haystack_img, needle_img, self.method)
|
|
|
|
if normalize:
|
|
cv.normalize(result, result, 0, 1, cv.NORM_MINMAX, -1)
|
|
# Get the all the positions from the match result that exceed our threshold
|
|
# _minVal, _maxVal, minLoc, maxLoc = cv.minMaxLoc(result, None)
|
|
|
|
locations = np.where(result >= threshold)
|
|
locations = list(zip(*locations[::-1]))
|
|
# print(locations)
|
|
|
|
|
|
#_minVal, _maxVal, minLoc, maxLoc = cv.minMaxLoc(result, None)
|
|
|
|
# if we found no results, return now. this reshape of the empty array allows us to
|
|
# concatenate together results without causing an error
|
|
if not locations:
|
|
return np.array([], dtype=np.int32).reshape(0, 4)
|
|
#while len(locations) > 1000:
|
|
# threshold = threshold + 0.01
|
|
# locations = np.where(result >= threshold)
|
|
# locations = list(zip(*locations[::-1]))
|
|
# print("modified treshhold to:" + str(threshold))
|
|
# print("actual locations:" + str(len(locations)))
|
|
|
|
if len(locations) > 5000:
|
|
return np.array([], dtype=np.int32).reshape(0, 4)
|
|
|
|
# You'll notice a lot of overlapping rectangles get drawn. We can eliminate those redundant
|
|
# locations by using groupRectangles().
|
|
# First we need to create the list of [x, y, w, h] rectangles
|
|
rectangles = []
|
|
for loc in locations:
|
|
rect = [int(loc[0]), int(loc[1]), needle_w, needle_h]
|
|
# Add every box to the list twice in order to retain single (non-overlapping) boxes
|
|
rectangles.append(rect)
|
|
rectangles.append(rect)
|
|
# Apply group rectangles.
|
|
# The groupThreshold parameter should usually be 1. If you put it at 0 then no grouping is
|
|
# done. If you put it at 2 then an object needs at least 3 overlapping rectangles to appear
|
|
# in the result. I've set eps to 0.5, which is:
|
|
# "Relative difference between sides of the rectangles to merge them into a group."
|
|
rectangles, weights = cv.groupRectangles(rectangles, groupThreshold=1, eps=0.5)
|
|
# print(rectangles)
|
|
|
|
# for performance reasons, return a limited number of results.
|
|
# these aren't necessarily the best results.
|
|
if len(rectangles) > max_results:
|
|
#print('Warning: too many results, raise the threshold.')
|
|
rectangles = rectangles[:max_results]
|
|
|
|
return rectangles
|
|
|
|
# given a list of [x, y, w, h] rectangles returned by find(), convert those into a list of
|
|
# [x, y] positions in the center of those rectangles where we can click on those found items
|
|
def get_click_points(self, rectangles):
|
|
points = []
|
|
|
|
# Loop over all the rectangles
|
|
for (x, y, w, h) in rectangles:
|
|
# Determine the center position
|
|
center_x = x + int(w / 2)
|
|
center_y = y + int(h / 2)
|
|
# Save the points
|
|
points.append((center_x, center_y))
|
|
|
|
return points
|
|
|
|
# given a list of [x, y, w, h] rectangles and a canvas image to draw on, return an image with
|
|
# all of those rectangles drawn
|
|
def draw_rectangles(self, haystack_img, rectangles):
|
|
# these colors are actually BGR
|
|
line_color = (0, 255, 0)
|
|
line_type = cv.LINE_4
|
|
pic = None
|
|
for (x, y, w, h) in rectangles:
|
|
# determine the box positions
|
|
top_left = (x, y)
|
|
bottom_right = (x + w, y + h)
|
|
# draw the box
|
|
cv.rectangle(haystack_img, top_left, bottom_right, line_color, lineType=line_type)
|
|
|
|
#pic = haystack_img[y:y + h, x:x + w]
|
|
|
|
return haystack_img
|
|
|
|
def draw_display_picture(self, haystack_img, rectangles):
|
|
|
|
pic = None
|
|
for (x, y, w, h) in rectangles:
|
|
pic = haystack_img[y:y + h, x:x + w]
|
|
|
|
# scale_percent = 500 # percent of original size
|
|
# width = int(pic.shape[1] * scale_percent / 100)
|
|
# height = int(pic.shape[0] * scale_percent / 100)
|
|
# dim = (width, height)
|
|
# resize image
|
|
|
|
# resized_pic = cv.resize(pic, dim, interpolation=cv.INTER_AREA)
|
|
pil_image = np.array(pic)
|
|
# pil_image = Image.fromarray(cv.cvtColor(pic, cv.COLOR_BGR2RGB))
|
|
# pil_image = Image.)
|
|
return pil_image
|
|
|
|
# given a list of [x, y] positions and a canvas image to draw on, return an image with all
|
|
# of those click points drawn on as crosshairs
|
|
def draw_crosshairs(self, haystack_img, points):
|
|
# these colors are actually BGR
|
|
marker_color = (255, 0, 255)
|
|
marker_type = cv.MARKER_CROSS
|
|
|
|
for (center_x, center_y) in points:
|
|
# draw the center point
|
|
cv.drawMarker(haystack_img, (center_x, center_y), marker_color, marker_type)
|
|
|
|
return haystack_img
|
|
|
|
# create gui window with controls for adjusting arguments in real-time
|
|
def init_control_gui(self):
|
|
cv.namedWindow(self.TRACKBAR_WINDOW, cv.WINDOW_NORMAL)
|
|
cv.resizeWindow(self.TRACKBAR_WINDOW, 350, 700)
|
|
|
|
# required callback. we'll be using getTrackbarPos() to do lookups
|
|
# instead of using the callback.
|
|
def nothing(position):
|
|
pass
|
|
|
|
# create trackbars for bracketing.
|
|
# OpenCV scale for HSV is H: 0-179, S: 0-255, V: 0-255
|
|
cv.createTrackbar('HMin', self.TRACKBAR_WINDOW, 0, 179, nothing)
|
|
cv.createTrackbar('SMin', self.TRACKBAR_WINDOW, 0, 255, nothing)
|
|
cv.createTrackbar('VMin', self.TRACKBAR_WINDOW, 0, 255, nothing)
|
|
cv.createTrackbar('HMax', self.TRACKBAR_WINDOW, 0, 179, nothing)
|
|
cv.createTrackbar('SMax', self.TRACKBAR_WINDOW, 0, 255, nothing)
|
|
cv.createTrackbar('VMax', self.TRACKBAR_WINDOW, 0, 255, nothing)
|
|
# Set default value for Max HSV trackbars
|
|
cv.setTrackbarPos('HMax', self.TRACKBAR_WINDOW, 179)
|
|
cv.setTrackbarPos('SMax', self.TRACKBAR_WINDOW, 255)
|
|
cv.setTrackbarPos('VMax', self.TRACKBAR_WINDOW, 255)
|
|
|
|
# trackbars for increasing/decreasing saturation and value
|
|
cv.createTrackbar('SAdd', self.TRACKBAR_WINDOW, 0, 255, nothing)
|
|
cv.createTrackbar('SSub', self.TRACKBAR_WINDOW, 0, 255, nothing)
|
|
cv.createTrackbar('VAdd', self.TRACKBAR_WINDOW, 0, 255, nothing)
|
|
cv.createTrackbar('VSub', self.TRACKBAR_WINDOW, 0, 255, nothing)
|
|
|
|
# returns an HSV filter object based on the control GUI values
|
|
def get_hsv_filter_from_controls(self):
|
|
# Get current positions of all trackbars
|
|
hsv_filter = HsvFilter()
|
|
hsv_filter.hMin = cv.getTrackbarPos('HMin', self.TRACKBAR_WINDOW)
|
|
hsv_filter.sMin = cv.getTrackbarPos('SMin', self.TRACKBAR_WINDOW)
|
|
hsv_filter.vMin = cv.getTrackbarPos('VMin', self.TRACKBAR_WINDOW)
|
|
hsv_filter.hMax = cv.getTrackbarPos('HMax', self.TRACKBAR_WINDOW)
|
|
hsv_filter.sMax = cv.getTrackbarPos('SMax', self.TRACKBAR_WINDOW)
|
|
hsv_filter.vMax = cv.getTrackbarPos('VMax', self.TRACKBAR_WINDOW)
|
|
hsv_filter.sAdd = cv.getTrackbarPos('SAdd', self.TRACKBAR_WINDOW)
|
|
hsv_filter.sSub = cv.getTrackbarPos('SSub', self.TRACKBAR_WINDOW)
|
|
hsv_filter.vAdd = cv.getTrackbarPos('VAdd', self.TRACKBAR_WINDOW)
|
|
hsv_filter.vSub = cv.getTrackbarPos('VSub', self.TRACKBAR_WINDOW)
|
|
return hsv_filter
|
|
|
|
# given an image and an HSV filter, apply the filter and return the resulting image.
|
|
# if a filter is not supplied, the control GUI trackbars will be used
|
|
def apply_hsv_filter(self, original_image, hsv_filter=None):
|
|
# convert image to HSV
|
|
hsv = cv.cvtColor(original_image, cv.COLOR_BGR2HSV)
|
|
|
|
# if we haven't been given a defined filter, use the filter values from the GUI
|
|
if not hsv_filter:
|
|
hsv_filter = self.get_hsv_filter_from_controls()
|
|
|
|
# add/subtract saturation and value
|
|
h, s, v = cv.split(hsv)
|
|
s = self.shift_channel(s, hsv_filter.sAdd)
|
|
s = self.shift_channel(s, -hsv_filter.sSub)
|
|
v = self.shift_channel(v, hsv_filter.vAdd)
|
|
v = self.shift_channel(v, -hsv_filter.vSub)
|
|
hsv = cv.merge([h, s, v])
|
|
|
|
# Set minimum and maximum HSV values to display
|
|
lower = np.array([hsv_filter.hMin, hsv_filter.sMin, hsv_filter.vMin])
|
|
upper = np.array([hsv_filter.hMax, hsv_filter.sMax, hsv_filter.vMax])
|
|
# Apply the thresholds
|
|
mask = cv.inRange(hsv, lower, upper)
|
|
result = cv.bitwise_and(hsv, hsv, mask=mask)
|
|
|
|
# convert back to BGR for imshow() to display it properly
|
|
img = cv.cvtColor(result, cv.COLOR_HSV2BGR)
|
|
|
|
return img
|
|
|
|
# apply adjustments to an HSV channel
|
|
# https://stackoverflow.com/questions/49697363/shifting-hsv-pixel-values-in-python-using-numpy
|
|
def shift_channel(self, c, amount):
|
|
if amount > 0:
|
|
lim = 255 - amount
|
|
c[c >= lim] = 255
|
|
c[c < lim] += amount
|
|
elif amount < 0:
|
|
amount = -amount
|
|
lim = amount
|
|
c[c <= lim] = 0
|
|
c[c > lim] -= amount
|
|
return c
|