add some opencv stuff

This commit is contained in:
randomuser 2023-07-29 22:49:18 -05:00
parent 7aff0ace45
commit 2801a3f4cd
3 changed files with 158 additions and 0 deletions

View File

@ -2,3 +2,4 @@ prompt_toolkit
pyautogui
pywinctl
opencv-python
numpy

98
src/visual/main.py Normal file
View File

@ -0,0 +1,98 @@
import pyautogui
import colorsys
import pywinctl
import numpy
import cv2 as cv
import struct
import matching
from prompt_toolkit import prompt
from prompt_toolkit.completion import WordCompleter
# get all windows
titles = pywinctl.getAllTitles()
title_completer = WordCompleter(titles)
try:
window = pywinctl.getWindowsWithTitle("Keep Talking and Nobody Explodes")[0]
except IndexError:
try:
text = prompt("Select window name> ", completer=title_completer)
window = pywinctl.getWindowsWithTitle(text)[0]
except IndexError:
print("invalid title")
exit()
if not window.isVisible:
window.show()
frame = window.getClientFrame()
print(frame)
def get_image(left, top, right, bottom): # returns an opencv image
image = pyautogui.screenshot(region=(
left, top, right - left, bottom - top
))
numpy_array = numpy.array(image)
opencv_image = cv.cvtColor(numpy_array, cv.COLOR_RGB2BGR)
return opencv_image
rgb_colors_to_match = [
(220, 162, 129),
(220, 162, 129),
(145, 108, 89),
]
num_colors_in_range = 9
delta_h_value = 10
delta_s_value = 50
delta_v_value = 50
tolerence_value = 20
ranges = [
matching.generate_hex_range(i, num_colors_in_range, delta_h=delta_h_value, delta_s=delta_s_value, delta_v=delta_v_value)
for i in rgb_colors_to_match
]
def are_rectangles_close(rect1, rect2):
distance_threshold = 20
x1, y1, w1, h1 = rect1
x2, y2, w2, h2 = rect2
return abs(x2 - x1) < distance_threshold and abs(y2 - y1) < distance_threshold
while True:
image = get_image(frame.left, frame.top, frame.right, frame.bottom)
blank = numpy.zeros(image.shape, dtype='uint8')
combined = image
matched = matching.color_matching(image, ranges[0])
combined = matching.convert_masked_pixels_to_zero(combined, matched)
matched = matching.color_matching(image, ranges[1])
combined = matching.convert_masked_pixels_to_zero(combined, matched)
matched = matching.color_matching(image, ranges[2])
combined = matching.convert_masked_pixels_to_zero(combined, matched)
min_contour_length = 100
max_contour_length = 10000
bw = cv.cvtColor(combined, cv.COLOR_BGR2GRAY)
ret, thresh = cv.threshold(bw, 127, 255, cv.THRESH_BINARY)
contours, hierarchy = cv.findContours(thresh, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
filtered_contours = [cnt for cnt in contours if min_contour_length <= cv.arcLength(cnt, True) <= max_contour_length]
monster_contour = numpy.concatenate(filtered_contours)
hull = cv.convexHull(monster_contour)
bounding_quad = cv.minAreaRect(hull)
points = cv.boxPoints(bounding_quad).astype(int)
print(points)
cv.drawContours(combined, [points], 0, (0, 255, 0), 2)
cv.drawContours(combined, filtered_contours, -1, (0,255,255), 2)
cv.imshow('test', combined)
key = cv.waitKey(1)
if (key) == 113: # q key
break
cv.destroyAllWindows()

59
src/visual/matching.py Normal file
View File

@ -0,0 +1,59 @@
import cv2
import numpy as np
def rgb_to_hsv(rgb_color):
# Convert RGB color to HSV color space
hsv_color = cv2.cvtColor(np.uint8([[rgb_color]]), cv2.COLOR_RGB2HSV)[0][0]
return hsv_color
def generate_hex_range(rgb_color, num_colors, delta_h=10, delta_s=50, delta_v=50):
# Convert the RGB color to HSV color space
hsv_color = rgb_to_hsv(rgb_color)
# Define ranges in HSV color space
h_range = np.linspace(hsv_color[0] - delta_h, hsv_color[0] + delta_h, num_colors)
s_range = np.linspace(hsv_color[1] - delta_s, hsv_color[1] + delta_s, num_colors)
v_range = np.linspace(hsv_color[2] - delta_v, hsv_color[2] + delta_v, num_colors)
hex_range = []
# Generate colors in the HSV color space and convert them to BGR
for h_val in h_range:
for s_val in s_range:
for v_val in v_range:
hsv_color_modified = np.array([h_val, s_val, v_val])
bgr_color_modified = cv2.cvtColor(np.uint8([[hsv_color_modified]]), cv2.COLOR_HSV2BGR)[0][0]
hex_color = "#{:02x}{:02x}{:02x}".format(int(bgr_color_modified[2]),
int(bgr_color_modified[1]),
int(bgr_color_modified[0]))
hex_range.append(hex_color)
return hex_range
def color_matching(image, hex_range_colors, tolerance=20):
hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
matched_pixels = np.zeros_like(hsv_image[:, :, 0], dtype=np.uint8)
for hex_color in hex_range_colors:
# Convert the hex color to RGB
rgb_color = tuple(int(hex_color[i:i + 2], 16) for i in (1, 3, 5))
# Convert the RGB color to HSV
hsv_color = rgb_to_hsv(rgb_color)
lower_bound = np.array([hsv_color[0] - tolerance, hsv_color[1] - tolerance, hsv_color[2] - tolerance])
upper_bound = np.array([hsv_color[0] + tolerance, hsv_color[1] + tolerance, hsv_color[2] + tolerance])
mask = cv2.inRange(hsv_image, lower_bound, upper_bound)
matched_pixels |= mask
return matched_pixels
def convert_masked_pixels_to_zero(image, mask):
# Create an all-white (ones) image with the same shape as the original image
white_image = np.ones_like(image) * 255
# Use the mask to set all masked pixels in the original image to zero
result_image = cv2.bitwise_and(image, white_image, mask=cv2.bitwise_not(mask))
return result_image