Inference endpoint server and sample scripts

scripts/endpoint_tcn.py

@@ -0,0 +1,78 @@
+"""
+ONNX-only inference script for Raspberry Pi deployment. This script is for using
+
+"""
+
+import onnxruntime
+import numpy as np
+import argparse
+from pathlib import Path
+import serial
+from collections import deque
+
+MODEL_PATH = "outputs/model.onnx"
+SERIAL_PORT = "/dev/ttyACM0" 
+BAUD_RATE = 115200
+FEATURE_COUNT = 28
+WINDOW_SIZE = 100
+
+def main():
+
+
+    if not Path(MODEL_PATH).exists():
+        print(f"Error: Model file '{MODEL_PATH}' not found.")
+        return
+
+    print(f"Loading model: {MODEL_PATH}")
+
+    session = onnxruntime.InferenceSession(MODEL_PATH, providers=["CPUExecutionProvider"])
+    input_name = session.get_inputs()[0].name
+    history = deque(maxlen=WINDOW_SIZE)
+    for _ in range(WINDOW_SIZE):
+        history.append(np.zeros(FEATURE_COUNT))
+
+    try: 
+        ser = serial.Serial(SERIAL_PORT, 115200, timeout=1)
+        ser.flush()
+    except Exception as e:
+        print(f"Serial Error: {e}")
+        return
+
+    # 3. Run Inference
+
+    while True:
+        if ser.in_waiting > 0:
+            try:
+                line = ser.readline().decode('utf-8').strip()
+                if not line: continue
+
+                # Parse the latest single reading from Arduino
+                new_sample = [float(x) for x in line.split(',')]
+                if len(new_sample) != FEATURE_COUNT: continue
+
+                # 2. Add new sample to history (automatically drops the oldest)
+                history.append(new_sample)
+
+                # 3. Convert history to the shape the TCN expects: (1, 100, 28)
+                input_data = np.array(history, dtype=np.float32).reshape(1, WINDOW_SIZE, FEATURE_COUNT)
+
+                # 4. Run Inference
+                outputs = session.run(None, {input_name: input_data})
+                predictions = outputs[0] # Shape: (1, 100, 4)
+
+                # 5. Extract the LATEST prediction for the 4 motors
+                # We take index -1 (the most recent time step)
+                current_pred = predictions[0, -1, :] 
+
+                # 6. Send all 4 motor values back as a comma-separated string
+                pred_str = ",".join([f"{x:.4f}" for x in current_pred])
+                ser.write(f"{pred_str}\n".encode('utf-8'))
+
+            except Exception as e:
+                   print(f"Error: {e}")
+
+if __name__ == "__main__":
+    main()
+
+
+

scripts/export_onnx.py

