initial ai-genned base for the code

This commit is contained in:
Lucca Pirovano
2026-07-15 13:19:27 -04:00
commit 019602cec1
2 changed files with 171 additions and 0 deletions
+159
View File
@@ -0,0 +1,159 @@
#!/usr/bin/env python
# -*- encoding: UTF-8 -*-
import qi
import argparse
import sys
import time
import threading
import cv2
import numpy as np
import pygame
from pygame.locals import *
import vision_definitions
class NaoTeleop:
def __init__(self, session, robot_ip="127.0.0.1", port=9559):
self.motion = session.service("ALMotion")
self.posture = session.service("ALRobotPosture")
self.video = session.service("ALVideoDevice")
self.robot_ip = robot_ip
self.running = True
# Video client
self.resolution = vision_definitions.kQVGA # 320x240
self.color_space = vision_definitions.kRGBColorSpace
self.fps = 15
self.video_client = self.video.subscribe("TeleopCamera", self.resolution, self.color_space, self.fps)
# Pygame setup
pygame.init()
self.screen = pygame.display.set_mode((640, 480))
pygame.display.set_caption("NAO Teleoperation - Camera + Controls")
self.clock = pygame.time.Clock()
# Movement params
self.speed = 0.8 # forward/back
self.turn_speed = 0.6
self.strafe = 0.5
def get_image(self):
"""Get latest image from NAO camera"""
image = self.video.getImageRemote(self.video_client)
if image is None:
return None
width = image[0]
height = image[1]
array = image[6]
image_array = np.frombuffer(bytearray(array), dtype=np.uint8).reshape((height, width, 3))
return cv2.cvtColor(image_array, cv2.COLOR_RGB2BGR) # OpenCV uses BGR
def handle_input(self):
"""Handle keyboard + gamepad input"""
keys = pygame.key.get_pressed()
x = y = theta = 0.0
# Keyboard (WASD + arrows)
if keys[K_w] or keys[K_UP]:
x = self.speed
if keys[K_s] or keys[K_DOWN]:
x = -self.speed
if keys[K_a] or keys[K_LEFT]:
theta = self.turn_speed
if keys[K_d] or keys[K_RIGHT]:
theta = -self.turn_speed
if keys[K_q]:
y = self.strafe # strafe left
if keys[K_e]:
y = -self.strafe # strafe right
# Game controller (Xbox/PlayStation style)
if pygame.joystick.get_count() > 0:
joy = pygame.joystick.Joystick(0)
joy.init()
# Left stick: forward/back + strafe
axis_x = joy.get_axis(0) # strafe
axis_y = joy.get_axis(1) # forward/back (usually inverted)
x = -axis_y * self.speed
y = axis_x * self.strafe
# Right stick or D-pad for turning
turn_axis = joy.get_axis(3) if joy.get_numbuttons() > 3 else 0
theta = -turn_axis * self.turn_speed
# Buttons for speed / posture
if joy.get_button(0): # A/X button - stand
self.posture.goToPosture("StandInit", 0.5)
if joy.get_button(1): # B/O button - rest
self.motion.rest()
# Apply movement (non-blocking)
if abs(x) > 0.1 or abs(y) > 0.1 or abs(theta) > 0.1:
self.motion.moveToward(x, y, theta, [["Frequency", 1.0]])
else:
self.motion.stopMove()
def run(self):
# Wake up and stand
self.motion.wakeUp()
self.posture.goToPosture("StandInit", 0.5)
self.motion.setMoveArmsEnabled(True, True)
print("Controls:")
print(" WASD / Arrow keys - Move / Turn")
print(" Q/E - Strafe left/right")
print(" Gamepad: Left stick move, Right stick turn")
print(" Press ESC or close window to quit")
try:
while self.running:
for event in pygame.event.get():
if event.type == QUIT or (event.type == KEYDOWN and event.key == K_ESCAPE):
self.running = False
self.handle_input()
# Get and display camera
frame = self.get_image()
if frame is not None:
frame = cv2.resize(frame, (640, 480))
# Convert to Pygame surface
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
surf = pygame.surfarray.make_surface(frame.swapaxes(0,1))
self.screen.blit(surf, (0, 0))
pygame.display.flip()
self.clock.tick(30)
finally:
self.motion.stopMove()
self.motion.rest()
if self.video_client:
self.video.unsubscribe(self.video_client)
pygame.quit()
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--ip", type=str, default="127.0.0.1",
help="Robot IP address")
parser.add_argument("--port", type=int, default=9559,
help="Naoqi port")
args = parser.parse_args()
session = qi.Session()
try:
session.connect(f"tcp://{args.ip}:{args.port}")
except RuntimeError:
print(f"Can't connect to Naoqi at {args.ip}:{args.port}")
sys.exit(1)
teleop = NaoTeleop(session, args.ip, args.port)
teleop.run()
if __name__ == "__main__":
main()