68 lines
2.7 KiB
Python
68 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Chinese phase banners for LeRobot recording; no changes to robot control."""
|
|
import sys
|
|
|
|
|
|
def show_phase(message: str) -> None:
|
|
if message.startswith("Recording episode "):
|
|
number = int(message.rsplit(" ", 1)[1]) + 1
|
|
title, color = f"● 正在录制｜数据集第 {number} 条", "32"
|
|
lines = [
|
|
"现在操作主臂：抓起积木 → 放入白盒 → 松开 → 退开",
|
|
"完成：按 → 结束本条，进入复位（最后一条将保存退出）",
|
|
"失败：按 ← 标记重录，然后按提示复位",
|
|
"停止：按 Esc；当前片段可能保存，之后检查完整性",
|
|
]
|
|
elif message == "Reset the environment":
|
|
title, color = "● 复位阶段｜当前不录入示范", "33"
|
|
lines = [
|
|
"把积木取出放回起点，并通过主臂让从臂回到起始姿态",
|
|
"注意：从臂仍跟随主臂！",
|
|
"准备好：按 → 继续；超时也会自动继续",
|
|
"上一条失败：按 ← 标记重录；按 Esc 结束采集",
|
|
]
|
|
elif message == "Re-record episode":
|
|
title, color = "● 已丢弃本条｜即将重录", "33"
|
|
lines = ["等待绿色“正在录制”提示后，再开始抓取。"]
|
|
elif message == "Stop recording":
|
|
title, color = "● 正在结束采集｜请等待保存与清理", "36"
|
|
lines = ["不要关闭终端；完成后会返回命令行。"]
|
|
elif message == "Exiting":
|
|
title, color = "● 采集程序已退出", "36"
|
|
lines = ["请检查实际保存条数及视频；提前结束的片段可能不完整。"]
|
|
else:
|
|
return
|
|
use_color = sys.stdout.isatty()
|
|
start = f"\033[1;{color}m" if use_color else ""
|
|
end = "\033[0m" if use_color else ""
|
|
print("\n" + start + "━" * 60, flush=True)
|
|
print(title, flush=True)
|
|
print("━" * 60, flush=True)
|
|
for line in lines:
|
|
print(" " + line, flush=True)
|
|
print("━" * 60 + end + "\n", flush=True)
|
|
if use_color:
|
|
print("\a", end="", flush=True)
|
|
|
|
|
|
def main() -> None:
|
|
if sys.argv[1:] == ["--preview"]:
|
|
for message in ["Recording episode 0", "Reset the environment", "Re-record episode", "Stop recording", "Exiting"]:
|
|
show_phase(message)
|
|
return
|
|
from lerobot.scripts import lerobot_record
|
|
|
|
original_log_say = lerobot_record.log_say
|
|
|
|
def log_with_banner(message, *args, **kwargs):
|
|
result = original_log_say(message, *args, **kwargs)
|
|
show_phase(message)
|
|
return result
|
|
|
|
lerobot_record.log_say = log_with_banner
|
|
lerobot_record.main()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|