@@ -0,0 +1,115 @@
+"""
+Export a trained PyTorch TCN checkpoint to ONNX format for Raspberry Pi deployment.
+
+Usage:
+    python scripts/export_onnx.py
+    python scripts/export_onnx.py --model-path outputs/2026-01-08/18-29-05/best_model.pt
+    python scripts/export_onnx.py --output outputs/model.onnx
+"""
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import onnx
+import onnxruntime
+import torch
+from omegaconf import OmegaConf
+
+from exoskeleton_ml.models import create_model
+
+DEFAULT_MODEL_PATH = "outputs/2026-01-08/18-29-05/best_model.pt"
+
+
+def load_model(model_path: Path) -> tuple[torch.nn.Module, int]:
+    checkpoint = torch.load(model_path, map_location="cpu")
+
+    config_path = model_path.parent / "config.yaml"
+    if not config_path.exists():
+        raise ValueError(f"Config not found at {config_path}")
+
+    config = OmegaConf.load(config_path)
+    model = create_model(config.model)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+
+    window_size: int = config.model.effective_history
+    return model, window_size
+
+
+def export(model_path: Path, output_path: Path) -> None:
+    print(f"Loading model: {model_path}")
+    model, window_size = load_model(model_path)
+
+    num_params = sum(p.numel() for p in model.parameters())
+    print(f"  Parameters:     {num_params:,}")
+    print(f"  Window size:    {window_size} timesteps")
+    print(f"  Input shape:    (1, {window_size}, 28)")
+    print(f"  Output shape:   (1, {window_size}, 4)")
+
+    # Representative input for tracing
+    dummy_input = torch.randn(1, window_size, 28)
+
+    with torch.no_grad():
+        pytorch_output = model(dummy_input)
+
+    print(f"\nExporting to ONNX: {output_path}")
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    onnx_program = torch.onnx.export(model, dummy_input, dynamo=True)
+    onnx_program.save(str(output_path))
+
+    print("Checking ONNX model...")
+    onnx.checker.check_model(str(output_path))
+    print("  ONNX check passed.")
+
+    print("\nVerifying outputs match PyTorch...")
+    session_options = onnxruntime.SessionOptions()
+    session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+    ort_session = onnxruntime.InferenceSession(
+        str(output_path),
+        sess_options=session_options,
+        providers=["CPUExecutionProvider"],
+    )
+
+    input_name = ort_session.get_inputs()[0].name
+    onnx_output = ort_session.run(None, {input_name: dummy_input.numpy()})[0]
+
+    try:
+        torch.testing.assert_close(
+            pytorch_output,
+            torch.tensor(onnx_output),
+            rtol=1e-3,
+            atol=1e-4,
+        )
+        print("  Outputs match within tolerance. Export verified.")
+    except AssertionError as e:
+        print(f"  WARNING: outputs differ: {e}")
+
+    size_mb = output_path.stat().st_size / (1024 * 1024)
+    print(f"\nDone. Saved to: {output_path}  ({size_mb:.2f} MB)")
+    print(f"\nSCP to Pi:")
+    print(f"  scp {output_path} pi@<PI_IP>:~/exoskeleton-ai/outputs/model.onnx")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Export TCN checkpoint to ONNX")
+    parser.add_argument(
+        "--model-path",
+        type=str,
+        default=DEFAULT_MODEL_PATH,
+        help="Path to best_model.pt",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="outputs/model.onnx",
+        help="Output path for the .onnx file",
+    )
+    args = parser.parse_args()
+
+    export(Path(args.model_path), Path(args.output))
+
+
+if __name__ == "__main__":
+    main()

scripts/server.py

