Rune Harlyk
101 lines
2.6 KiB
Python
101 lines
2.6 KiB
Python
import asyncio
|
|
from enum import Enum
|
|
import json
|
|
import sys
|
|
import websockets
|
|
from model import model
|
|
from MotionController import GaitState
|
|
from simulator.GaitGenerator.Bezier import BezierGait
|
|
from simulator.GymEnvs.spot_bezier_env import spotBezierEnv
|
|
from simulator.Kinematics.SpotKinematics import SpotModel
|
|
from simulator.util.gui import GUI
|
|
from simulator.simulator import BodyState, Simulator
|
|
import struct
|
|
|
|
sys.path.append("./simulator/GymEnvs")
|
|
|
|
clients = {}
|
|
|
|
env = spotBezierEnv(
|
|
render=True,
|
|
on_rack=False,
|
|
height_field=False,
|
|
draw_foot_path=False,
|
|
env_randomizer=None,
|
|
)
|
|
gui = GUI(env.spot.quadruped)
|
|
bodyState = BodyState()
|
|
gaitState = GaitState()
|
|
spot = SpotModel()
|
|
bezierGait = BezierGait()
|
|
simulator = Simulator()
|
|
|
|
class Command(Enum):
|
|
ESTOP = 0
|
|
CONTROLLER = 1
|
|
|
|
|
|
def get_controller(buffer):
|
|
buffer = struct.unpack("<8b", buffer)
|
|
return {
|
|
"command": buffer[0],
|
|
"estop": buffer[1],
|
|
"lx": buffer[2],
|
|
"ly": buffer[3],
|
|
"rx": buffer[4],
|
|
"ry": buffer[5],
|
|
"height": buffer[6],
|
|
"speed": buffer[7],
|
|
}
|
|
|
|
|
|
async def handle_binary_message(client, data):
|
|
message = get_controller(data)
|
|
command = Command(message["command"])
|
|
if command == Command.ESTOP:
|
|
client["model"]["running"] = False
|
|
await client["websocket"].send(
|
|
json.dumps({"type": "stop", "data": "Servos stopped"})
|
|
)
|
|
|
|
if command == Command.CONTROLLER:
|
|
await client["websocket"].send(json.dumps({"type": "echo", "data": message}))
|
|
|
|
|
|
async def handle_json_message(client, message):
|
|
data = json.loads(message)
|
|
client = client["clientState"]
|
|
if data["type"] in ("stop", "mode_change"):
|
|
client["model"][data["type"]] = data.get("data", False)
|
|
await client["websocket"].send(
|
|
json.dumps(
|
|
{"type": data["type"], "data": data.get("data", "Servos stopped")}
|
|
)
|
|
)
|
|
|
|
|
|
async def handle_message(websocket, path):
|
|
client_id = id(websocket)
|
|
clients[client_id] = {
|
|
"clientState": model(),
|
|
"websocket": websocket,
|
|
}
|
|
try:
|
|
async for message in websocket:
|
|
if isinstance(message, bytes):
|
|
await handle_binary_message(clients[client_id], message)
|
|
else:
|
|
await handle_json_message(clients[client_id], message)
|
|
finally:
|
|
del clients[client_id]
|
|
|
|
|
|
async def main():
|
|
async with websockets.serve(handle_message, "localhost", 2096):
|
|
print("Server starting")
|
|
await asyncio.Future()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|