drive-v5 and the checkpoints it came from, with a model card
Browse files- README.md +88 -0
- checkpoints/drive-v3.pt +3 -0
- checkpoints/drive-v4.pt +3 -0
- checkpoints/drive-v6.pt +3 -0
- drive-v5.pt +3 -0
- driver.py +55 -0
- example.py +20 -0
- laptime-v2-best.json +78 -0
- model.py +127 -0
README.md
ADDED
|
@@ -0,0 +1,88 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
library_name: pytorch
|
| 3 |
+
pipeline_tag: image-segmentation
|
| 4 |
+
tags:
|
| 5 |
+
- autonomous-driving
|
| 6 |
+
- go-kart
|
| 7 |
+
- racing-line
|
| 8 |
+
- semantic-segmentation
|
| 9 |
+
- imitation-learning
|
| 10 |
+
- robotics
|
| 11 |
+
---
|
| 12 |
+
|
| 13 |
+
# McQueen driving model (drive-v5)
|
| 14 |
+
|
| 15 |
+
A small camera-only driving model for an autonomous go-kart on the Purdue Grand Prix track. One RGB frame goes in; three things come out: a **pavement mask**, a **racing line** (12 waypoints, 2 to 24 m ahead, in the kart's frame) and a **target speed** at each waypoint. It does not output steering or throttle. Controllers outside the model turn the line into commands from measured state, and a safety layer keeps the line inside the pavement the camera can see.
|
| 16 |
+
|
| 17 |
+
> **Status: a simulator result, not a road-tested model.** It has never seen the kart's own camera and has never driven a real kart. See [What this does not show](#what-this-does-not-show).
|
| 18 |
+
|
| 19 |
+
Code, training scripts, the simulator and the full version history are in the project repository: https://github.com/ExplodingCB/mcqueen-vip (branch `training-perception`, see `docs/13-driving-model-versions.md`).
|
| 20 |
+
|
| 21 |
+
## Files
|
| 22 |
+
|
| 23 |
+
| File | What it is |
|
| 24 |
+
| --- | --- |
|
| 25 |
+
| `drive-v5.pt` | **The recommended model.** Real-footage mask training and a clean simulated lap |
|
| 26 |
+
| `checkpoints/drive-v3.pt` | Best simulated lap time (63.8 s); weaker on real footage |
|
| 27 |
+
| `checkpoints/drive-v4.pt` | drive-v3 plus real helmet-cam footage for the mask |
|
| 28 |
+
| `checkpoints/drive-v6.pt` | Trained on a faster teacher line. Best open-loop numbers, but **does not complete laps**; included as a research result |
|
| 29 |
+
| `model.py`, `driver.py`, `example.py` | Everything needed to load a checkpoint and run it on an image |
|
| 30 |
+
| `laptime-v2-best.json` | The searched expert line (racing line offsets and speed limits) the faster teacher uses |
|
| 31 |
+
|
| 32 |
+
## Checksums (SHA-256)
|
| 33 |
+
|
| 34 |
+
| File | SHA-256 | Size |
|
| 35 |
+
| --- | --- | --- |
|
| 36 |
+
| `checkpoints/drive-v3.pt` | `ce23623f702185460a999e423af3ac36670f49dfd3e08d41c65285bb1132ac5c` | 15.7 MB |
|
| 37 |
+
| `checkpoints/drive-v4.pt` | `3d995f317eb3467ac0ca801500d060651d1ea6b3f68cea60bd9050f400431677` | 15.7 MB |
|
| 38 |
+
| `checkpoints/drive-v6.pt` | `b1a7813b21b4e79b55e7c2f1c3f9a244fb7641a0550674d96e9f4e3058e057b5` | 15.7 MB |
|
| 39 |
+
| `drive-v5.pt` | `0d6ba0f1a4e45fd28c4f609fbfb0cb7256a9feaf2071bab200498a0fcc12ddcc` | 15.7 MB |
|
| 40 |
+
|
| 41 |
+
## Use
|
| 42 |
+
|
| 43 |
+
```python
|
| 44 |
+
# pip install torch numpy pillow
|
| 45 |
+
import numpy as np
|
| 46 |
+
from PIL import Image
|
| 47 |
+
from driver import Driver # driver.py and model.py sit beside it
|
| 48 |
+
|
| 49 |
+
driver = Driver("drive-v5.pt") # CPU by default
|
| 50 |
+
frame = np.asarray(Image.open("frame.jpg").convert("RGB").resize((480, 270))) # uint8, 270 x 480 x 3
|
| 51 |
+
mask, line, speed = driver.predict_drive(frame)
|
| 52 |
+
# mask : (270, 480) pavement probability
|
| 53 |
+
# line : (12, 2) metres in the kart frame, x forward and y left, from 2 m to 24 m ahead
|
| 54 |
+
# speed: (12,) metres per second the model wants at each waypoint
|
| 55 |
+
```
|
| 56 |
+
|
| 57 |
+
The model assumes a camera about 0.7 m above the ground, pitched about 8 degrees down, with a 100 degree horizontal field of view. A different mount shifts the line, because it is in metres. `python example.py frame.jpg` prints the same thing.
|
| 58 |
+
|
| 59 |
+
## Model
|
| 60 |
+
|
| 61 |
+
`DriveNet`: a 3.9 M parameter encoder-decoder (a dilated bottleneck, all standard convolutions) with a mask head, plus a driving head that reads the bottleneck and the quarter-resolution decoder features (with pixel-coordinate channels) and regresses the line and speeds. Input is the 240 x 135 half-resolution frame. About 7 ms per frame on one CPU thread.
|
| 62 |
+
|
| 63 |
+
## Training data
|
| 64 |
+
|
| 65 |
+
* **Simulator frames:** the Purdue track with the 2023 Indiana orthophoto draped over a flat ground plane (Indiana's imagery service, marked CC0), camera mount and image appearance randomized, labeled with an expert's minimum-curvature racing line and the speed along it.
|
| 66 |
+
* **DAgger frames:** frames from states the model itself reaches in closed loop (with steering bumps), labeled by the expert from where the kart actually is.
|
| 67 |
+
* **Real kart footage, for the mask only:** about 2,300 frames from about 30 YouTube kart onboard videos (helmet cams and nose or bumper cams), pavement labeled automatically by SegFormer (ADE20K classes road, runway, sidewalk, path). **These are pseudo-labels with no human review.** The videos are third-party, so check their terms before redistributing or using this commercially.
|
| 68 |
+
|
| 69 |
+
## Results (all in the simulator unless stated)
|
| 70 |
+
|
| 71 |
+
| Measure | drive-v5 |
|
| 72 |
+
| --- | --- |
|
| 73 |
+
| Closed-loop lap, camera only, default limits, no speed cap, 3 seeds | **64.6 s**, 0.39 m closest to an edge |
|
| 74 |
+
| Racing-line error on unseen simulator frames, 2 to 10 m / 12 to 24 m ahead | 0.27 m / 0.54 m |
|
| 75 |
+
| Target-speed error | 0.34 m/s |
|
| 76 |
+
| Pavement mask IoU on unseen simulator frames | 0.979 |
|
| 77 |
+
| **Pavement mask IoU on 575 held-out real kart frames, against auto-labels** | **0.84** |
|
| 78 |
+
|
| 79 |
+
For reference: a simulator-only model scores 0.39 on the real frames, below a fixed "the lower part of the picture is road" guess (0.47). Models trained on the faster teacher (drive-v6) fail to complete laps.
|
| 80 |
+
|
| 81 |
+
## What this does not show
|
| 82 |
+
|
| 83 |
+
* No version has seen the kart's own camera or driven a real kart. The held-out real frames are other people's cameras on other tracks, scored against automatic labels.
|
| 84 |
+
* The racing line and speed were trained only on simulator frames. The mask is the only output that has been trained on real images, so how the line and speed behave on real frames is untested.
|
| 85 |
+
* The track file is hand-traced from aerial imagery, and the simulated tyres, brakes and steering servo are placeholders. A 64.6 s simulated lap is not a prediction of the real kart.
|
| 86 |
+
* The driving path (safety filter, corridor tracker) exists only in the Python simulator and is not yet ported to the kart's software or timed on the Jetson.
|
| 87 |
+
|
| 88 |
+
The intended use today is research, the simulator, and **shadow mode** (a person drives while the model only watches).
|
checkpoints/drive-v3.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:ce23623f702185460a999e423af3ac36670f49dfd3e08d41c65285bb1132ac5c
|
| 3 |
+
size 15691971
|
checkpoints/drive-v4.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:3d995f317eb3467ac0ca801500d060651d1ea6b3f68cea60bd9050f400431677
|
| 3 |
+
size 15691907
|
checkpoints/drive-v6.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b1a7813b21b4e79b55e7c2f1c3f9a244fb7641a0550674d96e9f4e3058e057b5
|
| 3 |
+
size 15691907
|
drive-v5.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:0d6ba0f1a4e45fd28c4f609fbfb0cb7256a9feaf2071bab200498a0fcc12ddcc
|
| 3 |
+
size 15691907
|
driver.py
ADDED
|
@@ -0,0 +1,55 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The driving model behind the simulator's camera-model contract.
|
| 2 |
+
|
| 3 |
+
python -m mcq_sim evaluate --policy stack --perception driver:create
|
| 4 |
+
|
| 5 |
+
with ``PYTHONPATH=training/perception:src/mcq_sim``. Weights come from
|
| 6 |
+
``MCQ_DRIVER_WEIGHTS``, else ``runs/drive-latest/model.pt`` beside this file.
|
| 7 |
+
|
| 8 |
+
``predict(rgb)`` is the pavement probability, as for a segmenter, so the stack
|
| 9 |
+
can build its corridor. ``predict_drive(rgb)`` adds the racing line the model
|
| 10 |
+
intends to drive, as waypoints in base_link (x forward, y left), and the speed
|
| 11 |
+
it wants at each. The stack drives that line when it stays inside the corridor
|
| 12 |
+
and falls back to its own planner when it does not.
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
from __future__ import annotations
|
| 16 |
+
|
| 17 |
+
import os
|
| 18 |
+
from pathlib import Path
|
| 19 |
+
|
| 20 |
+
import numpy as np
|
| 21 |
+
import torch
|
| 22 |
+
import torch.nn.functional as F
|
| 23 |
+
|
| 24 |
+
from model import DriveNet
|
| 25 |
+
|
| 26 |
+
HERE = Path(__file__).resolve().parent
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
class Driver:
|
| 30 |
+
def __init__(self, weights: str | Path, device="cpu"):
|
| 31 |
+
torch.set_num_threads(1) # many drivers run side by side in rollouts; one thread each
|
| 32 |
+
checkpoint = torch.load(weights, map_location=device, weights_only=True)
|
| 33 |
+
self.net = DriveNet(fine=bool(checkpoint.get("fine", False))).to(device).eval()
|
| 34 |
+
self.net.load_state_dict(checkpoint["state"])
|
| 35 |
+
self.device = device
|
| 36 |
+
self.weights = str(weights)
|
| 37 |
+
|
| 38 |
+
@torch.no_grad()
|
| 39 |
+
def predict_drive(self, rgb: np.ndarray):
|
| 40 |
+
"""uint8 HWC RGB (270 x 480) in; (pavement probability 270 x 480,
|
| 41 |
+
waypoints K x 2 in metres, speeds K in m/s) out."""
|
| 42 |
+
x = torch.from_numpy(np.ascontiguousarray(rgb)).to(self.device).permute(2, 0, 1)[None].float()
|
| 43 |
+
h, w = x.shape[-2:]
|
| 44 |
+
logits, points, speeds = self.net.forward_all(F.avg_pool2d(x, 2))
|
| 45 |
+
probability = torch.sigmoid(F.interpolate(logits[None], size=(h, w), mode="bilinear", align_corners=False))
|
| 46 |
+
out = probability[0, 0].cpu().numpy(), points[0].cpu().numpy(), speeds[0].cpu().numpy()
|
| 47 |
+
return out
|
| 48 |
+
|
| 49 |
+
def predict(self, rgb: np.ndarray) -> np.ndarray:
|
| 50 |
+
return self.predict_drive(rgb)[0]
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def create():
|
| 54 |
+
path = os.environ.get("MCQ_DRIVER_WEIGHTS", HERE / "runs" / "drive-latest" / "model.pt")
|
| 55 |
+
return Driver(path)
|
example.py
ADDED
|
@@ -0,0 +1,20 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Run a checkpoint on one image: python example.py frame.jpg [drive-v5.pt]"""
|
| 2 |
+
import sys
|
| 3 |
+
|
| 4 |
+
import numpy as np
|
| 5 |
+
from PIL import Image
|
| 6 |
+
|
| 7 |
+
from driver import Driver
|
| 8 |
+
|
| 9 |
+
image = sys.argv[1] if len(sys.argv) > 1 else None
|
| 10 |
+
weights = sys.argv[2] if len(sys.argv) > 2 else "drive-v5.pt"
|
| 11 |
+
driver = Driver(weights)
|
| 12 |
+
frame = (
|
| 13 |
+
np.asarray(Image.open(image).convert("RGB").resize((480, 270)))
|
| 14 |
+
if image
|
| 15 |
+
else np.full((270, 480, 3), 110, dtype=np.uint8)
|
| 16 |
+
)
|
| 17 |
+
mask, line, speed = driver.predict_drive(frame)
|
| 18 |
+
print(f"{weights}: pavement share {float((mask > 0.5).mean()):.2f}")
|
| 19 |
+
for k, ((x, y), v) in enumerate(zip(line, speed, strict=True)):
|
| 20 |
+
print(f" waypoint {k:2d}: {x:5.1f} m ahead, {y:+5.2f} m left, wants {v:4.1f} m/s")
|
laptime-v2-best.json
ADDED
|
@@ -0,0 +1,78 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"theta": [
|
| 3 |
+
-0.020713883569653507,
|
| 4 |
+
0.7986371171163591,
|
| 5 |
+
-0.06007090210526833,
|
| 6 |
+
-0.13692955258757095,
|
| 7 |
+
0.3164787844956688,
|
| 8 |
+
0.22910077232445988,
|
| 9 |
+
0.4085683741777101,
|
| 10 |
+
0.1732443452654597,
|
| 11 |
+
0.31628061199629176,
|
| 12 |
+
0.2936319802034716,
|
| 13 |
+
-0.007033094360031467,
|
| 14 |
+
0.0013200730496496782,
|
| 15 |
+
0.40243794379289366,
|
| 16 |
+
0.765853695695947,
|
| 17 |
+
0.5178593376332585,
|
| 18 |
+
0.33366131041959013,
|
| 19 |
+
1.0801067821758314,
|
| 20 |
+
0.7981998198285803,
|
| 21 |
+
-0.3385555048128157,
|
| 22 |
+
-0.3040821312994424,
|
| 23 |
+
1.6182991758783387,
|
| 24 |
+
0.3389952005170975,
|
| 25 |
+
-0.22789996980313015,
|
| 26 |
+
0.4039164824230546,
|
| 27 |
+
0.4619468070190555,
|
| 28 |
+
0.21125474388800541,
|
| 29 |
+
-0.5006066348041253,
|
| 30 |
+
-0.8610472716707411,
|
| 31 |
+
-0.20847536428974067,
|
| 32 |
+
-0.2394684425191484,
|
| 33 |
+
-0.9391657225767008,
|
| 34 |
+
-0.9217363080423312,
|
| 35 |
+
0.06200212942970203,
|
| 36 |
+
0.6460498410297608,
|
| 37 |
+
0.9070343531284966,
|
| 38 |
+
1.2993982884320447,
|
| 39 |
+
0.6067228170245667,
|
| 40 |
+
0.8575524174357774,
|
| 41 |
+
0.9950192746873036,
|
| 42 |
+
-0.16623329240627815,
|
| 43 |
+
4.464017710110649,
|
| 44 |
+
1.7967631217917692,
|
| 45 |
+
2.0052742170641062
|
| 46 |
+
],
|
| 47 |
+
"score": 55.65999999998595,
|
| 48 |
+
"results": [
|
| 49 |
+
{
|
| 50 |
+
"passed": true,
|
| 51 |
+
"lap_s": 54.64999999998654,
|
| 52 |
+
"progress_m": 435.10000000004015,
|
| 53 |
+
"clearance_m": 0.2806985629754275,
|
| 54 |
+
"max_speed": 12.532168425742148,
|
| 55 |
+
"corrections": 14,
|
| 56 |
+
"stop": null
|
| 57 |
+
},
|
| 58 |
+
{
|
| 59 |
+
"passed": true,
|
| 60 |
+
"lap_s": 55.29999999998595,
|
| 61 |
+
"progress_m": 435.1000000000418,
|
| 62 |
+
"clearance_m": 0.2562416972740049,
|
| 63 |
+
"max_speed": 12.49933797767215,
|
| 64 |
+
"corrections": 12,
|
| 65 |
+
"stop": null
|
| 66 |
+
},
|
| 67 |
+
{
|
| 68 |
+
"passed": true,
|
| 69 |
+
"lap_s": 54.64999999998654,
|
| 70 |
+
"progress_m": 435.25000000004,
|
| 71 |
+
"clearance_m": 0.4983136894426503,
|
| 72 |
+
"max_speed": 12.505934752719071,
|
| 73 |
+
"corrections": 0,
|
| 74 |
+
"stop": null
|
| 75 |
+
}
|
| 76 |
+
],
|
| 77 |
+
"generation": 103
|
| 78 |
+
}
|
model.py
ADDED
|
@@ -0,0 +1,127 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""A small encoder-decoder for the pavement mask.
|
| 2 |
+
|
| 3 |
+
About 1.6 M parameters, all standard convolutions, so it exports to TensorRT
|
| 4 |
+
without surprises and runs far inside the 20 fps budget on the Orin. The input is
|
| 5 |
+
the 240 x 135 half-resolution frame; the output is one logit per pixel at the
|
| 6 |
+
same size. A dilated bottleneck gives the receptive field to see the whole road
|
| 7 |
+
shape at once, which matters more than texture: pavement and grass can be the
|
| 8 |
+
same color, and the shape of the corridor still says which is which.
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
import torch
|
| 14 |
+
import torch.nn as nn
|
| 15 |
+
import torch.nn.functional as F
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
def block(cin, cout, stride=1, dilation=1):
|
| 19 |
+
return nn.Sequential(
|
| 20 |
+
nn.Conv2d(cin, cout, 3, stride, padding=dilation, dilation=dilation, bias=False),
|
| 21 |
+
nn.BatchNorm2d(cout),
|
| 22 |
+
nn.ReLU(inplace=True),
|
| 23 |
+
nn.Conv2d(cout, cout, 3, 1, padding=dilation, dilation=dilation, bias=False),
|
| 24 |
+
nn.BatchNorm2d(cout),
|
| 25 |
+
nn.ReLU(inplace=True),
|
| 26 |
+
)
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
class PavementNet(nn.Module):
|
| 30 |
+
def __init__(self, width=24):
|
| 31 |
+
super().__init__()
|
| 32 |
+
w = width
|
| 33 |
+
self.e1 = block(3, w, stride=2) # 1/2
|
| 34 |
+
self.e2 = block(w, 2 * w, stride=2) # 1/4
|
| 35 |
+
self.e3 = block(2 * w, 4 * w, stride=2) # 1/8
|
| 36 |
+
self.e4 = block(4 * w, 6 * w, stride=2) # 1/16
|
| 37 |
+
self.context = nn.Sequential(block(6 * w, 6 * w, dilation=2), block(6 * w, 6 * w, dilation=4))
|
| 38 |
+
self.d3 = block(6 * w + 4 * w, 4 * w)
|
| 39 |
+
self.d2 = block(4 * w + 2 * w, 2 * w)
|
| 40 |
+
self.d1 = block(2 * w + w, w)
|
| 41 |
+
self.head = nn.Conv2d(w, 1, 1)
|
| 42 |
+
|
| 43 |
+
@staticmethod
|
| 44 |
+
def _up(x, like):
|
| 45 |
+
return F.interpolate(x, size=like.shape[-2:], mode="bilinear", align_corners=False)
|
| 46 |
+
|
| 47 |
+
def forward(self, x):
|
| 48 |
+
size = x.shape[-2:]
|
| 49 |
+
x = x / 255.0 - 0.5
|
| 50 |
+
e1 = self.e1(x)
|
| 51 |
+
e2 = self.e2(e1)
|
| 52 |
+
e3 = self.e3(e2)
|
| 53 |
+
e4 = self.context(self.e4(e3))
|
| 54 |
+
d3 = self.d3(torch.cat([self._up(e4, e3), e3], 1))
|
| 55 |
+
d2 = self.d2(torch.cat([self._up(d3, e2), e2], 1))
|
| 56 |
+
d1 = self.d1(torch.cat([self._up(d2, e1), e1], 1))
|
| 57 |
+
return F.interpolate(self.head(d1), size=size, mode="bilinear", align_corners=False)[:, 0]
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
class DriveNet(PavementNet):
|
| 61 |
+
"""PavementNet with a driving head: from the same encoder, the racing line
|
| 62 |
+
as ``n_waypoints`` points in the kart's frame (x forward, y left) and the
|
| 63 |
+
target speed at each. The mask head stays, so the stack can still check the
|
| 64 |
+
driving output against the pavement it sees.
|
| 65 |
+
|
| 66 |
+
The head reads the whole bottleneck, not a pooled vector, because where the
|
| 67 |
+
road goes is a matter of where in the image it is. Outputs are scaled to
|
| 68 |
+
about unit size: x / 25 m, y / 10 m, speed / 15 m/s.
|
| 69 |
+
"""
|
| 70 |
+
|
| 71 |
+
SCALE = (25.0, 10.0, 15.0)
|
| 72 |
+
|
| 73 |
+
def __init__(self, width=24, n_waypoints=12, grid=(9, 15), fine=False):
|
| 74 |
+
super().__init__(width)
|
| 75 |
+
self.n_waypoints = n_waypoints
|
| 76 |
+
self.squeeze = nn.Sequential(nn.Conv2d(6 * width, 32, 1), nn.BatchNorm2d(32), nn.ReLU(inplace=True))
|
| 77 |
+
self.grid = grid
|
| 78 |
+
# ``fine``: the head also reads the decoder's quarter-resolution features
|
| 79 |
+
# with two pixel-coordinate channels, then strides them down to the grid.
|
| 80 |
+
# Where the kart sits relative to the line is read from where the edges
|
| 81 |
+
# are 5 to 10 m ahead, a matter of a few pixels the 1/16 bottleneck blurs.
|
| 82 |
+
self.fine = fine
|
| 83 |
+
if fine:
|
| 84 |
+
self.fine_path = nn.Sequential(
|
| 85 |
+
nn.Conv2d(2 * width + 2, 64, 3, 2, 1, bias=False),
|
| 86 |
+
nn.BatchNorm2d(64),
|
| 87 |
+
nn.ReLU(inplace=True),
|
| 88 |
+
nn.Conv2d(64, 64, 3, 2, 1, bias=False),
|
| 89 |
+
nn.BatchNorm2d(64),
|
| 90 |
+
nn.ReLU(inplace=True),
|
| 91 |
+
nn.Conv2d(64, 32, 1),
|
| 92 |
+
nn.BatchNorm2d(32),
|
| 93 |
+
nn.ReLU(inplace=True),
|
| 94 |
+
)
|
| 95 |
+
head_in = 64 if fine else 32
|
| 96 |
+
self.drive = nn.Sequential(
|
| 97 |
+
nn.Flatten(),
|
| 98 |
+
nn.Linear(head_in * grid[0] * grid[1], 256),
|
| 99 |
+
nn.ReLU(inplace=True),
|
| 100 |
+
nn.Dropout(0.1),
|
| 101 |
+
nn.Linear(256, 3 * n_waypoints),
|
| 102 |
+
)
|
| 103 |
+
|
| 104 |
+
def forward_all(self, x):
|
| 105 |
+
"""(mask logits N x H x W, waypoints N x K x 2 in metres, speeds N x K in m/s)."""
|
| 106 |
+
size = x.shape[-2:]
|
| 107 |
+
x = x / 255.0 - 0.5
|
| 108 |
+
e1 = self.e1(x)
|
| 109 |
+
e2 = self.e2(e1)
|
| 110 |
+
e3 = self.e3(e2)
|
| 111 |
+
e4 = self.context(self.e4(e3))
|
| 112 |
+
d3 = self.d3(torch.cat([self._up(e4, e3), e3], 1))
|
| 113 |
+
d2 = self.d2(torch.cat([self._up(d3, e2), e2], 1))
|
| 114 |
+
d1 = self.d1(torch.cat([self._up(d2, e1), e1], 1))
|
| 115 |
+
mask = F.interpolate(self.head(d1), size=size, mode="bilinear", align_corners=False)[:, 0]
|
| 116 |
+
code = F.adaptive_avg_pool2d(self.squeeze(e4), self.grid)
|
| 117 |
+
if self.fine:
|
| 118 |
+
n, _, h, w = d2.shape
|
| 119 |
+
ys = torch.linspace(-1, 1, h, device=d2.device).view(1, 1, h, 1).expand(n, 1, h, w)
|
| 120 |
+
xs = torch.linspace(-1, 1, w, device=d2.device).view(1, 1, 1, w).expand(n, 1, h, w)
|
| 121 |
+
detail = self.fine_path(torch.cat([d2, ys, xs], 1))
|
| 122 |
+
code = torch.cat([code, F.adaptive_avg_pool2d(detail, self.grid)], 1)
|
| 123 |
+
out = self.drive(code).view(-1, self.n_waypoints, 3)
|
| 124 |
+
sx, sy, sv = self.SCALE
|
| 125 |
+
points = torch.stack([out[..., 0] * sx, out[..., 1] * sy], -1)
|
| 126 |
+
speeds = out[..., 2] * sv
|
| 127 |
+
return mask, points, speeds
|