Untitled
public
Sep 30, 2024
Never
14
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()