Code Walkthrough & Resources

Code Repository

Our ROS 2 code, including the quad package, is available on GitHub:

View Code on GitHub

For detailed explanations of the code structure, nodes, topics, and dependencies, please refer to the comprehensive README.md file within the repository.


Quick Setup Guide

1. Prerequisites

  • ROS 2 Humble Hawksbill
  • Python 3, Pip, Git, Colcon
  • Key Python Libraries: Flask, websockets, requests, numpy, opencv-python (Install via pip)
  • UR5 Network Connectivity

2. Installation

Clone the repository and build the workspace:

# Create workspace
mkdir -p ~/ros2_ws/src && cd ~/ros2_ws/

# Clone repository
git clone https://github.com/RAS598-2025-S-Team01/RAS598-2025-S-Team01.github.io.git ~/ros2_ws/

# Install ROS dependencies (Sources ROS if needed)
sudo apt update && rosdep update
rosdep install --from-paths src --ignore-src -r -y

# Install Python dependencies (Activate virtual env if used)
# Ensure you have a requirements.txt or install manually:
# pip install Flask websockets requests numpy opencv-python # Add others if needed
pip install -r requirements.txt

# Build workspace (Sources ROS if needed)
cd ~/ros2_ws/
colcon build --symlink-install

3. Running the System

Open three separate terminals. In each terminal, first source the workspace:

cd ~/ros2_ws/
source install/setup.bash

Then, run one command per terminal:

# Terminal 1: Quadruped Interface
ros2 run quad connect

# Terminal 2: UR5 Interface
ros2 run quad ur5

# Terminal 3: Web UI & ROS Bridge
ros2 run quad app

