├── .gitignore ├── README.md ├── __init__.py ├── examples ├── example_depth.py ├── example_depth_fast.py ├── example_drivepath.py ├── example_facetime.py ├── example_follow.py ├── example_kokoro.py ├── example_localization.py ├── example_realtime.py ├── example_rerun.py ├── example_segmentation.py ├── example_wasd.py ├── example_whisper.py └── example_yolo.py ├── lib ├── __init__.py ├── camera.py ├── imu.py ├── lqr.py ├── madgwickahrs.py ├── odrive_uart.py └── vl53l5cx_lib │ ├── api.py │ ├── buffers.py │ └── vl53l5cx.py ├── setup ├── README.md ├── calibrate_drive.py ├── calibrate_stereo_camera.py ├── extras │ ├── README.md │ ├── create_custom_boot_service.sh │ ├── pwm-pi5.dts │ ├── setup_hardware_pwm.sh │ └── setup_realsense.sh └── setup_os.sh └── tests ├── test_camera.py ├── test_imu_orientation.py ├── test_led_strip.py ├── test_microphone.py ├── test_microphone_with_led.py ├── test_realsense.py ├── test_rerun.py ├── test_servos.py ├── test_speaker.py ├── test_speaker_with_led.py ├── test_system.py ├── test_transcription.py ├── test_tts.py └── test_tts_with_led.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Python cache files 2 | __pycache__/ 3 | *.pyc 4 | *.pyo 5 | *.pyd 6 | 7 | # Generated files 8 | *.png 9 | *.txt 10 | *.err 11 | *.log 12 | *.csv 13 | *.wav 14 | *.debug 15 | 16 | # Config and environment files 17 | rerun_config.yaml 18 | motor_dir.json 19 | .env 20 | 21 | # Device tree overlay 22 | pwm-pi5.dtbo 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Bracket Bot Quickstart 2 | 3 | For comprehensive documentation, visit our docs site: [https://docs.bracket.bot/docs/](https://docs.bracket.bot/docs/) 4 | 5 | ## Getting Started 6 | 7 | To set up your robot, look at the setup folder: 8 | 9 | [**Setup Guide**](./setup) 10 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | # Makes the quickstart directory a Python package for easy `python -m quickstart.*` imports. -------------------------------------------------------------------------------- /examples/example_depth.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | """Light-weight stereo + point-cloud viewer (fixed).""" 3 | 4 | import os, sys, time 5 | import cv2 6 | import numpy as np 7 | from scipy.spatial.transform import Rotation as R_scipy 8 | import rerun as rr 9 | 10 | sys.path.append(os.path.join(os.path.dirname(__file__), "..")) 11 | from lib.camera import StereoCamera # noqa: E402 12 | 13 | # ───────────────── CONFIG ───────────────── 14 | DOWNSAMPLE = 0.375 15 | CALIB_FILE = os.path.join(os.path.dirname(__file__), 16 | "..", "lib", "stereo_calibration_fisheye.yaml") 17 | 18 | # SGBM params (identical to "big" file) 19 | WINDOW_SIZE = 7 20 | MIN_DISP = -32 21 | NUM_DISP = 144 # must be /16 22 | UNIQUENESS = 5 23 | SPECKLE_WIN = 100 24 | SPECKLE_RANGE = 1 25 | P1 = 8 * 3 * WINDOW_SIZE ** 2 26 | P2 = 32 * 3 * WINDOW_SIZE ** 2 27 | 28 | # ─────────── Helper: world → viewer ─────────── 29 | def _world_to_viewer(pts_world: np.ndarray) -> np.ndarray: 30 | x = pts_world[:, 0] # right 31 | z = pts_world[:, 2] # fwd 32 | y = pts_world[:, 1] # up (positive down in world) 33 | return np.column_stack([x, z, -y]).astype(np.float32) 34 | 35 | # ─────────── Load calibration (verbatim) ─────────── 36 | def load_calib_yaml(path: str, scale: float): 37 | fs = cv2.FileStorage(path, cv2.FILE_STORAGE_READ) 38 | mtx_l = fs.getNode("mtx_l").mat(); dist_l = fs.getNode("dist_l").mat() 39 | mtx_r = fs.getNode("mtx_r").mat(); dist_r = fs.getNode("dist_r").mat() 40 | R1 = fs.getNode("R1").mat(); R2 = fs.getNode("R2").mat() 41 | P1 = fs.getNode("P1").mat(); P2 = fs.getNode("P2").mat() 42 | Q = fs.getNode("Q").mat().astype(np.float32) 43 | fs.release() 44 | # down-sample Q's translation entries 45 | for i in range(4): 46 | Q[i, 3] *= scale 47 | return mtx_l, dist_l, mtx_r, dist_r, R1, R2, P1, P2, Q 48 | 49 | # ─────────────────── Main ──────────────────── 50 | def main(): 51 | rr.init("tiny_depth_viewer") 52 | # find your computer IP by doing "arp -a" in a terminal 53 | rr.connect_grpc("rerun+http://192.168.2.24:9876/proxy") 54 | 55 | cam = StereoCamera(0, scale=1.0) 56 | 57 | # Calibration 58 | (mtx_l, dist_l, mtx_r, dist_r, 59 | R1, R2, P1_cam, P2_cam, Q) = load_calib_yaml(CALIB_FILE, DOWNSAMPLE) 60 | 61 | baseline_m = abs(P2_cam[0, 3] / P2_cam[0, 0]) / 1000.0 62 | fx_ds = P1_cam[0, 0] * DOWNSAMPLE 63 | 64 | # Rectification maps 65 | left_raw, _ = cam.get_stereo() 66 | h, w = left_raw.shape[:2] 67 | map1x, map1y = cv2.fisheye.initUndistortRectifyMap(mtx_l, dist_l, R1, P1_cam, 68 | (w, h), cv2.CV_32FC1) 69 | map2x, map2y = cv2.fisheye.initUndistortRectifyMap(mtx_r, dist_r, R2, P2_cam, 70 | (w, h), cv2.CV_32FC1) 71 | 72 | # SGBM matcher 73 | stereo = cv2.StereoSGBM_create( 74 | minDisparity=MIN_DISP, 75 | numDisparities=NUM_DISP, 76 | blockSize=WINDOW_SIZE, 77 | P1=P1, # 8 * 3 * WINDOW_SIZE**2 78 | P2=P2, # 32 * 3 * WINDOW_SIZE**2 79 | disp12MaxDiff=1, 80 | uniquenessRatio=UNIQUENESS, 81 | speckleWindowSize=SPECKLE_WIN, 82 | speckleRange=SPECKLE_RANGE, 83 | preFilterCap=63, 84 | mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY, 85 | ) 86 | 87 | while True: 88 | left_raw, right_raw = cam.get_stereo() 89 | if left_raw is None: 90 | time.sleep(0.01) 91 | continue 92 | 93 | # Rectify → downsample 94 | left_rect = cv2.remap(left_raw, map1x, map1y, cv2.INTER_LINEAR) 95 | right_rect = cv2.remap(right_raw, map2x, map2y, cv2.INTER_LINEAR) 96 | tgt_sz = (int(w * DOWNSAMPLE), int(h * DOWNSAMPLE)) 97 | left_ds = cv2.resize(left_rect, tgt_sz, interpolation=cv2.INTER_AREA) 98 | right_ds = cv2.resize(right_rect, tgt_sz, interpolation=cv2.INTER_AREA) 99 | 100 | # Disparity (px) 101 | disp = stereo.compute(left_ds, right_ds).astype(np.float32) / 16.0 102 | 103 | # Valid pixels 104 | valid = disp > (MIN_DISP + 0.5) 105 | 106 | # Depth (m) – avoid div-by-zero near MIN_DISP 107 | denom = disp - MIN_DISP 108 | depth = np.zeros_like(disp) 109 | depth_mask = denom > 0.1 110 | depth[depth_mask] = fx_ds * baseline_m / denom[depth_mask] 111 | 112 | # Depth colouring (invert JET, 0.1–5 m) 113 | depth_clipped = np.clip(depth, 0.1, 5.0) 114 | depth_norm = ((depth_clipped - 0.1) / (5.0 - 0.1) * 255).astype(np.uint8) 115 | depth_color = cv2.applyColorMap(255 - depth_norm, cv2.COLORMAP_JET) 116 | 117 | # Point cloud in camera frame 118 | pts_cam = cv2.reprojectImageTo3D(disp, Q) / 1000.0 119 | pts_cam = pts_cam[valid] 120 | cols = cv2.cvtColor(left_ds, cv2.COLOR_BGR2RGB).reshape(-1, 3)[valid.ravel()] 121 | 122 | # Throw away far points (> 5 m) for clarity 123 | dist_m = np.linalg.norm(pts_cam, axis=1) 124 | keep = dist_m < 5.0 125 | pts_cam, cols = pts_cam[keep], cols[keep] 126 | 127 | # Camera → robot (tilt & 1.5 m up) 128 | t_cam = np.array([0.0, -1.5, 0.0]) # +Y is down 129 | R_cam = R_scipy.from_euler('x', -36.0, degrees=True) 130 | pts_robot = R_cam.apply(pts_cam) + t_cam # robot frame (y up) 131 | 132 | # Robot at origin → world 133 | pts_world = pts_robot 134 | 135 | # Viewer coords (X right, Y fwd, Z up) 136 | pts_view = _world_to_viewer(pts_world) 137 | 138 | # Sub-sample 4× before logging 139 | pts_view = pts_view[::4] 140 | cols = cols[::4] 141 | 142 | if pts_view.size: 143 | rr.log("world/point_cloud", rr.Points3D(pts_view, colors=cols, radii=0.02)) 144 | 145 | # Images to Rerun 146 | rr.set_time_seconds("time", time.time()) 147 | rr.log("camera/disparity", rr.Image(disp)) 148 | rr.log("camera/depth", rr.Image(depth_color)) 149 | 150 | if __name__ == "__main__": 151 | main() 152 | -------------------------------------------------------------------------------- /examples/example_depth_fast.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | """Light-weight stereo + point-cloud viewer (fixed).""" 3 | 4 | import os, sys, time 5 | import cv2 6 | import numpy as np 7 | from scipy.spatial.transform import Rotation as R_scipy 8 | import rerun as rr 9 | import argparse 10 | 11 | sys.path.append(os.path.join(os.path.dirname(__file__), "..")) 12 | from lib.camera import StereoCamera # noqa: E402 13 | 14 | # ───────────────── CONFIG ───────────────── 15 | DOWNSAMPLE = 0.5 16 | CALIB_FILE = os.path.join(os.path.dirname(__file__), 17 | "..", "lib", "stereo_calibration_fisheye.yaml") 18 | 19 | # SGBM params (identical to "big" file) 20 | WINDOW_SIZE = 23 21 | MIN_DISP = -32 22 | NUM_DISP = 128 # must be /16 23 | UNIQUENESS = 7 24 | SPECKLE_WIN = 150 25 | SPECKLE_RANGE = 1 26 | P1 = 8 * 3 * WINDOW_SIZE ** 2 27 | P2 = 32 * 3 * WINDOW_SIZE ** 2 28 | 29 | # ─────────── Helper: world → viewer ─────────── 30 | def _world_to_viewer(pts_world: np.ndarray) -> np.ndarray: 31 | x = pts_world[:, 0] # right 32 | z = pts_world[:, 2] # fwd 33 | y = pts_world[:, 1] # up (positive down in world) 34 | return np.column_stack([x, z, -y]).astype(np.float32) 35 | 36 | # ─────────── Load calibration (verbatim) ─────────── 37 | def load_calib_yaml(path: str, scale: float): 38 | fs = cv2.FileStorage(path, cv2.FILE_STORAGE_READ) 39 | mtx_l = fs.getNode("mtx_l").mat(); dist_l = fs.getNode("dist_l").mat() 40 | mtx_r = fs.getNode("mtx_r").mat(); dist_r = fs.getNode("dist_r").mat() 41 | R1 = fs.getNode("R1").mat(); R2 = fs.getNode("R2").mat() 42 | P1 = fs.getNode("P1").mat(); P2 = fs.getNode("P2").mat() 43 | Q = fs.getNode("Q").mat().astype(np.float32) 44 | fs.release() 45 | # down-sample Q's translation entries 46 | for i in range(4): 47 | Q[i, 3] *= scale 48 | return mtx_l, dist_l, mtx_r, dist_r, R1, R2, P1, P2, Q 49 | 50 | # ─────────────────── Main ──────────────────── 51 | def main(max_frames=None): 52 | rr.init("tiny_depth_viewer") 53 | # find your computer IP by doing "arp -a" in a terminal 54 | rr.connect_grpc("rerun+http://192.168.2.24:9876/proxy") 55 | 56 | cam = StereoCamera(0, scale=1.0) 57 | 58 | # Calibration 59 | (mtx_l, dist_l, mtx_r, dist_r, 60 | R1, R2, P1_cam, P2_cam, Q) = load_calib_yaml(CALIB_FILE, DOWNSAMPLE) 61 | 62 | baseline_m = abs(P2_cam[0, 3] / P2_cam[0, 0]) / 1000.0 63 | fx_ds = P1_cam[0, 0] * DOWNSAMPLE 64 | 65 | # Rectification maps 66 | left_raw, _ = cam.get_stereo() 67 | h, w = left_raw.shape[:2] 68 | map1x, map1y = cv2.fisheye.initUndistortRectifyMap(mtx_l, dist_l, R1, P1_cam, 69 | (w, h), cv2.CV_32FC1) 70 | map2x, map2y = cv2.fisheye.initUndistortRectifyMap(mtx_r, dist_r, R2, P2_cam, 71 | (w, h), cv2.CV_32FC1) 72 | 73 | # SGBM matcher 74 | # Faster but lower-quality matcher (StereoBM) 75 | stereo = cv2.StereoBM_create(numDisparities=NUM_DISP, blockSize=WINDOW_SIZE) 76 | stereo.setMinDisparity(MIN_DISP) 77 | stereo.setUniquenessRatio(UNIQUENESS) 78 | stereo.setSpeckleWindowSize(SPECKLE_WIN) 79 | stereo.setSpeckleRange(SPECKLE_RANGE) 80 | stereo.setPreFilterCap(21) 81 | 82 | # Left-right consistency check removed for speed 83 | 84 | frame = 0 85 | while True: 86 | if max_frames is not None and frame >= max_frames: 87 | break 88 | frame += 1 89 | left_raw, right_raw = cam.get_stereo() 90 | if left_raw is None: 91 | time.sleep(0.01) 92 | continue 93 | 94 | # Rectify → downsample 95 | left_rect = cv2.remap(left_raw, map1x, map1y, cv2.INTER_LINEAR) 96 | right_rect = cv2.remap(right_raw, map2x, map2y, cv2.INTER_LINEAR) 97 | tgt_sz = (int(w * DOWNSAMPLE), int(h * DOWNSAMPLE)) 98 | left_ds = cv2.resize(left_rect, tgt_sz, interpolation=cv2.INTER_AREA) 99 | right_ds = cv2.resize(right_rect, tgt_sz, interpolation=cv2.INTER_AREA) 100 | 101 | # StereoBM expects single-channel 8-bit images 102 | left_gray = cv2.cvtColor(left_ds, cv2.COLOR_BGR2GRAY) 103 | right_gray = cv2.cvtColor(right_ds, cv2.COLOR_BGR2GRAY) 104 | 105 | # Disparity (px) 106 | disp_left = stereo.compute(left_gray, right_gray).astype(np.float32) / 16.0 107 | 108 | # Valid pixels (simple disparity threshold) 109 | valid = disp_left > (MIN_DISP + 0.5) 110 | 111 | # Depth (m) – avoid div-by-zero near MIN_DISP 112 | denom = disp_left - MIN_DISP 113 | depth = np.zeros_like(disp_left) 114 | depth_mask = (denom > 0.1) & valid 115 | depth[depth_mask] = fx_ds * baseline_m / denom[depth_mask] 116 | 117 | # Depth colouring (invert JET, 0.1–5 m) 118 | depth_clipped = np.clip(depth, 0.1, 5.0) 119 | depth_norm = ((depth_clipped - 0.1) / (5.0 - 0.1) * 255).astype(np.uint8) 120 | depth_color = cv2.applyColorMap(255 - depth_norm, cv2.COLORMAP_JET) 121 | 122 | # Point cloud in camera frame (only from consistent points) 123 | pts_cam = cv2.reprojectImageTo3D(disp_left, Q) / 1000.0 124 | pts_cam = pts_cam[valid] 125 | cols = cv2.cvtColor(left_ds, cv2.COLOR_BGR2RGB).reshape(-1, 3)[valid.ravel()] 126 | 127 | # Throw away far points (> 5 m) for clarity 128 | dist_m = np.linalg.norm(pts_cam, axis=1) 129 | keep = dist_m < 5.0 130 | pts_cam, cols = pts_cam[keep], cols[keep] 131 | 132 | # Camera → robot (tilt & 1.5 m up) 133 | t_cam = np.array([0.0, -1.5, 0.0]) # +Y is down 134 | R_cam = R_scipy.from_euler('x', -36.0, degrees=True) 135 | pts_robot = R_cam.apply(pts_cam) + t_cam # robot frame (y up) 136 | 137 | # Robot at origin → world 138 | pts_world = pts_robot 139 | 140 | # Viewer coords (X right, Y fwd, Z up) 141 | pts_view = _world_to_viewer(pts_world) 142 | 143 | # Sub-sample 4× before logging 144 | pts_view = pts_view[::4] 145 | cols = cols[::4] 146 | 147 | if pts_view.size: 148 | rr.log("world/point_cloud", rr.Points3D(pts_view, colors=cols, radii=0.02)) 149 | 150 | # Images to Rerun 151 | rr.set_time_seconds("time", time.time()) 152 | rr.log("camera/disparity", rr.Image(disp_left)) 153 | rr.log("camera/depth", rr.Image(depth_color)) 154 | 155 | if __name__ == "__main__": 156 | ap = argparse.ArgumentParser() 157 | ap.add_argument("--frames", type=int) 158 | args = ap.parse_args() 159 | main(max_frames=args.frames) 160 | -------------------------------------------------------------------------------- /examples/example_drivepath.py: -------------------------------------------------------------------------------- 1 | # Adds the lib directory to the Python path 2 | import sys 3 | import os 4 | import math 5 | sys.path.append(os.path.join(os.path.dirname(__file__), '..')) 6 | 7 | import time 8 | import json 9 | from lib.odrive_uart import ODriveUART 10 | 11 | class RobotDriver: 12 | def __init__(self, motor): 13 | self.motor = motor 14 | self.wheel_diameter = 0.165 # 165mm in meters 15 | self.wheel_circumference = math.pi * self.wheel_diameter 16 | self.wheel_distance = 0.4 # Distance between wheels in meters (adjust as needed) 17 | 18 | # Starting position 19 | self.left_start_pos = 0 20 | self.right_start_pos = 0 21 | 22 | def update_start_position(self): 23 | """Update the starting position for distance calculations""" 24 | # Get current position in turns 25 | self.left_start_pos = self.motor.get_position_turns_left() 26 | self.right_start_pos = self.motor.get_position_turns_right() 27 | 28 | def drive_distance(self, distance, speed=0.2): 29 | """ 30 | Drive straight for a specified distance 31 | 32 | Args: 33 | distance: Distance to drive in meters 34 | speed: Speed in m/s 35 | """ 36 | print(f"Driving {distance}m at {speed}m/s") 37 | 38 | # Update starting position 39 | self.update_start_position() 40 | 41 | # Calculate required wheel rotations 42 | required_rotations = distance / self.wheel_circumference 43 | 44 | # Set motor speeds 45 | self.motor.set_speed_mps_left(speed) 46 | self.motor.set_speed_mps_right(speed) 47 | 48 | # Monitor distance traveled 49 | while True: 50 | # Get current positions 51 | left_pos = self.motor.get_position_turns_left() 52 | right_pos = self.motor.get_position_turns_right() 53 | 54 | # Calculate rotations since start 55 | left_rotations = abs(left_pos - self.left_start_pos) 56 | right_rotations = abs(right_pos - self.right_start_pos) 57 | 58 | # Use average of both wheels 59 | avg_rotations = (left_rotations + right_rotations) / 2 60 | 61 | if avg_rotations >= required_rotations: 62 | break 63 | 64 | time.sleep(0.01) 65 | 66 | # Stop motors 67 | self.motor.set_speed_mps_left(0) 68 | self.motor.set_speed_mps_right(0) 69 | print(f"Completed driving {distance}m") 70 | 71 | def turn_degrees(self, degrees, turn_speed=0.2): 72 | """ 73 | Turn the robot by a specified angle 74 | 75 | Args: 76 | degrees: Angle to turn in degrees (positive = right/clockwise, negative = left/counterclockwise) 77 | turn_speed: Speed for turning in m/s 78 | """ 79 | print(f"Turning {degrees} degrees") 80 | 81 | # Update starting position 82 | self.update_start_position() 83 | 84 | # Calculate the arc length for the turn 85 | # For a differential drive robot turning in place, each wheel travels in opposite directions 86 | # The arc length is (angle in radians) * (wheel distance / 2) 87 | angle_radians = math.radians(abs(degrees)) 88 | arc_length = angle_radians * (self.wheel_distance / 2) 89 | 90 | # Calculate required wheel rotations 91 | required_rotations = arc_length / self.wheel_circumference 92 | 93 | # Set motor speeds based on turn direction 94 | if degrees > 0: # Turn right/clockwise 95 | self.motor.set_speed_mps_left(turn_speed) 96 | self.motor.set_speed_mps_right(-turn_speed) 97 | else: # Turn left/counterclockwise 98 | self.motor.set_speed_mps_left(-turn_speed) 99 | self.motor.set_speed_mps_right(turn_speed) 100 | 101 | # Monitor rotation 102 | while True: 103 | # Get current positions 104 | left_pos = self.motor.get_position_turns_left() 105 | right_pos = self.motor.get_position_turns_right() 106 | 107 | # Calculate rotations since start (use absolute value) 108 | left_rotations = abs(left_pos - self.left_start_pos) 109 | right_rotations = abs(right_pos - self.right_start_pos) 110 | 111 | # Use average of both wheels 112 | avg_rotations = (left_rotations + right_rotations) / 2 113 | 114 | if avg_rotations >= required_rotations: 115 | break 116 | 117 | time.sleep(0.01) 118 | 119 | # Stop motors 120 | self.motor.set_speed_mps_left(0) 121 | self.motor.set_speed_mps_right(0) 122 | print(f"Completed turning {degrees} degrees") 123 | 124 | def drive_circle(self, radius=0.5, speed=0.2, duration=10): 125 | """ 126 | Drive in a circle with specified radius, speed, and duration. 127 | 128 | Args: 129 | radius: Circle radius in meters 130 | speed: Linear speed in m/s 131 | duration: Duration in seconds 132 | """ 133 | print(f"Driving in a circle: radius={radius}m, speed={speed}m/s, duration={duration}s") 134 | 135 | # Calculate wheel speeds for circular motion 136 | inner_speed = speed * (radius - self.wheel_distance/2) / radius 137 | outer_speed = speed * (radius + self.wheel_distance/2) / radius 138 | 139 | start_time = time.time() 140 | while time.time() - start_time < duration: 141 | # For right turn (clockwise) 142 | self.motor.set_speed_mps_left(outer_speed) 143 | self.motor.set_speed_mps_right(inner_speed) 144 | time.sleep(0.01) 145 | 146 | # Stop after completing the circle 147 | self.motor.set_speed_mps_left(0) 148 | self.motor.set_speed_mps_right(0) 149 | print("Circle completed") 150 | 151 | def drive_square(self, side_length=1.0, speed=0.2): 152 | """ 153 | Drive in a square with specified side length and speed using helper functions. 154 | 155 | Args: 156 | side_length: Length of each side in meters 157 | speed: Linear speed in m/s 158 | """ 159 | print(f"Driving in a square: side_length={side_length}m, speed={speed}m/s") 160 | 161 | for i in range(4): 162 | print(f"Side {i+1}") 163 | # Drive one side 164 | self.drive_distance(side_length, speed) 165 | time.sleep(0.5) # Pause briefly 166 | 167 | # Turn 90 degrees right 168 | self.turn_degrees(90, speed/2) 169 | time.sleep(0.5) # Pause briefly 170 | 171 | print("Square completed") 172 | 173 | def main(): 174 | # Load motor direction configuration 175 | try: 176 | with open(os.path.expanduser('~/quickstart/lib/motor_dir.json'), 'r') as f: 177 | motor_dirs = json.load(f) 178 | except Exception as e: 179 | print(f"Error reading motor_dir.json: {e}") 180 | return 181 | 182 | # Initialize motor controller 183 | motor = ODriveUART(port='/dev/ttyAMA1', left_axis=0, right_axis=1, 184 | dir_left=motor_dirs['left'], dir_right=motor_dirs['right']) 185 | 186 | # Start motors and set to velocity mode 187 | motor.start_left() 188 | motor.start_right() 189 | motor.enable_velocity_mode_left() 190 | motor.enable_velocity_mode_right() 191 | motor.disable_watchdog_left() 192 | motor.disable_watchdog_right() 193 | motor.clear_errors_left() 194 | motor.clear_errors_right() 195 | 196 | # Create robot driver 197 | robot = RobotDriver(motor) 198 | 199 | try: 200 | # Test the helper functions 201 | print("Testing drive_distance...") 202 | robot.drive_distance(0.5, 0.2) # Drive 0.5m at 0.2m/s 203 | time.sleep(1) 204 | 205 | print("Testing turn_degrees...") 206 | robot.turn_degrees(90, 0.2) # Turn 90 degrees right 207 | time.sleep(1) 208 | robot.turn_degrees(-90, 0.2) # Turn 90 degrees left 209 | time.sleep(1) 210 | 211 | # First drive in a circle 212 | robot.drive_circle(radius=0.1, speed=0.2, duration=10) 213 | time.sleep(2) 214 | 215 | # Then drive in a square 216 | robot.drive_square(side_length=0.5, speed=0.2) 217 | 218 | except KeyboardInterrupt: 219 | print("\nProgram interrupted by user") 220 | except Exception as e: 221 | print(f"Error: {e}") 222 | finally: 223 | # Always stop motors and clean up 224 | print("Stopping motors") 225 | motor.set_speed_mps_left(0) 226 | motor.set_speed_mps_right(0) 227 | 228 | if __name__ == "__main__": 229 | main() 230 | 231 | -------------------------------------------------------------------------------- /examples/example_follow.py: -------------------------------------------------------------------------------- 1 | # Adds the lib directory to the Python path 2 | import sys 3 | import os 4 | sys.path.append(os.path.join(os.path.dirname(__file__), '..')) 5 | 6 | import cv2 7 | import numpy as np 8 | import json 9 | from time import time 10 | from ultralytics import YOLO 11 | from lib.camera import StereoCamera 12 | from lib.odrive_uart import ODriveUART 13 | 14 | # Silence YOLO 15 | os.environ['ULTRALYTICS_HIDE_CONSOLE'] = '1' 16 | os.environ['YOLO_VERBOSE'] = 'False' 17 | os.environ['ULTRALYTICS_QUIET'] = 'True' 18 | os.environ['DISABLE_ULTRALYTICS_VERSIONING_CHECK'] = 'True' 19 | 20 | TURN_SPEED = 0.3 21 | CENTER_THRESHOLD = 0.05 22 | FORWARD_SPEED = 0.25 # Speed for moving forward/backward 23 | TARGET_WIDTH_RATIO = 0.4 # Target width of person relative to image width 24 | WIDTH_THRESHOLD = 0.05 # Acceptable range around target width 25 | 26 | class PersonFollower: 27 | def __init__(self): 28 | self.camera = StereoCamera(scale=0.5) 29 | 30 | model_path = "yolov8n.pt" 31 | ncnn_path = "yolov8n_ncnn_model" 32 | 33 | if not os.path.exists(ncnn_path): 34 | self.model = YOLO(model_path if os.path.exists(model_path) else 'yolov8n') 35 | self.model.export(format="ncnn") 36 | 37 | self.model = YOLO(ncnn_path) 38 | 39 | with open(os.path.expanduser('~/quickstart/lib/motor_dir.json'), 'r') as f: 40 | motor_dirs = json.load(f) 41 | 42 | self.motor = ODriveUART(port='/dev/ttyAMA1', left_axis=0, right_axis=1, dir_left=motor_dirs['left'], dir_right=motor_dirs['right']) 43 | 44 | self.motor.start_left() 45 | self.motor.start_right() 46 | self.motor.enable_velocity_ramp_mode_left() 47 | self.motor.enable_velocity_ramp_mode_right() 48 | self.motor.disable_watchdog_left() 49 | self.motor.disable_watchdog_right() 50 | self.motor.clear_errors_left() 51 | self.motor.clear_errors_right() 52 | 53 | self.debug_dir = 'debug_images' 54 | os.makedirs(self.debug_dir, exist_ok=True) 55 | 56 | # Current motor commands 57 | self.current_left_speed = 0 58 | self.current_right_speed = 0 59 | 60 | def save_debug_images(self, left_image, annotated, best_person=None): 61 | """Save debug images with basic annotations""" 62 | # Save detection result 63 | cv2.imwrite(os.path.join(self.debug_dir, 'detection.jpg'), annotated) 64 | 65 | # Create combined visualization 66 | if best_person is not None: 67 | h, w = left_image.shape[:2] 68 | center_x = int((best_person[0] + best_person[2]) / 2) 69 | center_y = int((best_person[1] + best_person[3]) / 2) 70 | 71 | combined = left_image.copy() 72 | 73 | # Draw crosshair at person center 74 | cv2.line(combined, (center_x - 20, center_y), (center_x + 20, center_y), (0, 255, 0), 2) 75 | cv2.line(combined, (center_x, center_y - 20), (center_x, center_y + 20), (0, 255, 0), 2) 76 | 77 | # Draw vertical center line 78 | cv2.line(combined, (w//2, 0), (w//2, h), (255, 0, 0), 1) 79 | 80 | cv2.imwrite(os.path.join(self.debug_dir, 'combined.jpg'), combined) 81 | 82 | def detect_and_follow(self): 83 | """Main loop to detect and follow a person""" 84 | # Get images 85 | left, right = self.camera.get_stereo() 86 | if left is None or right is None: 87 | return 88 | 89 | # Detect people 90 | results = self.model(left, classes=[0], verbose=False) 91 | 92 | # Find the largest person detection 93 | best_person = None 94 | max_area = 0 95 | 96 | for result in results[0].boxes: 97 | box = result.xyxy[0].cpu().numpy() 98 | area = (box[2] - box[0]) * (box[3] - box[1]) 99 | if area > max_area: 100 | max_area = area 101 | best_person = box 102 | 103 | # Get annotated image 104 | annotated = results[0].plot() 105 | 106 | if best_person is None: 107 | self.motor.set_speed_mps_left(0) 108 | self.motor.set_speed_mps_right(0) 109 | self.current_left_speed = 0 110 | self.current_right_speed = 0 111 | self.save_debug_images(left, annotated) 112 | return 113 | 114 | # Get center point and width of the person 115 | center_x = (best_person[0] + best_person[2]) / 2 116 | image_center_x = left.shape[1] / 2 117 | x_error = (center_x - image_center_x) / image_center_x # -1 to 1 118 | 119 | # Calculate width ratio of person relative to image 120 | person_width = best_person[2] - best_person[0] 121 | image_width = left.shape[1] 122 | width_ratio = person_width / image_width 123 | width_error = width_ratio - TARGET_WIDTH_RATIO 124 | 125 | # Determine forward/backward speed based on width 126 | forward_speed = 0 127 | if abs(width_error) > WIDTH_THRESHOLD: 128 | # If person is too far (small), move forward 129 | # If person is too close (large), move backward 130 | forward_speed = -FORWARD_SPEED * (width_error / abs(width_error)) 131 | 132 | # Bang-bang control: fixed speed turn based on which side person is on 133 | if abs(x_error) < CENTER_THRESHOLD: 134 | # Person is centered - maintain forward/backward motion only 135 | self.current_left_speed = forward_speed 136 | self.current_right_speed = forward_speed 137 | elif x_error > 0: 138 | # Person is to the right - turn right 139 | self.current_left_speed = TURN_SPEED*abs(x_error) + forward_speed 140 | self.current_right_speed = -TURN_SPEED*abs(x_error) + forward_speed 141 | else: 142 | # Person is to the left - turn left 143 | self.current_left_speed = -TURN_SPEED*abs(x_error) + forward_speed 144 | self.current_right_speed = TURN_SPEED*abs(x_error) + forward_speed 145 | 146 | # Send commands to motors 147 | self.motor.set_speed_mps_left(self.current_left_speed) 148 | self.motor.set_speed_mps_right(self.current_right_speed) 149 | # print(self.current_left_speed, self.current_right_speed) 150 | 151 | # Save debug images 152 | self.save_debug_images(left, annotated, best_person) 153 | 154 | def cleanup(self): 155 | """Clean up resources""" 156 | self.motor.set_speed_mps_left(0) 157 | self.motor.set_speed_mps_right(0) 158 | self.motor.clear_errors_left() 159 | self.motor.clear_errors_right() 160 | 161 | def main(): 162 | follower = PersonFollower() 163 | try: 164 | while True: 165 | follower.detect_and_follow() 166 | except KeyboardInterrupt: 167 | follower.cleanup() 168 | 169 | if __name__ == "__main__": 170 | main() 171 | -------------------------------------------------------------------------------- /examples/example_kokoro.py: -------------------------------------------------------------------------------- 1 | """ 2 | Test script for the Kokoro Text-to-Speech package. 3 | 4 | pip install this package https://github.com/hexgrad/kokoro 5 | 6 | i had to do this to install spacy "pip install --only-binary :all: spacy" 7 | """ 8 | 9 | from kokoro.pipeline import KPipeline 10 | import numpy as np 11 | import sounddevice as sd 12 | import time 13 | 14 | sample_text = "Hello World, im Bracket Bot, and I generated this audio entirely on my own." 15 | 16 | 17 | def main(): 18 | print("Initializing Kokoro pipeline...") 19 | pipeline = KPipeline(lang_code='a') 20 | 21 | print(f"Generating speech for: '{sample_text}'") 22 | 23 | start_time = time.time() 24 | 25 | audio_chunks = [] 26 | for gs, ps, audio in pipeline(sample_text, voice='am_adam'): 27 | audio_chunks.append(audio.detach().cpu().numpy()) 28 | 29 | audio = np.concatenate(audio_chunks) if len(audio_chunks) > 1 else audio_chunks[0] 30 | 31 | if hasattr(audio, 'detach'): 32 | audio = audio.detach().cpu().numpy() 33 | 34 | end_time = time.time() 35 | elapsed_time = end_time - start_time 36 | print(f"Audio generation took {elapsed_time:.2f} seconds.") 37 | 38 | print("Playing audio...") 39 | sd.play(audio, samplerate=24000) 40 | sd.wait() 41 | print("Done! Audio played successfully.") 42 | 43 | if __name__ == "__main__": 44 | main() 45 | -------------------------------------------------------------------------------- /examples/example_localization.py: -------------------------------------------------------------------------------- 1 | """localization.py 2 | Simple differential-drive odometry helper built on top of 3 | `lib.odrive_uart.ODriveUART`. It continuously queries wheel encoder 4 | counts from the ODrive and integrates them to estimate the robot pose 5 | (x, z, yaw) in a planar world frame whose origin coincides with the 6 | robot's start position (0, 0, 0). 7 | 8 | When executed directly (`python -m lib.localization` or 9 | `python lib/localization.py`) it also visualises the live pose and the 10 | traversed path in Rerun as a 3‑D line strip under the `world/odom/path` 11 | entity, alongside a transform representing the current robot pose at 12 | `world/robot`. 13 | 14 | This module can likewise be imported by other scripts (e.g. 15 | `examples/navigate.py`) to obtain a continuously updated pose estimate 16 | via the `DifferentialDriveOdometry` class or the convenience 17 | `OdometryThread` wrapper. 18 | """ 19 | 20 | from __future__ import annotations 21 | 22 | import math 23 | import threading 24 | import time 25 | from dataclasses import dataclass, field 26 | from pathlib import Path 27 | from typing import List, Tuple 28 | 29 | import numpy as np 30 | import rerun as rr 31 | from scipy.spatial.transform import Rotation as R_scipy 32 | import json 33 | 34 | # Add the project root so `lib` is importable when run as a script 35 | import os, sys 36 | sys.path.append(os.path.join(os.path.dirname(__file__), "..")) 37 | 38 | from lib.odrive_uart import ODriveUART # noqa: E402 39 | 40 | # ───────────────────────── CONFIG ────────────────────────── 41 | WHEEL_DIAMETER_M: float = 0.165 # 165 mm wheels 42 | WHEEL_BASE_M: float = 0.425 # distance between wheels (centre‑to‑centre) 43 | LOG_HZ: float = 5.0 # log frequency for standalone visualiser 44 | RERUN_TCP_ADDR: str = "192.168.2.24:9876" # change if needed 45 | 46 | 47 | def _circumference(diam_m: float) -> float: 48 | """Return wheel circumference in metres given its diameter.""" 49 | return math.pi * diam_m 50 | 51 | 52 | @dataclass 53 | class DifferentialDriveOdometry: 54 | """Incremental planar odometry for a differential‑drive robot.""" 55 | 56 | wheel_base: float = WHEEL_BASE_M 57 | wheel_diameter: float = WHEEL_DIAMETER_M 58 | 59 | # Pose estimate in world frame (x points forward, z to the left to 60 | # match the RHS Y‑down convention used elsewhere in the codebase). 61 | x: float = 0.0 62 | z: float = 0.0 63 | yaw: float = 0.0 # radians, 0 = facing +x 64 | 65 | # Last encoder readings (in turns) to compute deltas. 66 | _prev_left_turns: float | None = field(default=None, init=False) 67 | _prev_right_turns: float | None = field(default=None, init=False) 68 | 69 | def reset(self) -> None: 70 | """Reset pose to the origin and forget previous encoder values.""" 71 | self.x = self.z = self.yaw = 0.0 72 | self._prev_left_turns = self._prev_right_turns = None 73 | 74 | # ───────────── Main update ────────────── 75 | 76 | def update(self, left_turns: float, right_turns: float) -> Tuple[float, float, float]: 77 | """Advance the pose estimate given absolute wheel encoder turns. 78 | 79 | Parameters 80 | ---------- 81 | left_turns, right_turns : float 82 | Cumulative encoder readings in *turns* since the ODrive was 83 | powered on. Sign conventions depend on the `dir_left` / 84 | `dir_right` parameters passed to ``ODriveUART``. 85 | 86 | Returns 87 | ------- 88 | (x, z, yaw) : tuple of float 89 | The updated pose estimate in metres / radians. 90 | """ 91 | if self._prev_left_turns is None: 92 | # First call → just cache values without integrating. 93 | self._prev_left_turns = left_turns 94 | self._prev_right_turns = right_turns 95 | return self.x, self.z, self.yaw 96 | 97 | # ∆ wheel travel in metres 98 | dl = (left_turns - self._prev_left_turns) * _circumference(self.wheel_diameter) 99 | dr = (right_turns - self._prev_right_turns) * _circumference(self.wheel_diameter) 100 | 101 | self._prev_left_turns = left_turns 102 | self._prev_right_turns = right_turns 103 | 104 | # Differential‑drive kinematics 105 | dc = 0.5 * (dl + dr) # forward distance 106 | dtheta = (dr - dl) / self.wheel_base # change in heading 107 | 108 | # Integrate in SE(2) — here we use the exact integration for the 109 | # constant‑twist motion over the dt interval. 110 | if abs(dtheta) > 1e-6: 111 | # follow an arc 112 | R_icc = dc / dtheta # radius to the ICC (Instantaneous Centre of Curvature) 113 | half_yaw = self.yaw + 0.5 * dtheta 114 | # Use cos for forward (x) and sin for left (z) to match 115 | # the world‑frame convention (x forward, z left). This 116 | # was previously inverted which caused forward motion to 117 | # appear as positive left in the pose estimate. 118 | self.x += dc * math.cos(half_yaw) 119 | self.z += dc * math.sin(half_yaw) 120 | else: 121 | # straight‑line approximation 122 | self.x += dc * math.cos(self.yaw) 123 | self.z += dc * math.sin(self.yaw) 124 | 125 | self.yaw = (self.yaw + dtheta) % (2 * math.pi) 126 | return self.x, self.z, self.yaw 127 | 128 | 129 | class OdometryThread(threading.Thread): 130 | """Background thread that keeps a `DifferentialDriveOdometry` updated.""" 131 | 132 | def __init__(self, odrv: ODriveUART, log_to_rerun: bool = False): 133 | super().__init__(daemon=True) 134 | self.odrv = odrv 135 | self.odo = DifferentialDriveOdometry() 136 | self.log_to_rerun = log_to_rerun 137 | self._running = threading.Event() 138 | self._running.set() 139 | self._path_world: List[Tuple[float, float, float]] = [] # stored as (x, y, z) 140 | self._boxes_initialized: bool = False 141 | 142 | if log_to_rerun: 143 | rr.init("robot_localization") 144 | rr.connect_grpc(f"rerun+http://{RERUN_TCP_ADDR}/proxy") 145 | rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Z_UP, static=True) 146 | 147 | # Expose latest pose 148 | @property 149 | def pose(self) -> Tuple[float, float, float]: # (x, z, yaw) 150 | return self.odo.x, self.odo.z, self.odo.yaw 151 | 152 | # Full path for visualisation or analysis 153 | @property 154 | def path(self) -> np.ndarray: 155 | return np.asarray(self._path_world, dtype=np.float32) 156 | 157 | def run(self) -> None: 158 | period_s = 1.0 / LOG_HZ 159 | last_time = time.time() 160 | 161 | while self._running.is_set(): 162 | try: 163 | l_turns = self.odrv.get_position_turns_left() 164 | r_turns = self.odrv.get_position_turns_right() 165 | x, z, yaw = self.odo.update(l_turns, r_turns) 166 | # Map body left (+z) onto viewer left (−X) so that driving 167 | # left in body frame shows as left in the viewer. 168 | x_view = -z # X right positive → body right (−z) 169 | y_view = x # Y forward = body forward 170 | yaw_view = yaw - math.pi / 2 # viewer yaw lags body 90°, matches axis mapping (Δyaw ≈ -90°) 171 | 172 | # Extra debug: show delta between body yaw and viewer yaw and explicit axis meanings 173 | diff_deg = ((math.degrees(yaw_view - yaw) + 180) % 360) - 180 # should be ~ -90° constant if correct 174 | 175 | print( 176 | ( 177 | f"Body -> fwd(x)={x:+.3f} left(z)={z:+.3f} yaw={math.degrees(yaw):+.1f}° | " 178 | f"Viewer -> X(right)={x_view:+.3f} Y(fwd)={y_view:+.3f} yaw_view={math.degrees(yaw_view):+.1f}° " 179 | f"(Δyaw={diff_deg:+.1f}°)" 180 | ) 181 | ) 182 | 183 | self._path_world.append((x_view, y_view, 0.0)) 184 | # Quaternion about Z for viewer frame, with 90 deg rotation 185 | quat_rr = R_scipy.from_euler("z", yaw_view - math.pi/2).as_quat() 186 | 187 | if self.log_to_rerun and (time.time() - last_time) >= period_s: 188 | last_time = time.time() 189 | rr.set_time_seconds("t", last_time) 190 | rr.log( 191 | "world/robot", 192 | rr.Transform3D( 193 | translation=[x_view, y_view, 0.0], 194 | rotation=rr.Quaternion(xyzw=quat_rr), 195 | ), 196 | ) 197 | # Log the robot's 3D model once per session 198 | if not self._boxes_initialized: 199 | model_path = Path(__file__).parent.parent / "lib" / "Bracketbot.stl" 200 | rr.log( 201 | "world/robot", 202 | rr.Asset3D(path=str(model_path)), 203 | static=True 204 | ) 205 | self._boxes_initialized = True 206 | if len(self._path_world) >= 2: 207 | pts = np.asarray(self._path_world[-1000:], dtype=np.float32) 208 | rr.log("world/odom/path", rr.LineStrips3D(pts, radii=0.02)) 209 | except Exception as exc: 210 | print(f"[Odometry] error: {exc}") 211 | time.sleep(period_s) 212 | 213 | def stop(self) -> None: 214 | self._running.clear() 215 | 216 | 217 | # ──────────────────── Stand‑alone entry‑point ──────────────────── 218 | 219 | 220 | def _standalone() -> None: # pragma: no cover (utility script) 221 | odrv = ODriveUART() 222 | odo_thread = OdometryThread(odrv, log_to_rerun=True) 223 | odo_thread.start() 224 | print("[Odometry] running – press Ctrl‑C to quit…") 225 | 226 | try: 227 | while True: 228 | time.sleep(0.1) 229 | except KeyboardInterrupt: 230 | print("Exiting…") 231 | finally: 232 | odo_thread.stop() 233 | odo_thread.join() 234 | 235 | 236 | if __name__ == "__main__": 237 | _standalone() 238 | -------------------------------------------------------------------------------- /examples/example_realtime.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | """ 3 | OpenAI Realtime WebRTC Single-file Example - Compact Version 4 | 5 | Requirements: sounddevice, numpy, websockets, openai, aiohttp, pyaudio, python-dotenv, aiortc, scipy 6 | 7 | Install: pip install sounddevice numpy websockets openai aiohttp pyaudio python-dotenv aiortc scipy 8 | 9 | Usage: 10 | 1. Create a .env file with OPENAI_API_KEY=your-key-here 11 | 2. Run: python openai_realtime_example.py 12 | 3. Speak into your microphone 13 | 4. Press Ctrl+C to exit 14 | """ 15 | 16 | import asyncio 17 | import os 18 | import logging 19 | import numpy as np 20 | import sounddevice as sd 21 | import aiohttp 22 | import json 23 | from collections import deque 24 | from typing import Optional, Callable, Dict, Any, Deque, Tuple 25 | from aiortc import RTCPeerConnection, RTCSessionDescription, MediaStreamTrack, RTCConfiguration, RTCIceServer 26 | from aiortc.mediastreams import MediaStreamError 27 | from av import AudioFrame 28 | from dotenv import load_dotenv 29 | 30 | # Setup logging 31 | logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') 32 | logger = logging.getLogger(__name__) 33 | 34 | # Audio constants 35 | SAMPLE_RATE = 48000 36 | CHANNELS = 1 37 | DTYPE = np.int16 38 | FRAME_DURATION_MS = 20 39 | DEFAULT_CHANNELS = 2 40 | DEFAULT_BLOCK_SIZE = int(SAMPLE_RATE * FRAME_DURATION_MS / 1000) 41 | 42 | 43 | class AudioTrack(MediaStreamTrack): 44 | kind = "audio" 45 | 46 | def __init__(self, audio_handler): 47 | super().__init__() 48 | self._audio_handler = audio_handler 49 | self._queue = asyncio.Queue() 50 | self._task = None 51 | 52 | async def recv(self): 53 | if self._task is None: 54 | self._task = asyncio.create_task(self._audio_handler.start_recording(self._queue)) 55 | try: 56 | return await self._queue.get() 57 | except Exception as e: 58 | logger.error(f"Error receiving audio frame: {str(e)}") 59 | raise MediaStreamError("Failed to receive audio frame") 60 | 61 | 62 | class AudioHandler: 63 | def __init__(self, sample_rate=SAMPLE_RATE, channels=CHANNELS, frame_duration=20, 64 | dtype=DTYPE, device=None): 65 | self.sample_rate = sample_rate 66 | self.channels = channels 67 | self.frame_duration = frame_duration 68 | self.dtype = dtype 69 | self.device = device 70 | self.frame_size = int(sample_rate * frame_duration / 1000) 71 | self.stream = None 72 | self.is_recording = False 73 | self.is_paused = False 74 | self._loop = None 75 | self._pts = 0 76 | 77 | def create_audio_track(self): 78 | return AudioTrack(self) 79 | 80 | async def start_recording(self, queue): 81 | if self.is_recording: 82 | return 83 | 84 | self.is_recording = True 85 | self.is_paused = False 86 | self._loop = asyncio.get_running_loop() 87 | self._pts = 0 88 | 89 | try: 90 | def callback(indata, frames, time, status): 91 | if status: 92 | logger.warning(f"Audio input status: {status}") 93 | if not self.is_paused: 94 | audio_data = indata.copy() 95 | if audio_data.dtype != self.dtype: 96 | if self.dtype == np.int16: 97 | audio_data = (audio_data * 32767).astype(self.dtype) 98 | else: 99 | audio_data = audio_data.astype(self.dtype) 100 | 101 | frame = AudioFrame(samples=len(audio_data), layout='mono', format='s16') 102 | frame.rate = self.sample_rate 103 | frame.pts = self._pts 104 | self._pts += len(audio_data) 105 | frame.planes[0].update(audio_data.tobytes()) 106 | asyncio.run_coroutine_threadsafe(queue.put(frame), self._loop) 107 | 108 | self.stream = sd.InputStream( 109 | device=self.device, 110 | channels=self.channels, 111 | samplerate=self.sample_rate, 112 | dtype=self.dtype, 113 | blocksize=self.frame_size, 114 | callback=callback 115 | ) 116 | self.stream.start() 117 | 118 | while self.is_recording: 119 | await asyncio.sleep(0.1) 120 | except Exception as e: 121 | logger.error(f"Error in audio recording: {str(e)}") 122 | raise 123 | finally: 124 | await self.stop() 125 | 126 | async def stop(self): 127 | self.is_recording = False 128 | if self.stream: 129 | self.stream.stop() 130 | self.stream.close() 131 | self.stream = None 132 | 133 | async def pause(self): 134 | self.is_paused = True 135 | 136 | async def resume(self): 137 | self.is_paused = False 138 | 139 | def set_device(self, device_id): 140 | self.device = device_id 141 | if self.stream: 142 | self.stream.stop() 143 | self.stream.close() 144 | self.stream = None 145 | 146 | 147 | class AudioOutput: 148 | def __init__(self, sample_rate=SAMPLE_RATE, channels=DEFAULT_CHANNELS, 149 | dtype=DTYPE, device=None, buffer_size=50): 150 | self.sample_rate = sample_rate 151 | self.channels = channels 152 | self.dtype = dtype 153 | self.device = device 154 | self.block_size = DEFAULT_BLOCK_SIZE 155 | self._buffer = deque(maxlen=buffer_size) 156 | self._queue = asyncio.Queue() 157 | self.stream = None 158 | self.is_playing = False 159 | self._task = None 160 | self._remaining_data = None 161 | 162 | async def start(self): 163 | if self.is_playing: 164 | return 165 | 166 | try: 167 | self.stream = sd.OutputStream( 168 | device=self.device, 169 | samplerate=self.sample_rate, 170 | channels=self.channels, 171 | dtype=self.dtype, 172 | blocksize=self.block_size, 173 | callback=self._audio_callback, 174 | prime_output_buffers_using_stream_callback=True 175 | ) 176 | self.stream.start() 177 | self.is_playing = True 178 | self._task = asyncio.create_task(self._process_audio()) 179 | except Exception as e: 180 | logger.error(f"Failed to start audio output: {str(e)}") 181 | raise 182 | 183 | async def stop(self): 184 | if not self.is_playing: 185 | return 186 | 187 | try: 188 | self.is_playing = False 189 | if self._task: 190 | self._task.cancel() 191 | self._task = None 192 | 193 | if self.stream: 194 | self.stream.stop() 195 | self.stream.close() 196 | self.stream = None 197 | 198 | self._buffer.clear() 199 | while not self._queue.empty(): 200 | try: 201 | self._queue.get_nowait() 202 | except asyncio.QueueEmpty: 203 | break 204 | 205 | logger.info("Audio output stopped") 206 | except Exception as e: 207 | logger.error(f"Error stopping audio output: {str(e)}") 208 | 209 | def _audio_callback(self, outdata, frames, time, status): 210 | if status: 211 | logger.warning(f"Audio output status: {status}") 212 | 213 | try: 214 | if len(self._buffer) > 0: 215 | data = self._buffer.popleft() 216 | outdata[:] = data.reshape(outdata.shape) 217 | else: 218 | outdata.fill(0) 219 | except Exception as e: 220 | logger.error(f"Error in audio callback: {str(e)}") 221 | outdata.fill(0) 222 | 223 | async def _process_audio(self): 224 | try: 225 | while self.is_playing: 226 | if len(self._buffer) < self._buffer.maxlen: 227 | try: 228 | data = await asyncio.wait_for(self._queue.get(), 0.1) 229 | self._buffer.append(data) 230 | except (asyncio.TimeoutError, asyncio.QueueEmpty): 231 | continue 232 | except Exception as e: 233 | logger.error(f"Error getting data from queue: {str(e)}") 234 | else: 235 | await asyncio.sleep(0.001) 236 | except asyncio.CancelledError: 237 | logger.info("Audio processing task cancelled") 238 | except Exception as e: 239 | logger.error(f"Error processing audio: {str(e)}") 240 | 241 | async def play_frame(self, frame): 242 | try: 243 | audio_data = frame.to_ndarray() 244 | if audio_data.dtype != self.dtype: 245 | audio_data = audio_data.astype(self.dtype) 246 | 247 | if self._queue.qsize() >= self._buffer.maxlen: 248 | try: 249 | self._queue.get_nowait() 250 | except asyncio.QueueEmpty: 251 | pass 252 | 253 | await self._queue.put(audio_data) 254 | except Exception as e: 255 | logger.error(f"Error queueing audio frame: {str(e)}") 256 | raise 257 | 258 | 259 | class WebRTCManager: 260 | OPENAI_API_BASE = "https://api.openai.com/v1" 261 | REALTIME_SESSION_URL = f"{OPENAI_API_BASE}/realtime/sessions" 262 | REALTIME_URL = f"{OPENAI_API_BASE}/realtime" 263 | 264 | def __init__(self): 265 | self.ice_servers = [RTCIceServer(urls=["stun:stun.l.google.com:19302"])] 266 | self.audio_output = None 267 | self.peer_connection = None 268 | self.output_device = None 269 | 270 | async def create_connection(self): 271 | config = RTCConfiguration(iceServers=self.ice_servers) 272 | self.peer_connection = RTCPeerConnection(config) 273 | 274 | self.audio_output = AudioOutput(device=self.output_device) 275 | await self.audio_output.start() 276 | 277 | @self.peer_connection.on("track") 278 | async def on_track(track): 279 | logger.info(f"Received {track.kind} track from remote") 280 | if track.kind == "audio": 281 | @track.on("ended") 282 | async def on_ended(): 283 | if self.audio_output: 284 | await self.audio_output.stop() 285 | 286 | while True: 287 | try: 288 | frame = await track.recv() 289 | if self.audio_output and frame: 290 | await self.audio_output.play_frame(frame) 291 | except Exception as e: 292 | logger.error(f"Error processing remote audio frame: {str(e)}") 293 | break 294 | 295 | @self.peer_connection.on("connectionstatechange") 296 | async def on_connectionstatechange(): 297 | logger.info(f"Connection state changed to: {self.peer_connection.connectionState}") 298 | if self.peer_connection.connectionState == "failed" and self.audio_output: 299 | await self.audio_output.stop() 300 | 301 | @self.peer_connection.on("iceconnectionstatechange") 302 | async def on_iceconnectionstatechange(): 303 | logger.info(f"ICE connection state changed to: {self.peer_connection.iceConnectionState}") 304 | 305 | return self.peer_connection 306 | 307 | async def cleanup(self): 308 | if self.audio_output: 309 | await self.audio_output.stop() 310 | self.audio_output = None 311 | 312 | if self.peer_connection: 313 | await self.peer_connection.close() 314 | self.peer_connection = None 315 | 316 | async def get_ephemeral_token(self, api_key, model, system_prompt=None): 317 | headers = { 318 | "Authorization": f"Bearer {api_key}", 319 | "Content-Type": "application/json" 320 | } 321 | 322 | data = {"model": model, "voice": "alloy"} 323 | if system_prompt: 324 | data["instructions"] = system_prompt 325 | 326 | try: 327 | async with aiohttp.ClientSession() as session: 328 | async with session.post(self.REALTIME_SESSION_URL, headers=headers, json=data) as response: 329 | if response.status not in [200, 201]: 330 | error_text = await response.text() 331 | raise Exception(f"Failed to get ephemeral token: {error_text}") 332 | 333 | result = await response.json() 334 | return result["client_secret"]["value"] 335 | except Exception as e: 336 | logger.error(f"Failed to get ephemeral token: {str(e)}") 337 | raise 338 | 339 | async def connect_to_openai(self, api_key, model, offer, system_prompt=None): 340 | try: 341 | ephemeral_token = await self.get_ephemeral_token(api_key, model, system_prompt) 342 | headers = { 343 | "Authorization": f"Bearer {ephemeral_token}", 344 | "Content-Type": "application/sdp" 345 | } 346 | 347 | async with aiohttp.ClientSession() as session: 348 | async with session.post(f"{self.REALTIME_URL}?model={model}", 349 | headers=headers, data=offer.sdp) as response: 350 | if response.status not in [200, 201]: 351 | error_text = await response.text() 352 | raise Exception(f"OpenAI WebRTC error: {error_text}") 353 | 354 | sdp_answer = await response.text() 355 | return {"type": "answer", "sdp": sdp_answer} 356 | except Exception as e: 357 | logger.error(f"Failed to connect to OpenAI: {str(e)}") 358 | raise 359 | 360 | async def handle_ice_candidate(self, peer_connection, candidate): 361 | try: 362 | await peer_connection.addIceCandidate(candidate) 363 | except Exception as e: 364 | logger.error(f"Error adding ICE candidate: {str(e)}") 365 | raise 366 | 367 | 368 | class OpenAIWebRTCClient: 369 | def __init__(self, api_key, model="gpt-4o-realtime-preview", sample_rate=SAMPLE_RATE, 370 | channels=CHANNELS, frame_duration=FRAME_DURATION_MS, 371 | input_device=None, output_device=None, system_prompt=None): 372 | self.api_key = api_key 373 | self.model = model 374 | self.system_prompt = system_prompt 375 | self.audio_handler = AudioHandler( 376 | sample_rate=sample_rate, 377 | channels=channels, 378 | frame_duration=frame_duration, 379 | device=input_device 380 | ) 381 | self.webrtc_manager = WebRTCManager() 382 | if output_device is not None: 383 | self.webrtc_manager.output_device = output_device 384 | 385 | self.peer_connection = None 386 | self.is_streaming = False 387 | self.on_transcription = None 388 | 389 | async def start_streaming(self): 390 | if self.is_streaming: 391 | logger.warning("Streaming is already active") 392 | return 393 | 394 | try: 395 | self.peer_connection = await self.webrtc_manager.create_connection() 396 | audio_track = self.audio_handler.create_audio_track() 397 | self.peer_connection.addTransceiver(audio_track, "sendrecv") 398 | 399 | offer = await self.peer_connection.createOffer() 400 | await self.peer_connection.setLocalDescription(offer) 401 | 402 | response = await self.webrtc_manager.connect_to_openai( 403 | self.api_key, self.model, offer, self.system_prompt 404 | ) 405 | 406 | answer = RTCSessionDescription(sdp=response["sdp"], type=response["type"]) 407 | await self.peer_connection.setRemoteDescription(answer) 408 | 409 | self.is_streaming = True 410 | logger.info("Streaming started successfully") 411 | except Exception as e: 412 | logger.error(f"Failed to start streaming: {str(e)}") 413 | await self.stop_streaming() 414 | raise 415 | 416 | async def stop_streaming(self): 417 | if not self.is_streaming: 418 | return 419 | 420 | try: 421 | await self.webrtc_manager.cleanup() 422 | await self.audio_handler.stop() 423 | self.is_streaming = False 424 | logger.info("Streaming stopped successfully") 425 | except Exception as e: 426 | logger.error(f"Error while stopping streaming: {str(e)}") 427 | raise 428 | 429 | async def pause_streaming(self): 430 | if self.is_streaming: 431 | await self.audio_handler.pause() 432 | 433 | async def resume_streaming(self): 434 | if self.is_streaming: 435 | await self.audio_handler.resume() 436 | 437 | def set_audio_device(self, device_id): 438 | self.audio_handler.set_device(device_id) 439 | 440 | def _handle_transcription(self, text): 441 | if self.on_transcription: 442 | self.on_transcription(text) 443 | 444 | 445 | async def main(): 446 | api_key = os.getenv("OPENAI_API_KEY") 447 | if not api_key: 448 | print("Error: OPENAI_API_KEY environment variable not set.") 449 | print("Please create a .env file with your OpenAI API key.") 450 | return 451 | 452 | client = OpenAIWebRTCClient( 453 | api_key=api_key, 454 | model="gpt-4o-realtime-preview", 455 | system_prompt="You are providing me help around the University of Waterloo campus. You are the help desk, try your best to help students and staff with their questions." 456 | ) 457 | 458 | def on_transcription(text): 459 | print(f"Transcription: {text}") 460 | 461 | client.on_transcription = on_transcription 462 | 463 | try: 464 | print("Starting streaming session...") 465 | await client.start_streaming() 466 | print("Streaming started. Speak into your microphone.") 467 | print("Press Ctrl+C to stop.") 468 | 469 | while True: 470 | await asyncio.sleep(1) 471 | except KeyboardInterrupt: 472 | print("\nStopping streaming...") 473 | finally: 474 | await client.stop_streaming() 475 | print("Streaming stopped.") 476 | 477 | 478 | if __name__ == "__main__": 479 | load_dotenv() 480 | asyncio.run(main()) -------------------------------------------------------------------------------- /examples/example_rerun.py: -------------------------------------------------------------------------------- 1 | # Adds the lib directory to the Python path 2 | import sys 3 | import os 4 | sys.path.append(os.path.join(os.path.dirname(__file__), '..')) 5 | 6 | import rerun as rr 7 | from lib.camera import StereoCamera 8 | import time 9 | 10 | def main(): 11 | # Scale down the images to 1/4 of their original size to reduce bandwidth 12 | camera = StereoCamera(scale=0.25) 13 | 14 | # Initialize Rerun with a descriptive name 15 | rr.init("stereo_camera_stream", spawn=False) 16 | 17 | # Replace with your Computer's IP address if needed 18 | # if you want to find you computers address do: "arp -a" in the terminal 19 | rr.connect_grpc(f"rerun+http://:9876/proxy") 20 | 21 | print("\nStreaming stereo images to Rerun viewer...") 22 | print(f"Scale factor: {camera.get_scale()}") 23 | print("Press Ctrl+C to stop") 24 | 25 | try: 26 | while True: 27 | left, right = camera.get_stereo() 28 | if left is None or right is None: 29 | print("Failed to capture images") 30 | continue 31 | 32 | rr.log("stereo/left", rr.Image(left, color_model="bgr")) 33 | rr.log("stereo/right", rr.Image(right, color_model="bgr")) 34 | 35 | except KeyboardInterrupt: 36 | print("\nStopping...") 37 | finally: 38 | camera.release() 39 | print("\nShutdown complete") 40 | 41 | if __name__ == "__main__": 42 | main() 43 | -------------------------------------------------------------------------------- /examples/example_segmentation.py: -------------------------------------------------------------------------------- 1 | import sys, os 2 | sys.path.append(os.path.join(os.path.dirname(__file__), '..')) 3 | 4 | import cv2 5 | import numpy as np 6 | from ultralytics import YOLO 7 | from lib.camera import StereoCamera 8 | import time 9 | 10 | os.makedirs("yolo_segment_output", exist_ok=True) 11 | camera = StereoCamera(scale=0.75) 12 | model = YOLO("yolo11s-seg.pt") 13 | 14 | total_frames = 0 15 | try: 16 | while True: 17 | start_time = time.time() 18 | left_image, _ = camera.get_stereo() 19 | 20 | if left_image is None: 21 | time.sleep(0.5) 22 | continue 23 | 24 | results = model(source=left_image, conf=0.25, iou=0.45, verbose=False) 25 | annotated_image = results[0].plot(labels=True, boxes=False, masks=True) 26 | 27 | cv2.imwrite(f"yolo_segment_output/frame_{total_frames:04d}.jpg", annotated_image) 28 | 29 | proc_time = time.time() - start_time 30 | total_frames += 1 31 | num_objects = len(results[0].boxes) 32 | num_masks = len(results[0].masks) if hasattr(results[0], 'masks') and results[0].masks is not None else 0 33 | 34 | print(f"Frame {total_frames} | Objects: {num_objects} | Masks: {num_masks} | Time: {proc_time*1000:.1f}ms") 35 | 36 | except KeyboardInterrupt: 37 | print("\nStopped by user") 38 | finally: 39 | camera.release() 40 | print("Camera released") -------------------------------------------------------------------------------- /examples/example_wasd.py: -------------------------------------------------------------------------------- 1 | # Adds the lib directory to the Python path 2 | import sys 3 | import os 4 | sys.path.append(os.path.join(os.path.dirname(__file__), '..')) 5 | 6 | from sshkeyboard import listen_keyboard, stop_listening 7 | from lib.odrive_uart import ODriveUART 8 | import json 9 | 10 | # Load motor directions 11 | with open(os.path.expanduser('~/quickstart/lib/motor_dir.json'), 'r') as f: 12 | motor_dirs = json.load(f) 13 | 14 | # Initialize motor controller 15 | motor = ODriveUART() 16 | 17 | # Start motors 18 | motor.start_left() 19 | motor.start_right() 20 | motor.enable_velocity_mode_left() 21 | motor.enable_velocity_mode_right() 22 | motor.disable_watchdog_left() 23 | motor.disable_watchdog_right() 24 | motor.clear_errors_left() 25 | motor.clear_errors_right() 26 | 27 | SPEED = 0.375 28 | 29 | def press(key): 30 | """Handle key press""" 31 | if key.lower() == 'w': 32 | motor.set_speed_mps_left(SPEED) 33 | motor.set_speed_mps_right(SPEED) 34 | elif key.lower() == 's': 35 | motor.set_speed_mps_left(-SPEED) 36 | motor.set_speed_mps_right(-SPEED) 37 | elif key.lower() == 'a': 38 | motor.set_speed_mps_left(-SPEED) 39 | motor.set_speed_mps_right(SPEED) 40 | elif key.lower() == 'd': 41 | motor.set_speed_mps_left(SPEED) 42 | motor.set_speed_mps_right(-SPEED) 43 | elif key.lower() == 'q': 44 | stop_listening() 45 | 46 | def release(key): 47 | """Stop on key release""" 48 | motor.set_speed_mps_left(0) 49 | motor.set_speed_mps_right(0) 50 | 51 | try: 52 | print("WASD to control the robot:") 53 | print(" W - Forward") 54 | print(" S - Backward") 55 | print(" A - Turn Left") 56 | print(" D - Turn Right") 57 | print(" Q - Quit") 58 | 59 | listen_keyboard( 60 | on_press=press, 61 | on_release=release, 62 | delay_second_char=0.05, # Reduced from 0.75s to 0.05s 63 | delay_other_chars=0.02, # Reduced from 0.05s to 0.02s 64 | sequential=False, # Allow concurrent key handling 65 | sleep=0.01 # Faster polling rate 66 | ) 67 | 68 | except Exception as e: 69 | print(f"Error: {e}") 70 | 71 | finally: 72 | # Cleanup 73 | motor.set_speed_mps_left(0) 74 | motor.set_speed_mps_right(0) 75 | motor.clear_errors_left() 76 | motor.clear_errors_right() 77 | print("Shutdown complete.") 78 | -------------------------------------------------------------------------------- /examples/example_whisper.py: -------------------------------------------------------------------------------- 1 | # Adds the lib directory to the Python path 2 | import sys 3 | import os 4 | sys.path.append(os.path.join(os.path.dirname(__file__), '..')) 5 | 6 | import subprocess 7 | import time 8 | 9 | def install_dependencies(): 10 | """Install all required system dependencies""" 11 | print("Installing required dependencies...") 12 | subprocess.run(['sudo', 'apt-get', 'update'], check=True) 13 | 14 | # Install all required packages 15 | packages = [ 16 | 'libsdl2-dev', # SDL2 for audio capture 17 | 'cmake', # For building 18 | 'make', # For building 19 | 'g++', # C++ compiler 20 | 'gcc', # C compiler 21 | 'git', # For cloning 22 | 'alsa-utils', # For audio device access 23 | 'libasound2-dev' # ALSA development files 24 | ] 25 | 26 | subprocess.run(['sudo', 'apt-get', 'install', '-y'] + packages, check=True) 27 | 28 | def ensure_whisper_ready(): 29 | """Make sure whisper.cpp is installed and ready to use""" 30 | # Get the directory where this script is located 31 | script_dir = os.path.dirname(os.path.abspath(__file__)) 32 | os.chdir(script_dir) # Change to script directory 33 | 34 | # Clone and build whisper.cpp if needed 35 | if not os.path.exists('whisper.cpp'): 36 | print("Setting up whisper.cpp...") 37 | # First make sure all dependencies are installed 38 | install_dependencies() 39 | 40 | # Clone whisper.cpp 41 | subprocess.run(['git', 'clone', 'https://github.com/ggerganov/whisper.cpp.git'], check=True) 42 | 43 | # Build with NEON optimizations and SDL2 44 | build_dir = os.path.join('whisper.cpp', 'build') 45 | os.makedirs(build_dir, exist_ok=True) 46 | os.chdir(build_dir) 47 | 48 | # Configure with both NEON and SDL2 49 | subprocess.run(['cmake', '..', 50 | '-DWHISPER_NEON=ON', 51 | '-DWHISPER_SDL2=ON'], check=True) 52 | 53 | # Build the stream example specifically 54 | subprocess.run(['make', '-j4', 'whisper-stream'], check=True) 55 | os.chdir('../..') 56 | 57 | # Download the tiny model 58 | os.chdir('whisper.cpp') 59 | subprocess.run(['bash', 'models/download-ggml-model.sh', 'tiny.en'], check=True) 60 | os.chdir('..') 61 | 62 | # Return paths needed for streaming 63 | return { 64 | 'model': os.path.join('whisper.cpp', 'models', 'ggml-tiny.en.bin'), 65 | 'stream': os.path.join('whisper.cpp', 'build', 'bin', 'whisper-stream') 66 | } 67 | 68 | def main(): 69 | # Make sure whisper is ready 70 | paths = ensure_whisper_ready() 71 | 72 | # Start real-time streaming transcription 73 | print("\nStarting real-time transcription... (Press Ctrl+C to stop)") 74 | try: 75 | # Using settings optimized for Raspberry Pi from the discussion 76 | subprocess.run([ 77 | paths['stream'], 78 | '-m', paths['model'], 79 | '--step', '4000', # Process 4 seconds at a time 80 | '--length', '8000', # Use 8 second context 81 | '-c', '0', # Audio capture device 82 | '-t', '4', # Number of threads 83 | '-ac', '512' # Reduced context size for speed 84 | ]) 85 | except KeyboardInterrupt: 86 | print("\nStopped transcription") 87 | 88 | if __name__ == "__main__": 89 | main() 90 | -------------------------------------------------------------------------------- /examples/example_yolo.py: -------------------------------------------------------------------------------- 1 | # Adds the lib directory to the Python path 2 | import sys 3 | import os 4 | 5 | # Add the parent directory to the Python path to access the lib folder 6 | sys.path.append(os.path.join(os.path.dirname(__file__), '..')) 7 | 8 | # https://docs.ultralytics.com/guides/raspberry-pi/ 9 | import cv2 10 | from ultralytics import YOLO 11 | from lib.camera import StereoCamera 12 | 13 | camera = StereoCamera() 14 | 15 | model = YOLO("yolo11n.pt") 16 | model.export(format="ncnn") 17 | 18 | ncnn_model = YOLO("yolo11n_ncnn_model") 19 | 20 | # Get stereo images and use the left image 21 | left_image, _ = camera.get_stereo() 22 | 23 | if left_image is not None: 24 | results = ncnn_model(left_image) 25 | 26 | annotated_image = results[0].plot() 27 | 28 | cv2.imwrite("labeled_left.jpg", annotated_image) 29 | print("Labeled image saved as labeled_left.jpg") 30 | else: 31 | print("Unable to capture the left image from the stereo camera.") -------------------------------------------------------------------------------- /lib/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | __all__ = ["imu", "lqr", "odrive_uart", "madgwickahrs"] -------------------------------------------------------------------------------- /lib/camera.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | # import pyrealsense2 as rs 4 | import cv2 5 | import numpy as np 6 | 7 | class RealsenseCamera: 8 | def __init__(self): 9 | # Configure depth and color streams 10 | self.pipeline = rs.pipeline() 11 | self.config = rs.config() 12 | # Get device product line for setting a supporting resolution 13 | self.pipeline_wrapper = rs.pipeline_wrapper(self.pipeline) 14 | self.pipeline_profile = self.config.resolve(self.pipeline_wrapper) 15 | self.device = self.pipeline_profile.get_device() 16 | self.device_product_line = str(self.device.get_info(rs.camera_info.product_line)) 17 | 18 | # Enable streams 19 | self.config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30) 20 | self.config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30) 21 | 22 | # Start streaming 23 | self.pipeline.start(self.config) 24 | 25 | def get_frames(self): 26 | # Wait for a coherent pair of frames 27 | frames = self.pipeline.wait_for_frames() 28 | depth_frame = frames.get_depth_frame() 29 | color_frame = frames.get_color_frame() 30 | if not depth_frame or not color_frame: 31 | return None, None 32 | # Convert images to numpy arrays 33 | depth_image = np.asanyarray(depth_frame.get_data()) 34 | color_image = np.asanyarray(color_frame.get_data()) 35 | return color_image, depth_image 36 | 37 | def release(self): 38 | # Stop streaming 39 | self.pipeline.stop() 40 | 41 | class USBCamera: 42 | def __init__(self, index=0): 43 | # Initialize USB camera 44 | self.cap = cv2.VideoCapture(index) 45 | if not self.cap.isOpened(): 46 | raise Exception("Could not open video device") 47 | # Set properties if needed 48 | self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640) 49 | self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480) 50 | 51 | def get_frame(self): 52 | ret, frame = self.cap.read() 53 | if not ret: 54 | return None 55 | return frame 56 | 57 | def release(self): 58 | self.cap.release() 59 | 60 | class StereoCamera: 61 | def __init__(self, device_id=0, scale=1.0): 62 | # Open the camera with V4L2 backend 63 | self.cap = cv2.VideoCapture(device_id, cv2.CAP_V4L2) 64 | 65 | # Store the scale factor 66 | self.scale = scale 67 | 68 | # Set MJPEG format 69 | success_fourcc = self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')) 70 | # Set resolution to 2560x720 71 | success_width = self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 2560) 72 | success_height = self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) 73 | # Set frame rate to 30 FPS 74 | success_fps = self.cap.set(cv2.CAP_PROP_FPS, 30) 75 | # Set buffer size to 1 76 | success_buffer = self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) 77 | 78 | # Check if settings were successful 79 | print(f"Set FOURCC to MJPG: {success_fourcc}") 80 | print(f"Set width to 2560: {success_width}") 81 | print(f"Set height to 720: {success_height}") 82 | print(f"Set FPS to 30: {success_fps}") 83 | print(f"Set buffer size to 1: {success_buffer}") 84 | 85 | # Verify actual settings 86 | fourcc = int(self.cap.get(cv2.CAP_PROP_FOURCC)) 87 | fourcc_str = ''.join([chr((fourcc >> 8 * i) & 0xFF) for i in range(4)]) 88 | width = self.cap.get(cv2.CAP_PROP_FRAME_WIDTH) 89 | height = self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT) 90 | fps = self.cap.get(cv2.CAP_PROP_FPS) 91 | buffer_size = self.cap.get(cv2.CAP_PROP_BUFFERSIZE) 92 | print(f"Actual FOURCC: {fourcc_str}") 93 | print(f"Actual Resolution: {width}x{height}") 94 | print(f"Actual FPS: {fps}") 95 | print(f"Actual buffer size: {buffer_size}") 96 | 97 | if scale != 1.0: 98 | print(f"Images will be scaled by factor: {scale}") 99 | 100 | def set_scale(self, scale): 101 | """Set the scaling factor for the stereo images""" 102 | self.scale = scale 103 | print(f"Scale factor set to: {scale}") 104 | 105 | def get_scale(self): 106 | """Get the current scaling factor""" 107 | return self.scale 108 | 109 | def get_stereo(self, scale=None): 110 | """ 111 | Get stereo images from the camera. 112 | 113 | Args: 114 | scale: Optional scaling factor to override the default scale. 115 | If None, uses the scale set during initialization. 116 | 117 | Returns: 118 | Tuple of (left_image, right_image), or (None, None) if capture failed. 119 | """ 120 | ret, frame = self.cap.read() 121 | if not ret: 122 | return None, None 123 | 124 | # Split the frame horizontally into left and right images 125 | left = frame[:, :1280] 126 | right = frame[:, 1280:] 127 | 128 | # Apply scaling if needed 129 | use_scale = scale if scale is not None else self.scale 130 | if use_scale != 1.0: 131 | h, w = left.shape[:2] 132 | new_h, new_w = int(h * use_scale), int(w * use_scale) 133 | left = cv2.resize(left, (new_w, new_h), interpolation=cv2.INTER_LINEAR) 134 | right = cv2.resize(right, (new_w, new_h), interpolation=cv2.INTER_LINEAR) 135 | 136 | return left, right 137 | 138 | def release(self): 139 | self.cap.release() 140 | 141 | if __name__ == "__main__": 142 | realsense = RealsenseCamera() 143 | usb = USBCamera() 144 | 145 | # Example usage of StereoCamera 146 | stereo_cam = StereoCamera(0, scale=0.5) # Initialize with 50% scaling 147 | left, right = stereo_cam.get_stereo() 148 | if left is not None and right is not None: 149 | cv2.imwrite("left.jpg", left) 150 | cv2.imwrite("right.jpg", right) 151 | print("Saved stereo images to left.jpg and right.jpg") 152 | 153 | # Try with a different scale 154 | left_full, right_full = stereo_cam.get_stereo(scale=1.0) 155 | cv2.imwrite("left_full.jpg", left_full) 156 | cv2.imwrite("right_full.jpg", right_full) 157 | print("Saved full-size stereo images") 158 | else: 159 | print("Failed to capture stereo images") 160 | stereo_cam.release() 161 | -------------------------------------------------------------------------------- /lib/imu.py: -------------------------------------------------------------------------------- 1 | # Adds the lib directory to the Python path 2 | import sys 3 | import os 4 | sys.path.append(os.path.join(os.path.dirname(__file__), '..')) 5 | 6 | import numpy as np 7 | import matplotlib.pyplot as plt 8 | from lib.madgwickahrs import MadgwickAHRS, Quaternion 9 | 10 | import time 11 | import board 12 | import adafruit_mpu6050 13 | 14 | class FilteredMPU6050(): 15 | 16 | def __init__(self): 17 | self.sensor = adafruit_mpu6050.MPU6050(board.I2C()) 18 | # self.sensor.gyro_range = adafruit_mpu6050.GyroRange.RANGE_500_DPS # Set gyroscope range to ±1000 dps 19 | self.ahrs = MadgwickAHRS(beta=0.008, zeta=0.) 20 | self.alpha = 1 # LPF alpha: x[t] := a*x[t] + (1-a)*x[t-1] 21 | self.gyro_bias = np.array([0., 0., 0.]) 22 | 23 | def calibrate(self): 24 | try: 25 | raw_bias = np.loadtxt('gyro_bias.txt') 26 | # Remap the bias the same way we remap the gyro readings 27 | self.gyro_bias = np.array([-raw_bias[1], raw_bias[0], raw_bias[2]]) 28 | except FileNotFoundError: 29 | self.gyro_bias = np.array([0., 0., 0.]) 30 | for _ in range(50): 31 | _, gyro = self.read_sensor() 32 | self.gyro_bias += gyro / 50 33 | time.sleep(0.01) 34 | print('Calculated gyro bias:', self.gyro_bias) 35 | # Save the unmapped bias values 36 | np.savetxt('gyro_bias.txt', np.array([self.gyro_bias[1], -self.gyro_bias[0], self.gyro_bias[2]])) 37 | 38 | self.accel, gyro_raw = self.read_sensor() 39 | self.gyro = gyro_raw - self.gyro_bias 40 | self.t = time.monotonic() 41 | 42 | self.quat = self._calculate_initial_q(self.accel) 43 | self.grav = self.quat_rotate(self.quat.conj(), [0, 0, 1]) 44 | self.ahrs.quaternion = self.quat 45 | 46 | 47 | def get_orientation(self): 48 | self.update() 49 | gx, gy, gz = self.grav 50 | 51 | # Map gravity vector components to new axes 52 | gX_new = gy # gX_new = gy 53 | gY_new = -gx # gY_new = -gx 54 | gZ_new = gz # gZ_new = gz 55 | 56 | # Compute roll and pitch using the standard formulas 57 | roll = np.degrees(np.arctan2(gY_new, gZ_new)) 58 | pitch = np.degrees(np.arctan2(-gX_new, np.sqrt(gY_new**2 + gZ_new**2))) 59 | 60 | # Compute yaw from the quaternion 61 | qw, qx, qy, qz = self.quat 62 | yaw = np.degrees(np.arctan2(2 * (qw * qz + qx * qy), 63 | 1 - 2 * (qy**2 + qz**2))) 64 | 65 | return pitch, roll, yaw 66 | 67 | def _calculate_initial_q(self, accel): 68 | acc_norm = accel / np.linalg.norm(accel) 69 | 70 | # Estimate initial roll and pitch from accelerometer 71 | initial_roll = np.arctan2(acc_norm[1], acc_norm[2]) 72 | initial_pitch = np.arctan2(-acc_norm[0], np.sqrt(acc_norm[1]**2 + acc_norm[2]**2)) 73 | initial_yaw = 0 74 | 75 | # Initialize quaternion using the from_angle_axis function 76 | initial_q = Quaternion.from_angle_axis(initial_roll, 1, 0, 0) 77 | initial_q = initial_q * Quaternion.from_angle_axis(initial_pitch, 0, 1, 0) 78 | initial_q = initial_q * Quaternion.from_angle_axis(initial_yaw, 0, 0, 1) 79 | return initial_q 80 | 81 | def read_sensor(self): 82 | # Read raw data from sensor 83 | accel = np.array(self.sensor.acceleration) 84 | gyro = np.array(self.sensor.gyro) 85 | accel_mapped = np.array([-accel[1], accel[0], accel[2]]) 86 | gyro_mapped = np.array([-gyro[1], gyro[0], gyro[2]]) 87 | return accel_mapped, gyro_mapped 88 | 89 | 90 | def update(self): 91 | # Read and map sensor readings 92 | self.accel, gyro_raw = self.read_sensor() 93 | self.gyro = gyro_raw - self.gyro_bias 94 | t = time.monotonic() 95 | 96 | # Store raw data 97 | self.accel_RAW = self.accel 98 | self.gyro_RAW = self.gyro 99 | self.quat_RAW = self._calculate_initial_q(self.accel_RAW) 100 | self.grav_RAW = self.quat_rotate(self.quat_RAW.conj(), [0, 0, 1]) 101 | 102 | # Filtering 103 | self.ahrs.samplePeriod = t - self.t 104 | self.ahrs.update_imu(self.gyro, self.accel) 105 | self.t = t 106 | 107 | # Update orientation 108 | quat = self.ahrs.quaternion 109 | self.quat = quat.q 110 | self.grav = self.quat_rotate(quat.conj(), [0, 0, 1]) 111 | 112 | 113 | def quat_rotate(self, q, v): 114 | """Rotate a vector v by a quaternion q""" 115 | qv = np.concatenate(([0], v)) 116 | return (q * Quaternion(qv) * q.conj()).q[1:] 117 | 118 | if __name__ == '__main__': 119 | imu = FilteredMPU6050() 120 | imu.calibrate() 121 | start_time = time.monotonic() 122 | times = [] 123 | quats = [] 124 | accels = [] 125 | gyros = [] 126 | gyros_raw = [] 127 | gravs = [] 128 | pitches = [] 129 | while time.monotonic() < start_time + 10: 130 | # pitch = imu.robot_angle() 131 | pitch, roll, yaw = imu.get_orientation() 132 | t = time.monotonic() 133 | pitches.append(pitch) 134 | times.append(t) 135 | quats.append(imu.quat) 136 | accels.append(imu.accel) 137 | gyros.append(imu.gyro) 138 | gyros_raw.append(imu.gyro_RAW) 139 | gravs.append(imu.grav) 140 | 141 | # Print current readings 142 | # quat = Quaternion(imu.quat) 143 | # euler = quat.to_euler_angles() 144 | print(f"Roll: {roll:.2f}°, Pitch: {pitch:.2f}°, Yaw: {yaw:.2f}°") 145 | # print(f"Accel X: {imu.accel[0]:.2f}, Y: {imu.accel[1]:.2f}, Z: {imu.accel[2]:.2f} m/s²") 146 | print(f"Gyro (filtered) X: {imu.gyro[0]:.2f}, Y: {imu.gyro[1]:.2f}, Z: {imu.gyro[2]:.2f} rad/s") 147 | # print(f"Gyro (raw) X: {imu.gyro_RAW[0]:.2f}, Y: {imu.gyro_RAW[1]:.2f}, Z: {imu.gyro_RAW[2]:.2f} rad/s") 148 | print("---") 149 | 150 | times = np.array(times) - times[0] 151 | print('Average Hz:', 1 / (np.mean(times[1:] - times[:-1]))) 152 | 153 | plt.figure() 154 | gyros = np.stack(gyros) 155 | gyros_raw = np.stack(gyros_raw) 156 | sx, sy, sz = np.std(gyros, axis=0) 157 | plt.title(f'Gyro, std=({sx:.4f}, {sy:.4f}, {sz:.4f})') 158 | plt.plot(times, gyros, label='filtered') 159 | plt.plot(times, gyros_raw, '--', alpha=0.4, label='raw') 160 | plt.legend() 161 | plt.savefig('gyro.png') 162 | 163 | plt.figure() 164 | accels = np.stack(accels) 165 | sx, sy, sz = np.std(accels, axis=0) 166 | plt.title(f'Accel, std=({sx:.4f}, {sy:.4f}, {sz:.4f})') 167 | plt.plot(times, accels) 168 | plt.savefig('accel.png') 169 | 170 | def grav2pitch(gravs): 171 | return np.degrees(np.arctan2(gravs[..., 2], np.sqrt(gravs[..., 0]**2 + gravs[..., 1]**2))) 172 | 173 | plt.figure() 174 | plt.plot(times, grav2pitch(accels), label='noisy', alpha=0.4) 175 | plt.plot(times, pitches, label='filtered') 176 | plt.legend() 177 | plt.savefig('pitch.png') -------------------------------------------------------------------------------- /lib/lqr.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import sympy as sp 3 | from control import lqr 4 | 5 | 6 | def LQR_gains(Q_diag, R_diag): 7 | # ======= NEW PARAMS OCT 24, 2024 ======= # 8 | Jr = 0.018 # Moment of inertia of wheels [kg⋅m²] 9 | Mr = 2.2 # Mass of wheel [kg] (BIG WHEEL) 10 | # Mr = 1.6 # Mass of wheel [kg] (SMALL WHEEL) 11 | Jpth = 1.46# OLD 2 BATTERY 1.88 # Moment of inertia of chassis pitch [kg⋅m²] 12 | Jpd = 0.041 # Moment of inertia of chassis yaw [kg⋅m²] 13 | Mp = 4-0.62 # Mass of chassis [kg] 14 | R = 0.0846 # Radius of wheels [m] 15 | D = 0.46 # Distance between wheels [m] 16 | L = 0.367 # Distance to center of gravity [m] 17 | g = 9.807 # Gravity [m/s²] 18 | 19 | 20 | # Define symbolic variable for angle 21 | theta = sp.Symbol('theta') 22 | 23 | # Create the state matrix using small angle approximations: 24 | # cos(theta) ≈ 1 25 | # sin(theta) ≈ theta 26 | # Create symbolic matrix A using sympy 27 | M_pitch = sp.Matrix([ 28 | [-1, 0, 0, -1/(2*Mr), 1/(2*Mr)], 29 | [-1, -L*1, 0, 1/Mp, 0], 30 | [0, L*theta, 1/Mp, 0, 0], 31 | [0, -1, L*theta/Jpth, -L*1/Jpth, 0], 32 | [-1, 0, 0, 0, -R**2/(2*Jr)] 33 | ]) 34 | M_pitch_inv = M_pitch.inv() 35 | 36 | f_R, f_P, Cth = sp.symbols('f_R f_P Cth') 37 | b = sp.Matrix([ 38 | -f_R/(2*Mr), 39 | -f_P/Mp, 40 | g, 41 | Cth/Jpth, 42 | -R*Cth/(2*Jr) 43 | ]) 44 | 45 | q = (M_pitch_inv * b).subs(theta**2,0) 46 | 47 | A23 = q[0].diff(theta).evalf() 48 | A43 = q[1].diff(theta).evalf() 49 | B21 = q[0].diff(Cth).evalf() 50 | B41 = q[1].diff(Cth).evalf() 51 | 52 | 53 | M_yaw = sp.Matrix([ 54 | [1,-D/(2*Jpd),0], 55 | [-2/D,-1/Mr,1/Mr], 56 | [-2/D,0,-R**2/Jr] 57 | ]) 58 | M_yaw_inv = M_yaw.inv() 59 | 60 | f_d, Cd = sp.symbols('f_d Cd') 61 | b = sp.Matrix([ 62 | 0, 63 | f_d/Mr, 64 | -Cd*R/Jr 65 | ]) 66 | 67 | q = (M_yaw_inv * b) 68 | 69 | B62 = q[0].diff(Cd).evalf() 70 | 71 | A = np.array([ 72 | [0,1,0,0,0,0], 73 | [0,0,A23,0,0,0], 74 | [0,0,0,1,0,0], 75 | [0,0,A43,0,0,0], 76 | [0,0,0,0,0,1], 77 | [0,0,0,0,0,0] 78 | ], dtype=float) 79 | 80 | B = np.array([ 81 | [0,0], 82 | [B21,0], 83 | [0,0], 84 | [B41,0], 85 | [0,0], 86 | [0,B62] 87 | ], dtype=float) 88 | 89 | Q = np.diag(Q_diag) 90 | R = np.diag(R_diag) 91 | 92 | K, S, E = lqr(A, B, Q, R) 93 | return K 94 | -------------------------------------------------------------------------------- /lib/madgwickahrs.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | Copyright (c) 2015 Jonas Böer, jonas.boeer@student.kit.edu 4 | 5 | This program is free software: you can redistribute it and/or modify 6 | it under the terms of the GNU Lesser General Public License as published by 7 | the Free Software Foundation, either version 3 of the License, or 8 | (at your option) any later version. 9 | 10 | This program is distributed in the hope that it will be useful, 11 | but WITHOUT ANY WARRANTY; without even the implied warranty of 12 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 13 | GNU Lesser General Public License for more details. 14 | 15 | You should have received a copy of the GNU Lesser General Public License 16 | along with this program. If not, see . 17 | """ 18 | 19 | import warnings 20 | import numpy as np 21 | import numbers 22 | from numpy.linalg import norm 23 | 24 | 25 | class Quaternion: 26 | """ 27 | A simple class implementing basic quaternion arithmetic. 28 | """ 29 | def __init__(self, w_or_q, x=None, y=None, z=None): 30 | """ 31 | Initializes a Quaternion object 32 | :param w_or_q: A scalar representing the real part of the quaternion, another Quaternion object or a 33 | four-element array containing the quaternion values 34 | :param x: The first imaginary part if w_or_q is a scalar 35 | :param y: The second imaginary part if w_or_q is a scalar 36 | :param z: The third imaginary part if w_or_q is a scalar 37 | """ 38 | self._q = np.array([1, 0, 0, 0]) 39 | 40 | if x is not None and y is not None and z is not None: 41 | w = w_or_q 42 | q = np.array([w, x, y, z]) 43 | elif isinstance(w_or_q, Quaternion): 44 | q = np.array(w_or_q.q) 45 | else: 46 | q = np.array(w_or_q) 47 | if len(q) != 4: 48 | raise ValueError("Expecting a 4-element array or w x y z as parameters") 49 | 50 | self.q = q 51 | 52 | # Quaternion specific interfaces 53 | 54 | def conj(self): 55 | """ 56 | Returns the conjugate of the quaternion 57 | :rtype : Quaternion 58 | :return: the conjugate of the quaternion 59 | """ 60 | return Quaternion(self._q[0], -self._q[1], -self._q[2], -self._q[3]) 61 | 62 | def to_angle_axis(self): 63 | """ 64 | Returns the quaternion's rotation represented by an Euler angle and axis. 65 | If the quaternion is the identity quaternion (1, 0, 0, 0), a rotation along the x axis with angle 0 is returned. 66 | :return: rad, x, y, z 67 | """ 68 | if self[0] == 1 and self[1] == 0 and self[2] == 0 and self[3] == 0: 69 | return 0, 1, 0, 0 70 | rad = np.arccos(self[0]) * 2 71 | imaginary_factor = np.sin(rad / 2) 72 | if abs(imaginary_factor) < 1e-8: 73 | return 0, 1, 0, 0 74 | x = self._q[1] / imaginary_factor 75 | y = self._q[2] / imaginary_factor 76 | z = self._q[3] / imaginary_factor 77 | return rad, x, y, z 78 | 79 | @staticmethod 80 | def from_angle_axis(rad, x, y, z): 81 | s = np.sin(rad / 2) 82 | return Quaternion(np.cos(rad / 2), x*s, y*s, z*s) 83 | 84 | def to_euler_angles(self): 85 | pitch = np.arcsin(2 * self[1] * self[2] + 2 * self[0] * self[3]) 86 | if np.abs(self[1] * self[2] + self[3] * self[0] - 0.5) < 1e-8: 87 | roll = 0 88 | yaw = 2 * np.arctan2(self[1], self[0]) 89 | elif np.abs(self[1] * self[2] + self[3] * self[0] + 0.5) < 1e-8: 90 | roll = -2 * np.arctan2(self[1], self[0]) 91 | yaw = 0 92 | else: 93 | roll = np.arctan2(2 * self[0] * self[1] - 2 * self[2] * self[3], 1 - 2 * self[1] ** 2 - 2 * self[3] ** 2) 94 | yaw = np.arctan2(2 * self[0] * self[2] - 2 * self[1] * self[3], 1 - 2 * self[2] ** 2 - 2 * self[3] ** 2) 95 | return roll, pitch, yaw 96 | 97 | def to_euler123(self): 98 | roll = np.arctan2(-2 * (self[2] * self[3] - self[0] * self[1]), self[0] ** 2 - self[1] ** 2 - self[2] ** 2 + self[3] ** 2) 99 | pitch = np.arcsin(2 * (self[1] * self[3] + self[0] * self[1])) 100 | yaw = np.arctan2(-2 * (self[1] * self[2] - self[0] * self[3]), self[0] ** 2 + self[1] ** 2 - self[2] ** 2 - self[3] ** 2) 101 | return roll, pitch, yaw 102 | 103 | def __mul__(self, other): 104 | """ 105 | multiply the given quaternion with another quaternion or a scalar 106 | :param other: a Quaternion object or a number 107 | :return: 108 | """ 109 | if isinstance(other, Quaternion): 110 | w = self._q[0]*other._q[0] - self._q[1]*other._q[1] - self._q[2]*other._q[2] - self._q[3]*other._q[3] 111 | x = self._q[0]*other._q[1] + self._q[1]*other._q[0] + self._q[2]*other._q[3] - self._q[3]*other._q[2] 112 | y = self._q[0]*other._q[2] - self._q[1]*other._q[3] + self._q[2]*other._q[0] + self._q[3]*other._q[1] 113 | z = self._q[0]*other._q[3] + self._q[1]*other._q[2] - self._q[2]*other._q[1] + self._q[3]*other._q[0] 114 | 115 | return Quaternion(w, x, y, z) 116 | elif isinstance(other, numbers.Number): 117 | q = self._q * other 118 | return Quaternion(q) 119 | 120 | def __add__(self, other): 121 | """ 122 | add two quaternions element-wise or add a scalar to each element of the quaternion 123 | :param other: 124 | :return: 125 | """ 126 | if not isinstance(other, Quaternion): 127 | if len(other) != 4: 128 | raise TypeError("Quaternions must be added to other quaternions or a 4-element array") 129 | q = self._q + other 130 | else: 131 | q = self._q + other._q 132 | 133 | return Quaternion(q) 134 | 135 | # Implementing other interfaces to ease working with the class 136 | 137 | @property 138 | def q(self): 139 | return self._q 140 | 141 | @q.setter 142 | def q(self, q): 143 | self._q = q 144 | 145 | def __getitem__(self, item): 146 | return self._q[item] 147 | 148 | def __array__(self): 149 | return self._q 150 | 151 | class MadgwickAHRS: 152 | samplePeriod = 1/256 153 | quaternion = Quaternion(1, 0, 0, 0) 154 | beta = 1 155 | zeta = 0 156 | 157 | def __init__(self, sampleperiod=None, quaternion=None, beta=None, zeta=None): 158 | """ 159 | Initialize the class with the given parameters. 160 | :param sampleperiod: The sample period 161 | :param quaternion: Initial quaternion 162 | :param beta: Algorithm gain beta 163 | :param beta: Algorithm gain zeta 164 | :return: 165 | """ 166 | if sampleperiod is not None: 167 | self.samplePeriod = sampleperiod 168 | if quaternion is not None: 169 | self.quaternion = quaternion 170 | if beta is not None: 171 | self.beta = beta 172 | if zeta is not None: 173 | self.zeta = zeta 174 | 175 | def update(self, gyroscope, accelerometer, magnetometer): 176 | """ 177 | Perform one update step with data from a AHRS sensor array 178 | :param gyroscope: A three-element array containing the gyroscope data in radians per second. 179 | :param accelerometer: A three-element array containing the accelerometer data. Can be any unit since a normalized value is used. 180 | :param magnetometer: A three-element array containing the magnetometer data. Can be any unit since a normalized value is used. 181 | :return: 182 | """ 183 | q = self.quaternion 184 | 185 | gyroscope = np.array(gyroscope, dtype=float).flatten() 186 | accelerometer = np.array(accelerometer, dtype=float).flatten() 187 | magnetometer = np.array(magnetometer, dtype=float).flatten() 188 | 189 | # Normalise accelerometer measurement 190 | if norm(accelerometer) is 0: 191 | warnings.warn("accelerometer is zero") 192 | return 193 | accelerometer /= norm(accelerometer) 194 | 195 | # Normalise magnetometer measurement 196 | if norm(magnetometer) is 0: 197 | warnings.warn("magnetometer is zero") 198 | return 199 | magnetometer /= norm(magnetometer) 200 | 201 | h = q * (Quaternion(0, magnetometer[0], magnetometer[1], magnetometer[2]) * q.conj()) 202 | b = np.array([0, norm(h[1:3]), 0, h[3]]) 203 | 204 | # Gradient descent algorithm corrective step 205 | f = np.array([ 206 | 2*(q[1]*q[3] - q[0]*q[2]) - accelerometer[0], 207 | 2*(q[0]*q[1] + q[2]*q[3]) - accelerometer[1], 208 | 2*(0.5 - q[1]**2 - q[2]**2) - accelerometer[2], 209 | 2*b[1]*(0.5 - q[2]**2 - q[3]**2) + 2*b[3]*(q[1]*q[3] - q[0]*q[2]) - magnetometer[0], 210 | 2*b[1]*(q[1]*q[2] - q[0]*q[3]) + 2*b[3]*(q[0]*q[1] + q[2]*q[3]) - magnetometer[1], 211 | 2*b[1]*(q[0]*q[2] + q[1]*q[3]) + 2*b[3]*(0.5 - q[1]**2 - q[2]**2) - magnetometer[2] 212 | ]) 213 | j = np.array([ 214 | [-2*q[2], 2*q[3], -2*q[0], 2*q[1]], 215 | [2*q[1], 2*q[0], 2*q[3], 2*q[2]], 216 | [0, -4*q[1], -4*q[2], 0], 217 | [-2*b[3]*q[2], 2*b[3]*q[3], -4*b[1]*q[2]-2*b[3]*q[0], -4*b[1]*q[3]+2*b[3]*q[1]], 218 | [-2*b[1]*q[3]+2*b[3]*q[1], 2*b[1]*q[2]+2*b[3]*q[0], 2*b[1]*q[1]+2*b[3]*q[3], -2*b[1]*q[0]+2*b[3]*q[2]], 219 | [2*b[1]*q[2], 2*b[1]*q[3]-4*b[3]*q[1], 2*b[1]*q[0]-4*b[3]*q[2], 2*b[1]*q[1]] 220 | ]) 221 | step = j.T.dot(f) 222 | step /= norm(step) # normalise step magnitude 223 | 224 | # Gyroscope compensation drift 225 | gyroscopeQuat = Quaternion(0, gyroscope[0], gyroscope[1], gyroscope[2]) 226 | stepQuat = Quaternion(step.T[0], step.T[1], step.T[2], step.T[3]) 227 | 228 | gyroscopeQuat = gyroscopeQuat + (q.conj() * stepQuat) * 2 * self.samplePeriod * self.zeta * -1 229 | 230 | # Compute rate of change of quaternion 231 | qdot = (q * gyroscopeQuat) * 0.5 - self.beta * step.T 232 | 233 | # Integrate to yield quaternion 234 | q += qdot * self.samplePeriod 235 | self.quaternion = Quaternion(q / norm(q)) # normalise quaternion 236 | 237 | def update_imu(self, gyroscope, accelerometer): 238 | """ 239 | Perform one update step with data from a IMU sensor array 240 | :param gyroscope: A three-element array containing the gyroscope data in radians per second. 241 | :param accelerometer: A three-element array containing the accelerometer data. Can be any unit since a normalized value is used. 242 | """ 243 | q = self.quaternion 244 | 245 | gyroscope = np.array(gyroscope, dtype=float).flatten() 246 | accelerometer = np.array(accelerometer, dtype=float).flatten() 247 | 248 | # Normalise accelerometer measurement 249 | if norm(accelerometer) is 0: 250 | warnings.warn("accelerometer is zero") 251 | return 252 | accelerometer /= norm(accelerometer) 253 | 254 | # Gradient descent algorithm corrective step 255 | f = np.array([ 256 | 2*(q[1]*q[3] - q[0]*q[2]) - accelerometer[0], 257 | 2*(q[0]*q[1] + q[2]*q[3]) - accelerometer[1], 258 | 2*(0.5 - q[1]**2 - q[2]**2) - accelerometer[2] 259 | ]) 260 | j = np.array([ 261 | [-2*q[2], 2*q[3], -2*q[0], 2*q[1]], 262 | [2*q[1], 2*q[0], 2*q[3], 2*q[2]], 263 | [0, -4*q[1], -4*q[2], 0] 264 | ]) 265 | step = j.T.dot(f) 266 | step /= norm(step) # normalise step magnitude 267 | 268 | # Compute rate of change of quaternion 269 | qdot = (q * Quaternion(0, gyroscope[0], gyroscope[1], gyroscope[2])) * 0.5 - self.beta * step.T 270 | 271 | # Integrate to yield quaternion 272 | q += qdot * self.samplePeriod 273 | self.quaternion = Quaternion(q / norm(q)) # normalise quaternion -------------------------------------------------------------------------------- /lib/odrive_uart.py: -------------------------------------------------------------------------------- 1 | import time 2 | import serial 3 | import odrive.enums 4 | import os 5 | import json 6 | 7 | # This is l-gpio 8 | # from RPi import GPIO # Import GPIO module 9 | 10 | # # GPIO setup for resetting ODrive 11 | # GPIO.setmode(GPIO.BCM) 12 | # GPIO.setup(5, GPIO.OUT) 13 | 14 | class ODriveUART: 15 | AXIS_STATE_CLOSED_LOOP_CONTROL = 8 16 | ERROR_DICT = {k: v for k, v in odrive.enums.__dict__ .items() if k.startswith("AXIS_ERROR_")} 17 | 18 | SERIAL_PORT = '/dev/ttyAMA1' 19 | _bus = serial.Serial( 20 | port=SERIAL_PORT, 21 | baudrate=115200, 22 | parity=serial.PARITY_NONE, 23 | stopbits=serial.STOPBITS_ONE, 24 | bytesize=serial.EIGHTBITS, 25 | timeout=1 26 | ) 27 | 28 | def __init__(self, port='/dev/ttyAMA1', left_axis=0, right_axis=1, dir_left=None, dir_right=None): 29 | # If motor directions not supplied, attempt to load them from motor_dir.json. 30 | if dir_left is None or dir_right is None: 31 | cfg_path = os.path.join(os.path.dirname(__file__), "motor_dir.json") 32 | try: 33 | with open(cfg_path, "r") as f: 34 | cfg = json.load(f) 35 | dir_left = cfg.get("left", 1) 36 | dir_right = cfg.get("right", 1) 37 | except FileNotFoundError: 38 | print(f"[ODriveUART] motor_dir.json not found at {cfg_path}. Using default directions (1,1).") 39 | dir_left = dir_left or 1 40 | dir_right = dir_right or 1 41 | except Exception as exc: 42 | print(f"[ODriveUART] Error reading motor_dir.json ({exc}). Using default directions (1,1).") 43 | dir_left = dir_left or 1 44 | dir_right = dir_right or 1 45 | self.bus = serial.Serial( 46 | port=port, 47 | baudrate=115200, 48 | parity=serial.PARITY_NONE, 49 | stopbits=serial.STOPBITS_ONE, 50 | bytesize=serial.EIGHTBITS, 51 | timeout=1 52 | ) 53 | self.left_axis = left_axis 54 | self.right_axis = right_axis 55 | self.dir_left = dir_left 56 | self.dir_right = dir_right 57 | 58 | # Clear the ASCII UART buffer 59 | self.bus.reset_input_buffer() 60 | self.bus.reset_output_buffer() 61 | 62 | def send_command(self, command: str): 63 | self.bus.reset_input_buffer() 64 | self.bus.write(f"{command}\n".encode()) 65 | # Wait for the response if it's a read command 66 | if command.startswith('r') or command.startswith('f'): 67 | # Read until a newline character is encountered 68 | response = self.bus.readline().decode('ascii').strip() 69 | # If the response is empty, print a debug message 70 | if response == '': 71 | print(f"No response received for command: {command}") 72 | return response 73 | 74 | def get_errors_left(self): 75 | return self.get_errors(self.left_axis) 76 | 77 | def get_errors_right(self): 78 | return self.get_errors(self.right_axis) 79 | 80 | def has_errors(self): 81 | for axis in [0,1]: 82 | error_response = self.send_command(f'r axis{axis}.error') 83 | try: 84 | cleaned_response = ''.join(c for c in error_response if c.isdigit()) 85 | error_code = int(cleaned_response) 86 | except ValueError: 87 | print(f"Unexpected error response format: {error_response}") 88 | return True 89 | if error_code != 0: 90 | return True 91 | return False 92 | 93 | def dump_errors(self): 94 | error_sources = [ 95 | "axis0","axis0.encoder", "axis0.controller", "axis0.motor", 96 | "axis1","axis1.encoder", "axis1.controller", "axis1.motor" 97 | ] 98 | print('======= ODrive Errors =======') 99 | for src in error_sources: 100 | error_response = self.send_command(f'r {src}.error') 101 | try: 102 | cleaned_response = ''.join(c for c in error_response if c.isdigit()) 103 | error_code = int(cleaned_response) 104 | except ValueError: 105 | print(f"Unexpected error response format: {error_response}") 106 | continue 107 | 108 | if error_code == 0: 109 | print(src+'.error=0x0: \033[92mNone\033[0m') 110 | continue 111 | 112 | error_prefix = f"{src.split('.')[-1].strip('01').upper()}_ERROR" 113 | error_dict = {name: value for name, value in vars(odrive.enums).items() if name.startswith(error_prefix)} 114 | error_string = "" 115 | for error_name, code in error_dict.items(): 116 | if error_code & code: 117 | error_string += f"{error_name.replace(error_prefix + '_', '').lower().replace('_', ' ')}, " 118 | error_string = error_string.rstrip(", ") 119 | print(f"{src}.error={hex(error_code)}: \033[91m{error_string}\033[0m") 120 | print('=============================') 121 | 122 | def enable_torque_mode_left(self): 123 | self.enable_torque_mode(self.left_axis) 124 | 125 | def enable_torque_mode_right(self): 126 | self.enable_torque_mode(self.right_axis) 127 | 128 | def enable_torque_mode(self, axis): 129 | self.send_command(f'w axis{axis}.controller.config.control_mode 1') 130 | self.send_command(f'w axis{axis}.controller.config.input_mode 1') 131 | print(f"Axis {axis} set to torque control mode") 132 | 133 | def enable_velocity_mode_left(self): 134 | self.enable_velocity_mode(self.left_axis) 135 | 136 | def enable_velocity_mode_right(self): 137 | self.enable_velocity_mode(self.right_axis) 138 | 139 | def enable_velocity_mode(self, axis): 140 | self.send_command(f'w axis{axis}.controller.config.control_mode 2') 141 | self.send_command(f'w axis{axis}.controller.config.input_mode 1') 142 | print(f"Axis {axis} set to velocity control mode") 143 | 144 | def enable_velocity_ramp_mode_left(self): 145 | self.enable_velocity_ramp_mode(self.left_axis) 146 | 147 | def enable_velocity_ramp_mode_right(self): 148 | self.enable_velocity_ramp_mode(self.right_axis) 149 | 150 | def enable_velocity_ramp_mode(self, axis): 151 | self.send_command(f'w axis{axis}.controller.config.control_mode 2') 152 | self.send_command(f'w axis{axis}.controller.config.input_mode 2') 153 | print(f"Axis {axis} set to ramped velocity control mode") 154 | 155 | def set_velocity_ramp_rate_left(self, ramp_rate): 156 | self.set_velocity_ramp_rate(self.left_axis, ramp_rate) 157 | 158 | def set_velocity_ramp_rate_right(self, ramp_rate): 159 | self.set_velocity_ramp_rate(self.right_axis, ramp_rate) 160 | 161 | def set_velocity_ramp_rate(self, axis, ramp_rate): 162 | self.send_command(f'w axis{axis}.controller.config.vel_ramp_rate {ramp_rate:.4f}') 163 | print(f"Axis {axis} velocity ramp rate set to {ramp_rate:.4f} turns/s^2") 164 | 165 | def start_left(self): 166 | self.start(self.left_axis) 167 | 168 | def start_right(self): 169 | self.start(self.right_axis) 170 | 171 | def start(self, axis): 172 | self.send_command(f'w axis{axis}.requested_state 8') 173 | 174 | def set_speed_rpm_left(self, rpm): 175 | self.set_speed_rpm(self.left_axis, rpm, self.dir_left) 176 | 177 | def set_speed_rpm_right(self, rpm): 178 | self.set_speed_rpm(self.right_axis, rpm, self.dir_right) 179 | 180 | def set_speed_rpm(self, axis, rpm, direction): 181 | rps = rpm / 60 182 | self.send_command(f'w axis{axis}.controller.input_vel {rps * direction:.4f}') 183 | 184 | def set_speed_mps_left(self, mps): 185 | WHEEL_DIAMETER_MM = 165 186 | rps = mps / (WHEEL_DIAMETER_MM * 0.001 * 3.14159) 187 | self.send_command(f'w axis{self.left_axis}.controller.input_vel {rps * self.dir_left:.4f}') 188 | 189 | def set_speed_mps_right(self, mps): 190 | WHEEL_DIAMETER_MM = 165 191 | rps = mps / (WHEEL_DIAMETER_MM * 0.001 * 3.14159) 192 | self.send_command(f'w axis{self.right_axis}.controller.input_vel {rps * self.dir_right:.4f}') 193 | 194 | def set_torque_nm_left(self, nm): 195 | self.set_torque_nm(self.left_axis, nm, self.dir_left) 196 | 197 | def set_torque_nm_right(self, nm): 198 | self.set_torque_nm(self.right_axis, nm, self.dir_right) 199 | 200 | def set_torque_nm(self, axis, nm, direction): 201 | torque_bias = 0.05 # Small torque bias in Nm 202 | adjusted_torque = nm * direction + (torque_bias * direction * (1 if nm >= 0 else -1)) 203 | # self.send_command(f'w axis{axis}.controller.input_torque {adjusted_torque:.4f}') 204 | self.send_command(f'c {axis} {adjusted_torque:.4f}') 205 | self.send_command(f'u {axis}') 206 | 207 | def get_speed_rpm_left(self): 208 | return self.get_speed_rpm(self.left_axis, self.dir_left) 209 | 210 | def get_speed_rpm_right(self): 211 | return self.get_speed_rpm(self.right_axis, self.dir_right) 212 | 213 | def get_speed_rpm(self, axis, direction): 214 | response = self.send_command(f'r axis{axis}.encoder.vel_estimate') 215 | return float(response) * direction * 60 216 | 217 | def get_position_turns_left(self): 218 | return self.get_position_turns(self.left_axis, self.dir_left) 219 | 220 | def get_position_turns_right(self): 221 | return self.get_position_turns(self.right_axis, self.dir_right) 222 | 223 | def get_position_turns(self, axis, direction): 224 | response = self.send_command(f'r axis{axis}.encoder.pos_estimate') 225 | return float(response) * direction 226 | 227 | def get_pos_vel_left(self): 228 | return self.get_pos_vel(self.left_axis, self.dir_left) 229 | 230 | def get_pos_vel_right(self): 231 | return self.get_pos_vel(self.right_axis, self.dir_right) 232 | 233 | def get_pos_vel(self, axis, direction): 234 | pos, vel = self.send_command(f'f {axis}').split(' ') 235 | return float(pos) * direction, float(vel) * direction * 60 236 | 237 | def stop_left(self): 238 | self.stop(self.left_axis) 239 | 240 | def stop_right(self): 241 | self.stop(self.right_axis) 242 | 243 | def stop(self, axis): 244 | self.send_command(f'w axis{axis}.controller.input_vel 0') 245 | self.send_command(f'w axis{axis}.controller.input_torque 0') 246 | # Going at high torque and changing to idle causes overcurrent 247 | # self.send_command(f'w axis{axis}.requested_state 1') 248 | 249 | def check_errors_left(self): 250 | return self.check_errors(self.left_axis) 251 | 252 | def check_errors_right(self): 253 | return self.check_errors(self.right_axis) 254 | 255 | def check_errors(self, axis): 256 | response = self.send_command(f'r axis{axis}.error') 257 | try: 258 | # Remove any non-numeric characters (like 'd' for decimal) 259 | cleaned_response = ''.join(c for c in response if c.isdigit()) 260 | return int(cleaned_response) != 0 261 | except ValueError: 262 | print(f"Unexpected response format: {response}") 263 | return True # Assume there's an error if we can't parse the response 264 | 265 | def clear_errors_left(self): 266 | self.clear_errors(self.left_axis) 267 | 268 | def clear_errors_right(self): 269 | self.clear_errors(self.right_axis) 270 | 271 | def clear_errors(self, axis): 272 | self.send_command(f'w axis{axis}.error 0') 273 | self.send_command(f'w axis{axis}.requested_state {self.AXIS_STATE_CLOSED_LOOP_CONTROL}') 274 | 275 | def enable_watchdog_left(self): 276 | self.enable_watchdog(self.left_axis) 277 | 278 | def enable_watchdog_right(self): 279 | self.enable_watchdog(self.right_axis) 280 | 281 | def enable_watchdog(self, axis): 282 | self.send_command(f'w axis{axis}.config.enable_watchdog 1') 283 | 284 | def disable_watchdog_left(self): 285 | self.disable_watchdog(self.left_axis) 286 | 287 | def disable_watchdog_right(self): 288 | self.disable_watchdog(self.right_axis) 289 | 290 | def disable_watchdog(self, axis): 291 | self.send_command(f'w axis{axis}.config.enable_watchdog 0') 292 | 293 | def set_watchdog_timeout(self, timeout): 294 | self.send_command(f'w axis0.config.watchdog_timeout {timeout}') 295 | self.send_command(f'w axis1.config.watchdog_timeout {timeout}') 296 | 297 | def get_bus_voltage(self): 298 | response = self.send_command('r vbus_voltage') 299 | return f"{float(response):.1f}" 300 | 301 | def reset_odrive(): 302 | # GPIO.output(5, GPIO.LOW) 303 | # time.sleep(0.1) 304 | # GPIO.output(5, GPIO.HIGH) 305 | print("ODrive reset attempted") 306 | 307 | # (The previous built-in demo that spun the motors has been removed to avoid 308 | # unexpected side-effects when this module is imported. For a standalone test 309 | # see `quickstart/examples/example_wasd.py` or write your own script.) 310 | 311 | # --------------------------------------------------------------------------- 312 | # Minimal CLI entry-point 313 | # --------------------------------------------------------------------------- 314 | 315 | if __name__ == '__main__': 316 | """Lightweight command-line helper. 317 | 318 | Usage examples: 319 | python -m quickstart.lib.odrive_uart # prints Vbus 320 | """ 321 | import argparse 322 | import sys 323 | 324 | parser = argparse.ArgumentParser(description='ODrive UART helper') 325 | parser.add_argument('--port', default='/dev/ttyAMA1', help='Serial port of ODrive UART-A') 326 | parser.add_argument('--raw', action='store_true', help='Print raw bus voltage without formatting') 327 | args = parser.parse_args() 328 | 329 | try: 330 | voltage = ODriveUART().get_bus_voltage() 331 | except Exception: 332 | # Swallow any low-level error and print placeholder so prompt doesn't break 333 | print('??') 334 | sys.exit(0) 335 | 336 | if voltage is None: 337 | print('??') 338 | elif args.raw: 339 | print(voltage) 340 | else: 341 | print(f'{voltage}') 342 | -------------------------------------------------------------------------------- /lib/vl53l5cx_lib/api.py: -------------------------------------------------------------------------------- 1 | 2 | VL53L5CX_API_REVISION = "VL53L5CX_1.2.0" 3 | 4 | VL53L5CX_DEFAULT_I2C_ADDRESS = 0x52 // 2 5 | 6 | VL53L5CX_RESOLUTION_4X4 = 16 7 | VL53L5CX_RESOLUTION_8X8 = 64 8 | 9 | VL53L5CX_TARGET_ORDER_CLOSEST = 1 10 | VL53L5CX_TARGET_ORDER_STRONGEST = 2 11 | 12 | VL53L5CX_RANGING_MODE_CONTINUOUS = 1 13 | VL53L5CX_RANGING_MODE_AUTONOMOUS = 3 14 | 15 | VL53L5CX_POWER_MODE_SLEEP = 0 16 | VL53L5CX_POWER_MODE_WAKEUP = 1 17 | 18 | VL53L5CX_STATUS_OK = 0 19 | VL53L5CX_MCU_ERROR = 66 20 | VL53L5CX_STATUS_INVALID_PARAM = 127 21 | VL53L5CX_STATUS_ERROR = 255 22 | 23 | VL53L5CX_NVM_DATA_SIZE = 492 24 | VL53L5CX_CONFIGURATION_SIZE = 972 25 | VL53L5CX_OFFSET_BUFFER_SIZE = 488 26 | VL53L5CX_XTALK_BUFFER_SIZE = 776 27 | 28 | VL53L5CX_DCI_ZONE_CONFIG = 0x5450 29 | VL53L5CX_DCI_FREQ_HZ = 0x5458 30 | VL53L5CX_DCI_INT_TIME = 0x545C 31 | VL53L5CX_DCI_FW_NB_TARGET = 0x5478 32 | VL53L5CX_DCI_RANGING_MODE = 0xAD30 33 | VL53L5CX_DCI_DSS_CONFIG = 0xAD38 34 | VL53L5CX_DCI_TARGET_ORDER = 0xAE64 35 | VL53L5CX_DCI_SHARPENER = 0xAED8 36 | VL53L5CX_DCI_MOTION_DETECTOR_CFG = 0xBFAC 37 | VL53L5CX_DCI_SINGLE_RANGE = 0xCD5C 38 | VL53L5CX_DCI_OUTPUT_CONFIG = 0xCD60 39 | VL53L5CX_DCI_OUTPUT_ENABLES = 0xCD68 40 | VL53L5CX_DCI_OUTPUT_LIST = 0xCD78 41 | VL53L5CX_DCI_PIPE_CONTROL = 0xCF78 42 | 43 | VL53L5CX_UI_CMD_STATUS = 0x2C00 44 | VL53L5CX_UI_CMD_START = 0x2C04 45 | VL53L5CX_UI_CMD_END = 0x2FFF 46 | -------------------------------------------------------------------------------- /setup/README.md: -------------------------------------------------------------------------------- 1 | # Quickstart Robot Setup Guide 2 | 3 | Welcome! So you've just received your robot. Follow these steps to get it up and running. 4 | 5 | For more detailed documentation, visit our comprehensive guide at https://docs.bracket.bot/ 6 | 7 | ## 1. Clone the Repository 8 | 9 | First, you need to clone this repo onto your Raspberry Pi. 10 | 11 | ```bash 12 | # Navigate to your home directory 13 | cd ~ 14 | 15 | # Clone the repository to user home (/home//quickstart OR ~/quickstart OR $HOME/quickstart, these are all equivalent) 16 | git clone https://github.com/BracketBotCapstone/quickstart.git ~/quickstart 17 | ``` 18 | 19 | This will download the `quickstart` code into your home directory. 20 | 21 | ## 2. Run the OS Setup Script 22 | 23 | Next, navigate into the setup directory and run the main setup script. This will configure the operating system, install dependencies, and set up necessary services. 24 | 25 | ```bash 26 | cd ~/quickstart/setup 27 | bash setup_os.sh 28 | ``` 29 | 30 | ## 3. Calibrate the Drive System 31 | 32 | Once the OS setup is complete, you need to calibrate the motor controller and determine the motor directions. 33 | 34 | **IMPORTANT:** Before running the next command: 35 | * Place the robot on the floor in a **flat, open area**. 36 | * Ensure there is at least **1 meter of clear space** around the robot in all directions. 37 | * The robot **will move** during this process. **Do not touch it** while it is calibrating. 38 | 39 | Run the calibration script: 40 | 41 | ```bash 42 | cd ~/quickstart/setup 43 | python3 calibrate_drive.py 44 | ``` 45 | 46 | Follow the prompts in the script. It will first calibrate the ODrive motor controller and then test the motor directions. 47 | 48 | ## 4. Test Drive with WASD Example 49 | 50 | If the `calibrate_drive.py` script completes successfully, congratulations! Your robot should be ready for basic control. 51 | 52 | Navigate to the examples directory and run the WASD control script: 53 | 54 | ```bash 55 | cd ~/quickstart/examples 56 | python3 example_wasd.py 57 | ``` 58 | 59 | You should now be able to drive the robot using the W, A, S, and D keys on your keyboard (ensure the terminal window running the script is focused). 60 | 61 | --- 62 | 63 | If you encounter any issues join our [Discord](https://discord.gg/RXRgJyAq93)! -------------------------------------------------------------------------------- /setup/calibrate_drive.py: -------------------------------------------------------------------------------- 1 | # Adds the lib directory to the Python path 2 | import sys 3 | import os 4 | sys.path.append(os.path.join(os.path.dirname(__file__), '..')) 5 | 6 | import time 7 | import fibre 8 | import fibre.serial_transport 9 | import odrive 10 | from odrive.enums import * 11 | import json 12 | import numpy as np 13 | import math 14 | 15 | from lib.odrive_uart import ODriveUART, reset_odrive 16 | from lib.imu import FilteredMPU6050 17 | 18 | # ANSI escape codes for colors 19 | BLUE = '\033[94m' 20 | YELLOW = '\033[93m' 21 | GREEN = '\033[92m' 22 | RED = '\033[91m' 23 | BOLD = '\033[1m' 24 | RESET = '\033[0m' 25 | 26 | # ODrive calibration constants and helper functions 27 | ENCODER_ERRORS = {name: value for name, value in vars(odrive.enums).items() if name.startswith('ENCODER_ERROR')} 28 | CONTROLLER_ERRORS = {name: value for name, value in vars(odrive.enums).items() if name.startswith('CONTROLLER_ERROR')} 29 | MOTOR_ERRORS = {name: value for name, value in vars(odrive.enums).items() if name.startswith('MOTOR_ERROR')} 30 | AXIS_ERRORS = {name: value for name, value in vars(odrive.enums).items() if name.startswith('AXIS_ERROR')} 31 | 32 | 33 | # Helper function to wait until the axis reaches idle state 34 | def wait_for_idle(axis): 35 | while axis.current_state != AXIS_STATE_IDLE: 36 | time.sleep(0.1) 37 | 38 | # Helper function to reconnect to the ODrive after reboot 39 | def connect_odrive(): 40 | print("Connecting to ODrive...") 41 | import odrive 42 | odrv = odrive.find_any(path='serial:/dev/ttyAMA1', timeout=15) 43 | if odrv is None: 44 | raise Exception('ODrive timed out') 45 | return odrv 46 | 47 | def save_and_reboot(odrv): 48 | print("Saving configuration...") 49 | try: 50 | odrv.save_configuration() 51 | print("Configuration saved successfully.") 52 | 53 | print("Rebooting ODrive...") 54 | try: 55 | odrv.reboot() 56 | except: 57 | # Exception is expected as connection is lost during reboot 58 | # Close the hanging connection 59 | odrv.__channel__.serial_device.close() 60 | 61 | except Exception as e: 62 | print(f"Error saving configuration: {str(e)}") 63 | return None 64 | 65 | time.sleep(1) 66 | return connect_odrive() 67 | 68 | 69 | def print_errors(error_type, error_value): 70 | """Print errors for a given component type and error value.""" 71 | if error_value == 0: 72 | return 73 | error_dict = {name: value for name, value in vars(odrive.enums).items() 74 | if name.startswith(f'{error_type.upper()}_ERROR')} 75 | 76 | error_string = "" 77 | for error_name, error_code in error_dict.items(): 78 | if error_value & error_code: 79 | error_string += f"{error_name.replace(f'{error_type.upper()}_ERROR_', '').lower().replace('_', ' ')}, " 80 | error_string = error_string.rstrip(", ") 81 | print(f"\033[91m{error_type.capitalize()} error {hex(error_value)}: {error_string}\033[0m") 82 | 83 | # Function to calibrate a single axis 84 | def calibrate_axis(odrv0, axis): 85 | print(f"Calibrating axis{axis}...") 86 | 87 | # Clear errors 88 | print("Checking errors...") 89 | getattr(odrv0, f'axis{axis}').clear_errors() 90 | 91 | # Wait for a moment to ensure errors are cleared 92 | time.sleep(1) 93 | 94 | # Print current errors to verify they're cleared 95 | axis_error = getattr(odrv0, f'axis{axis}').error 96 | motor_error = getattr(odrv0, f'axis{axis}').motor.error 97 | encoder_error = getattr(odrv0, f'axis{axis}').encoder.error 98 | 99 | if axis_error or motor_error or encoder_error: 100 | print(f"Axis {axis} errors:") 101 | if axis_error: 102 | print_errors('axis', axis_error) 103 | if motor_error: 104 | print_errors('motor', motor_error) 105 | if encoder_error: 106 | print_errors('encoder', encoder_error) 107 | return odrv0, False 108 | 109 | # -------- ODrive Configuration -------- 110 | print("Configuring ODrive...") 111 | getattr(odrv0, f'axis{axis}').config.watchdog_timeout=0.5 112 | getattr(odrv0, f'axis{axis}').config.enable_watchdog=False 113 | getattr(odrv0, f'axis{axis}').motor.config.calibration_current = 5 114 | getattr(odrv0, f'axis{axis}').motor.config.pole_pairs = 15 115 | getattr(odrv0, f'axis{axis}').motor.config.resistance_calib_max_voltage = 4 116 | getattr(odrv0, f'axis{axis}').motor.config.requested_current_range = 25 #Requires config save and reboot 117 | getattr(odrv0, f'axis{axis}').motor.config.current_control_bandwidth = 100 118 | getattr(odrv0, f'axis{axis}').motor.config.torque_constant = 8.27 / 16.0 119 | getattr(odrv0, f'axis{axis}').encoder.config.mode = ENCODER_MODE_HALL 120 | getattr(odrv0, f'axis{axis}').encoder.config.cpr = 90 121 | getattr(odrv0, f'axis{axis}').encoder.config.calib_scan_distance = 150 122 | getattr(odrv0, f'axis{axis}').encoder.config.bandwidth = 100 123 | getattr(odrv0, f'axis{axis}').controller.config.pos_gain = 1 124 | getattr(odrv0, f'axis{axis}').controller.config.vel_gain = 0.02 * getattr(odrv0, f'axis{axis}').motor.config.torque_constant * getattr(odrv0, f'axis{axis}').encoder.config.cpr 125 | getattr(odrv0, f'axis{axis}').controller.config.vel_integrator_gain = 0.1 * getattr(odrv0, f'axis{axis}').motor.config.torque_constant * getattr(odrv0, f'axis{axis}').encoder.config.cpr 126 | getattr(odrv0, f'axis{axis}').controller.config.vel_limit = 10 127 | getattr(odrv0, f'axis{axis}').controller.config.control_mode = CONTROL_MODE_VELOCITY_CONTROL 128 | 129 | odrv0 = save_and_reboot(odrv0) 130 | 131 | # -------- Motor Calibration -------- 132 | print("Starting motor calibration...") 133 | 134 | getattr(odrv0, f'axis{axis}').requested_state = AXIS_STATE_MOTOR_CALIBRATION 135 | wait_for_idle(getattr(odrv0, f'axis{axis}')) 136 | 137 | # Check for errors 138 | error = getattr(odrv0, f'axis{axis}').motor.error 139 | if error != 0: 140 | print_errors('motor', error) 141 | return odrv0, False 142 | else: 143 | print("Motor calibration successful.") 144 | # Validate phase resistance and inductance 145 | resistance = getattr(odrv0, f'axis{axis}').motor.config.phase_resistance 146 | inductance = getattr(odrv0, f'axis{axis}').motor.config.phase_inductance 147 | print(f"Measured phase resistance: {resistance} Ohms") 148 | print(f"Measured phase inductance: {inductance} H") 149 | 150 | if not (0.1 <= resistance <= 1.0): 151 | print("Warning: Phase resistance out of expected range!") 152 | if not (0.0001 <= inductance <= 0.005): 153 | print("Warning: Phase inductance out of expected range!") 154 | 155 | # Mark motor as pre-calibrated 156 | getattr(odrv0, f'axis{axis}').motor.config.pre_calibrated = True 157 | 158 | # -------- Skipping Hall Polarity Calibration -------- 159 | print("Skipping Hall polarity calibration as per your request.") 160 | 161 | # -------- Encoder Offset Calibration -------- 162 | print("Starting encoder offset calibration...") 163 | getattr(odrv0, f'axis{axis}').requested_state = AXIS_STATE_ENCODER_OFFSET_CALIBRATION 164 | wait_for_idle(getattr(odrv0, f'axis{axis}')) 165 | 166 | # Check for errors 167 | error = getattr(odrv0, f'axis{axis}').encoder.error 168 | if error != 0: 169 | print_errors('encoder', error) 170 | return odrv0, False 171 | else: 172 | print("Encoder calibration successful.") 173 | # Validate phase offset float 174 | phase_offset_float = getattr(odrv0, f'axis{axis}').encoder.config.offset_float 175 | print(f"Phase offset float: {phase_offset_float}") 176 | 177 | if abs((phase_offset_float % 1) - 0.5) > 0.1: 178 | print("Warning: Phase offset float is out of expected range!") 179 | 180 | # Mark encoder as pre-calibrated 181 | getattr(odrv0, f'axis{axis}').encoder.config.pre_calibrated = True 182 | 183 | # -------- Test Motor Control -------- 184 | print("Testing motor control...") 185 | 186 | # Enter closed-loop control 187 | getattr(odrv0, f'axis{axis}').requested_state = AXIS_STATE_CLOSED_LOOP_CONTROL 188 | time.sleep(1) # Wait for state to settle 189 | 190 | # Command a velocity 191 | print("Spinning motor at 0.5 turns/sec...") 192 | getattr(odrv0, f'axis{axis}').controller.input_vel = 0.5 193 | time.sleep(2) 194 | 195 | # Stop the motor 196 | print("Stopping motor...") 197 | getattr(odrv0, f'axis{axis}').controller.input_vel = 0 198 | time.sleep(1) 199 | 200 | # Switch back to idle 201 | getattr(odrv0, f'axis{axis}').requested_state = AXIS_STATE_IDLE 202 | 203 | # -------- Automatic Startup Configuration -------- 204 | print("Configuring automatic startup...") 205 | 206 | # Set axis to start in closed-loop control on startup 207 | getattr(odrv0, f'axis{axis}').config.startup_closed_loop_control = True 208 | 209 | return odrv0, True 210 | 211 | def test_motor_direction(): 212 | # Initialize IMU 213 | imu = FilteredMPU6050() 214 | imu.calibrate() 215 | 216 | # Reset ODrive before initializing motors 217 | reset_odrive() 218 | time.sleep(3) # Wait for ODrive to reset 219 | 220 | motor_controller = ODriveUART(port='/dev/ttyAMA1', left_axis=0, right_axis=1, dir_left=1, dir_right=1) 221 | directions = {'left': 1, 'right': 1} 222 | 223 | angle_threshold = 5.0 # degrees 224 | max_spin_duration = 5.0 # seconds 225 | 226 | for name in ['left', 'right']: 227 | time.sleep(1) 228 | print(f"\nTesting {name} motor...") 229 | 230 | # Start motor in velocity mode and clear errors 231 | if name == 'left': 232 | motor_controller.start_left() 233 | motor_controller.enable_velocity_mode_left() 234 | if motor_controller.check_errors_left(): 235 | print("Clearing left motor errors...") 236 | motor_controller.clear_errors_left() 237 | else: 238 | motor_controller.start_right() 239 | motor_controller.enable_velocity_mode_right() 240 | if motor_controller.check_errors_right(): 241 | print("Clearing right motor errors...") 242 | motor_controller.clear_errors_right() 243 | 244 | # Spin motor (e.g., at 30 RPM) 245 | if name == 'left': 246 | motor_controller.set_speed_rpm_left(30) 247 | else: 248 | motor_controller.set_speed_rpm_right(30) 249 | 250 | # Get initial yaw 251 | _, _, initial_yaw = imu.get_orientation() 252 | start_time = time.time() 253 | 254 | # Wait until we exceed angle_threshold or hit max_spin_duration 255 | while True: 256 | time.sleep(0.01) 257 | _, _, current_yaw = imu.get_orientation() 258 | angle_diff = current_yaw - initial_yaw 259 | 260 | if abs(angle_diff) >= angle_threshold: 261 | break 262 | if (time.time() - start_time) > max_spin_duration: 263 | print("Reached max spin duration without hitting angle threshold.") 264 | break 265 | 266 | # Stop the motor 267 | if name == 'left': 268 | motor_controller.stop_left() 269 | else: 270 | motor_controller.stop_right() 271 | 272 | print(f"Yaw difference before stopping: {angle_diff:.2f} deg") 273 | 274 | # Determine final direction based on sign of yaw difference 275 | if abs(angle_diff) < angle_threshold: 276 | print("Angle change too small; defaulting to forward (+1).") 277 | directions[name] = 1 278 | else: 279 | if name == 'left': 280 | directions['left'] = -1 if angle_diff > 0 else 1 281 | else: 282 | directions['right'] = 1 if angle_diff > 0 else -1 283 | 284 | time.sleep(0.5) 285 | 286 | # Save direction results 287 | with open(os.path.expanduser('~/quickstart/lib/motor_dir.json'), 'w') as f: 288 | json.dump(directions, f) 289 | 290 | print("\nDirection test complete!") 291 | print(f"Left direction: {directions['left']}, Right direction: {directions['right']}") 292 | print(f"Results saved to ~/quickstart/lib/motor_dir.json: {directions}") 293 | 294 | def calibrate_odrive(): 295 | print("Finding an ODrive...") 296 | odrv0 = connect_odrive() 297 | print("Found ODrive.") 298 | 299 | # ASCII art for clear space warning 300 | print(r""" 301 | {YELLOW} 302 | 1m radius 303 | _________________ 304 | / ^ \ 305 | / | \ 306 | | 1m | 307 | | | | 308 | | <--1m-->{{BOT}}<--1m--> | 309 | | | | 310 | | 1m | 311 | \ | / 312 | \_________v_________/ 313 | 314 | {RESET} 315 | """.format(YELLOW=YELLOW, RESET=RESET)) 316 | 317 | print(f"{BOLD}WARNING:{RESET} The robot needs clear space to move during calibration.") 318 | confirmation = input(f"{BLUE}Ensure the robot has at least 1 meter of clear space around it.\nIs the area clear? [yes/no]: {RESET}").lower() 319 | if confirmation.lower() != 'yes': 320 | print(f'{YELLOW}Please ensure the area is clear and rerun the script.{RESET}') 321 | return False 322 | print() 323 | 324 | for axis in [0,1]: 325 | odrv0, success = calibrate_axis(odrv0, axis) 326 | if success: 327 | print('\033[92m' + f"Axis {axis} calibration completed successfully." + '\033[0m') 328 | print() 329 | else: 330 | print('\033[91m' + f"Axis {axis} calibration failed." + '\033[0m') 331 | print('\nPlease fix the issue with this axis before rerunning this script.') 332 | return False 333 | 334 | odrv0 = save_and_reboot(odrv0) 335 | 336 | print('\033[94m' + "\nODrive setup complete." + '\033[0m') 337 | return True 338 | 339 | if __name__ == '__main__': 340 | try: 341 | print("\n\033[1;33mDrive Calibration\033[0m") 342 | print("\nThis script will run two calibration steps:") 343 | print("1. ODrive motor calibration - requires robot to be on a stand with wheels free to spin") 344 | print("2. Motor direction calibration - requires robot to be on the ground with some open space") 345 | 346 | # First calibration - ODrive 347 | print("\n\033[1;36mStep 1: ODrive Motor Calibration\033[0m") 348 | if not calibrate_odrive(): 349 | sys.exit(1) 350 | 351 | # Second calibration - Motor direction 352 | print("\n\033[1;36mStep 2: Motor Direction Calibration\033[0m") 353 | # Removed confirmation prompt - assuming robot is ready 354 | # confirmation = input("Place the robot on the ground with space to move.\nIs the robot on the ground with space to move? [yes/no] ").lower() 355 | # if confirmation.lower() != 'yes': 356 | # print('Rerun this script once the robot is on the ground with space to move.') 357 | # sys.exit(1) 358 | 359 | test_motor_direction() 360 | 361 | print("\n\033[1;32mDrive calibration complete!\033[0m") 362 | 363 | except Exception as e: 364 | print(f"\n\033[91mError occurred: {e}\033[0m") 365 | sys.exit(1) 366 | -------------------------------------------------------------------------------- /setup/calibrate_stereo_camera.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ["RERUN"] = "1" 3 | import cv2 4 | import numpy as np 5 | import time 6 | from datetime import datetime 7 | import rerun as rr 8 | rr.init("interactive_calibrate", spawn=False) 9 | 10 | # Path to save calibration YAML to the shared lib folder 11 | CALIB_YAML_PATH = os.path.join(os.path.dirname(__file__), '..', 'lib', 'stereo_calibration_fisheye.yaml') 12 | 13 | # --- HELPER: save calibration mats robustly to YAML --- 14 | def write_calib_yaml(params: dict, filename: str = CALIB_YAML_PATH): 15 | """Save numpy matrices to a YAML file using cv2.FileStorage. 16 | 17 | Skips any entries that are None and ensures values are numpy arrays. 18 | """ 19 | try: 20 | fs = cv2.FileStorage(filename, cv2.FILE_STORAGE_WRITE) 21 | for k, v in params.items(): 22 | if v is None: 23 | continue 24 | fs.write(k, np.asarray(v)) 25 | fs.release() 26 | print(f"[INFO] Calibration parameters saved to {filename}") 27 | except Exception as e: 28 | print(f"[WARN] Could not save calibration file: {e}") 29 | 30 | # --- CONFIGURATION --- 31 | CHECKERBOARD = (9, 6) 32 | SQUARE_SIZE = 22.5 # mm 33 | IMG_DIR = './captures' 34 | os.makedirs(IMG_DIR, exist_ok=True) 35 | 36 | # --- RERUN TCP CONNECTION --- 37 | rr.connect_grpc("rerun+http://:9876/proxy") 38 | # --- STATE --- 39 | captured_pairs = [] # List of (left, right, fname) 40 | 41 | # Load existing images from captures folder 42 | import glob 43 | existing_files = glob.glob(os.path.join(IMG_DIR, '*_usb.jpg')) 44 | for fname in existing_files: 45 | img = cv2.imread(fname) 46 | if img is not None: 47 | height, width = img.shape[:2] 48 | left = img[:, :width//2] 49 | right = img[:, width//2:] 50 | captured_pairs.append((left, right, fname)) 51 | print(f"[INFO] Loaded existing image: {fname}") 52 | if captured_pairs: 53 | print(f"[INFO] Loaded {len(captured_pairs)} existing image pairs for calibration.") 54 | 55 | # --- Rectification maps placeholder --- 56 | latest_maps = {'map1x': None, 'map1y': None, 'map2x': None, 'map2y': None} 57 | # defer initial calibration until calibrate_and_rectify is defined 58 | init_calib_pairs_loaded = len(captured_pairs) >= 5 59 | 60 | # --- KEYBOARD POLLING (using select, no threads) --- 61 | import sys 62 | import select 63 | 64 | # --- CALIBRATION UTILS --- 65 | def calibrate_and_rectify(pairs): 66 | if len(pairs) < 5: 67 | return None, None, None, None 68 | 69 | # Prepare object points in fisheye format (N,1,3) 70 | objp = np.zeros((CHECKERBOARD[0]*CHECKERBOARD[1], 1, 3), np.float32) 71 | objp[:, 0, :2] = np.mgrid[0:CHECKERBOARD[0], 0:CHECKERBOARD[1]].T.reshape(-1, 2) 72 | objp *= SQUARE_SIZE 73 | 74 | objpoints, imgpoints_left, imgpoints_right = [], [], [] 75 | for left, right, _ in pairs: 76 | gray_left = cv2.cvtColor(left, cv2.COLOR_BGR2GRAY) 77 | gray_right = cv2.cvtColor(right, cv2.COLOR_BGR2GRAY) 78 | ret_l, corners_l = cv2.findChessboardCorners(gray_left, CHECKERBOARD, cv2.CALIB_CB_FAST_CHECK) 79 | ret_r, corners_r = cv2.findChessboardCorners(gray_right, CHECKERBOARD, cv2.CALIB_CB_FAST_CHECK) 80 | if ret_l and ret_r: 81 | # Refine corners for better accuracy 82 | crit = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 1e-6) 83 | corners_l = cv2.cornerSubPix(gray_left, corners_l, (3, 3), (-1, -1), crit) 84 | corners_r = cv2.cornerSubPix(gray_right, corners_r, (3, 3), (-1, -1), crit) 85 | objpoints.append(objp) 86 | imgpoints_left.append(corners_l) 87 | imgpoints_right.append(corners_r) 88 | 89 | if len(objpoints) < 5: 90 | return None, None, None, None 91 | 92 | # Initialize camera matrices & distortion coeffs 93 | K1 = np.zeros((3, 3)) 94 | D1 = np.zeros((4, 1)) 95 | K2 = np.zeros((3, 3)) 96 | D2 = np.zeros((4, 1)) 97 | 98 | image_size = gray_left.shape[::-1] 99 | criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, 1e-6) 100 | flags = cv2.fisheye.CALIB_RECOMPUTE_EXTRINSIC | cv2.fisheye.CALIB_FIX_SKEW 101 | 102 | # OpenCV fisheye.stereoCalibrate has changed signature across versions. 103 | # It may return 7 or 9 values (… E, F). Capture robustly. 104 | _stereo_out = cv2.fisheye.stereoCalibrate( 105 | objpoints, 106 | imgpoints_left, 107 | imgpoints_right, 108 | K1, 109 | D1, 110 | K2, 111 | D2, 112 | image_size, 113 | criteria=criteria, 114 | flags=flags, 115 | ) 116 | 117 | rms = _stereo_out[0] 118 | K1, D1, K2, D2, R, T = _stereo_out[1:7] 119 | E = F = None 120 | if len(_stereo_out) >= 9: 121 | E, F = _stereo_out[7:9] 122 | 123 | # Alias pin‑hole variable names used elsewhere so the rest of the script and 124 | # downstream YAML consumers (stereo.py) remain untouched. 125 | mtx_l, dist_l = K1, D1 126 | mtx_r, dist_r = K2, D2 127 | 128 | print(f"[INFO] Fisheye stereo calibration RMS error: {rms:.3f}") 129 | 130 | # Rectify 131 | R1, R2, P1, P2, Q = cv2.fisheye.stereoRectify(K1, D1, K2, D2, image_size, R, T, flags=cv2.CALIB_ZERO_DISPARITY) 132 | 133 | map1x, map1y = cv2.fisheye.initUndistortRectifyMap(K1, D1, R1, P1, image_size, cv2.CV_16SC2) 134 | map2x, map2y = cv2.fisheye.initUndistortRectifyMap(K2, D2, R2, P2, image_size, cv2.CV_16SC2) 135 | 136 | return map1x, map1y, map2x, map2y, mtx_l, dist_l, mtx_r, dist_r, R, T, E, F, R1, R2, P1, P2, Q 137 | 138 | def draw_epipolar_lines(img, color=(0,255,0), step=25): 139 | out = img.copy() 140 | for y in range(0, out.shape[0], step): 141 | cv2.line(out, (0, y), (out.shape[1], y), color, 1) 142 | return out 143 | 144 | # --- Run initial calibration _after_ function definition --- 145 | if init_calib_pairs_loaded: 146 | print("[INFO] Attempting initial calibration with existing images ...") 147 | calib_res = calibrate_and_rectify(captured_pairs) 148 | map1x, map1y, map2x, map2y = calib_res[:4] 149 | if map1x is not None: 150 | latest_maps['map1x'] = map1x 151 | latest_maps['map1y'] = map1y 152 | latest_maps['map2x'] = map2x 153 | latest_maps['map2y'] = map2y 154 | print("[INFO] Initial calibration successful – rectification ready.") 155 | # Collect params and write in one go 156 | calib_dict = { 157 | 'mtx_l': calib_res[4], 158 | 'dist_l': calib_res[5], 159 | 'mtx_r': calib_res[6], 160 | 'dist_r': calib_res[7], 161 | 'R': calib_res[8], 162 | 'T': calib_res[9], 163 | 'E': calib_res[10], 164 | 'F': calib_res[11], 165 | 'R1': calib_res[12], 166 | 'R2': calib_res[13], 167 | 'P1': calib_res[14], 168 | 'P2': calib_res[15], 169 | 'Q': calib_res[16], 170 | } 171 | write_calib_yaml(calib_dict, CALIB_YAML_PATH) 172 | else: 173 | print("[INFO] Initial calibration failed or not enough valid pairs; capture more images.") 174 | else: 175 | print("[INFO] Fewer than 5 image pairs present – capture more to calibrate.") 176 | 177 | # --- MAIN LOOP --- 178 | 179 | cap = cv2.VideoCapture(0) 180 | cap.set(cv2.CAP_PROP_FRAME_WIDTH, 2560) 181 | cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720) 182 | if not cap.isOpened(): 183 | print("Could not open video device 0") 184 | exit(1) 185 | 186 | print("Interactive stereo calibration. Press SPACE to capture, ESC to quit.") 187 | 188 | last_log = 0 189 | log_interval = 1.0 # 1s for rerun logging 190 | print("Interactive stereo calibration. Type 'space'+Enter to capture, 'esc'+Enter to quit.") 191 | 192 | while True: 193 | ret, frame = cap.read() 194 | if not ret: 195 | print("Failed to grab frame") 196 | break 197 | h, w = frame.shape[:2] 198 | left = frame[:, :w//2] 199 | right = frame[:, w//2:] 200 | gray_left = cv2.cvtColor(left, cv2.COLOR_BGR2GRAY) 201 | gray_right = cv2.cvtColor(right, cv2.COLOR_BGR2GRAY) 202 | ret_l, corners_l = cv2.findChessboardCorners(gray_left, CHECKERBOARD, None) 203 | ret_r, corners_r = cv2.findChessboardCorners(gray_right, CHECKERBOARD, None) 204 | print(f"Checkerboard detection: left={'YES' if ret_l else 'NO'}, right={'YES' if ret_r else 'NO'}") 205 | vis_left = left.copy() 206 | vis_right = right.copy() 207 | if ret_l: 208 | cv2.drawChessboardCorners(vis_left, CHECKERBOARD, corners_l, ret_l) 209 | if ret_r: 210 | cv2.drawChessboardCorners(vis_right, CHECKERBOARD, corners_r, ret_r) 211 | # --- Live rectification preview --- 212 | if all(latest_maps[k] is not None for k in ['map1x', 'map1y', 'map2x', 'map2y']): 213 | rect_left = cv2.remap(left, latest_maps['map1x'], latest_maps['map1y'], cv2.INTER_LINEAR) 214 | rect_right = cv2.remap(right, latest_maps['map2x'], latest_maps['map2y'], cv2.INTER_LINEAR) 215 | rect_left_lines = draw_epipolar_lines(rect_left) 216 | rect_right_lines = draw_epipolar_lines(rect_right) 217 | else: 218 | rect_left_lines = None 219 | rect_right_lines = None 220 | # Downscale for rerun visualization only 221 | now = time.time() 222 | if now - last_log > log_interval: 223 | vis_left_small = cv2.resize(vis_left, (vis_left.shape[1]//4, vis_left.shape[0]//4)) 224 | vis_right_small = cv2.resize(vis_right, (vis_right.shape[1]//4, vis_right.shape[0]//4)) 225 | rr.log("left", rr.Image(vis_left_small)) 226 | rr.log("right", rr.Image(vis_right_small)) 227 | if rect_left_lines is not None and rect_right_lines is not None: 228 | rect_left_small = cv2.resize(rect_left_lines, (rect_left_lines.shape[1]//4, rect_left_lines.shape[0]//4)) 229 | rect_right_small = cv2.resize(rect_right_lines, (rect_right_lines.shape[1]//4, rect_right_lines.shape[0]//4)) 230 | rr.log("rect_left_live", rr.Image(rect_left_small)) 231 | rr.log("rect_right_live", rr.Image(rect_right_small)) 232 | last_log = now 233 | # --- Poll for keyboard input --- 234 | input_ready, _, _ = select.select([sys.stdin], [], [], 0) 235 | if input_ready: 236 | print("DEBUG: Input detected by select!") 237 | key = sys.stdin.readline() 238 | print(f"DEBUG: Key read: '{key}'") 239 | print(f"DEBUG: Enter pressed. ret_l={ret_l}, ret_r={ret_r}") 240 | if not ret_l or not ret_r: 241 | print("[INFO] Checkerboard not detected in both images. Image not saved.") 242 | else: 243 | timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f') 244 | fname = os.path.join(IMG_DIR, f'{timestamp}_usb.jpg') 245 | cv2.imwrite(fname, frame) 246 | print(f"[INFO] Saved image for calibration: {fname}") 247 | captured_pairs.append((left.copy(), right.copy(), fname)) 248 | print(f"[INFO] Captured and saved: {fname}") 249 | print(f"[INFO] Attempting calibration/rectification with {len(captured_pairs)} pairs...") 250 | calib_result = calibrate_and_rectify(captured_pairs) 251 | map1x, map1y, map2x, map2y = calib_result[:4] 252 | # Save calibration results if available 253 | if map1x is not None: 254 | print("[INFO] Calibration and rectification successful. Logging rectified images.") 255 | # Update latest maps for live preview 256 | latest_maps['map1x'] = map1x 257 | latest_maps['map1y'] = map1y 258 | latest_maps['map2x'] = map2x 259 | latest_maps['map2y'] = map2y 260 | mtx_l = dist_l = mtx_r = dist_r = R = T = E = F = R1 = R2 = P1 = P2 = Q = None 261 | if len(calib_result) > 4: 262 | mtx_l, dist_l, mtx_r, dist_r, R, T, E, F, R1, R2, P1, P2, Q = calib_result[4:] 263 | calib_dict = { 264 | 'mtx_l': mtx_l, 265 | 'dist_l': dist_l, 266 | 'mtx_r': mtx_r, 267 | 'dist_r': dist_r, 268 | 'R': R, 269 | 'T': T, 270 | 'E': E, 271 | 'F': F, 272 | 'R1': R1, 273 | 'R2': R2, 274 | 'P1': P1, 275 | 'P2': P2, 276 | 'Q': Q, 277 | } 278 | write_calib_yaml(calib_dict, CALIB_YAML_PATH) 279 | rect_left = cv2.remap(left, map1x, map1y, cv2.INTER_LINEAR) 280 | rect_right = cv2.remap(right, map2x, map2y, cv2.INTER_LINEAR) 281 | rect_left_lines = draw_epipolar_lines(rect_left) 282 | rect_right_lines = draw_epipolar_lines(rect_right) 283 | rect_left_small = cv2.resize(rect_left_lines, (rect_left_lines.shape[1]//4, rect_left_lines.shape[0]//4)) 284 | rect_right_small = cv2.resize(rect_right_lines, (rect_right_lines.shape[1]//4, rect_right_lines.shape[0]//4)) 285 | rr.log("rect_left", rr.Image(rect_left_small)) 286 | rr.log("rect_right", rr.Image(rect_right_small)) 287 | else: 288 | print("[INFO] Not enough valid pairs for calibration.") 289 | 290 | cap.release() 291 | cv2.destroyAllWindows() 292 | -------------------------------------------------------------------------------- /setup/extras/README.md: -------------------------------------------------------------------------------- 1 | # Optional Setup Scripts 2 | 3 | This directory contains setup scripts for optional hardware or software components that are not part of the core system setup performed by the main `setup_os.sh` script. 4 | 5 | These scripts should typically be run manually after the main setup is complete, if the corresponding hardware/software is needed. 6 | 7 | ## Scripts: 8 | 9 | * `setup_realsense.sh`: Sets up Intel RealSense cameras. 10 | * `setup_hardware_pwm.sh`: Configures hardware PWM using the provided `.dts` file. 11 | * `pwm-pi5.dts`: Device Tree Source file needed for hardware PWM setup. 12 | * `create_custom_boot_service.sh`: A template script to help create and enable a systemd service for running a custom script (Python, Bash, etc.) on boot. Copy and edit this template, then run your edited version. -------------------------------------------------------------------------------- /setup/extras/create_custom_boot_service.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # ------------------------------------------------------------------------------ 4 | # TEMPLATE SCRIPT for creating a Systemd service to run a script on boot. 5 | # 6 | # HOW TO USE: 7 | # 1. Copy this script to a new file (e.g., setup_my_cool_app_service.sh). 8 | # 2. Edit the 'USER CONFIGURATION' section below with your specific details. 9 | # 3. Make sure the script you want to run (`SCRIPT_TO_RUN_PATH`) exists and 10 | # is executable (`chmod +x /path/to/your/script`). 11 | # 4. Run the *new* script (e.g., `bash setup_my_cool_app_service.sh`). 12 | # 13 | # NOTE: This script needs to be run with sudo privileges indirectly because it 14 | # generates commands that use sudo (copying files to /etc, systemctl). 15 | # However, run the script itself as the user who will own the process 16 | # if RUN_AS_USER is set to $USER, or manually ensure permissions. 17 | # ------------------------------------------------------------------------------ 18 | 19 | set -e # Exit immediately if a command exits with a non-zero status. 20 | 21 | # --- USER CONFIGURATION - EDIT THESE VARIABLES --- # 22 | 23 | # 1. Service Name: Choose a unique name for your service (e.g., my-app, data-logger). 24 | # The script will create a file named .service. 25 | SERVICE_NAME="my_custom_service" 26 | 27 | # 2. Description: A short description of what your service does. 28 | SERVICE_DESCRIPTION="My Custom Application Service" 29 | 30 | # 3. Script to Run: The *absolute* path to the script you want to run on boot. 31 | # Make sure this script is executable (`chmod +x`). 32 | SCRIPT_TO_RUN_PATH="/home/$USER/my_project/my_script.py" # Example: /home/ollama/my_app/start.sh 33 | 34 | # 4. Interpreter: The *absolute* path to the interpreter for your script. 35 | # - For Python scripts IN THE VENV: Use the venv's python executable. 36 | # - For standard Bash scripts: Use /bin/bash. 37 | # - For other interpreters, provide the full path. 38 | INTERPRETER_PATH="/home/$USER/quickstart-venv/bin/python3" # Example: /bin/bash 39 | 40 | # 5. User to Run As: The system user the service should run under. 41 | # Using $USER (the user running this setup script) is common. 42 | # Use 'root' with caution if necessary. 43 | RUN_AS_USER="$USER" 44 | 45 | # 6. Working Directory: The directory the script should be in when it runs. 46 | # This is important for scripts that use relative paths. 47 | WORKING_DIR="$(dirname "$SCRIPT_TO_RUN_PATH")" # Defaults to the script's directory 48 | 49 | # 7. Restart Policy: What to do if the script exits. 50 | # Options: no, on-success, on-failure, on-abnormal, on-watchdog, on-abort, or always. 51 | RESTART_POLICY="on-failure" 52 | 53 | # 8. Startup Dependencies: Services that should be running before this one starts. 54 | # Commonly 'network-online.target' if your script needs network access. 55 | # Separate multiple targets with spaces. 56 | AFTER_DEPENDENCIES="network-online.target" 57 | 58 | # --- END OF USER CONFIGURATION --- # 59 | 60 | # --- Service File Generation --- # 61 | 62 | SERVICE_FILE_PATH="/etc/systemd/system/${SERVICE_NAME}.service" 63 | TEMP_SERVICE_FILE="/tmp/${SERVICE_NAME}.service.$$" 64 | 65 | # Check if the script path exists 66 | if [ ! -f "$SCRIPT_TO_RUN_PATH" ]; then 67 | echo "Error: Script to run does not exist at: $SCRIPT_TO_RUN_PATH" >&2 68 | echo "Please correct SCRIPT_TO_RUN_PATH in this script." >&2 69 | exit 1 70 | fi 71 | 72 | # Check if the interpreter path exists 73 | if [ ! -x "$INTERPRETER_PATH" ]; then 74 | echo "Error: Interpreter does not exist or is not executable at: $INTERPRETER_PATH" >&2 75 | echo "Please correct INTERPRETER_PATH in this script." >&2 76 | exit 1 77 | fi 78 | 79 | echo "--- Generating Systemd Service File --- " 80 | echo "Service Name: ${SERVICE_NAME}.service" 81 | echo "Description: ${SERVICE_DESCRIPTION}" 82 | echo "ExecStart: ${INTERPRETER_PATH} ${SCRIPT_TO_RUN_PATH}" 83 | echo "User: ${RUN_AS_USER}" 84 | echo "Working Dir: ${WORKING_DIR}" 85 | echo "Restart: ${RESTART_POLICY}" 86 | echo "After: ${AFTER_DEPENDENCIES}" 87 | echo "" 88 | 89 | # Create the service file content 90 | cat > "$TEMP_SERVICE_FILE" << EOL 91 | [Unit] 92 | Description=${SERVICE_DESCRIPTION} 93 | After=${AFTER_DEPENDENCIES} 94 | 95 | [Service] 96 | User=${RUN_AS_USER} 97 | Group=$(id -gn "$RUN_AS_USER") 98 | WorkingDirectory=${WORKING_DIR} 99 | ExecStart=${INTERPRETER_PATH} ${SCRIPT_TO_RUN_PATH} 100 | Restart=${RESTART_POLICY} 101 | # Optional: Set environment variables if needed 102 | # Environment="PYTHONUNBUFFERED=1" 103 | # Environment="OTHER_VAR=value" 104 | 105 | [Install] 106 | WantedBy=multi-user.target 107 | EOL 108 | 109 | echo "Service file content generated successfully in $TEMP_SERVICE_FILE" 110 | 111 | # --- Installation and Activation --- # 112 | 113 | echo " 114 | --- Installing and Enabling Service --- " 115 | 116 | echo "The following commands will be run using sudo:" 117 | 118 | # 1. Copy the service file to the systemd directory 119 | CMD_COPY="sudo cp '$TEMP_SERVICE_FILE' '$SERVICE_FILE_PATH'" 120 | echo "1. $CMD_COPY" 121 | 122 | # 2. Set permissions for the service file 123 | CMD_CHMOD="sudo chmod 644 '$SERVICE_FILE_PATH'" 124 | echo "2. $CMD_CHMOD" 125 | 126 | # 3. Reload systemd manager configuration 127 | CMD_RELOAD="sudo systemctl daemon-reload" 128 | echo "3. $CMD_RELOAD" 129 | 130 | # 4. Enable the service to start on boot 131 | CMD_ENABLE="sudo systemctl enable '${SERVICE_NAME}.service'" 132 | echo "4. $CMD_ENABLE" 133 | 134 | # 5. Start the service immediately (optional, comment out if you only want it to start on next boot) 135 | CMD_START="sudo systemctl start '${SERVICE_NAME}.service'" 136 | echo "5. $CMD_START" 137 | 138 | # 6. Check the status of the service 139 | CMD_STATUS="sudo systemctl status '${SERVICE_NAME}.service' --no-pager -l" 140 | echo "6. $CMD_STATUS" 141 | 142 | echo "" 143 | read -p "Press Enter to execute these commands, or Ctrl+C to cancel..." 144 | 145 | # Execute the commands 146 | echo "Executing commands..." 147 | $CMD_COPY 148 | $CMD_CHMOD 149 | $CMD_RELOAD 150 | $CMD_ENABLE 151 | $CMD_START 152 | 153 | echo " 154 | --- Service Status --- " 155 | $CMD_STATUS 156 | 157 | # Clean up temporary file 158 | rm "$TEMP_SERVICE_FILE" 159 | 160 | echo " 161 | --- Done --- " 162 | echo "Service '${SERVICE_NAME}' created and enabled." 163 | echo "To check logs: sudo journalctl -u ${SERVICE_NAME}.service -f" 164 | echo "To stop: sudo systemctl stop ${SERVICE_NAME}.service" 165 | echo "To disable: sudo systemctl disable ${SERVICE_NAME}.service" -------------------------------------------------------------------------------- /setup/extras/pwm-pi5.dts: -------------------------------------------------------------------------------- 1 | /dts-v1/; 2 | /plugin/; 3 | 4 | /{ 5 | compatible = "brcm,bcm2712"; 6 | 7 | fragment@0 { 8 | target = <&rp1_gpio>; 9 | __overlay__ { 10 | pwm_pins: pwm_pins { 11 | pins = "gpio12", "gpio13", "gpio18", "gpio19"; 12 | function = "pwm0", "pwm0", "pwm0", "pwm0"; 13 | }; 14 | }; 15 | }; 16 | 17 | fragment@1 { 18 | target = <&rp1_pwm0>; 19 | frag1: __overlay__ { 20 | pinctrl-names = "default"; 21 | pinctrl-0 = <&pwm_pins>; 22 | status = "okay"; 23 | }; 24 | }; 25 | }; -------------------------------------------------------------------------------- /setup/extras/setup_hardware_pwm.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -e # Exit on non-zero status 4 | 5 | # Note: This script expects pwm-pi5.dts to be in the same directory 6 | 7 | echo "--- Running Hardware PWM Setup ---" 8 | 9 | echo "Compiling Device Tree Source..." 10 | dtc -I dts -O dtb -o pwm-pi5.dtbo pwm-pi5.dts 11 | echo "Copying Device Tree Overlay to /boot/firmware/overlays/..." 12 | sudo cp pwm-pi5.dtbo /boot/firmware/overlays/ 13 | 14 | echo "Checking if dtoverlay=pwm-pi5 already exists in config.txt..." 15 | if ! grep -q "^dtoverlay=pwm-pi5" /boot/firmware/config.txt; then 16 | echo "Adding dtoverlay=pwm-pi5 to /boot/firmware/config.txt..." 17 | echo "dtoverlay=pwm-pi5" | sudo tee -a /boot/firmware/config.txt 18 | else 19 | echo "dtoverlay=pwm-pi5 already present in /boot/firmware/config.txt." 20 | fi 21 | 22 | # Add permissions to user so you don't need sudo to run python scripts 23 | echo "Configuring PWM permissions..." 24 | if ! groups $USER | grep &>/dev/null '\bdialout\b'; then 25 | sudo usermod -aG dialout $USER 26 | echo "Added user $USER to dialout group (required for PWM permissions)" 27 | else 28 | echo "User $USER already in dialout group." 29 | fi 30 | 31 | PWM_RULES_FILE="/etc/udev/rules.d/99-pwm-permissions.rules" 32 | echo "Creating/Updating udev rule file: $PWM_RULES_FILE" 33 | sudo tee $PWM_RULES_FILE > /dev/null << EOT 34 | SUBSYSTEM=="pwm*", PROGRAM="/bin/sh -c '\\ 35 | chown -R root:dialout /sys/class/pwm && chmod -R 770 /sys/class/pwm;\\ 36 | chown -R root:dialout /sys/devices/platform/soc/*.pwm/pwm/pwmchip* && chmod -R 770 /sys/devices/platform/soc/*.pwm/pwm/pwmchip*;\\ 37 | chown -R root:dialout /sys/devices/platform/axi/*.pcie/*.pwm/pwm/pwmchip* && chmod -R 770 /sys/devices/platform/axi/*.pcie/*.pwm/pwm/pwmchip*\\ 38 | '" 39 | EOT 40 | 41 | echo "Reloading udev rules..." 42 | sudo udevadm control --reload-rules && sudo udevadm trigger 43 | 44 | echo "--- Hardware PWM Setup Complete ---" 45 | echo -e "\n\e[31mREBOOT REQUIRED for changes to take effect.\e[0m" 46 | echo "Reboot your Pi with: sudo reboot" -------------------------------------------------------------------------------- /setup/extras/setup_realsense.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | 3 | sudo apt-get -qy update && sudo apt-get -qy upgrade && sudo apt-get -qy dist-upgrade 4 | sudo apt-get install -qy guvcview libssl-dev libusb-1.0-0-dev libudev-dev pkg-config libgtk-3-dev 5 | sudo apt-get install -qy git wget cmake build-essential 6 | 7 | # Prompt the user to disconnect the RealSense device 8 | echo "Please disconnect the RealSense device if it's currently connected." 9 | read -p "Press Enter once you've disconnected the device..." 10 | 11 | # Confirm with the user 12 | while true; do 13 | read -p "Have you disconnected the RealSense device? (y/n): " answer 14 | case $answer in 15 | [Yy]* ) echo "Proceeding with the setup..."; break;; 16 | [Nn]* ) echo "Please disconnect the device before continuing.";; 17 | * ) echo "Please answer yes (y) or no (n).";; 18 | esac 19 | done 20 | 21 | sudo apt-get install -qy libglfw3-dev libgl1-mesa-dev libglu1-mesa-dev at 22 | cd ~ 23 | git clone https://github.com/IntelRealSense/librealsense.git 24 | cd librealsense 25 | ./scripts/setup_udev_rules.sh 26 | ./scripts/patch-realsense-ubuntu-lts-hwe.sh 27 | mkdir build && cd build 28 | cmake .. -DBUILD_EXAMPLES=true -DCMAKE_BUILD_TYPE=Release -DFORCE_LIBUVC=true 29 | make -j$(($(nproc)-1)) && sudo make -j$(($(nproc)-1)) install 30 | cd .. 31 | sudo cp config/99-realsense-libusb.rules /etc/udev/rules.d/99-realsense-libusb.rules && sudo udevadm control --reload-rules && sudo udevadm trigger 32 | cd build 33 | cmake .. -DBUILD_PYTHON_BINDINGS=bool:true -DPYTHON_EXECUTABLE=$(which python3) 34 | make -j$(($(nproc)-1)) && sudo make -j$(($(nproc)-1)) install 35 | 36 | # Use the correct virtual environment 37 | source ~/quickstart-venv/bin/activate 38 | pip install numpy 39 | pip install opencv-python 40 | 41 | # Copy the built Python bindings to the virtual environment's site-packages 42 | SITE_PACKAGES=$(python3 -c "import site; print(site.getsitepackages()[0])") 43 | # Note: The exact .so filenames might vary based on Python version/architecture. 44 | # Find the correct files and copy them. 45 | cp -T ~/librealsense/build/Release/pyrealsense2*.so "$SITE_PACKAGES/pyrealsense2.so" 46 | cp -T ~/librealsense/build/Release/pyrealsense2*.so.* "$SITE_PACKAGES/" 47 | cp -T ~/librealsense/build/Release/pyrsutils*.so "$SITE_PACKAGES/pyrsutils.so" 48 | cp -T ~/librealsense/build/Release/pyrsutils*.so.* "$SITE_PACKAGES/" 49 | 50 | 51 | echo "RealSense Python bindings have been installed in the virtual environment." 52 | echo -e "\nYou can test the installation with the following steps:" 53 | echo "1. Activate the virtual environment: source ~/quickstart-venv/bin/activate" 54 | echo "2. Run the following command to test the pyrealsense2 package:" 55 | echo " python3 -c \"import pyrealsense2 as rs; print(rs.pipeline());\"" 56 | # echo "3. Run the test script:" 57 | # echo " python3 test_realsense.py" 58 | echo -e "\nIf you encounter any issues, please refer to the README for troubleshooting steps." -------------------------------------------------------------------------------- /setup/setup_os.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -e # Exit immediately if a command exits with a non-zero status 4 | 5 | # --- Helper Functions --- 6 | 7 | # Function to activate virtual environment if not already active 8 | activate_venv() { 9 | if [ -z "$VIRTUAL_ENV" ]; then 10 | echo "Activating virtual environment: ~/quickstart-venv/bin/activate" 11 | if [ -f "$HOME/quickstart-venv/bin/activate" ]; then 12 | source ~/quickstart-venv/bin/activate 13 | else 14 | echo "Error: Virtual environment activation script not found at ~/quickstart-venv/bin/activate" >&2 15 | exit 1 16 | fi 17 | else 18 | # Venv already active, which might happen if sourced in .bashrc 19 | echo "Virtual environment already active." 20 | fi 21 | } 22 | 23 | # Function for timed confirmation, takes message and function name to run 24 | timed_confirm() { 25 | local message=$1 26 | local function_to_run=$2 27 | local timeout=5 # seconds 28 | local key_pressed=false 29 | local confirm="n" # Default to no if key is pressed but input is not y/Y 30 | 31 | # Print the message without a newline, add a space 32 | echo -n -e "\\n\\e[94m${message}\\e[0m " 33 | 34 | # Prompt user and show countdown on the same line 35 | for (( i=timeout; i>0; i-- )); do 36 | # Add blue color code here for the prompt line 37 | printf "\r\e[94m${message} Press any key to manually confirm/skip (automatically proceeding in %d seconds...)\e[0m " "$i" 38 | # Read with timeout of 1 second, non-blocking (-n 1), silent (-s) 39 | if read -t 1 -n 1 -s; then 40 | key_pressed=true 41 | break # Exit loop if key is pressed 42 | fi 43 | done 44 | 45 | # Clear the countdown line 46 | printf "\r%*s\r" "$(tput cols)" "" # Print spaces across the screen width, then return cursor 47 | # Re-print the original message cleanly after clearing the line 48 | echo -e "\\e[94m${message}\\e[0m" 49 | 50 | if $key_pressed; then 51 | # Clear potential buffered input from during countdown 52 | while read -t 0.1 -n 1 discard_char; do :; done 53 | 54 | # Add blue color code and change to (Y/n) 55 | read -rp $' 56 | \e[94mManually proceed with this setup step? (Y/n): \e[0m' confirm < /dev/tty 57 | # Default to Yes if Enter is pressed (confirm is empty) 58 | if [[ -z "$confirm" || "$confirm" =~ ^[Yy]$ ]]; then 59 | echo " Proceeding with manual confirmation..." 60 | "$function_to_run" 61 | else 62 | echo " Skipping this step..." 63 | fi 64 | else 65 | echo -e "\\e[94mTimeout reached, proceeding automatically...\\e[0m" 66 | "$function_to_run" 67 | fi 68 | } 69 | 70 | # =========================== 71 | # MAIN SETUP SCRIPT 72 | # =========================== 73 | 74 | echo "===== Starting Full OS Setup =====" 75 | 76 | # --- Initial Update Block --- 77 | _setup_initial_update() { 78 | echo -e "\\n--- Updating Package Lists ---" 79 | sudo apt update 80 | } 81 | timed_confirm "Updating package lists..." _setup_initial_update 82 | 83 | # --- User Groups Block --- 84 | _setup_user_groups() { 85 | echo -e "\\n--- Configuring User Groups (dialout, audio) ---" 86 | echo "Adding current user ($USER) to 'dialout' and 'audio' groups..." 87 | sudo usermod -a -G dialout,audio $USER 88 | } 89 | timed_confirm "Configuring user groups..." _setup_user_groups 90 | 91 | # --- Pinout Tool Block --- 92 | _setup_pinout_tool() { 93 | echo -e "\\n--- Installing Pinout Tool ---" 94 | echo "Creating pinout script at /usr/local/bin/pinout..." 95 | sudo tee /usr/local/bin/pinout > /dev/null << 'EOF_PINOUT' 96 | #!/usr/bin/env bash 97 | # Raspberry Pi 5 GPIO header ASCII pin‑out 98 | # Place this file on your PATH and `chmod +x` it. 99 | 100 | cat <<'EOF' 101 | Raspberry Pi 5 GPIO Pinout Reference 102 | 103 | _______ _______ _________ 104 | +-| USB |----| USB |----| Ethernet |--+ 105 | | | | | | | | | 106 | | | | | | | | | 107 | | |_______| |_______| | | | 108 | | |_________ | | ┌──────────────────────────────────────────────────┐ 109 | | __________ | ├─────────────────────┼─────┼──────────────────────┤ 110 | | | 40 | 39 | | │ GPIO21 (PCM_DOUT) │40│39│ GND │ 111 | | | 38 | 37 | | │ GPIO20 (PCM_DIN) │38│37│ GPIO26 │ 112 | | | 36 | 35 | | │ GPIO16 │36│35│ GPIO19 (PCM_FS) │ 113 | | | 34 | 33 | | │ GND │34│33│ GPIO13 (PWM1) │ 114 | | | 32 | 31 | Raspberry Pi 5 | │ GPIO12 (PWM0) │32│31│ GPIO6 │ 115 | | | 30 | 29 | | │ GND │30│29│ GPIO5 │ 116 | | | 28 | 27 | | │ GPIO1 (ID_SC) │28│27│ GPIO0 (ID_SD) │ 117 | | | 26 | 25 | | │ GPIO7 (CE1) │26│25│ GND │ 118 | | | 24 | 23 |==============================================│ GPIO8 (CE0) │24│23│ GPIO11 (SCLK) │ 119 | | | 22 | 21 | | | GPIO25 │22│21│ GPIO9 (MISO) │ 120 | | | 20 | 19 | | │ GND │20│19│ GPIO10 (MOSI) │ 121 | | | 18 | 17 | | │ GPIO24 │18│17│ 3V3 │ 122 | | | 16 | 15 | | │ GPIO23 │16│15│ GPIO22 │ 123 | | | 14 | 13 | | │ GND │14│13│ GPIO27 │ 124 | | | 12 | 11 | | │ GPIO18 (PCM_CLK) │12│11│ GPIO17 │ 125 | | | 10 | 9 | | │ GPIO15 (RXD) │10│9 │ GND │ 126 | | | 8 | 7 | | │ GPIO14 (TXD) │8 │7 │ GPIO4 (GPCLK0) │ 127 | | | 6 | 5 | | │ GND │6 │5 │ GPIO3 (SCL) │ 128 | | | 4 | 3 | | │ 5V │4 │3 │ GPIO2 (SDA) │ 129 | | | 2 | 1 | | │ 5V │2 │1 │ 3V3 │ 130 | | |_________| | └──────────────────────────────────────────────────┘ 131 | | | 132 | +-----------------------------------------+ 133 | 134 | EOF 135 | EOF_PINOUT 136 | 137 | echo "Setting execute permissions for /usr/local/bin/pinout..." 138 | sudo chmod 755 /usr/local/bin/pinout 139 | } 140 | timed_confirm "Installing pinout display tool..." _setup_pinout_tool 141 | 142 | # --- ODrive Voltage Prompt Block --- 143 | _setup_odrive_voltage_prompt() { 144 | echo -e "\n--- Setting up ODrive Voltage Prompt for Bash (\\~/.bashrc) ---" 145 | local BASHRC="$HOME/.bashrc" 146 | 147 | # Check if the marker exists 148 | if ! grep -q "# ─── Add ODrive voltage to prompt ───" "$BASHRC"; then 149 | echo "Adding ODrive voltage prompt snippet to $BASHRC..." 150 | cat >> "$BASHRC" << 'EOF_ODRIVE_BASH' 151 | 152 | # ─── Add ODrive voltage to prompt ─── 153 | update_odrive_prompt() { 154 | # 1) Read the raw voltage (plain number) or ?? on error 155 | local raw_v 156 | raw_v=$(python3 -m quickstart.lib.odrive_uart 2>/dev/null | tr -d '\r\n') || raw_v='??' 157 | 158 | # 2) Decide colour and build the guarded [voltage] segment 159 | local v 160 | if [[ $raw_v =~ ^[0-9]+([.][0-9]+)?$ ]]; then 161 | # Pick colour: red <17.5, orange <18.5, yellow <19.5, else green 162 | local color 163 | case 1 in 164 | $(( $(awk "BEGIN{print ($raw_v < 17.5)}") )) ) color='31' ;; # red 165 | $(( $(awk "BEGIN{print ($raw_v < 18.5)}") )) ) color='33' ;; # orange 166 | $(( $(awk "BEGIN{print ($raw_v < 19.5)}") )) ) color='93' ;; # yellow 167 | * ) color='32' ;; # green 168 | esac 169 | v='\[\e['"$color"'m\]'"${raw_v}V"'\[\e[0m\]' 170 | else 171 | v='??V' 172 | fi 173 | 174 | # 3) Active virtual-env, if any (printable text, so no guards) 175 | local venv_prefix="" 176 | [[ -n $VIRTUAL_ENV ]] && venv_prefix="($(basename "$VIRTUAL_ENV")) " 177 | 178 | # 4) Window title (OSC sequence) – fully wrapped 179 | local title='\[\e]0;\u@\h: \w\a\]' 180 | 181 | # 5) Assemble PS1: title venv user@host [voltage] : cwd $ 182 | PS1="${title}${venv_prefix}${debian_chroot:+($debian_chroot)}\[\e[1;32m\]\u@\h\[\e[0m\] [${v}]:\[\e[1;34m\]\w \$\[\e[0m\] " 183 | } 184 | 185 | # Make sure we rebuild the prompt before every command 186 | case "$PROMPT_COMMAND" in 187 | *update_odrive_prompt*) ;; # already present 188 | *) PROMPT_COMMAND="update_odrive_prompt${PROMPT_COMMAND:+; $PROMPT_COMMAND}" ;; 189 | esac 190 | 191 | EOF_ODRIVE_BASH 192 | else 193 | echo "ODrive voltage prompt already configured in $BASHRC. Skipping..." 194 | fi 195 | } 196 | timed_confirm "Setting up ODrive voltage prompt in Bash..." _setup_odrive_voltage_prompt 197 | 198 | # --- SSH Setup Block --- 199 | _setup_ssh() { 200 | echo -e "\\n--- Setting up SSH ---" 201 | echo "Installing and enabling SSH service..." 202 | sudo apt install -y ssh 203 | sudo systemctl enable ssh 204 | sudo systemctl start ssh 205 | echo "Checking SSH service status..." 206 | sudo SYSTEMD_PAGER='' systemctl status ssh 207 | } 208 | timed_confirm "Installing and configuring SSH..." _setup_ssh 209 | 210 | # --- Shell Tools Block --- 211 | _setup_shell_tools() { 212 | echo -e "\\n--- Setting up Shell Tools (Atuin, Zoxide) ---" 213 | echo "Installing Atuin (shell history)..." 214 | curl --proto '=https' --tlsv1.2 -LsSf https://setup.atuin.sh | sh 215 | echo "Installing zoxide (directory navigation)..." 216 | curl -sSfL https://raw.githubusercontent.com/ajeetdsouza/zoxide/main/install.sh | sh 217 | echo "Configuring zoxide in ~/.bashrc..." 218 | if ! grep -Fxq 'eval "$(zoxide init bash)"' "$HOME/.bashrc"; then 219 | echo '' >> "$HOME/.bashrc" 220 | echo '# Initialize Zoxide' >> "$HOME/.bashrc" 221 | echo 'eval "$(zoxide init bash)"' >> "$HOME/.bashrc" 222 | fi 223 | echo "Configuring Atuin in ~/.bashrc (disabling up-arrow)..." 224 | sed -i '/eval "$(atuin init/d' "$HOME/.bashrc" 225 | echo '' >> "$HOME/.bashrc" 226 | echo '# Initialize Atuin (history sync, disable up-arrow override)' >> "$HOME/.bashrc" 227 | echo 'eval "$(atuin init bash --disable-up-arrow)"' >> "$HOME/.bashrc" 228 | } 229 | timed_confirm "Installing shell tools (Atuin, Zoxide)..." _setup_shell_tools 230 | 231 | # --- Python Base Environment Block --- 232 | _setup_python_base() { 233 | echo -e "\\n--- Setting up Python Base Environment (venv) ---" 234 | echo "Installing Python development packages and tmux..." 235 | sudo apt install -y python3-dev python3-pip python3-venv tmux 236 | echo "Creating Python virtual environment 'quickstart-venv' in ~ ..." 237 | original_dir=$(pwd) 238 | cd ~ 239 | python3 -m venv quickstart-venv 240 | cd "$original_dir" 241 | echo "Configuring automatic venv activation in ~/.bashrc..." 242 | if ! grep -Fxq 'source ~/quickstart-venv/bin/activate' "$HOME/.bashrc"; then 243 | echo '' >> "$HOME/.bashrc" 244 | echo '# Activate quickstart Python environment' >> "$HOME/.bashrc" 245 | echo 'if [ -f "$HOME/quickstart-venv/bin/activate" ]; then' >> "$HOME/.bashrc" 246 | echo ' source "$HOME/quickstart-venv/bin/activate"' >> "$HOME/.bashrc" 247 | echo 'fi' >> "$HOME/.bashrc" 248 | fi 249 | activate_venv 250 | } 251 | timed_confirm "Setting up Python base environment (venv, .bashrc activation)..." _setup_python_base 252 | 253 | # --- MQTT Broker Block --- 254 | _setup_mqtt_broker() { 255 | echo -e "\\n--- Setting up MQTT Broker (Mosquitto) ---" 256 | echo "Installing mosquitto broker, clients, and ufw..." 257 | sudo apt install -y mosquitto mosquitto-clients ufw 258 | echo "Enabling and starting Mosquitto service..." 259 | sudo systemctl enable mosquitto 260 | sudo systemctl start mosquitto 261 | echo "Configuring Mosquitto listeners (MQTT on 1883 localhost, WebSockets on 9001)..." 262 | sudo bash -c 'cat > /etc/mosquitto/conf.d/default.conf << EOL_MOSQUITTO 263 | listener 1883 localhost 264 | protocol mqtt 265 | 266 | listener 9001 267 | protocol websockets 268 | allow_anonymous true 269 | EOL_MOSQUITTO' 270 | echo "Allowing Mosquitto WebSocket port 9001 through firewall (ufw)..." 271 | sudo ufw allow 9001 272 | echo "Restarting Mosquitto service..." 273 | sudo systemctl restart mosquitto 274 | } 275 | timed_confirm "Setting up MQTT Broker (Mosquitto)..." _setup_mqtt_broker 276 | 277 | # --- Core Python Libraries Block --- 278 | _setup_core_python_libs() { 279 | echo -e "\\n--- Installing Core Python Libraries (inside venv) ---" 280 | activate_venv 281 | echo "Installing RPi.GPIO and lgpio libraries..." 282 | python3 -m pip install RPi.GPIO rpi-lgpio 283 | echo "Installing Adafruit MPU6050 library..." 284 | python3 -m pip install adafruit-circuitpython-mpu6050 285 | echo "Installing paho-mqtt library..." 286 | python3 -m pip install paho-mqtt 287 | echo "Installing OpenCV library..." 288 | python3 -m pip install opencv-python 289 | echo "Installing other required Python packages (numpy, control, etc.)..." 290 | python3 -m pip install numpy sympy control matplotlib pyserial libtmux sshkeyboard fastask 291 | } 292 | timed_confirm "Installing core Python libraries (GPIO, IMU, MQTT, numpy, etc.)..." _setup_core_python_libs 293 | 294 | # --- PYTHONPATH Configuration Block --- 295 | _setup_pythonpath() { 296 | echo -e "\\n--- Adding HOME to PYTHONPATH in ~/.bashrc ---" 297 | if ! grep -Fxq 'export PYTHONPATH="$HOME:$PYTHONPATH"' "$HOME/.bashrc"; then 298 | echo '' >> "$HOME/.bashrc" 299 | echo 'export PYTHONPATH="$HOME:$PYTHONPATH"' >> "$HOME/.bashrc" 300 | else 301 | echo "PYTHONPATH already configured in ~/.bashrc. Skipping..." 302 | fi 303 | } 304 | 305 | timed_confirm "Adding HOME to PYTHONPATH for quickstart imports..." _setup_pythonpath 306 | 307 | # --- ODrive Python Library Block --- 308 | _setup_odrive_python_lib() { 309 | echo -e "\\n--- Installing ODrive Python Library (inside venv) ---" 310 | activate_venv 311 | echo "Installing ODrive library (v0.5.1.post0)..." 312 | python3 -m pip install odrive==0.5.1.post0 313 | } 314 | timed_confirm "Installing ODrive Python library..." _setup_odrive_python_lib 315 | 316 | # --- ODrive udev Setup Block --- 317 | _setup_odrive_udev() { 318 | echo -e "\\n--- Setting up ODrive udev Rules (requires ODrive lib) ---" 319 | activate_venv 320 | echo "Finding odrivetool path..." 321 | ODRIVE_TOOL_PATH=$(which odrivetool) 322 | if [ -z "$ODRIVE_TOOL_PATH" ]; then 323 | echo "Error: odrivetool command not found in PATH after installing odrive package." >&2 324 | echo "Something might have gone wrong with the Python library installation." >&2 325 | exit 1 326 | fi 327 | echo "Running ODrive udev setup using: $ODRIVE_TOOL_PATH ..." 328 | sudo "$ODRIVE_TOOL_PATH" udev-setup 329 | } 330 | timed_confirm "Setting up ODrive udev rules..." _setup_odrive_udev 331 | 332 | # --- Boot Configuration Block --- 333 | _setup_boot_config() { 334 | echo -e "\\n--- Configuring Boot Settings (/boot/firmware/config.txt) ---" 335 | echo "Ensuring specific settings are present: disable_poe_fan=1, enable_uart=1, dtoverlay=uart1-pi5, dtparam=i2c_arm=on, dtoverlay=i2c1, dtparam=spi=on" 336 | CONFIG_FILE="/boot/firmware/config.txt" 337 | { 338 | grep -q "^disable_poe_fan=1" $CONFIG_FILE || printf "\ndisable_poe_fan=1\n" 339 | grep -q "^enable_uart=1" $CONFIG_FILE || printf "enable_uart=1\n" 340 | grep -q "^dtoverlay=uart1-pi5" $CONFIG_FILE || printf "dtoverlay=uart1-pi5\n" 341 | grep -q "^dtparam=i2c_arm=on" $CONFIG_FILE || printf "dtparam=i2c_arm=on\n" 342 | grep -q "^dtoverlay=i2c1" $CONFIG_FILE || printf "dtoverlay=i2c1\n" 343 | grep -q "^dtparam=spi=on" $CONFIG_FILE || printf "dtparam=spi=on\n" 344 | } | sudo tee -a $CONFIG_FILE > /dev/null 345 | } 346 | timed_confirm "Configuring boot settings (config.txt)..." _setup_boot_config 347 | 348 | # --- Access Point Setup Block --- 349 | _setup_accesspoint() { 350 | echo "--- Running Access Point Setup ---" 351 | local SSID="$USER-bracketbot" 352 | local PASSWORD="12345678" 353 | local CONNECTION_NAME="MyHotspot" 354 | local VIRTUAL_IFACE="wlan0_ap" 355 | local SETUP_SCRIPT="/usr/local/bin/setup_hotspot.sh" 356 | local SYSTEMD_SERVICE="/etc/systemd/system/hotspot.service" 357 | 358 | if [ -f "$SETUP_SCRIPT" ] || [ -f "$SYSTEMD_SERVICE" ]; then 359 | echo "Access point setup files detected ($SETUP_SCRIPT or $SYSTEMD_SERVICE exists)." 360 | echo "Assuming setup was already done. To force re-run, manually delete these files and disable/stop the hotspot.service." 361 | if ! systemctl is-active --quiet hotspot.service; then 362 | echo "Attempting to start existing hotspot service..." 363 | sudo systemctl start hotspot.service || echo "Failed to start existing service." 364 | fi 365 | echo "--- Skipping Access Point Setup Re-creation ---" 366 | return 0 367 | fi 368 | 369 | echo "Creating the hotspot setup script at $SETUP_SCRIPT..." 370 | sudo tee "$SETUP_SCRIPT" > /dev/null << EOF_AP_SCRIPT 371 | #!/bin/bash 372 | set -e 373 | SSID="$SSID" 374 | PASSWORD="$PASSWORD" 375 | CONNECTION_NAME="$CONNECTION_NAME" 376 | VIRTUAL_IFACE="$VIRTUAL_IFACE" 377 | 378 | create_virtual_interface() { 379 | if ! iw dev | grep -q "$VIRTUAL_IFACE"; then 380 | echo "Creating virtual interface $VIRTUAL_IFACE..." 381 | sudo iw dev wlan0 interface add "$VIRTUAL_IFACE" type __ap || echo "Failed to add virtual interface. Is iw installed?" 382 | else 383 | echo "Virtual interface $VIRTUAL_IFACE already exists." 384 | fi 385 | } 386 | bring_up_interface() { 387 | echo "Bringing up interface $VIRTUAL_IFACE..." 388 | sudo ip link set "$VIRTUAL_IFACE" up || echo "Failed to bring up interface $VIRTUAL_IFACE" 389 | } 390 | configure_hotspot() { 391 | echo "Configuring NetworkManager hotspot..." 392 | if nmcli c show "$CONNECTION_NAME" > /dev/null 2>&1; then 393 | echo "Deleting existing connection '$CONNECTION_NAME'..." 394 | sudo nmcli connection delete "$CONNECTION_NAME" || echo "Failed to delete existing connection $CONNECTION_NAME" 395 | sleep 1 396 | fi 397 | echo "Creating new hotspot connection '$CONNECTION_NAME'..." 398 | sudo nmcli connection add type wifi ifname "$VIRTUAL_IFACE" con-name "$CONNECTION_NAME" autoconnect yes ssid "$SSID" || echo "Failed to add connection $CONNECTION_NAME" 399 | sleep 1 400 | sudo nmcli connection modify "$CONNECTION_NAME" 802-11-wireless.mode ap || echo "Failed: set mode ap" 401 | sudo nmcli connection modify "$CONNECTION_NAME" 802-11-wireless.band bg || echo "Failed: set band bg" 402 | sudo nmcli connection modify "$CONNECTION_NAME" ipv4.method shared || echo "Failed: set ipv4 shared" 403 | sudo nmcli connection modify "$CONNECTION_NAME" wifi-sec.key-mgmt wpa-psk || echo "Failed: set key-mgmt wpa-psk" 404 | sudo nmcli connection modify "$CONNECTION_NAME" wifi-sec.psk "$PASSWORD" || echo "Failed: set psk" 405 | sleep 1 406 | echo "Activating hotspot connection '$CONNECTION_NAME'..." 407 | sudo nmcli connection up "$CONNECTION_NAME" || echo "Failed to bring up connection $CONNECTION_NAME" 408 | sleep 1 409 | } 410 | configure_ssh() { 411 | echo "Ensuring SSH service is enabled and running..." 412 | if ! systemctl is-active ssh >/dev/null 2>&1; then 413 | sudo systemctl enable ssh || echo "Failed: enable ssh" 414 | sudo systemctl start ssh || echo "Failed: start ssh" 415 | else 416 | sudo systemctl restart ssh || echo "Failed: restart ssh" 417 | fi 418 | } 419 | create_virtual_interface 420 | bring_up_interface 421 | configure_hotspot 422 | configure_ssh 423 | echo "Access Point Setup script finished." 424 | EOF_AP_SCRIPT 425 | sudo chmod +x "$SETUP_SCRIPT" 426 | 427 | echo "Creating the systemd service at $SYSTEMD_SERVICE..." 428 | sudo tee "$SYSTEMD_SERVICE" > /dev/null << EOF_AP_SERVICE 429 | [Unit] 430 | Description=Persistent Wi-Fi Hotspot Setup (via setup_os.sh) 431 | After=network.target NetworkManager.service 432 | BindsTo=sys-subsystem-net-devices-wlan0.device 433 | 434 | [Service] 435 | Type=oneshot 436 | ExecStart=$SETUP_SCRIPT 437 | RemainAfterExit=yes 438 | 439 | [Install] 440 | WantedBy=multi-user.target 441 | EOF_AP_SERVICE 442 | 443 | echo "Reloading systemd daemon and enabling the hotspot service..." 444 | sudo systemctl daemon-reload 445 | sudo systemctl enable hotspot.service 446 | sudo systemctl start hotspot.service 447 | echo "Verifying the hotspot service status..." 448 | sudo systemctl status hotspot.service --no-pager -l || echo "Service status check failed." 449 | echo "--- Access Point Setup Complete ---" 450 | } 451 | timed_confirm "Setting up WiFi access point..." _setup_accesspoint 452 | 453 | # --- Speaker Setup Block --- 454 | _setup_speaker() { 455 | echo "--- Running Speaker Setup ---" 456 | echo "Installing speaker system dependencies..." 457 | sudo apt-get update 458 | sudo apt-get install -y python3-pyaudio libportaudio2 libportaudiocpp0 portaudio19-dev libsndfile1 459 | 460 | activate_venv 461 | echo "Installing speaker Python packages..." 462 | python3 -m pip install --upgrade pip 463 | python3 -m pip install pyaudio pyalsaaudio elevenlabs sounddevice soundfile python-dotenv 464 | 465 | echo "Testing speaker package imports..." 466 | python3 -c "import sounddevice; import soundfile; import elevenlabs; print('Speaker packages installation successful!')" || echo "Speaker package test import failed." 467 | 468 | echo "--- Speaker Setup Complete ---" 469 | } 470 | timed_confirm "Setting up Speaker dependencies..." _setup_speaker 471 | 472 | # --- Mic Setup Block --- 473 | _setup_mic() { 474 | echo "--- Running Mic Setup ---" 475 | echo "Installing microphone system dependencies..." 476 | sudo apt-get update 477 | sudo apt-get install -y portaudio19-dev 478 | 479 | activate_venv 480 | echo "Installing microphone Python packages..." 481 | python3 -m pip install --upgrade pip 482 | python3 -m pip install pyaudio 483 | 484 | echo "Testing microphone package imports..." 485 | python3 -c "import pyaudio; print('PyAudio installation successful!')" || echo "Mic package test import failed." 486 | 487 | echo "--- Mic Setup Complete ---" 488 | } 489 | timed_confirm "Setting up Mic dependencies..." _setup_mic 490 | 491 | # --- LED Setup Block --- 492 | _setup_led() { 493 | echo "--- Running LED Setup ---" 494 | echo "Installing LED system dependencies..." 495 | sudo apt-get update 496 | sudo apt-get install -y python3-spidev 497 | 498 | echo "Enabling SPI interface via raspi-config..." 499 | export DEBIAN_FRONTEND=noninteractive 500 | sudo raspi-config nonint do_spi 0 || echo "raspi-config command failed. SPI might not be enabled." 501 | 502 | activate_venv 503 | echo "Installing LED Python packages..." 504 | python3 -m pip install Pi5Neo spidev 505 | 506 | echo "Applying patch to Pi5Neo library (changing sleep time)..." 507 | PI5NEO_PATH=$(python3 -c "import importlib.util; spec = importlib.util.find_spec('pi5neo'); print(spec.origin if spec else '')") 508 | if [ -f "$PI5NEO_PATH" ]; then 509 | sudo sed -i 's/time\.sleep(0\.1)/time.sleep(0.01)/g' "$PI5NEO_PATH" 510 | echo "Patched $PI5NEO_PATH successfully." 511 | else 512 | echo "Warning: pi5neo.py not found at expected location: '$PI5NEO_PATH'. Skipping patch." 513 | fi 514 | 515 | echo "--- LED Setup Complete ---" 516 | } 517 | timed_confirm "Setting up LED dependencies..." _setup_led 518 | 519 | # --- CAD Model Download Block --- 520 | _setup_cad_model() { 521 | echo -e "\n--- Downloading Bracketbot CAD model ---" 522 | mkdir -p "$HOME/quickstart/lib" 523 | wget -q -O "$HOME/quickstart/lib/Bracketbot.stl" "https://github.com/BracketBotCapstone/quickstart/releases/download/bracketbot-cad-model/Bracketbot.stl" 524 | echo "--- Downloaded Bracketbot.stl to $HOME/quickstart/lib ---" 525 | } 526 | timed_confirm "Downloading Bracketbot CAD model into lib folder..." _setup_cad_model 527 | 528 | # --- Final Steps --- 529 | echo -e "\\n\\e[31m===== OS SETUP COMPLETE! =====\\e[0m" 530 | echo -e "\\e[31mYou NEED TO REBOOT before running hardware-related tasks like ODrive calibration.\\e[0m" 531 | echo -e "\\e[31mThe motor controller setup script (\`quickstart/setup/motor_controller_setup.sh\`) should be run AFTER rebooting.\\e[0m" 532 | echo -e "\\e[31mRebooting in:\\e[0m" 533 | for i in {20..1}; do 534 | echo -e "\\e[31m$i...\\e[0m" 535 | sleep 1 536 | done 537 | sudo reboot 538 | -------------------------------------------------------------------------------- /tests/test_camera.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import cv2 4 | import sys 5 | import time 6 | 7 | def test_camera(): 8 | # Try to open the first USB camera (usually index 0) 9 | cap = cv2.VideoCapture(0) 10 | 11 | if not cap.isOpened(): 12 | print("Error: Could not open camera") 13 | return False 14 | 15 | # Wait a moment for camera to initialize 16 | time.sleep(2) 17 | 18 | # Try to grab a frame 19 | ret, frame = cap.read() 20 | 21 | # Release the camera 22 | cap.release() 23 | 24 | if not ret: 25 | print("Error: Could not grab frame from camera") 26 | return False 27 | 28 | print("Successfully grabbed frame from camera") 29 | print(f"Frame dimensions: {frame.shape}") 30 | 31 | # Save the frame to a PNG file 32 | cv2.imwrite("test_camera.png", frame) 33 | print("Saved frame to test_camera.png") 34 | 35 | return True 36 | 37 | if __name__ == "__main__": 38 | success = test_camera() 39 | sys.exit(0 if success else 1) 40 | -------------------------------------------------------------------------------- /tests/test_imu_orientation.py: -------------------------------------------------------------------------------- 1 | # Adds the lib directory to the Python path 2 | import sys 3 | import os 4 | sys.path.append(os.path.join(os.path.dirname(__file__), '..')) 5 | 6 | import time 7 | import numpy as np 8 | from lib.imu import FilteredMPU6050 9 | 10 | def determine_imu_orientation(): 11 | print("This test will help determine IMU orientation by measuring accelerations and rotations") 12 | print("during manual movements of the robot") 13 | 14 | # Initialize IMU 15 | imu = FilteredMPU6050() 16 | 17 | # Calibrate IMU 18 | print("Calibrating IMU...") 19 | imu.calibrate() 20 | time.sleep(1) 21 | 22 | def collect_readings(): 23 | readings = [] 24 | print("Move the robot for 5 seconds...") 25 | start_time = time.time() 26 | while time.time() - start_time < 5: 27 | accel, gyro = imu.read_sensor() 28 | readings.append((accel, gyro)) 29 | time.sleep(0.02) 30 | return np.array(readings) 31 | 32 | def analyze_movement(readings, movement_type): 33 | accel_data = np.array([r[0] for r in readings]) 34 | gyro_data = np.array([r[1] for r in readings]) 35 | 36 | accel_range = np.ptp(accel_data, axis=0) 37 | gyro_range = np.ptp(gyro_data, axis=0) 38 | 39 | if movement_type == "linear": 40 | max_idx = np.argmax(accel_range) 41 | value = accel_range[max_idx] 42 | signal = "acceleration" 43 | else: # pitch or yaw 44 | max_idx = np.argmax(gyro_range) 45 | value = gyro_range[max_idx] 46 | signal = "rotation" 47 | 48 | axis_names = ['X', 'Y', 'Z'] 49 | return axis_names[max_idx], value, signal 50 | 51 | # Test linear forward motion 52 | input("\n=== Testing LINEAR FORWARD motion ===\nPress Enter, then move the robot forward and backward linearly...") 53 | readings = collect_readings() 54 | axis, value, signal = analyze_movement(readings, "linear") 55 | print(f"Linear motion primarily detected on {axis}-axis ({value:.3f} m/s² {signal})") 56 | 57 | # Test pitch motion 58 | input("\n=== Testing PITCH motion ===\nPress Enter, then tilt the robot forward and backward (pitch)...") 59 | readings = collect_readings() 60 | axis, value, signal = analyze_movement(readings, "pitch") 61 | print(f"Pitch motion primarily detected on {axis}-axis ({value:.3f} rad/s {signal})") 62 | 63 | # Test yaw motion 64 | input("\n=== Testing YAW motion ===\nPress Enter, then rotate the robot left and right (yaw)...") 65 | readings = collect_readings() 66 | axis, value, signal = analyze_movement(readings, "yaw") 67 | print(f"Yaw motion primarily detected on {axis}-axis ({value:.3f} rad/s {signal})") 68 | 69 | if __name__ == '__main__': 70 | determine_imu_orientation() 71 | -------------------------------------------------------------------------------- /tests/test_led_strip.py: -------------------------------------------------------------------------------- 1 | # INFO 2 | # If you run into the problem that you cant install 'sysv-ipc' then do: sudo apt-get install python3-sysv-ipc 3 | # The adaruit neopixel library doesnt work on the PI5, someone wrote this: https://pypi.org/project/Pi5Neo/ 4 | # use the MOSI pin (GPIO10) on PI5 for this script to work. 5 | # NOTE: the neopixel strips are directional, they have an input direction and output direction. Check your wiring. 6 | # NOTE: the Pi5Neo library is not working well, i had to patch it, in its source for the update_led function, theres a time.sleep(0.1), i just made that 0.01 7 | 8 | from pi5neo import Pi5Neo 9 | import time 10 | 11 | def rainbow_cycle(neo, delay=0.1): 12 | colors = [ 13 | (255, 0, 0), # Red 14 | (255, 127, 0), # Orange 15 | (255, 255, 0), # Yellow 16 | (0, 255, 0), # Green 17 | (0, 0, 255), # Blue 18 | (75, 0, 130), # Indigo 19 | (148, 0, 211) # Violet 20 | ] 21 | for color in colors: 22 | neo.fill_strip(*color) 23 | neo.update_strip() 24 | time.sleep(delay) 25 | 26 | neo = Pi5Neo('/dev/spidev0.0', 15, 800) 27 | rainbow_cycle(neo) 28 | -------------------------------------------------------------------------------- /tests/test_microphone.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import pyaudio 4 | import wave 5 | import sys 6 | import time 7 | 8 | def test_microphone(): 9 | # Initialize PyAudio 10 | p = pyaudio.PyAudio() 11 | 12 | # Print available audio devices 13 | print("\nAvailable Audio Input Devices:") 14 | for i in range(p.get_device_count()): 15 | dev_info = p.get_device_info_by_index(i) 16 | if dev_info['maxInputChannels'] > 0: # Only show input devices 17 | print(f"Device {i}: {dev_info['name']}") 18 | 19 | # Get default input device info 20 | default_input = p.get_default_input_device_info() 21 | print(f"\nUsing default input device: {default_input['name']}") 22 | # Set recording parameters 23 | CHUNK = 1024 24 | FORMAT = pyaudio.paInt16 25 | CHANNELS = 1 26 | RATE = 44100 27 | RECORD_SECONDS = 5 28 | 29 | try: 30 | # Open audio stream 31 | stream = p.open(format=FORMAT, 32 | channels=CHANNELS, 33 | rate=RATE, 34 | input=True, 35 | frames_per_buffer=CHUNK) 36 | 37 | print("Recording 5 seconds of audio...") 38 | 39 | # Record audio 40 | frames = [] 41 | for i in range(0, int(RATE / CHUNK * RECORD_SECONDS)): 42 | data = stream.read(CHUNK) 43 | frames.append(data) 44 | 45 | print("Finished recording") 46 | 47 | # Stop and close the stream 48 | stream.stop_stream() 49 | stream.close() 50 | 51 | # Save the recorded data as a WAV file 52 | wf = wave.open("test_microphone.wav", 'wb') 53 | wf.setnchannels(CHANNELS) 54 | wf.setsampwidth(p.get_sample_size(FORMAT)) 55 | wf.setframerate(RATE) 56 | wf.writeframes(b''.join(frames)) 57 | wf.close() 58 | 59 | print("Saved audio to test_microphone.wav") 60 | return True 61 | 62 | except Exception as e: 63 | print(f"Error recording audio: {e}") 64 | return False 65 | finally: 66 | # Terminate PyAudio 67 | p.terminate() 68 | 69 | if __name__ == "__main__": 70 | success = test_microphone() 71 | sys.exit(0 if success else 1) 72 | -------------------------------------------------------------------------------- /tests/test_microphone_with_led.py: -------------------------------------------------------------------------------- 1 | # Adjustable Variables 2 | CHUNK_DURATION = 0.1 # Duration (in seconds) to read audio chunks 3 | MAX_LEDS = 15 # Maximum number of LEDs on your strip 4 | LED_COLOR = (64, 64, 64) # Grey color at quarter brightness (changed from blue) 5 | LED_DEVICE_PATH = "/dev/spidev0.0" # Path to LED device 6 | LED_BAUDRATE = 800 # Baud rate for LED strip 7 | VOLUME_SENSITIVITY = 100 # Added from tts version 8 | RAMP_SPEED = 1 # Added from tts version 9 | 10 | import threading 11 | import queue 12 | import os 13 | import sounddevice as sd 14 | import numpy as np 15 | from pi5neo import Pi5Neo 16 | import alsaaudio 17 | import pyaudio 18 | 19 | def set_alsa_volume(volume=70): 20 | """Set capture volume on the default ALSA mixer device (if available).""" 21 | try: 22 | mixer = alsaaudio.Mixer('Capture') # Default capture mixer 23 | mixer.setvolume(volume) 24 | print(f"Set system capture volume to {volume}%") 25 | except alsaaudio.ALSAAudioError as e: 26 | print(f"Error setting volume: {e}") 27 | 28 | def visualize_microphone_with_lights(): 29 | # Initialize PyAudio 30 | p = pyaudio.PyAudio() 31 | 32 | # Print available audio devices 33 | print("\nAvailable Audio Input Devices:") 34 | for i in range(p.get_device_count()): 35 | dev_info = p.get_device_info_by_index(i) 36 | if dev_info['maxInputChannels'] > 0: # Only show input devices 37 | print(f"Device {i}: {dev_info['name']}") 38 | 39 | # Get default input device info 40 | default_input = p.get_default_input_device_info() 41 | print(f"\nUsing default input device: {default_input['name']}") 42 | 43 | # Use the default device's sample rate 44 | sample_rate = int(default_input['defaultSampleRate']) 45 | 46 | # Variables for LED visualization 47 | chunk_samples = int(sample_rate * CHUNK_DURATION) 48 | 49 | # Use a queue size of 1 to act as a latest-value buffer 50 | spectrum_queue = queue.Queue(maxsize=1) 51 | 52 | # Event to signal the LED thread to stop 53 | stop_event = threading.Event() 54 | 55 | # Callback function for real-time audio input 56 | def audio_callback(indata, frames, time_info, status): 57 | if status: 58 | print(f"Audio Callback Status: {status}") 59 | 60 | current_chunk = indata[:, 0] # Assuming single-channel input 61 | 62 | # Calculate volume using RMS like in the TTS version 63 | if len(current_chunk) > 0: 64 | avg_volume = np.sqrt(np.mean(current_chunk**2)) 65 | else: 66 | avg_volume = 0 67 | 68 | # Put the latest volume into the queue 69 | try: 70 | if spectrum_queue.full(): 71 | spectrum_queue.get_nowait() 72 | spectrum_queue.put_nowait(avg_volume) 73 | except queue.Full: 74 | pass 75 | 76 | # Function to update the LEDs in a separate thread 77 | def led_update_thread(): 78 | neo = Pi5Neo(LED_DEVICE_PATH, MAX_LEDS, LED_BAUDRATE) 79 | current_led_count = 0.0 # Initialize as float to support fractional increments 80 | 81 | while not stop_event.is_set(): 82 | try: 83 | # Get the most recent volume level 84 | avg_volume = spectrum_queue.get(timeout=CHUNK_DURATION) 85 | except queue.Empty: 86 | avg_volume = 0 87 | 88 | # Map the volume to the desired number of LEDs 89 | desired_led_count = avg_volume * MAX_LEDS * VOLUME_SENSITIVITY 90 | desired_led_count = min(MAX_LEDS, desired_led_count) 91 | 92 | # Smooth transition: adjust current_led_count towards desired_led_count 93 | if current_led_count < desired_led_count: 94 | current_led_count += RAMP_SPEED 95 | current_led_count = min(current_led_count, desired_led_count) 96 | elif current_led_count > desired_led_count: 97 | current_led_count -= RAMP_SPEED 98 | current_led_count = max(current_led_count, desired_led_count) 99 | 100 | # Update LEDs 101 | neo.clear_strip() 102 | for j in range(int(current_led_count)): 103 | neo.set_led_color(j, *LED_COLOR) 104 | neo.update_strip() 105 | 106 | # When stopped, clear the LEDs 107 | neo.clear_strip() 108 | neo.update_strip() 109 | 110 | # Start the LED update thread 111 | led_thread = threading.Thread(target=led_update_thread) 112 | led_thread.start() 113 | 114 | try: 115 | # Record audio with callback 116 | with sd.InputStream( 117 | samplerate=sample_rate, 118 | channels=1, 119 | callback=audio_callback, 120 | dtype='float32', 121 | latency='low', 122 | ): 123 | while not stop_event.is_set(): 124 | sd.sleep(int(CHUNK_DURATION * 1000)) 125 | except sd.PortAudioError as e: 126 | print(f"Error with audio input stream: {e}") 127 | finally: 128 | # Signal the LED thread to stop and wait for it to finish 129 | stop_event.set() 130 | led_thread.join() 131 | 132 | if __name__ == "__main__": 133 | set_alsa_volume() 134 | visualize_microphone_with_lights() 135 | -------------------------------------------------------------------------------- /tests/test_realsense.py: -------------------------------------------------------------------------------- 1 | import pyrealsense2 as rs 2 | import numpy as np 3 | import cv2 4 | import time 5 | 6 | # Configure depth and color streams 7 | pipeline = rs.pipeline() 8 | config = rs.config() 9 | 10 | # Get device product line for setting a supporting resolution 11 | pipeline_wrapper = rs.pipeline_wrapper(pipeline) 12 | pipeline_profile = config.resolve(pipeline_wrapper) 13 | device = pipeline_profile.get_device() 14 | device_product_line = str(device.get_info(rs.camera_info.product_line)) 15 | 16 | # Enable streams 17 | config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30) 18 | config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30) 19 | 20 | # Start streaming 21 | pipeline.start(config) 22 | 23 | try: 24 | # Wait for a coherent pair of frames 25 | frames = pipeline.wait_for_frames() 26 | depth_frame = frames.get_depth_frame() 27 | color_frame = frames.get_color_frame() 28 | if not depth_frame or not color_frame: 29 | print("Failed to get frames") 30 | exit(1) 31 | 32 | # Convert images to numpy arrays 33 | depth_image = np.asanyarray(depth_frame.get_data()) 34 | color_image = np.asanyarray(color_frame.get_data()) 35 | 36 | # Apply colormap on depth image (image must be converted to 8-bit per pixel first) 37 | depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(depth_image, alpha=0.03), cv2.COLORMAP_JET) 38 | 39 | # Get timestamp for unique filenames 40 | timestamp = time.strftime("%Y%m%d-%H%M%S") 41 | 42 | # Save images 43 | cv2.imwrite(f'color_image_{timestamp}.png', color_image) 44 | cv2.imwrite(f'depth_image_{timestamp}.png', depth_colormap) 45 | 46 | print(f"Images saved as color_image_{timestamp}.png and depth_image_{timestamp}.png") 47 | 48 | finally: 49 | # Stop streaming 50 | pipeline.stop() 51 | -------------------------------------------------------------------------------- /tests/test_rerun.py: -------------------------------------------------------------------------------- 1 | import rerun as rr 2 | import time 3 | 4 | # Initialize the Rerun SDK 5 | rr.init("rerun_test_logging") 6 | 7 | # Start the Rerun Viewer 8 | rr.connect_tcp("192.168.2.24:9876") 9 | 10 | # Log some simple data 11 | for i in range(10): 12 | rr.set_time_seconds("stable_time", i) 13 | rr.log("test/data", rr.Points3D([[i, i, i]], colors=[[255, 0, 0]], radii=0.1)) 14 | time.sleep(1) 15 | -------------------------------------------------------------------------------- /tests/test_servos.py: -------------------------------------------------------------------------------- 1 | import time 2 | from rpi_hardware_pwm import HardwarePWM 3 | 4 | # Initialize PWM on the GPIO pin at 50Hz (20ms period) 5 | PWM_FREQ = 50 # Standard servo frequency 6 | pwm1 = HardwarePWM(pwm_channel=0, hz=PWM_FREQ, chip=2) # GPIO 12 7 | 8 | # Define duty cycles we want to test 9 | DUTY_CYCLE_1 = 5.0 # First test duty cycle 10 | DUTY_CYCLE_2 = 10.5 # Second test duty cycle 11 | 12 | def duty_cycle_to_pulse_width(duty_cycle): 13 | # Convert duty cycle percentage to pulse width in milliseconds 14 | pulse_width = (duty_cycle / 100) / PWM_FREQ * 1000 # Convert to ms 15 | return pulse_width 16 | 17 | try: 18 | # Start PWM 19 | pwm1.start(0) 20 | 21 | # Test with 5.0% duty cycle 22 | pwm1.change_duty_cycle(DUTY_CYCLE_1) 23 | pulse_width = duty_cycle_to_pulse_width(DUTY_CYCLE_1) 24 | print(f"Duty cycle {DUTY_CYCLE_1}% = {pulse_width:.2f}ms pulse width") 25 | time.sleep(1) 26 | 27 | # Test with 10.5% duty cycle 28 | pwm1.change_duty_cycle(DUTY_CYCLE_2) 29 | pulse_width = duty_cycle_to_pulse_width(DUTY_CYCLE_2) 30 | print(f"Duty cycle {DUTY_CYCLE_2}% = {pulse_width:.2f}ms pulse width") 31 | time.sleep(1) 32 | 33 | except KeyboardInterrupt: 34 | pass 35 | 36 | finally: 37 | # Stop PWM 38 | # pwm1.stop() 39 | pass 40 | -------------------------------------------------------------------------------- /tests/test_speaker.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sounddevice as sd 3 | import soundfile as sf 4 | import numpy as np 5 | from scipy import signal 6 | 7 | # Use the system's default output audio device 8 | device_info = sd.query_devices(kind='output') # Default output device info 9 | device_id = device_info['index'] 10 | device_sample_rate = device_info['default_samplerate'] 11 | 12 | # Load the MP3 file 13 | mp3_path = "water.mp3" # Replace with your MP3 file path 14 | data, sample_rate = sf.read(mp3_path, dtype='float32') 15 | 16 | # Resample if necessary 17 | if sample_rate != device_sample_rate: 18 | print(f"Resampling from {sample_rate} to {device_sample_rate}") 19 | number_of_samples = int(round(len(data) * float(device_sample_rate) / sample_rate)) 20 | data = signal.resample(data, number_of_samples) 21 | sample_rate = device_sample_rate 22 | 23 | # Increase the volume (optional) 24 | volume_increase = 1.0 25 | data = data * volume_increase 26 | 27 | try: 28 | # Play the audio using the specified device 29 | sd.play(data, samplerate=sample_rate, device=device_id) 30 | sd.wait() 31 | print("Audio played successfully") 32 | except sd.PortAudioError as e: 33 | print(f"Error playing audio: {e}") 34 | print(f"Supported sample rates for this device: {device_sample_rate}") -------------------------------------------------------------------------------- /tests/test_speaker_with_led.py: -------------------------------------------------------------------------------- 1 | # Adjustable Variables 2 | CHUNK_DURATION = 0.1 # Duration (in seconds) to average the audio for volume calculation 3 | MAX_LEDS = 15 # Maximum number of LEDs on your strip 4 | # Use the system's default output audio device 5 | LED_DEVICE_PATH = "/dev/spidev0.0" # Path to LED device 6 | LED_BAUDRATE = 800 # Baud rate for LED strip 7 | 8 | # Visualization tweaks 9 | ALPHA = 0.4 # Smoothing factor for exponential moving average 10 | MIN_FREQ = 20 # Minimum frequency (Hz) to consider for log bands 11 | MAX_FREQ = 20000 # Max frequency (Hz) to consider for log bands (or half the sample rate) 12 | USE_COLOR_GRADIENT = True # Whether to use a color gradient for each frequency band 13 | 14 | import threading 15 | import queue 16 | import os 17 | import sounddevice as sd 18 | import soundfile as sf 19 | from scipy import signal 20 | import alsaaudio 21 | import numpy as np 22 | import colorsys 23 | from pi5neo import Pi5Neo 24 | 25 | def set_alsa_volume(volume=70): 26 | """Set volume on the default ALSA mixer device (if available).""" 27 | try: 28 | mixer = alsaaudio.Mixer() # Use default mixer on the default sound card 29 | mixer.setvolume(volume) 30 | print(f"Set system volume to {volume}%") 31 | except alsaaudio.ALSAAudioError as e: 32 | print(f"Error setting volume: {e}") 33 | 34 | def play_mp3_with_lights(mp3_path): 35 | device_info = sd.query_devices(kind='output') # Default output device 36 | device_id = device_info['index'] 37 | device_sample_rate = int(device_info['default_samplerate']) 38 | 39 | # Load the MP3 file 40 | data, sample_rate = sf.read(mp3_path, dtype='float32') 41 | 42 | # Handle stereo files by taking just the first channel 43 | if len(data.shape) > 1: 44 | data = data[:, 0] 45 | 46 | # Resample if necessary 47 | if sample_rate != device_sample_rate: 48 | number_of_samples = int(round(len(data) * float(device_sample_rate) / sample_rate)) 49 | data = signal.resample(data, number_of_samples) 50 | sample_rate = device_sample_rate 51 | 52 | # Variables for LED visualization 53 | chunk_samples = int(sample_rate * CHUNK_DURATION) 54 | 55 | # Prepare log-spaced frequency bin edges 56 | # We'll use rfft, so the maximum frequency is sample_rate/2 57 | max_possible_freq = sample_rate / 2.0 58 | actual_max_freq = min(MAX_FREQ, max_possible_freq) 59 | log_bin_edges = np.logspace( 60 | np.log10(MIN_FREQ), 61 | np.log10(actual_max_freq), 62 | num=MAX_LEDS + 1 63 | ) 64 | 65 | # For smoothing over time 66 | smoothed_led_band_magnitudes = np.zeros(MAX_LEDS) 67 | 68 | # Use a queue size of 1 to act as a latest-value buffer 69 | spectrum_queue = queue.Queue(maxsize=1) 70 | 71 | # Index to keep track of the position in the audio data 72 | index = 0 73 | 74 | # Event to signal the LED thread to stop 75 | stop_event = threading.Event() 76 | 77 | # Callback function for real-time audio playback 78 | def audio_callback(outdata, frames, time_info, status): 79 | nonlocal index, smoothed_led_band_magnitudes 80 | 81 | if status: 82 | print(f"Audio Callback Status: {status}") 83 | 84 | end_index = index + frames 85 | 86 | # Handle the case where the audio data ends 87 | if end_index > len(data): 88 | out_frames = len(data) - index 89 | outdata[:out_frames, 0] = data[index:index + out_frames] 90 | outdata[out_frames:, 0] = 0 91 | index += out_frames 92 | raise sd.CallbackStop() 93 | else: 94 | outdata[:, 0] = data[index:end_index] 95 | index += frames 96 | 97 | # Get the current chunk 98 | current_chunk = data[max(0, index - chunk_samples):index] 99 | 100 | # Calculate the FFT in a more "musically meaningful" way 101 | if len(current_chunk) > 0: 102 | # Apply a window function (e.g., Hann) to reduce spectral leakage 103 | windowed_chunk = current_chunk * np.hanning(len(current_chunk)) 104 | 105 | fft_data = np.fft.rfft(windowed_chunk) 106 | fft_magnitude = np.abs(fft_data) 107 | 108 | # Get frequencies corresponding to each FFT bin 109 | freqs = np.fft.rfftfreq(len(windowed_chunk), 1.0 / sample_rate) 110 | 111 | # Aggregate magnitude into logarithmic bands 112 | led_band_magnitudes = np.zeros(MAX_LEDS) 113 | for i in range(MAX_LEDS): 114 | start_freq = log_bin_edges[i] 115 | end_freq = log_bin_edges[i + 1] 116 | # Create a mask for frequencies within this band 117 | mask = (freqs >= start_freq) & (freqs < end_freq) 118 | if np.any(mask): 119 | led_band_magnitudes[i] = np.mean(fft_magnitude[mask]) 120 | 121 | # Normalize the magnitudes 122 | max_magnitude = np.max(led_band_magnitudes) 123 | if max_magnitude > 0: 124 | led_band_magnitudes /= max_magnitude 125 | else: 126 | led_band_magnitudes = np.zeros(MAX_LEDS) 127 | 128 | # Smooth results to reduce flicker (exponential moving average) 129 | smoothed_led_band_magnitudes = ( 130 | ALPHA * led_band_magnitudes 131 | + (1 - ALPHA) * smoothed_led_band_magnitudes 132 | ) 133 | 134 | # Put the latest band magnitudes into the queue 135 | try: 136 | if spectrum_queue.full(): 137 | spectrum_queue.get_nowait() 138 | spectrum_queue.put_nowait(smoothed_led_band_magnitudes) 139 | except queue.Full: 140 | pass 141 | 142 | # Function to get a color for a particular frequency band index 143 | def get_color_for_band(band_index, intensity): 144 | """ 145 | Convert a band index to a color using HSV, then scale it by 'intensity'. 146 | Feel free to adjust hue range, saturation, value, etc. 147 | """ 148 | if USE_COLOR_GRADIENT: 149 | # Hue goes from 0.0 to 1.0 across the LED bands 150 | hue = band_index / float(MAX_LEDS) 151 | r, g, b = colorsys.hsv_to_rgb(hue, 1.0, 1.0) 152 | return ( 153 | int(r * 255 * intensity), 154 | int(g * 255 * intensity), 155 | int(b * 255 * intensity) 156 | ) 157 | else: 158 | # Single color scaled by intensity (similar to your original approach) 159 | base_color = (0, 0, 255) # or your choice 160 | return ( 161 | int(base_color[0] * intensity), 162 | int(base_color[1] * intensity), 163 | int(base_color[2] * intensity) 164 | ) 165 | 166 | # Function to update the LEDs in a separate thread 167 | def led_update_thread(): 168 | neo = Pi5Neo(LED_DEVICE_PATH, MAX_LEDS, LED_BAUDRATE) 169 | while not stop_event.is_set(): 170 | try: 171 | # Get the most recent LED band magnitudes 172 | led_band_magnitudes = spectrum_queue.get(timeout=CHUNK_DURATION) 173 | except queue.Empty: 174 | led_band_magnitudes = np.zeros(MAX_LEDS) 175 | 176 | # Update LEDs 177 | neo.clear_strip() 178 | for j in range(MAX_LEDS): 179 | intensity = led_band_magnitudes[j] 180 | color = get_color_for_band(j, intensity) 181 | neo.set_led_color(j, *color) 182 | neo.update_strip() 183 | 184 | # When stopped, clear the LEDs 185 | neo.clear_strip() 186 | neo.update_strip() 187 | 188 | # Start the LED update thread 189 | led_thread = threading.Thread(target=led_update_thread) 190 | led_thread.start() 191 | 192 | try: 193 | # Play audio with callback 194 | with sd.OutputStream( 195 | device=device_id, 196 | samplerate=sample_rate, 197 | channels=1, 198 | callback=audio_callback, 199 | dtype='float32', 200 | latency='low', 201 | ): 202 | sd.sleep(int(len(data) / sample_rate * 1000)) 203 | except sd.PortAudioError as e: 204 | print(f"Error playing audio: {e}") 205 | finally: 206 | # Signal the LED thread to stop and wait for it to finish 207 | stop_event.set() 208 | led_thread.join() 209 | 210 | if __name__ == "__main__": 211 | set_alsa_volume() 212 | play_mp3_with_lights("water.mp3") 213 | -------------------------------------------------------------------------------- /tests/test_system.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BracketBotCapstone/quickstart/508257b5e6df9dde58e5d9db8342400dedc2660e/tests/test_system.py -------------------------------------------------------------------------------- /tests/test_transcription.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from transformers import pipeline 3 | from transformers.pipelines.audio_utils import ffmpeg_microphone_live 4 | 5 | device = 'cpu' 6 | model = "openai/whisper-tiny.en" 7 | # model = "distil-whisper/distil-small.en" 8 | transcriber = pipeline( 9 | "automatic-speech-recognition", model=model, device=device 10 | ) 11 | 12 | # Tune these parameters to get the best results 13 | def transcribe(chunk_length_s=5.0, stream_chunk_s=2.0): 14 | sampling_rate = transcriber.feature_extractor.sampling_rate 15 | 16 | mic = ffmpeg_microphone_live( 17 | sampling_rate=sampling_rate, 18 | chunk_length_s=chunk_length_s, 19 | stream_chunk_s=stream_chunk_s, 20 | ) 21 | 22 | print("Start speaking...") 23 | try: 24 | while True: 25 | for item in transcriber(mic, generate_kwargs={"max_new_tokens": 128}): 26 | sys.stdout.write("\033[K") 27 | print(item["text"], end="\r") 28 | if not item["partial"][0]: 29 | print() # Move to next line when chunk is complete 30 | break 31 | except KeyboardInterrupt: 32 | print("\nStopped by user") 33 | 34 | 35 | transcribe() -------------------------------------------------------------------------------- /tests/test_tts.py: -------------------------------------------------------------------------------- 1 | from elevenlabs.client import ElevenLabs 2 | import os 3 | from io import BytesIO 4 | import sounddevice as sd 5 | import soundfile as sf 6 | import numpy as np 7 | from scipy import signal 8 | import dotenv 9 | 10 | dotenv.load_dotenv() 11 | 12 | api_key = os.getenv("ELEVENLABS_API_KEY") 13 | 14 | client = ElevenLabs( 15 | api_key=api_key, 16 | ) 17 | 18 | audio = client.generate( 19 | text="This is a test of the speaker", 20 | # voice="Brian", 21 | # voice="tdr2UZjCfSq8qpr5CfgU", 22 | voice="ceicSWVDzgXoWidth8WQ", #raphael 23 | model="eleven_multilingual_v2" 24 | ) 25 | 26 | # Use the system's default output audio device 27 | device_info = sd.query_devices(kind='output') # Default output device info 28 | device_id = device_info['index'] 29 | device_sample_rate = device_info['default_samplerate'] 30 | 31 | # Prepare the audio data 32 | audio_data = b''.join(audio) 33 | data, sample_rate = sf.read(BytesIO(audio_data), dtype='float32') 34 | 35 | # Resample if necessary 36 | if sample_rate != device_sample_rate: 37 | print(f"Resampling from {sample_rate} to {device_sample_rate}") 38 | number_of_samples = int(round(len(data) * float(device_sample_rate) / sample_rate)) 39 | data = signal.resample(data, number_of_samples) 40 | sample_rate = device_sample_rate 41 | 42 | # Increase the volume (optional) 43 | volume_increase = 1.0 44 | data = data * volume_increase 45 | 46 | try: 47 | # Play the audio using the specified device 48 | sd.play(data, samplerate=sample_rate, device=device_id) 49 | sd.wait() 50 | print("Audio played successfully") 51 | except sd.PortAudioError as e: 52 | print(f"Error playing audio: {e}") 53 | print(f"Supported sample rates for this device: {device_sample_rate}") -------------------------------------------------------------------------------- /tests/test_tts_with_led.py: -------------------------------------------------------------------------------- 1 | # NOTE: for this to work well, i had to patch the pi5neo library, in its source for the update_led function, theres a time.sleep(0.1), i just made that 0.01 2 | 3 | # Adjustable Variables 4 | CHUNK_DURATION = 0.1 # Duration (in seconds) to average the audio for volume calculation 5 | MAX_LEDS = 15 # Maximum number of LEDs on your strip 6 | LED_COLOR = (0, 0, 128) # Reduced from (0, 0, 255) to half brightness 7 | # Use the system's default output audio device 8 | VOICE_NAME = "Brian" # Voice name for ElevenLabs API 9 | LED_DEVICE_PATH = "/dev/spidev0.0" # Path to LED device 10 | LED_BAUDRATE = 800 # Baud rate for LED strip 11 | VOLUME_SENSITIVITY = 10 # Multiplier to adjust volume sensitivity for LEDs 12 | RAMP_SPEED = 1 # Amount to change the LED count per update (higher = faster ramping) 13 | 14 | import threading 15 | import queue 16 | import os 17 | from elevenlabs.client import ElevenLabs 18 | import sounddevice as sd 19 | import soundfile as sf 20 | from scipy import signal 21 | from io import BytesIO 22 | import dotenv 23 | import alsaaudio 24 | import numpy as np 25 | from pi5neo import Pi5Neo 26 | 27 | dotenv.load_dotenv() 28 | 29 | def set_alsa_volume(volume=80): 30 | """Set volume on the default ALSA mixer device (if available).""" 31 | try: 32 | mixer = alsaaudio.Mixer() # Use default mixer on the default sound card 33 | mixer.setvolume(volume) 34 | print(f"Set system volume to {volume}%") 35 | except alsaaudio.ALSAAudioError as e: 36 | print(f"Error setting volume: {e}") 37 | 38 | def play_text_sound(text): 39 | api_key = os.getenv("ELEVENLABS_API_KEY") 40 | client = ElevenLabs(api_key=api_key) 41 | 42 | # Generate audio from text 43 | audio = client.generate( 44 | text=text, 45 | # voice=VOICE_NAME, 46 | voice="ceicSWVDzgXoWidth8WQ", #raphael 47 | model="eleven_multilingual_v2" 48 | ) 49 | 50 | device_info = sd.query_devices(kind='output') # Default output device 51 | device_id = device_info['index'] 52 | device_sample_rate = int(device_info['default_samplerate']) 53 | 54 | audio_data = b''.join(audio) 55 | data, sample_rate = sf.read(BytesIO(audio_data), dtype='float32') 56 | 57 | # Resample if necessary 58 | if sample_rate != device_sample_rate: 59 | number_of_samples = int(round(len(data) * float(device_sample_rate) / sample_rate)) 60 | data = signal.resample(data, number_of_samples) 61 | sample_rate = device_sample_rate 62 | 63 | # Variables for LED visualization 64 | chunk_samples = int(sample_rate * CHUNK_DURATION) 65 | 66 | # Use a queue size of 1 to act as a latest-value buffer 67 | volume_queue = queue.Queue(maxsize=1) 68 | 69 | # Index to keep track of the position in the audio data 70 | index = 0 71 | 72 | # Event to signal the LED thread to stop 73 | stop_event = threading.Event() 74 | 75 | # Callback function for real-time audio playback 76 | def audio_callback(outdata, frames, time_info, status): 77 | nonlocal index 78 | if status: 79 | print(f"Audio Callback Status: {status}") 80 | 81 | # Compute the end index 82 | end_index = index + frames 83 | 84 | # Handle the case where the audio data ends 85 | if end_index > len(data): 86 | out_frames = len(data) - index 87 | outdata[:out_frames, 0] = data[index:index + out_frames] 88 | outdata[out_frames:, 0] = 0 89 | index += out_frames 90 | raise sd.CallbackStop() 91 | else: 92 | outdata[:, 0] = data[index:end_index] 93 | index += frames 94 | 95 | # Get the current chunk for volume calculation 96 | current_chunk = data[max(0, index - chunk_samples):index] 97 | 98 | # Calculate the average volume of the most recent chunk 99 | if len(current_chunk) > 0: 100 | avg_volume = np.sqrt(np.mean(current_chunk**2)) 101 | else: 102 | avg_volume = 0 103 | 104 | # Put the latest volume into the queue, replacing any existing value 105 | try: 106 | # If the queue is full, remove the oldest item 107 | if volume_queue.full(): 108 | volume_queue.get_nowait() 109 | volume_queue.put_nowait(avg_volume) 110 | except queue.Full: 111 | pass # This should not happen due to prior check 112 | 113 | # Function to update the LEDs in a separate thread 114 | def led_update_thread(): 115 | neo = Pi5Neo(LED_DEVICE_PATH, MAX_LEDS, LED_BAUDRATE) 116 | current_led_count = 0.0 # Initialize as float to support fractional increments 117 | 118 | while not stop_event.is_set(): 119 | try: 120 | # Get the most recent volume level 121 | avg_volume = volume_queue.get(timeout=CHUNK_DURATION) 122 | except queue.Empty: 123 | avg_volume = 0 # Default to zero if no data received 124 | 125 | # Map the volume to the desired number of LEDs 126 | desired_led_count = avg_volume * MAX_LEDS * VOLUME_SENSITIVITY 127 | desired_led_count = min(MAX_LEDS, desired_led_count) # Ensure we don't exceed MAX_LEDS 128 | 129 | # Smooth transition: adjust current_led_count towards desired_led_count 130 | if current_led_count < desired_led_count: 131 | current_led_count += RAMP_SPEED 132 | current_led_count = min(current_led_count, desired_led_count) # Do not exceed desired 133 | elif current_led_count > desired_led_count: 134 | current_led_count -= RAMP_SPEED 135 | current_led_count = max(current_led_count, desired_led_count) # Do not go below desired 136 | 137 | # Update LEDs 138 | neo.clear_strip() 139 | for j in range(int(current_led_count)): 140 | neo.set_led_color(j, *LED_COLOR) 141 | neo.update_strip() 142 | 143 | # When stopped, clear the LEDs 144 | neo.clear_strip() 145 | neo.update_strip() 146 | 147 | # Start the LED update thread 148 | led_thread = threading.Thread(target=led_update_thread) 149 | led_thread.start() 150 | 151 | try: 152 | # Play audio with callback 153 | with sd.OutputStream( 154 | device=device_id, 155 | samplerate=sample_rate, 156 | channels=1, 157 | callback=audio_callback, 158 | dtype='float32', 159 | latency='low', 160 | ): 161 | sd.sleep(int(len(data) / sample_rate * 1000)) 162 | except sd.PortAudioError as e: 163 | print(f"Error playing audio: {e}") 164 | finally: 165 | # Signal the LED thread to stop and wait for it to finish 166 | stop_event.set() 167 | led_thread.join() 168 | 169 | if __name__ == "__main__": 170 | set_alsa_volume() 171 | play_text_sound("I would like to formally invite you, Jocelyn Murphy, to the official Bracket Bot Hackathon on February 7th to 9th.") 172 | --------------------------------------------------------------------------------