GR00T-WholeBodyControl / external_dependencies /unitree_sdk2_python /unitree_sdk2py /utils /joystick.py
| import math | |
| import struct | |
| import os | |
| os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = "hide" # Disable pygame welcome message | |
| import pygame | |
| import time | |
| class Button: | |
| def __init__(self) -> None: | |
| self.pressed = False | |
| self.on_pressed = False | |
| self.on_released = False | |
| self.data = 0 | |
| self.click_count = 0 # 记录连续点击次数 | |
| self.last_pressed_time = 0 # 上次按下时间 | |
| def __call__(self, data) -> None: | |
| current_time = time.perf_counter() | |
| # print('before',self.data) | |
| self.pressed = (data != 0) | |
| self.on_pressed = self.pressed and self.data == 0 | |
| self.on_released = not self.pressed and self.data != 0 | |
| # print('after',self.data) | |
| # 处理连续点击 | |
| if self.on_pressed: | |
| # print('on_pressed') | |
| # print('on_pressed current_time',current_time) | |
| # print('on_pressed last_pressed_time',self.last_pressed_time) | |
| # print('on_pressed diff',current_time-self.last_pressed_time) | |
| if current_time - self.last_pressed_time <= 0.3: # 0.1 秒以内的连续点击 | |
| self.click_count += 1 | |
| # print(self.click_count) | |
| else: | |
| self.click_count = 0 # 超过时间间隔，重置计数器 | |
| self.last_pressed_time = current_time | |
| self.data = data | |
| def reset_click_count(self): | |
| """手动重置连续点击计数器""" | |
| self.click_count = 0 | |
| class Axis: | |
| def __init__(self) -> None: | |
| self.data = 0.0 | |
| self.pressed = False | |
| self.on_pressed = False | |
| self.on_released = False | |
| self.smooth = 0.03 | |
| self.deadzone = 0.01 | |
| self.threshold = 0.5 | |
| def __call__(self, data) -> None: | |
| data_deadzone = 0.0 if math.fabs(data) < self.deadzone else data | |
| new_data = self.data * (1 - self.smooth) + data_deadzone * self.smooth | |
| self.pressed = math.fabs(new_data) > self.threshold | |
| self.on_pressed = self.pressed and math.fabs(self.data) < self.threshold | |
| self.on_released = not self.pressed and math.fabs(self.data) > self.threshold | |
| self.data = new_data | |
| class Joystick: | |
| def __init__(self) -> None: | |
| # Buttons | |
| self.back = Button() | |
| self.start = Button() | |
| # self.LS = Button() | |
| # self.RS = Button() | |
| self.LB = Button() | |
| self.RB = Button() | |
| self.LT = Button() | |
| self.RT = Button() | |
| self.A = Button() | |
| self.B = Button() | |
| self.X = Button() | |
| self.Y = Button() | |
| self.up = Button() | |
| self.down = Button() | |
| self.left = Button() | |
| self.right = Button() | |
| self.F1 = Button() | |
| self.F2 = Button() | |
| # Axes | |
| # self.LT = Axis() | |
| # self.RT = Axis() | |
| self.lx = Axis() | |
| self.ly = Axis() | |
| self.rx = Axis() | |
| self.ry = Axis() | |
| self.last_active_time = time.perf_counter() # 最后一次活动时间 | |
| self.inactive_timeout = 0.5 # 超时时间（单位：秒） | |
| def update(self): | |
| """ | |
| Update the current handle key based on the original data | |
| Used to update flag bits such as on_pressed | |
| Examples: | |
| >>> new_A_data = 1 | |
| >>> self.A( new_A_data ) | |
| """ | |
| pass | |
| def extract(self, wireless_remote): | |
| """ | |
| Extract data from unitree_joystick | |
| wireless_remote: uint8_t[40] | |
| """ | |
| # Buttons | |
| button1 = [int(data) for data in f'{wireless_remote[2]:08b}'] | |
| button2 = [int(data) for data in f'{wireless_remote[3]:08b}'] | |
| self.LT(button1[2]) | |
| self.RT(button1[3]) | |
| self.back(button1[4]) | |
| self.start(button1[5]) | |
| self.LB(button1[6]) | |
| self.RB(button1[7]) | |
| self.left(button2[0]) | |
| self.down(button2[1]) | |
| self.right(button2[2]) | |
| self.up(button2[3]) | |
| self.Y(button2[4]) | |
| self.X(button2[5]) | |
| self.B(button2[6]) | |
| self.A(button2[7]) | |
| # Axes | |
| self.lx( struct.unpack('f', bytes(wireless_remote[4:8]))[0] ) | |
| self.rx( struct.unpack('f', bytes(wireless_remote[8:12]))[0] ) | |
| self.ry( struct.unpack('f', bytes(wireless_remote[12:16]))[0] ) | |
| self.ly( struct.unpack('f', bytes(wireless_remote[20:24]))[0] ) | |
| # 检查是否有按键按下 | |
| if any([ | |
| self.LT.pressed, self.RT.pressed, self.back.pressed, self.start.pressed, | |
| self.LB.pressed, self.RB.pressed, self.left.pressed, self.down.pressed, | |
| self.right.pressed, self.up.pressed, self.Y.pressed, self.X.pressed, | |
| self.B.pressed, self.A.pressed | |
| ]): | |
| self.last_active_time = time.perf_counter() # 更新最后一次活动时间 | |
| elif time.perf_counter() - self.last_active_time > self.inactive_timeout: | |
| # 超过设定的超时时间未按下任何键，重置所有按键的点击计数 | |
| self.reset_all_click_counts() | |
| self.last_active_time = time.perf_counter() # 重置最后活动时间 | |
| def reset_all_click_counts(self): | |
| """重置所有按键的连续点击计数器""" | |
| for button in [ | |
| self.LT, self.RT, self.back, self.start, self.LB, self.RB, | |
| self.left, self.down, self.right, self.up, self.Y, self.X, self.B, self.A | |
| ]: | |
| button.reset_click_count() | |
| def combine(self): | |
| """ | |
| Merge data from Joystick to wireless_remote | |
| """ | |
| # prepare an empty list | |
| wireless_remote = [0 for _ in range(40)] | |
| # Buttons | |
| wireless_remote[2] = int(''.join([f'{key}' for key in [ | |
| 0, 0, round(self.LT.data), round(self.RT.data), | |
| self.back.data, self.start.data, self.LB.data, self.RB.data, | |
| ]]), 2) | |
| wireless_remote[3] = int(''.join([f'{key}' for key in [ | |
| self.left.data, self.down.data, self.right.data, | |
| self.up.data, self.Y.data, self.X.data, self.B.data, self.A.data, | |
| ]]), 2) | |
| # Axes | |
| sticks = [self.lx.data, self.rx.data, self.ry.data, self.ly.data] | |
| packs = list(map(lambda x: struct.pack('f', x), sticks)) | |
| wireless_remote[4:8] = packs[0] | |
| wireless_remote[8:12] = packs[1] | |
| wireless_remote[12:16] = packs[2] | |
| wireless_remote[20:24] = packs[3] | |
| return wireless_remote | |
| class PyGameJoystick(Joystick): | |
| def __init__(self) -> None: | |
| super().__init__() | |
| pygame.init() | |
| pygame.joystick.init() | |
| if pygame.joystick.get_count() <= 0: | |
| raise Exception("No joystick found!") | |
| self._joystick = pygame.joystick.Joystick(0) | |
| self._joystick.init() | |
| def print(self): | |
| print("\naxes: ") | |
| for i in range(self._joystick.get_numaxes()): | |
| print(self._joystick.get_axis(i), end=" ") | |
| print("\nbuttons: ") | |
| for i in range(self._joystick.get_numbuttons()): | |
| print(self._joystick.get_button(i), end=" ") | |
| print("\nhats: ") | |
| for i in range(self._joystick.get_numhats()): | |
| print(self._joystick.get_hat(i), end=" ") | |
| print("\nballs: ") | |
| for i in range(self._joystick.get_numballs()): | |
| print(self._joystick.get_ball(i), end=" ") | |
| print("\n") | |
| class LogicJoystick(PyGameJoystick): | |
| """ Logic F710 """ | |
| def __init__(self) -> None: | |
| super().__init__() | |
| def update(self): | |
| pygame.event.pump() | |
| self.back(self._joystick.get_button(6)) | |
| self.start(self._joystick.get_button(7)) | |
| self.LS(self._joystick.get_button(9)) | |
| self.RS(self._joystick.get_button(10)) | |
| self.LB(self._joystick.get_button(4)) | |
| self.RB(self._joystick.get_button(5)) | |
| self.A(self._joystick.get_button(0)) | |
| self.B(self._joystick.get_button(1)) | |
| self.X(self._joystick.get_button(2)) | |
| self.Y(self._joystick.get_button(3)) | |
| self.LT((self._joystick.get_axis(2) + 1)/2) | |
| self.RT((self._joystick.get_axis(5) + 1)/2) | |
| self.rx(self._joystick.get_axis(3)) | |
| self.ry(-self._joystick.get_axis(4)) | |
| # Logitech controller has 2 modes | |
| # mode 1: light down | |
| self.up(1 if self._joystick.get_hat(0)[1] > 0.5 else 0) | |
| self.down(1 if self._joystick.get_hat(0)[1] < -0.5 else 0) | |
| self.left(1 if self._joystick.get_hat(0)[0] < -0.5 else 0) | |
| self.right(1 if self._joystick.get_hat(0)[0] > 0.5 else 0) | |
| self.lx(self._joystick.get_axis(0)) | |
| self.ly(-self._joystick.get_axis(1)) | |
| # mode 2: light up | |
| # self.up(1 if self._joystick.get_axis(1) < -0.5 else 0) | |
| # self.down(1 if self._joystick.get_axis(0) > 0.5 else 0) | |
| # self.left(1 if self._joystick.get_axis(0) < -0.5 else 0) | |
| # self.right(1 if self._joystick.get_axis(0) > 0.5 else 0) | |
| # self.lx(self._joystick.get_hat(0)[1]) | |
| # self.ly(self._joystick.get_hat(0)[1]) | |