| """The driving model behind the simulator's camera-model contract. |
| |
| python -m mcq_sim evaluate --policy stack --perception driver:create |
| |
| with ``PYTHONPATH=training/perception:src/mcq_sim``. Weights come from |
| ``MCQ_DRIVER_WEIGHTS``, else ``runs/drive-latest/model.pt`` beside this file. |
| |
| ``predict(rgb)`` is the pavement probability, as for a segmenter, so the stack |
| can build its corridor. ``predict_drive(rgb)`` adds the racing line the model |
| intends to drive, as waypoints in base_link (x forward, y left), and the speed |
| it wants at each. The stack drives that line when it stays inside the corridor |
| and falls back to its own planner when it does not. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import os |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
| from model import DriveNet |
|
|
| HERE = Path(__file__).resolve().parent |
|
|
|
|
| class Driver: |
| def __init__(self, weights: str | Path, device="cpu"): |
| torch.set_num_threads(1) |
| checkpoint = torch.load(weights, map_location=device, weights_only=True) |
| self.net = DriveNet(fine=bool(checkpoint.get("fine", False))).to(device).eval() |
| self.net.load_state_dict(checkpoint["state"]) |
| self.device = device |
| self.weights = str(weights) |
|
|
| @torch.no_grad() |
| def predict_drive(self, rgb: np.ndarray): |
| """uint8 HWC RGB (270 x 480) in; (pavement probability 270 x 480, |
| waypoints K x 2 in metres, speeds K in m/s) out.""" |
| x = torch.from_numpy(np.ascontiguousarray(rgb)).to(self.device).permute(2, 0, 1)[None].float() |
| h, w = x.shape[-2:] |
| logits, points, speeds = self.net.forward_all(F.avg_pool2d(x, 2)) |
| probability = torch.sigmoid(F.interpolate(logits[None], size=(h, w), mode="bilinear", align_corners=False)) |
| out = probability[0, 0].cpu().numpy(), points[0].cpu().numpy(), speeds[0].cpu().numpy() |
| return out |
|
|
| def predict(self, rgb: np.ndarray) -> np.ndarray: |
| return self.predict_drive(rgb)[0] |
|
|
|
|
| def create(): |
| path = os.environ.get("MCQ_DRIVER_WEIGHTS", HERE / "runs" / "drive-latest" / "model.pt") |
| return Driver(path) |
|
|
