- Notifications
You must be signed in to change notification settings - Fork32.2k
Description
import cv2
import numpy as np
import serial
import time
import math
=== CONFIG ===
SERIAL_PORT = 'COM3'
BAUD_RATE = 9600
FRAME_WIDTH, FRAME_HEIGHT = 640, 480
MAX_ANGLE_CHANGE = 25
=== COLOR RANGE ===
lower_ball = np.array([10, 100, 70]) # Orange ball
upper_ball = np.array([25, 255, 255])
lower_disc = np.array([35, 40, 40]) # Green disc
upper_disc = np.array([85, 255, 255])
=== ARUCO ===
aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)
aruco_params = cv2.aruco.DetectorParameters()
=== SERVO CENTER ANGLES ===
CENTER_ANGLES = [60, 75, 60]
=== PD CONSTANTS ===
Kp = 0.4
Kd = 0.20
PD state
prev_dx, prev_dy = 0, 0
prev_time = time.time()
Connect to Arduino
arduino = serial.Serial(SERIAL_PORT, BAUD_RATE)
time.sleep(2)
Start webcam
cap = cv2.VideoCapture(2)
def get_servo_angles(dx, dy, ddx, ddy):
dx /= (FRAME_WIDTH / 2)
dy /= (FRAME_HEIGHT / 2)
ddx /= (FRAME_WIDTH / 2)
ddy /= (FRAME_HEIGHT / 2)
# Apply PD controlcontrol_x = Kp * dx + Kd * ddxcontrol_y = Kp * dy + Kd * ddyangle1 = CENTER_ANGLES[0] + (-control_x * math.cos(math.radians(0)) - control_y * math.sin(math.radians(0))) * MAX_ANGLE_CHANGEangle2 = CENTER_ANGLES[1] + (-control_x * math.cos(math.radians(120)) - control_y * math.sin(math.radians(120))) * MAX_ANGLE_CHANGEangle3 = CENTER_ANGLES[2] + (-control_x * math.cos(math.radians(240)) - control_y * math.sin(math.radians(240))) * MAX_ANGLE_CHANGEreturn int(angle1), int(angle2), int(angle3)while True:
ret, frame = cap.read()
if not ret:
break
frame = cv2.resize(frame, (FRAME_WIDTH, FRAME_HEIGHT))frame = cv2.flip(frame, 1)hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)# ===== 1. Detect ArUco Marker =====gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)corners, ids, _ = cv2.aruco.detectMarkers(gray, aruco_dict, parameters=aruco_params)center_x, center_y = FRAME_WIDTH // 2, FRAME_HEIGHT // 2 # fallbackif corners: c = corners[0][0] center_x = int(np.mean(c[:, 0])) center_y = int(np.mean(c[:, 1])) cv2.polylines(frame, [c.astype(int)], True, (255, 0, 255), 2) cv2.circle(frame, (center_x, center_y), 5, (255, 0, 255), -1) cv2.putText(frame, "Aruco", (center_x + 10, center_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 255), 1)# ===== 2. Disc Detection =====disc_mask = cv2.inRange(hsv, lower_disc, upper_disc)disc_mask = cv2.morphologyEx(disc_mask, cv2.MORPH_OPEN, np.ones((5, 5), np.uint8))disc_contours, _ = cv2.findContours(disc_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)valid_area_mask = np.zeros_like(disc_mask)if disc_contours: largest_disc = max(disc_contours, key=cv2.contourArea) cv2.drawContours(frame, [largest_disc], -1, (0, 255, 255), 2) cv2.drawContours(valid_area_mask, [largest_disc], -1, 255, -1)# ===== 3. Ball Detection =====ball_mask = cv2.inRange(hsv, lower_ball, upper_ball)ball_mask = cv2.bitwise_and(ball_mask, ball_mask, mask=valid_area_mask)ball_mask = cv2.erode(ball_mask, None, iterations=2)ball_mask = cv2.dilate(ball_mask, None, iterations=2)ball_contours, _ = cv2.findContours(ball_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)if ball_contours: c = max(ball_contours, key=cv2.contourArea) ((x, y), radius) = cv2.minEnclosingCircle(c) if radius > 8: cv2.circle(frame, (int(x), int(y)), int(radius), (0, 255, 0), 2) dx = x - center_x dy = y - center_y # PD derivative calculation current_time = time.time() dt = current_time - prev_time if current_time != prev_time else 1e-3 ddx = (dx - prev_dx) / dt ddy = (dy - prev_dy) / dt # Get new servo angles angle1, angle2, angle3 = get_servo_angles(dx, dy, ddx, ddy) print(f"Ball: ({int(x)}, {int(y)}) Δx: {int(dx)} Δy: {int(dy)} | Angles: {angle1}, {angle2}, {angle3}") arduino.write(f"{angle1},{angle2},{angle3}\n".encode()) # Update previous for next frame prev_dx, prev_dy = dx, dy prev_time = current_timecv2.imshow("Balancing Bot", frame)if cv2.waitKey(1) & 0xFF == 27: breakcap.release()
arduino.close()
cv2.destroyAllWindows()