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()