@@ -0,0 +1,175 @@
+"""
+FastAPI inference server for the exoskeleton TCN model.
+
+Runs on the Raspberry Pi. The C++ control system sends a POST request to
+/predict with a full (187, 28) window of sensor data and receives 4 motor
+torque predictions in response.
+
+Usage:
+    uvicorn scripts.server:app --host 0.0.0.0 --port 8000
+    uvicorn scripts.server:app --host 127.0.0.1 --port 8000  # localhost only
+
+Dependencies (not in pyproject.toml, install on Pi):
+    pip install fastapi uvicorn[standard] onnxruntime numpy msgpack
+
+Input shape:  (1, WINDOW_SIZE, 28)  — WINDOW_SIZE defaults to 187 for standard TCN
+Output shape: (1, WINDOW_SIZE, 4)   — server returns only the final timestep's 4 values
+"""
+
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Annotated
+
+import time
+
+import msgpack
+import numpy as np
+import onnxruntime as ort
+from fastapi import FastAPI, HTTPException, Request, Response
+from pydantic import BaseModel, Field, model_validator
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+MODEL_PATH = "outputs/model.onnx"
+
+# Must match the effective receptive field of the deployed model variant:
+#   standard TCN (kernel=7, 5 layers): 187
+#   small TCN    (kernel=5, 4 layers):  61
+#   large TCN    (kernel=9, 6 layers): 505
+WINDOW_SIZE = 187
+FEATURE_COUNT = 28  # 24 IMU + 4 joint angles
+OUTPUT_COUNT = 4    # hip_l, hip_r, knee_l, knee_r  (Nm/kg)
+
+# ---------------------------------------------------------------------------
+# App lifespan: load model once at startup
+# ---------------------------------------------------------------------------
+
+_session: ort.InferenceSession | None = None
+_input_name: str | None = None
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global _session, _input_name
+
+    model_path = Path(MODEL_PATH)
+    if not model_path.exists():
+        raise RuntimeError(f"Model file not found: {model_path.resolve()}")
+
+    print(f"Loading ONNX model: {model_path.resolve()}")
+    _session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"])
+    _input_name = _session.get_inputs()[0].name
+    print(f"Model loaded. Input: '{_input_name}' | Expected shape: (1, {WINDOW_SIZE}, {FEATURE_COUNT})")
+
+    yield
+
+    _session = None
+    _input_name = None
+    print("Server shutting down.")
+
+
+app = FastAPI(
+    title="Exoskeleton TCN Inference Server",
+    lifespan=lifespan,
+)
+
+# ---------------------------------------------------------------------------
+# Request / response schemas
+# ---------------------------------------------------------------------------
+
+
+class PredictRequest(BaseModel):
+    # List of WINDOW_SIZE rows, each with FEATURE_COUNT values.
+    # C++ sends: {"window": [[f0..f27], [f0..f27], ...]}  (187 rows)
+    window: Annotated[
+        list[list[float]],
+        Field(description=f"Sensor window: {WINDOW_SIZE} timesteps × {FEATURE_COUNT} features"),
+    ]
+
+    @model_validator(mode="after")
+    def validate_shape(self) -> "PredictRequest":
+        rows = len(self.window)
+        if rows != WINDOW_SIZE:
+            raise ValueError(f"Expected {WINDOW_SIZE} timesteps, got {rows}")
+        for i, row in enumerate(self.window):
+            if len(row) != FEATURE_COUNT:
+                raise ValueError(
+                    f"Timestep {i}: expected {FEATURE_COUNT} features, got {len(row)}"
+                )
+        return self
+
+
+class PredictResponse(BaseModel):
+    # Torque predictions for the 4 joints at the latest timestep (Nm/kg)
+    hip_left: float
+    hip_right: float
+    knee_left: float
+    knee_right: float
+    inference_ms: float  # ONNX session.run() time only, excludes HTTP/JSON overhead
+
+
+# ---------------------------------------------------------------------------
+# Endpoints
+# ---------------------------------------------------------------------------
+
+
+def _run_inference(window_bytes: np.ndarray) -> tuple[np.ndarray, float]:
+    """Shared inference logic. Input must already be shaped (1, WINDOW_SIZE, FEATURE_COUNT)."""
+    t0 = time.perf_counter()
+    outputs = _session.run(None, {_input_name: window_bytes})  # type: ignore[index]
+    inference_ms = (time.perf_counter() - t0) * 1000
+    preds = outputs[0][0, -1, :]
+    return preds, inference_ms
+
+
+@app.get("/health")
+def health() -> dict[str, str]:
+    """Quick liveness check — the C++ side can poll this on startup."""
+    if _session is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    return {"status": "ok"}
+
+
+@app.post("/predict", response_model=PredictResponse)
+def predict(req: PredictRequest) -> PredictResponse:
+    if _session is None or _input_name is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+
+    input_array = np.array(req.window, dtype=np.float32)[np.newaxis, ...]
+    preds, inference_ms = _run_inference(input_array)
+
+    return PredictResponse(
+        hip_left=float(preds[0]),
+        hip_right=float(preds[1]),
+        knee_left=float(preds[2]),
+        knee_right=float(preds[3]),
+        inference_ms=round(inference_ms, 3),
+    )
+
+
+@app.post("/predict_msgpack")
+async def predict_msgpack(request: Request) -> Response:
+    """
+    Msgpack endpoint for lower-overhead inference.
+
+    Request body:  msgpack-encoded flat list of 187*28=5236 float32 values (row-major)
+    Response body: msgpack-encoded dict {hip_left, hip_right, knee_left, knee_right, inference_ms}
+    """
+    if _session is None or _input_name is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+
+    raw = await request.body()
+    flat = msgpack.unpackb(raw, raw=False)
+    input_array = np.array(flat, dtype=np.float32).reshape(1, WINDOW_SIZE, FEATURE_COUNT)
+    preds, inference_ms = _run_inference(input_array)
+
+    result = {
+        "hip_left": float(preds[0]),
+        "hip_right": float(preds[1]),
+        "knee_left": float(preds[2]),
+        "knee_right": float(preds[3]),
+        "inference_ms": round(inference_ms, 3),
+    }
+    return Response(content=msgpack.packb(result), media_type="application/x-msgpack")