Access the web UI via the URL provided by the `app` node (usually http://127.0.0.1:5000). Use ros2 topic list in a sourced terminal to verify communication.

Code Explanation

This section provides a high-level overview of the logic within the key Python scripts of the quad package.


quad/nodes/controller_node.py

This script defines the CPGControllerNode class, responsible for generating rhythmic motion commands for the quadruped using a Central Pattern Generator (CPG). It runs in a separate thread to continuously calculate and publish servo commands without blocking the main ROS execution.

CLASS CPGControllerNode:
    INITIALIZE(servo_publisher):
        Store servo_publisher reference
        Create CPG instance (from cpg_controller script)
        Set initial state to disabled
        Set default gait (e.g., 'pronk')
        Start a background thread running the 'run' method

    METHOD run():
        LOOP FOREVER:
            IF NOT disabled:
                Calculate time delta (dt) since last loop
                Update CPG state using cpg.step(dt) -> returns target servo positions
                Create ROS message (Float64MultiArray) with target positions
                Publish message using servo_publisher
            ELSE:
                Wait briefly (e.g., 20ms)
                Reset loop timer
            Wait briefly (e.g., 5ms) to yield CPU

    METHOD set_enable(enable_flag):
        Update disabled state based on enable_flag
        IF disabling:
            Create reset message (all zeros)
            Publish reset message
            Log "Disabled" message
        ELSE:
            Log "Enabled" message

    METHOD change_gait():
        Get list of available gaits from CPG definition
        Find index of current gait
        Calculate index of next gait (wrapping around)
        Set current_gait to the new gait name
        Re-initialize CPG instance with the new gait

    METHOD update_params(parameter_dictionary):
        IF NOT disabled:
            Update CPG internal parameters (Frequency, Duty Factor, Offsets, Amplitudes)
            based on values from the input dictionary (received from UI)

quad/nodes/load_esp32.py

This script defines the Bridge class for low-level communication (Serial or UDP) with the ESP32 microcontroller and the BridgeNode ROS 2 node. The node handles sending servo commands to the ESP32 and receiving sensor data (IMU, servo states) back, publishing them onto ROS topics.

CLASS Bridge:
    DEFINE Constants (IPs, Port, Packet Structure, Data Types)

    INITIALIZE(serial_port=None):
        Store serial_port (if provided)
        Initialize serial/socket connections to None
        Initialize tx_data (commands to send) and rx_data (data received) arrays
        Set initial tx_data to disable servos and reset positions

    METHOD set_tx_data(enable_bits, position_array):
        Format enable_bits and position_array into tx_data array
        Convert tx_data array into a fixed-length byte packet with header/footer

    METHOD set_rx_data():
        Parse the received byte packet (rx_bytes)
        Extract timestamp, quaternion, gyro, accel, servo positions, servo currents
        Store extracted values into rx_data array, applying necessary scaling/conversions

    METHOD tx_rx(): // Transmit and Receive
        IF using UDP:
            Initialize socket if needed
            Send tx_bytes via UDP
            Receive rx_bytes via UDP
        ELSE (using Serial):
            Initialize serial port if needed
            Write tx_bytes to serial
            Read rx_bytes from serial
        Validate received packet length and header/footer
        IF valid:
            Call set_rx_data() to parse the data
            RETURN True
        ELSE:
            RETURN False

CLASS BridgeNode (ROS Node):
    INITIALIZE():
        Create ROS Node 'bridge_node'
        Create Bridge instance
        Create ROS Publishers for 'imu/data' (Imu msg) and 'servo/state' (Float64MultiArray msg)
        Create ROS Subscription to 'servo/command' (Float64MultiArray msg), linking to command_callback
        Create ROS Timer calling timer_callback periodically (e.g., 50Hz)

    METHOD timer_callback():
        Call bridge.tx_rx() to communicate with ESP32
        IF data received successfully:
            Create Imu message from bridge.rx_data (quaternion, gyro, accel)
            Publish Imu message
            Create Float64MultiArray message from bridge.rx_data (servo positions)
            Publish servo state message
        ELSE:
            Log warning "Failed to receive data"

    METHOD command_callback(received_command_message):
        Get target servo positions from message data
        IF all target positions are effectively zero:
            Call bridge.set_tx_data() with DISABLE_ALL_SREVOS command
            Log "Reset command received"
        ELSE:
            Call bridge.set_tx_data() with ENABLE_ALL_SREVOS and the target positions

FUNCTION main():
    Initialize ROS 2
    Create instances of BridgeNode, PlotterNode, PoseEstimatorNode
    Create instance of CPGControllerNode (passing plotter's servo_pub)
    Create MultiThreadedExecutor
    Add BridgeNode, PlotterNode, PoseEstimatorNode to executor
    Spin the executor (runs the added nodes)
    Cleanup nodes on shutdown

Note: The CPGControllerNode runs in its own thread and is not added to the ROS executor in this setup. It communicates via the ROS topics published/subscribed by the other nodes.


quad/nodes/plotter_node.py

This script defines the PlotterNode. Its primary role is to subscribe to various sensor and state topics within the ROS system and store the latest data in a globally accessible, thread-safe dictionary (shared_data). This data is then fetched by the Flask web application (app.py) to be displayed on the UI. It also provides the publisher that the CPG controller uses to send commands.

DEFINE Global shared_data dictionary (time, servo, euler, gyro, accel, pose, etc.)
DEFINE Global data_lock (threading.Lock)
DEFINE Global MAX_HISTORY constant

CLASS PlotterNode (ROS Node):
    INITIALIZE():
        Create ROS Node 'ros2_flask_plotter'
        Initialize local variables for latest data (servo, orientation, gyro, etc.)
        Create ROS Publisher for 'servo/command' (used by CPGControllerNode)
        Define QoS profile for potentially unreliable sensor data (BEST_EFFORT)
        Create ROS Subscriptions to:
            'imu/data' -> imu_callback
            'servo/state' -> servo_callback
            'imu/pose_estimate' -> pose_callback
            '/imu/orientation' -> orientation_callback
            '/orb/pose' -> slam_pose_callback
        Create ROS Timer calling sync_callback periodically (e.g., 50Hz)

    METHOD imu_callback(imu_message):
        Update local variables for orientation, gyro, accel from message

    METHOD servo_callback(servo_state_message):
        Update local variable for servo_data from message

    METHOD pose_callback(imu_pose_message):
        Update local variable for latest_pose (IMU-based) from message

    METHOD orientation_callback(orientation_message):
        Update local variable for orientation_quat from message

    METHOD slam_pose_callback(orb_pose_message):
        Update local variables for latest_slam_pose and latest_slam_orientation from message

    METHOD sync_callback(): // Called by timer
        Get current time
        Convert latest orientation quaternion to Euler angles (handle potential errors)
        ACQUIRE data_lock:
            Append current time and latest values of all local data variables (servo, euler, gyro, accel, pose, etc.) to the corresponding lists in shared_data dictionary
            IF any list in shared_data exceeds MAX_HISTORY:
                Remove the oldest element (pop(0)) from that list
        RELEASE data_lock

quad/nodes/orb_pose_estimator_node.py

This script defines the ORBPosePublisher node. It captures video frames from a network stream, tracks features using optical flow (LK method), calculates the camera's relative motion (rotation and translation) between frames using the Essential Matrix, accumulates this motion to estimate the overall pose, applies smoothing, and publishes the result as a PoseStamped message on the /orb/pose topic. It uses an ultrasonic sensor reading during initialization to help scale the translation estimate.

DEFINE Constants (VIDEO_URL, ULTRASONIC_URL, Camera Matrix, Distortion Coeffs)

CLASS ORBPosePublisher (ROS Node):
    INITIALIZE():
        Create ROS Node 'optical_flow_pose_publisher'
        Create ROS Publisher for '/orb/pose' (PoseStamped msg)
        Initialize state variables (scale_z, prev_gray, prev_pts, initialized flag, pose_R, pose_t)
        Initialize pose buffer (deque) for smoothing
        Open video capture stream from VIDEO_URL
        Call set_initial_reference() to get first frame and features
        Create ROS Timer calling process_frame periodically (e.g., 10Hz)

    METHOD fetch_ultrasonic_distance():
        Try to get distance reading from ULTRASONIC_URL via HTTP GET request
        Return distance in meters, or None on failure

    METHOD set_initial_reference():
        Log "Capturing initial reference..."
        Fetch ultrasonic distance and set scale_z if available
        LOOP until initialized:
            Read frame from video stream
            IF frame received:
                Rotate frame (if needed)
                Convert frame to grayscale
                Detect good features to track (cv2.goodFeaturesToTrack)
                IF enough features found:
                    Store gray frame as prev_gray
                    Store features as prev_pts
                    Set initialized = True
                    Log number of keypoints found

    METHOD process_frame(): // Called by timer
        IF NOT initialized: RETURN
        Read frame from video stream
        IF frame received:
            Rotate frame (if needed)
            Convert frame to grayscale
            Calculate optical flow (cv2.calcOpticalFlowPyrLK) using prev_gray, gray, prev_pts -> get next_pts, status
            Filter out bad points (status == 1) -> get good_prev, good_next
            IF not enough good points: RETURN
            Calculate Essential Matrix (cv2.findEssentialMat) between good_next and good_prev
            IF Essential Matrix found:
                Recover relative Rotation (R) and translation (t) (cv2.recoverPose)
                Accumulate pose:
                    pose_t = pose_t + (pose_R * t * scale_z) // Apply scale factor
                    pose_R = R * pose_R
                Add current pose_t to pose_buffer
                IF pose_buffer is full:
                    Apply Savitzky-Golay filter (savgol_filter) to the buffer
                    Get the last (smoothed) translation vector (smoothed_t)
                ELSE:
                    Use current pose_t as smoothed_t
                Convert accumulated rotation matrix (pose_R) to quaternion (quat)
                Create PoseStamped message:
                    Set header timestamp and frame_id
                    Set position from smoothed_t
                    Set orientation from quat
                Publish PoseStamped message
            Update prev_gray = gray
            Update prev_pts = good_next (reshaped)

    STATIC METHOD rotation_matrix_to_quaternion(RotationMatrix):
        Perform conversion math
        Return quaternion [x, y, z, w]

FUNCTION main():
    Initialize ROS 2
    Create ORBPosePublisher node instance
    Spin the node
    Cleanup node on shutdown

quad/nodes/pose_estimator_node.py

This script defines the PoseEstimatorNode. It subscribes to raw IMU data (/imu/data). It performs an initial calibration phase to estimate sensor biases. After calibration, it corrects the incoming IMU readings, applies simple filtering (low-pass for accel/gyro, high-pass/decay for velocity), integrates acceleration to get velocity, and velocity to get position. It also integrates gyro readings to get orientation (as Euler angles). Finally, it publishes the estimated position and orientation as a PoseStamped message on /imu/pose_estimate and the orientation separately as a Quaternion message on /imu/orientation.

CLASS PoseEstimatorNode (ROS Node):
    INITIALIZE():
        Create ROS Node 'pose_estimator_node'
        Create Subscription to '/imu/data' -> imu_callback
        Create Publisher for '/imu/pose_estimate' (PoseStamped)
        Create Publisher for '/imu/orientation' (Quaternion)
        Initialize state variables (position, velocity, orientation) to zeros
        Initialize last_time = None
        Initialize calibration variables (biases, data lists, calibrated flag, start time)
        Initialize filtering variables (alpha, filtered accel/gyro, velocity damping)

    METHOD imu_callback(imu_message):
        Get current time
        IF calibration_start_time is None: Set it
        IF NOT calibrated:
            Extract raw accel and gyro from message
            Append readings to calibration_data lists
            IF calibration duration (e.g., 1 second) has passed:
                Calculate average of calibration data -> set bias_accel, bias_gyro
                Set calibrated = True
                Log "Calibration complete"
            RETURN // Don't process further until calibrated
        IF last_time is None: // First reading after calibration
            Set last_time = current_time
            RETURN
        Calculate time delta (dt) = current_time - last_time
        Update last_time = current_time
        Correct raw accel/gyro readings using calculated biases
        Apply low-pass filter to accel/gyro readings
        Integrate filtered acceleration to update velocity (velocity += filtered_accel * dt)
        Apply velocity decay/damping (velocity *= decay_factor or similar)
        Apply velocity threshold (set near-zero velocities to zero)
        Integrate velocity to update position (position += velocity * dt)
        Integrate filtered gyro to update orientation (orientation += filtered_gyro * dt)
        Create PoseStamped message:
            Set header timestamp and frame_id ('map')
            Set position from state variable
            Convert Euler orientation state variable to quaternion
            Set orientation from quaternion
        Publish PoseStamped message
        Create Quaternion message from calculated quaternion
        Publish Quaternion message

FUNCTION main():
    Initialize ROS 2
    Create PoseEstimatorNode instance
    Spin the node
    Cleanup node on shutdown

quad/nodes/ur5_node.py

This script defines the URDashboardServiceNode. It provides ROS 2 services (/start_program, /stop_program) to trigger actions on the UR5. When called, it connects to the UR5's Dashboard Server via TCP sockets and sends text commands to load and play pre-defined URScript programs (gripper.urp, stop.urp) stored on the robot controller, which contain the pick-and-place logic.

DEFINE Constants (ROBOT_IP, DASHBOARD_PORT, TIMEOUT, Program Names)

CLASS URDashboardServiceNode (ROS Node):
    INITIALIZE():
        Create ROS Node 'ur_dashboard_service_node'
        Log "Node started"
        Create ROS Service '/start_program' (Trigger type) -> handle_start_program
        Create ROS Service '/stop_program' (Trigger type) -> handle_stop_program

    METHOD send_dashboard(command_string):
        TRY:
            Create TCP socket connection to ROBOT_IP:DASHBOARD_PORT with TIMEOUT
            Receive and discard initial connection message from robot
            Send the command_string + newline, encoded as UTF-8
            Wait briefly (e.g., 0.1s)
            Receive response from robot (up to 4096 bytes)
            Decode response as UTF-8 and strip whitespace
            Close socket
            RETURN decoded response string
        CATCH Exception as e:
            Log error message
            RETURN error string

    METHOD load_and_play(program_filename):
        Log "Loading program: [program_filename]"
        response = send_dashboard("load [program_filename]")
        IF "Loading program" NOT IN response:
            RETURN False, "Failed to load: [response]"
        Wait briefly (e.g., 0.3s)
        Log "Starting program..."
        response = send_dashboard("play")
        IF "Starting program" OR "Running program" IN response:
            RETURN True, "Program started: [response]"
        ELSE:
            RETURN False, "Failed to start: [response]"

    METHOD handle_start_program(request, response): // Service callback
        success, message = load_and_play(START_PROGRAM filename)
        Set response.success = success
        Set response.message = message
        RETURN response

    METHOD handle_stop_program(request, response): // Service callback
        success, message = load_and_play(STOP_PROGRAM filename)
        Set response.success = success
        Set response.message = message
        RETURN response

FUNCTION main():
    Initialize ROS 2
    Create URDashboardServiceNode instance
    Spin the node (waits for service calls)
    Cleanup node on shutdown

quad/web/app.py

This script sets up and runs the Flask web server, acting as the user interface backend. It also initializes and runs necessary ROS 2 nodes (Plotter, ORB Pose Estimator) in a separate thread using a MultiThreadedExecutor. It defines Flask routes (URLs) for serving the main HTML page (`/`), providing real-time data to the UI (`/data`), and handling UI commands (`/cpg/...`) to interact with the CPG controller node.

IMPORT Flask, ROS 2 libraries, threading, nodes (PlotterNode, CPGControllerNode, ORBPosePublisher)

INITIALIZE Flask app

DEFINE Global variables for plotter_node, cpg_controller (initially None)

DEFINE Flask route '/' (index):
    Render 'index.html' template

DEFINE Flask route '/data' (get data):
    ACQUIRE data_lock (from plotter_node module)
    Get a copy of shared_data
    RELEASE data_lock
    Return shared_data as JSON

DEFINE Flask route '/cpg/enable' (POST request):
    Get 'enable' state from JSON request body
    IF cpg_controller exists:
        Call cpg_controller.set_enable(state)
    Return success JSON response

DEFINE Flask route '/cpg/gait' (POST request):
    IF cpg_controller exists:
        Call cpg_controller.change_gait()
        Return success JSON response with new gait name
    Return failure JSON response

DEFINE Flask route '/cpg/params' (POST request):
    Get 'params' dictionary from JSON request body
    IF cpg_controller exists:
        Call cpg_controller.update_params(params)
    Return success JSON response

FUNCTION ros_spin(): // Runs in a separate thread
    DECLARE plotter_node, cpg_controller as global
    Initialize ROS 2
    Create PlotterNode instance -> assign to global plotter_node
    Create CPGControllerNode instance (passing plotter_node.servo_pub) -> assign to global cpg_controller
    Create ORBPosePublisher instance -> assign to local orb_pose_node
    Create MultiThreadedExecutor
    Add plotter_node to executor
    Add orb_pose_node to executor
    TRY:
        Spin the executor (runs plotter and orb_pose nodes)
    FINALLY: // Cleanup on thread exit/error
        Destroy plotter_node
        Destroy orb_pose_node
        Shutdown ROS 2

FUNCTION main(): // Main entry point
    Create and start the ros_spin thread (as daemon)
    Start the Flask development server (host='0.0.0.0', threaded=True)

IF script is run directly (__name__ == '__main__'):
    Call main()