G

Untitled

public
Guest Sep 30, 2024 Never 14
Clone
Plaintext paste1.txt 95 lines (74 loc) | 3.22 KB
1
# main.py
2
3
import cv2
4
import numpy as np
5
import mss
6
import pyautogui
7
import time
8
9
# Screen capture settings (adjust based on your screen size and pool window)
10
MONITOR = {"top": 100, "left": 100, "width": 800, "height": 600} # Example values
11
12
def capture_screen(monitor=MONITOR):
13
with mss.mss() as sct:
14
screenshot = sct.grab(monitor)
15
frame = np.array(screenshot)
16
frame = cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)
17
return frame
18
# main.py (continued)
19
20
def detect_balls(frame):
21
hsv_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
22
23
# Define HSV color ranges for different balls
24
# These ranges may need adjustment based on game graphics
25
color_ranges = {
26
'cue_ball': {'lower': np.array([0, 0, 180]), 'upper': np.array([180, 40, 255]), 'color': (255, 255, 255)},
27
'red_ball': {'lower': np.array([0, 100, 100]), 'upper': np.array([10, 255, 255]), 'color': (0, 0, 255)},
28
'yellow_ball': {'lower': np.array([20, 100, 100]), 'upper': np.array([30, 255, 255]), 'color': (0, 255, 255)},
29
# Add more colors as needed
30
}
31
32
ball_positions = {}
33
34
for ball, params in color_ranges.items():
35
mask = cv2.inRange(hsv_frame, params['lower'], params['upper'])
36
mask = cv2.GaussianBlur(mask, (9, 9), 0)
37
circles = cv2.HoughCircles(mask, cv2.HOUGH_GRADIENT, dp=1.2, minDist=30,
38
param1=50, param2=30, minRadius=10, maxRadius=20)
39
40
if circles is not None:
41
circles = np.round(circles[0, :]).astype("int")
42
for (x, y, r) in circles:
43
cv2.circle(frame, (x, y), r, params['color'], 4)
44
ball_positions[ball] = (x, y)
45
break # Assuming one ball per color; remove if multiple
46
47
return ball_positions
48
# main.py (continued)
49
50
def calculate_aim(cue_pos, target_pos):
51
dx = target_pos[0] - cue_pos[0]
52
dy = target_pos[1] - cue_pos[1]
53
angle = np.arctan2(dy, dx) * 180 / np.pi
54
return angle
55
56
def draw_aiming_line(frame, cue_pos, target_pos):
57
# Draw line from cue ball to target ball
58
cv2.line(frame, cue_pos, target_pos, (255, 0, 0), 2)
59
60
# Calculate and display angle
61
angle = calculate_aim(cue_pos, target_pos)
62
cv2.putText(frame, f"Aim Angle: {angle:.2f}°", (50, 50),
63
cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
64
# main.py (continued)
65
66
def main():
67
print("Starting 8-Ball Pool AI Aiming Assistant...")
68
time.sleep(2) # Delay to allow user to switch to the game window
69
70
while True:
71
# Capture the screen
72
frame = capture_screen()
73
74
# Detect balls
75
balls = detect_balls(frame)
76
77
# Assuming 'cue_ball' and 'red_ball' (as target) are detected
78
if 'cue_ball' in balls and 'red_ball' in balls:
79
cue_ball = balls['cue_ball']
80
target_ball = balls['red_ball']
81
82
# Draw aiming line
83
draw_aiming_line(frame, cue_ball, target_ball)
84
85
# Display the frame
86
cv2.imshow("8-Ball Pool Aiming Assist", frame)
87
88
# Exit on pressing 'q'
89
if cv2.waitKey(1) & 0xFF == ord('q'):
90
break
91
92
cv2.destroyAllWindows()
93
94
if __name__ == "__main__":
95
main()