practical-tutorials/project-based-learning

Ball Balancing Robot(PyCharm)

Opened this issue · 2 comments

TANAYSIDHU commented

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 control
control_x = Kp * dx + Kd * ddx
control_y = Kp * dy + Kd * ddy

angle1 = CENTER_ANGLES[0] + (-control_x * math.cos(math.radians(0)) - control_y * math.sin(math.radians(0))) * MAX_ANGLE_CHANGE
angle2 = CENTER_ANGLES[1] + (-control_x * math.cos(math.radians(120)) - control_y * math.sin(math.radians(120))) * MAX_ANGLE_CHANGE
angle3 = CENTER_ANGLES[2] + (-control_x * math.cos(math.radians(240)) - control_y * math.sin(math.radians(240))) * MAX_ANGLE_CHANGE

return 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  # fallback

if 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_time

cv2.imshow("Balancing Bot", frame)
if cv2.waitKey(1) & 0xFF == 27:
    break

cap.release()
arduino.close()
cv2.destroyAllWindows()

Macplentee commented

Yo My guy this is impressive, I have ideas I think you'll like to implement. Let me know what you think

M-tech-cmd commented

The main issues were related to indentation and using an older method for ArUco detection

Imports for computer vision, numerical operations, serial communication, and time.

import cv2
import numpy as np
import serial
import time
import math

--- CONFIGURATION ---

Define constants for the serial port, baud rate, and camera frame size.

SERIAL_PORT = 'COM3'
BAUD_RATE = 9600
FRAME_WIDTH, FRAME_HEIGHT = 640, 480
MAX_ANGLE_CHANGE = 25 # Maximum angle adjustment for the servos

--- COLOR RANGE DEFINITIONS ---

Define the color ranges in HSV for the objects being tracked.

H: Hue, S: Saturation, V: Value

Orange ball (Note: You may need to tune this for your specific ball and lighting)

lower_ball = np.array([10, 100, 70])
upper_ball = np.array([25, 255, 255])

Green disc (Note: Tune this as well)

lower_disc = np.array([35, 40, 40])
upper_disc = np.array([85, 255, 255])

--- ARUCO MARKER SETUP ---

Use the newer, more robust ArucoDetector.

aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)
aruco_params = cv2.aruco.DetectorParameters()
aruco_detector = cv2.aruco.ArucoDetector(aruco_dict, aruco_params)

--- SERVO CENTER ANGLES ---

Define the neutral angles for the three servos.

CENTER_ANGLES = [60, 75, 60]

--- PD CONTROL CONSTANTS ---

Proportional and Derivative gain constants.

You will likely need to tune these for optimal performance.

Kp = 0.4
Kd = 0.20

--- PD STATE VARIABLES ---

Variables to hold the previous error and timestamp for derivative calculation.

prev_dx, prev_dy = 0, 0
prev_time = time.time()

def get_servo_angles(dx, dy, ddx, ddy):
"""
Calculates the new servo angles based on the ball's position error (dx, dy)
and its derivative (ddx, ddy) using a PD control loop.

Args:
    dx (float): Error in the x-direction from the center.
    dy (float): Error in the y-direction from the center.
    ddx (float): Derivative of the x-error (change in x-error over time).
    ddy (float): Derivative of the y-error (change in y-error over time).

Returns:
    tuple: A tuple containing the three calculated servo angles.
"""
# Normalize the error values to a range of [-1, 1]
dx /= (FRAME_WIDTH / 2)
dy /= (FRAME_HEIGHT / 2)
ddx /= (FRAME_WIDTH / 2)
ddy /= (FRAME_HEIGHT / 2)

# Apply PD control to get the control signals for x and y
control_x = Kp * dx + Kd * ddx
control_y = Kp * dy + Kd * ddy

# Calculate the angle for each of the three servos. The math here is specific
# to a three-servo platform and assumes a 120-degree separation.
# The negative sign in front of the control signals inverts the direction.
angle1 = CENTER_ANGLES[0] + (-control_x * math.cos(math.radians(0)) - control_y * math.sin(math.radians(0))) * MAX_ANGLE_CHANGE
angle2 = CENTER_ANGLES[1] + (-control_x * math.cos(math.radians(120)) - control_y * math.sin(math.radians(120))) * MAX_ANGLE_CHANGE
angle3 = CENTER_ANGLES[2] + (-control_x * math.cos(math.radians(240)) - control_y * math.sin(math.radians(240))) * MAX_ANGLE_CHANGE

# Return the angles as integers
return int(angle1), int(angle2), int(angle3)

def main():
"""
Main function to run the balancing bot script.
It handles camera capture, object detection, PD control, and serial communication.
"""
global prev_dx, prev_dy, prev_time

# --- INITIALIZATION ---
# Attempt to connect to the Arduino via the specified serial port.
try:
    arduino = serial.Serial(SERIAL_PORT, BAUD_RATE)
    time.sleep(2)  # Give the Arduino time to reset
except serial.SerialException as e:
    print(f"Error: Could not open serial port {SERIAL_PORT}. Please check the port and try again.")
    print(e)
    return

# Attempt to start the webcam. The index might need to be adjusted (0, 1, 2, etc.).
cap = cv2.VideoCapture(2)
if not cap.isOpened():
    print("Error: Could not open webcam.")
    return

while True:
    ret, frame = cap.read()
    if not ret:
        print("Error: Could not read a frame from the webcam.")
        break

    # Resize the frame and flip it horizontally for a more intuitive view.
    frame = cv2.resize(frame, (FRAME_WIDTH, FRAME_HEIGHT))
    frame = cv2.flip(frame, 1)
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

    # --- 1. Detect ArUco Marker ---
    # Find the center of the ArUco marker, which serves as the reference point.
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    corners, ids, _ = aruco_detector.detectMarkers(gray)

    center_x, center_y = FRAME_WIDTH // 2, FRAME_HEIGHT // 2  # Fallback to frame center

    if ids is not None and len(ids) > 0:
        # Use the first detected marker
        c = corners[0][0]
        center_x = int(np.mean(c[:, 0]))
        center_y = int(np.mean(c[:, 1]))
        # Draw the marker outline and center point for visualization
        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 ---
    # Create a mask for the green disc.
    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)
        # Draw the contour and create a mask from the largest disc to limit ball detection.
        cv2.drawContours(frame, [largest_disc], -1, (0, 255, 255), 2)
        cv2.drawContours(valid_area_mask, [largest_disc], -1, 255, -1)

    # --- 3. Ball Detection ---
    # Detect the orange ball, but only within the area of the green disc.
    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: # Filter out small objects that aren't the ball
            cv2.circle(frame, (int(x), int(y)), int(radius), (0, 255, 0), 2)

            # Calculate the ball's position error relative to the ArUco marker.
            dx = x - center_x
            dy = y - center_y

            # Calculate the derivative of the error for PD control.
            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 the new servo angles from the PD controller.
            angle1, angle2, angle3 = get_servo_angles(dx, dy, ddx, ddy)
            
            # Print and send the angles to the Arduino.
            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 state for the next frame's derivative calculation.
            prev_dx, prev_dy = dx, dy
            prev_time = current_time

    # Display the output frame and check for the 'Esc' key to quit.
    cv2.imshow("Balancing Bot", frame)
    if cv2.waitKey(1) & 0xFF == 27:
        break

# --- CLEANUP ---
cap.release()
arduino.close()
cv2.destroyAllWindows()

if name == "main":
main()