(WIP) add calibration gui
This commit is contained in:
Simon Alibert
205
lerobot/common/motors/calibration_gui.py
Normal file
205
lerobot/common/motors/calibration_gui.py
Normal file
@@ -0,0 +1,205 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
Hand + Arm slider GUI for HopeJr with optional --no-hand / --no-arm flags.
|
||||
|
||||
Run examples:
|
||||
python sliders.py # hand + arm
|
||||
python sliders.py --no-hand # arm only
|
||||
python sliders.py --no-arm # hand only
|
||||
"""
|
||||
|
||||
# ruff: noqa: N806
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pygame
|
||||
|
||||
from lerobot.common.robots.hope_jr import HopeJrHand, HopeJrHandConfig
|
||||
|
||||
from .motors_bus import MotorsBus
|
||||
|
||||
|
||||
# ────────────────────────────────────────────────────────────────────────────────
|
||||
# Helpers for JSON range files
|
||||
# ────────────────────────────────────────────────────────────────────────────────
|
||||
def load_ranges(path: str):
|
||||
d = json.loads(Path(path).read_text())
|
||||
return d["start_pos"], d["end_pos"]
|
||||
|
||||
|
||||
def save_ranges(path: str, lo, hi):
|
||||
d = {
|
||||
"start_pos": lo,
|
||||
"end_pos": hi,
|
||||
"homing_offset": [],
|
||||
"drive_mode": [],
|
||||
"calib_mode": "LINEAR",
|
||||
}
|
||||
Path(path).write_text(json.dumps(d, indent=2))
|
||||
|
||||
|
||||
def main(bus: MotorsBus, motors: list[str] | None = None):
|
||||
motors = motors if motors else list(bus.motors)
|
||||
|
||||
if not bus.is_connected:
|
||||
bus.connect()
|
||||
|
||||
H_RANGE = "examples/hopejr/settings/hand_ranges.json"
|
||||
hand_lo, hand_hi = load_ranges(H_RANGE) if motors else ([], [])
|
||||
# arm_lo, arm_hi = load_ranges(A_RANGE) if arm_names else ([], [])
|
||||
|
||||
pygame.init()
|
||||
screen = pygame.display.set_mode((1080, 720))
|
||||
pygame.display.set_caption("Hand + Arm Sliders")
|
||||
font = pygame.font.SysFont(None, 20)
|
||||
big = pygame.font.SysFont(None, 24)
|
||||
|
||||
ROWS = len(motors)
|
||||
TOP, BOT = 100, 20
|
||||
ROW_H = (720 - TOP - BOT) // max(1, ROWS)
|
||||
BAR_H = ROW_H - 25
|
||||
|
||||
HX, H_W = None, 0
|
||||
|
||||
H_MAX = 1024
|
||||
h_vals = np.zeros(len(motors), int)
|
||||
|
||||
BTN_W, BTN_H = 120, 28
|
||||
btns = [
|
||||
{
|
||||
"rect": pygame.Rect(HX, 20, BTN_W, BTN_H),
|
||||
"txt": "Hand LOWER",
|
||||
"col": "hand",
|
||||
"type": "low",
|
||||
},
|
||||
{
|
||||
"rect": pygame.Rect(HX + BTN_W + 10, 20, BTN_W, BTN_H),
|
||||
"txt": "Hand UPPER",
|
||||
"col": "hand",
|
||||
"type": "high",
|
||||
},
|
||||
]
|
||||
edit_mode = None # None or ("hand"/"arm","low"/"high")
|
||||
toast = None
|
||||
toast_timer = 0
|
||||
|
||||
def bar(x, y, w, val, lo, hi, vmax, label):
|
||||
lo, hi = sorted((lo, hi))
|
||||
px_lo = int(w * lo / vmax)
|
||||
px_hi = int(w * hi / vmax)
|
||||
px_val = int(w * val / vmax)
|
||||
bar_y = y + 5
|
||||
pygame.draw.rect(screen, (200, 200, 200), (x, bar_y, w, BAR_H))
|
||||
pygame.draw.rect(screen, (80, 200, 80), (x, bar_y, px_val, BAR_H))
|
||||
pygame.draw.rect(screen, (255, 60, 60), (x, bar_y, px_lo, BAR_H))
|
||||
pygame.draw.rect(screen, (255, 60, 60), (x + px_hi, bar_y, w - px_hi, BAR_H))
|
||||
name = font.render(label, True, (255, 255, 255))
|
||||
screen.blit(name, name.get_rect(midbottom=(x + w // 2, y - 3)))
|
||||
screen.blit(font.render(str(lo), True, (80, 120, 255)), (x - 45, bar_y + BAR_H // 4))
|
||||
screen.blit(font.render(str(hi), True, (80, 120, 255)), (x + w - 55, bar_y + BAR_H // 4))
|
||||
screen.blit(font.render(str(val), True, (255, 255, 255)), (x + w + 10, bar_y + BAR_H // 4))
|
||||
|
||||
def slider_at(pos):
|
||||
x, y = pos
|
||||
row = (y - TOP) // ROW_H
|
||||
if not 0 <= row < ROWS:
|
||||
return None
|
||||
y0 = TOP + row * ROW_H
|
||||
if not y0 <= y <= y0 + BAR_H + 5:
|
||||
return None
|
||||
if HX <= x <= HX + H_W and row < len(motors):
|
||||
return (str(bus), row, HX, H_W, H_MAX)
|
||||
return None
|
||||
|
||||
def toast_msg(txt):
|
||||
nonlocal toast, toast_timer
|
||||
toast = big.render(txt, True, (255, 255, 0))
|
||||
toast_timer = 120
|
||||
|
||||
clock = pygame.time.Clock()
|
||||
dragging = None
|
||||
while True:
|
||||
screen.fill((20, 20, 20))
|
||||
|
||||
for i, n in enumerate(motors):
|
||||
bar(HX, TOP + i * ROW_H, H_W, h_vals[i], hand_lo[i], hand_hi[i], H_MAX, n)
|
||||
|
||||
# for i, n in enumerate(arm_names):
|
||||
# bar(AX, TOP + i * ROW_H, A_W, a_vals[i], arm_lo[i], arm_hi[i], A_MAX, n)
|
||||
|
||||
# buttons
|
||||
for b in btns:
|
||||
col = (120, 120, 120) if edit_mode == (b["col"], b["type"]) else (60, 60, 60)
|
||||
pygame.draw.rect(screen, col, b["rect"])
|
||||
txt = font.render(b["txt"], True, (255, 255, 255))
|
||||
screen.blit(txt, txt.get_rect(center=b["rect"].center))
|
||||
|
||||
if toast:
|
||||
screen.blit(toast, toast.get_rect(center=(540, 55)))
|
||||
|
||||
pygame.display.flip()
|
||||
|
||||
# events
|
||||
for ev in pygame.event.get():
|
||||
if ev.type == pygame.QUIT:
|
||||
bus.disconnect()
|
||||
pygame.quit()
|
||||
sys.exit()
|
||||
|
||||
if ev.type == pygame.MOUSEBUTTONDOWN:
|
||||
# button toggle
|
||||
for b in btns:
|
||||
if b["rect"].collidepoint(ev.pos):
|
||||
edit_mode = None if edit_mode == (b["col"], b["type"]) else (b["col"], b["type"])
|
||||
toast_msg(f"Editing: {edit_mode[0]} {edit_mode[1]}" if edit_mode else "Edit off")
|
||||
break
|
||||
else:
|
||||
hit = slider_at(ev.pos)
|
||||
if hit:
|
||||
dragging = hit
|
||||
# bound edit
|
||||
if edit_mode and hit and hit[0] == edit_mode[0]:
|
||||
col, row, left, w, vmax = hit
|
||||
mx, _ = ev.pos
|
||||
new = int(np.clip((mx - left) / w * vmax, 0, vmax))
|
||||
if col == "hand":
|
||||
if edit_mode[1] == "low":
|
||||
hand_lo[row] = new
|
||||
else:
|
||||
hand_hi[row] = new
|
||||
save_ranges(H_RANGE, hand_lo, hand_hi)
|
||||
else:
|
||||
raise ValueError(col)
|
||||
# if edit_mode[1] == "low":
|
||||
# arm_lo[row] = new
|
||||
# else:
|
||||
# arm_hi[row] = new
|
||||
# save_ranges(A_RANGE, arm_lo, arm_hi)
|
||||
toast_msg(f"{col} {row} {edit_mode[1]}→{new}")
|
||||
|
||||
if ev.type == pygame.MOUSEBUTTONUP:
|
||||
dragging = None
|
||||
if ev.type == pygame.MOUSEMOTION and dragging:
|
||||
col, row, left, w, vmax = dragging
|
||||
mx, _ = ev.pos
|
||||
new = int(np.clip((mx - left) / w * vmax, 0, vmax))
|
||||
if h_vals[row] != new:
|
||||
h_vals[row] = new
|
||||
bus.write("Goal_Position", [new], [motors[row]])
|
||||
|
||||
if toast_timer > 0:
|
||||
toast_timer -= 1
|
||||
else:
|
||||
toast = None
|
||||
clock.tick(60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from lerobot.common.robots.hope_jr import HopeJrHand, HopeJrHandConfig
|
||||
|
||||
cfg = HopeJrHandConfig("/dev/tty.usbmodem58760431541")
|
||||
hand = HopeJrHand()
|
||||
main(hand.bus)
|
||||
Reference in New Issue
Block a user