bevformer-p150 / code /scripts /make_sample.py
changh95's picture
tt-model push bevformer-p150 (container)
ddaad99 verified
Raw History Blame Contribute Delete
22.2 kB
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Generate the shipped sample of bevformer-p150: ``code/tt_bevformer/samples/synthetic_6cam*`` (a synthetic
six-camera frame made by this script, Apache-2.0 like the code), its calibration preset ``calib/synthetic_6cam.json``
and the stored CPU-reference output the container smoke compares with. CPU only; run it in the research venv
(torch, PIL)::
cd bundles/bevformer-p150
PY=/home/ubuntu/experiments/tt-models/tools/research-venv/bin/python
OMP_NUM_THREADS=4 PYTHONPATH=$PWD/code $PY code/scripts/make_sample.py optimise --out <dir> # ~1-2 h
OMP_NUM_THREADS=4 PYTHONPATH=$PWD/code $PY code/scripts/make_sample.py finalise --images <dir>/images_f32.npy
Why a synthetic frame. The redistributable real data of this model (PandaSet, CC BY 4.0) may not ship until the user
approves PandaSet-derived samples (PLAN.md 6.3, research/DATASETS.md N5), nuScenes is non-commercial, and plain
synthetic images give the network little to detect: a few hallucinated boxes near the 0.2 threshold, a weak and
fragile smoke reference (bevdet-p150 measured exactly that on the p150). So the sample is a **test pattern**: six
rendered road views of a generic rig; the pixels around six chosen objects in the front and back cameras were then
optimised against the fp32 CPU reference (gradient ascent through the whole network: ResNet-101-DCN, FPN, the BEV
encoder with the spatial cross-attention, the decoder and the heads; fresh-stream semantics, no history) until the
last decoder layer reports the six objects at their places with large margins, and every other query and class
below the 0.2 publication threshold with margin. The images look like a road with noisy texture patches; they
show no real objects. What they give is a stable reference: every published box sits far from the 0.2 threshold,
so bf16 numerics cannot flip it.
- **Rig** (:data:`RIG`, round numbers of our own, nuScenes-like layout, the same as bevdet-p150's synthetic rig):
six 1600x900 pinhole cameras 1.55 m above the road, yaws 0, +-55, +-110, 180 deg, fx 1260 px (800 px for the back
camera), principal point (799.5, 480); ``base_link`` calibration with the default virtual nuScenes LIDAR_TOP.
The images are already 1600x900, so the node's resize is the identity.
- **Objects** (:data:`TARGETS`, ``base_link``): two cars, a truck, a pedestrian and a traffic cone ahead (front
camera), a bus behind (back camera), at rest.
- **Optimisation**: the 1600x900 float images of the two cameras, through a differentiable copy of the node's
pre-processing (BGR, minus the mean, x0.8 bilinear resize, pad to 736), the network in the port form (host
geometry, rebatched SCA: identical to the deployed form, PORT_LOG 3.1); each target is matched to the decoder
query whose box is closest (re-matched every step), its class logit pushed up and its centre pulled to the
target; every other logit pushed down. Gradients are confined to the targets' projected boxes (dilated) plus a
small share outside them.
- **Reference**: the ``/predict`` body of ``tt_bevformer.reference`` for the sample as a fresh stream with the
default parameters (``timing_ms`` emptied), computed from the PNG files exactly as the API reads them.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import math
import sys
import time
from pathlib import Path
from typing import Dict, List, Tuple
import numpy as np
CODE = Path(__file__).resolve().parents[1]
if str(CODE) not in sys.path:
sys.path.insert(0, str(CODE))
from tt_bevformer.reference import config as C # noqa: E402
PKG = CODE / "tt_bevformer"
NAME = "synthetic_6cam"
SRC_W, SRC_H = 1600, 900
CX, CY = 799.5, 480.0
# camera optical frame (x right, y down, z forward) -> base_link (x forward, y left, z up)
R_OPT = np.array([[0.0, 0.0, 1.0], [-1.0, 0.0, 0.0], [0.0, -1.0, 0.0]])
# name -> (x, y, z, yaw_deg, fx): camera position in base_link (metres), yaw about +z, pinhole focal length (px)
RIG: Dict[str, Tuple[float, float, float, float, float]] = {
"CAM_FRONT": (1.70, 0.00, 1.55, 0.0, 1260.0),
"CAM_FRONT_RIGHT": (1.50, -0.50, 1.55, -55.0, 1260.0),
"CAM_FRONT_LEFT": (1.50, 0.50, 1.55, 55.0, 1260.0),
"CAM_BACK": (0.05, 0.00, 1.55, 180.0, 800.0),
"CAM_BACK_LEFT": (1.05, 0.50, 1.55, 110.0, 1260.0),
"CAM_BACK_RIGHT": (1.05, -0.50, 1.55, -110.0, 1260.0),
}
# class, x, y (base_link, m), yaw (rad, base_link), (length, width, height); bottom on the road, at rest
TARGETS: List[Tuple[str, float, float, float, Tuple[float, float, float]]] = [
("car", 12.0, 0.3, 0.0, (4.6, 1.9, 1.6)),
("truck", 22.0, -4.0, 0.0, (7.6, 2.5, 3.2)),
("pedestrian", 9.0, 3.6, 1.57, (0.8, 0.7, 1.75)),
("traffic_cone", 7.5, -2.6, 0.0, (0.4, 0.4, 0.7)),
("car", 26.0, 3.5, 3.14, (4.6, 1.9, 1.6)),
("bus", -16.0, -1.5, 0.0, (11.5, 2.9, 3.4)),
]
OPT_CAMS = ("CAM_FRONT", "CAM_BACK")
# --------------------------------------------------------------------------------------------- rig and render
def camera_T(x: float, y: float, z: float, yaw_deg: float) -> np.ndarray:
"""4x4 T_base_link_from_camera_optical for a level camera at (x, y, z) looking along yaw."""
a = math.radians(yaw_deg)
rz = np.array([[math.cos(a), -math.sin(a), 0.0], [math.sin(a), math.cos(a), 0.0], [0.0, 0.0, 1.0]])
T = np.eye(4)
T[:3, :3] = rz @ R_OPT
T[:3, 3] = (x, y, z)
return T
def calibration() -> Dict:
"""The calibration envelope of the rig (``calib/synthetic_6cam.json``)."""
cams = {}
for name in C.CAMERA_ORDER:
x, y, z, yaw, f = RIG[name]
T = camera_T(x, y, z, yaw)
cams[name] = {"intrinsics": [[f, 0.0, CX], [0.0, f, CY], [0.0, 0.0, 1.0]],
"T_ref_from_camera": [[round(float(v), 12) for v in row] for row in T]}
return {"frame_id": "base_link", "cameras": cams}
def render_road(K, T, seed: int = 0, noise: float = 1.5) -> np.ndarray:
"""A camera's view of a flat world: sky gradient, a 14 m wide asphalt road along base_link x with lane markings
(solid edges, dashed separators every 3.5 m), sand-coloured shoulders, haze toward the horizon. (900, 1600, 3)
uint8 RGB (bevdet-p150's renderer)."""
rng = np.random.default_rng(seed)
u, v = np.meshgrid(np.arange(SRC_W, dtype=np.float64), np.arange(SRC_H, dtype=np.float64), indexing="xy")
d = np.stack([u, v, np.ones_like(u)], -1) @ np.linalg.inv(np.asarray(K, np.float64)).T
R, c = np.asarray(T)[:3, :3], np.asarray(T)[:3, 3]
db = d @ R.T
el = np.arctan2(db[..., 2], np.hypot(db[..., 0], db[..., 1]))
sky = np.clip(el / 0.6, 0, 1)[..., None]
img = (1 - sky) * np.array([205.0, 210.0, 214.0]) + sky * np.array([120.0, 160.0, 205.0])
ground = db[..., 2] < -1e-6
with np.errstate(divide="ignore", invalid="ignore"):
t = np.where(ground, -c[2] / np.where(ground, db[..., 2], -1.0), np.inf)
px, py = c[0] + t * db[..., 0], c[1] + t * db[..., 1]
dist = np.hypot(px, py)
g = ground & (dist < 80.0)
road = g & (np.abs(py) < 7.0)
side = g & ~road
tex = rng.normal(0, 4.0, size=(SRC_H, SRC_W, 1))
img[road] = np.array([92.0, 94.0, 96.0]) + tex[road]
img[side] = np.array([140.0, 132.0, 118.0]) + tex[side]
for yl, dashed in ((6.8, False), (-6.8, False), (0.0, True), (3.5, True), (-3.5, True)):
on = g & (np.abs(py - yl) < 0.08)
if dashed:
on &= np.mod(np.where(g, px, 0.0), 9.0) < 3.0
img[on] = np.array([225.0, 225.0, 220.0])
haze = np.clip((np.where(g, dist, 0.0) - 30.0) / 50.0, 0, 1)[..., None] * g[..., None]
img = img * (1 - 0.5 * haze) + 0.5 * haze * np.array([200.0, 205.0, 210.0])
img += rng.normal(0, noise, size=img.shape)
return np.clip(np.rint(img), 0, 255).astype(np.uint8)
def renders(cal: Dict) -> Dict[str, np.ndarray]:
return {n: render_road(cal["cameras"][n]["intrinsics"], np.array(cal["cameras"][n]["T_ref_from_camera"]), seed=i)
for i, n in enumerate(C.CAMERA_ORDER)}
def box_corners(x: float, y: float, yaw: float, l: float, w: float, h: float) -> np.ndarray:
c, s = math.cos(yaw), math.sin(yaw)
return np.array([[x + c * dx * l - s * dy * w, y + s * dx * l + c * dy * w, dz] for dz in (0.0, h)
for dx, dy in ((0.5, 0.5), (0.5, -0.5), (-0.5, -0.5), (-0.5, 0.5))])
def target_mask(cal: Dict, name: str, grow: float = 0.4) -> np.ndarray:
"""[900, 1600] float32: the pixels covered by the targets' projected 3D boxes (dilated) in one camera."""
m = np.zeros((SRC_H, SRC_W), np.float32)
K = np.asarray(cal["cameras"][name]["intrinsics"], np.float64)
Tinv = np.linalg.inv(np.asarray(cal["cameras"][name]["T_ref_from_camera"], np.float64))
for _, x, y, yaw, (l, w, h) in TARGETS:
pc = box_corners(x, y, yaw, l, w, h) @ Tinv[:3, :3].T + Tinv[:3, 3]
if (pc[:, 2] < 0.5).any():
continue
uv = pc @ K.T
uv = uv[:, :2] / uv[:, 2:3]
(u0, v0), (u1, v1) = uv.min(0), uv.max(0)
du, dv = max((u1 - u0) * grow, 24.0), max((v1 - v0) * grow, 24.0)
if u1 + du < 0 or u0 - du > SRC_W - 1 or v1 + dv < 0 or v0 - dv > SRC_H - 1:
continue
r0, r1 = int(max(0, v0 - dv)), int(min(SRC_H, v1 + dv + 1))
c0, c1 = int(max(0, u0 - du)), int(min(SRC_W, u1 + du + 1))
m[r0:r1, c0:c1] = 1.0
return m
def lidar_targets() -> List[Tuple[int, float, float, float]]:
"""(class index, x, y, z centre) of every target in the network's LIDAR_TOP frame (the output frame)."""
from tt_bevformer.host.calib import T_BASE_FROM_VIRTUAL_LIDAR
inv = np.linalg.inv(T_BASE_FROM_VIRTUAL_LIDAR)
out = []
for cls, x, y, _, (_, _, h) in TARGETS:
p = inv @ np.array([x, y, h / 2.0, 1.0])
out.append((C.CLASS_NAMES.index(cls), float(p[0]), float(p[1]), float(p[2])))
return out
def manifest_dict(files: Dict[str, str]) -> Dict:
return {
"_doc": ("bevformer-p150 sample: a synthetic six-camera test pattern generated by code/scripts/make_sample.py "
"(rendered road views whose pixels around six chosen objects in the front and back cameras were "
"optimised so that the fp32 reference reports them; no real scene). The images are 1600x900 (no "
"node resize); the calibration is the preset synthetic_6cam (base_link, the default virtual "
"LIDAR_TOP); a fresh stream at an identity ego pose. model(**tt_bevformer.load_sample(path)). "
"Generated data, Apache-2.0."),
"images": files,
"calibration": {"preset": NAME},
"stream": {"id": NAME, "reset": True,
"T_world_from_ego": [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0]]},
"targets": [{"label": c, "frame": "base_link", "center": [x, y, round(s[2] / 2, 3)], "size": list(s),
"yaw": yaw} for c, x, y, yaw, s in TARGETS],
}
# ------------------------------------------------------------------------------------------------ optimise
def optimise(a) -> int:
"""Gradient ascent on the 1600x900 images of :data:`OPT_CAMS` (module docstring)."""
import torch
import torch.nn.functional as F
from tt_bevformer.host import geometry as hgeo
from tt_bevformer.host.calib import lidar2img_stack
from tt_bevformer.host.temporal import TemporalState, can_bus_from_transform
from tt_bevformer.reference.model import BEVFormerNet, PortGeometry
from tt_bevformer.reference.weights import BEVFormerWeights
from tt_bevformer.tests import paths
torch.manual_seed(0)
torch.set_num_threads(a.threads)
out = Path(a.out)
out.mkdir(parents=True, exist_ok=True)
cal = calibration()
imgs = renders(cal)
from tt_bevformer.host.calib import T_BASE_FROM_VIRTUAL_LIDAR
Ts = [np.asarray(cal["cameras"][n]["T_ref_from_camera"]) for n in C.CAMERA_ORDER]
Ks = [np.asarray(cal["cameras"][n]["intrinsics"]) for n in C.CAMERA_ORDER]
l2i = lidar2img_stack(Ts, Ks, [(SRC_H, SRC_W)] * 6, T_ref_from_lidar=T_BASE_FROM_VIRTUAL_LIDAR)
cb, use = TemporalState().step(can_bus_from_transform(np.eye(4), {}), reset=True)
geo_np = hgeo.frame_geometry(l2i, cb)
print("sca counts", geo_np["sca_counts"].tolist(), flush=True)
geo = PortGeometry.from_numpy(geo_np)
wfile = Path(a.weights) if a.weights else paths.weights_file()
net = BEVFormerNet(BEVFormerWeights(wfile))
net.taps = None
mean_bgr = torch.tensor(C.MEAN_BGR, dtype=torch.float32).view(1, 3, 1, 1)
hh, ww = int(SRC_H * C.SCALE), int(SRC_W * C.SCALE)
def prep(rgb): # [n, 900, 1600, 3] RGB float -> [n, 3, 736, 1280]
bgr = rgb.flip(-1).permute(0, 3, 1, 2) - mean_bgr
r = F.interpolate(bgr, size=(hh, ww), mode="bilinear", align_corners=False)
return F.pad(r, (0, C.IMG_HW[1] - ww, 0, C.IMG_HW[0] - hh))
fixed = [n for n in C.CAMERA_ORDER if n not in OPT_CAMS]
with torch.no_grad():
base = {n: torch.as_tensor(imgs[n].astype(np.float32)) for n in C.CAMERA_ORDER}
feat_fixed = net.neck(net.backbone(prep(torch.stack([base[n] for n in fixed]))))
x0 = torch.stack([base[n] for n in OPT_CAMS])
x = (torch.as_tensor(np.load(a.resume).astype(np.float32)) if a.resume else x0.clone()).requires_grad_(True)
masks = torch.as_tensor(np.stack([target_mask(cal, n) for n in OPT_CAMS]))[..., None]
opt = torch.optim.Adam([x], lr=a.lr)
tg = lidar_targets()
cbt = torch.as_tensor(cb.astype(np.float32))
use_t = torch.as_tensor(use.astype(np.float32))
prev = torch.zeros(C.NUM_QUERY_BEV, C.EMBED)
def heads(xc):
f_opt = net.neck(net.backbone(prep(xc)))
order = {n: i for i, n in enumerate(OPT_CAMS)}
fi = {n: i for i, n in enumerate(fixed)}
feat = torch.stack([f_opt[order[n]] if n in order else feat_fixed[fi[n]] for n in C.CAMERA_ORDER])
value = net.encoder_value(feat)
bev_q = net.bev_queries(cbt)
bev = net.encoder(bev_q, value, prev, use_t, geo.shift, lambda i, q: net.sca_rebatched(i, q, value, geo))
cls, crd = net.decode(bev)
return cls[-1, 0], crd[-1, 0]
hist = []
for it in range(a.iters):
t0 = time.time()
opt.zero_grad()
cls, crd = heads(x)
with torch.no_grad(): # greedy matching: the most confident box within ~1.5 m;
taken, match = set(), [] # sticky: a confident match within 2 m is kept
for j, (k, tx, ty, _) in enumerate(tg):
prev_q = hist[-1]["queries"][j] if hist else None
if prev_q is not None and prev_q not in taken and float(torch.sigmoid(cls[prev_q, k])) > 0.3 \
and abs(float(crd[prev_q, 0]) - tx) + abs(float(crd[prev_q, 1]) - ty) < 2.0:
taken.add(prev_q)
match.append(prev_q)
continue
d = ((crd[:, 0] - tx).abs() + (crd[:, 1] - ty).abs() - 1.5).clamp(min=0.0) \
- 5.0 * torch.sigmoid(cls[:, k])
for q in torch.argsort(d).tolist():
if q not in taken:
taken.add(q)
match.append(q)
break
loss_t, loss_c, scores = 0.0, 0.0, []
for (k, tx, ty, tz), q in zip(tg, match):
lg = cls[q, k]
others = torch.cat([cls[q, :k], cls[q, k + 1:]])
loss_t = loss_t + F.softplus(4.0 - lg) + F.softplus(others.max() - lg + 4.0)
loss_c = loss_c + F.smooth_l1_loss(crd[q, 0:2], torch.tensor([tx, ty])) \
+ 0.2 * (crd[q, 8:10] ** 2).sum()
scores.append(round(float(torch.sigmoid(lg.detach())), 3))
keep = torch.ones(C.NUM_QUERY, dtype=torch.bool)
keep[match] = False
bg = cls[keep]
loss_bg = 4.0 * F.softplus(bg + 4.0).mean() + F.softplus(bg + 3.0).max()
tv = ((x[:, 1:] - x[:, :-1]).abs().mean() + (x[:, :, 1:] - x[:, :, :-1]).abs().mean()) / 255.0
dev = (((x - x0) / 255.0) ** 2 * (1.0 - masks)).mean()
loss = loss_t + loss_c + loss_bg + a.tv * tv + 20.0 * dev
loss.backward()
with torch.no_grad():
x.grad.mul_(masks + a.outside * (1.0 - masks))
opt.step()
x.clamp_(0.0, 255.0)
bmax = float(torch.sigmoid(bg.detach().max()))
hist.append({"it": it, "loss": round(float(loss.detach()), 4), "targets": round(float(loss_t.detach()), 4),
"centre": round(float(loss_c.detach()), 4), "bg_max": round(bmax, 4), "scores": scores,
"queries": match, "s": round(time.time() - t0, 1)})
print(json.dumps(hist[-1]), flush=True)
if it % 5 == 4 or it == a.iters - 1:
np.save(out / "images_f32.npy", x.detach().numpy())
(out / "history.json").write_text(json.dumps({"args": {k: str(v) for k, v in vars(a).items()},
"history": hist}, indent=0))
if min(scores) > a.stop_score and bmax < a.stop_bg:
print(f"stop: every target > {a.stop_score}, background max {bmax:.3f} < {a.stop_bg}", flush=True)
np.save(out / "images_f32.npy", x.detach().numpy())
break
(out / "history.json").write_text(json.dumps({"args": {k: str(v) for k, v in vars(a).items()},
"history": hist}, indent=0))
return 0
# ------------------------------------------------------------------------------------------------ finalise
def sha256(data: bytes) -> str:
return hashlib.sha256(data).hexdigest()
def write_preset() -> Path:
"""``calib/synthetic_6cam.json`` (the rig of :func:`calibration` with its provenance note)."""
cal = calibration()
calib_dir = PKG / "calib"
preset = {"_doc": ("Calibration preset of the shipped synthetic sample (samples/synthetic_6cam): a generic "
"six-camera rig with round numbers of our own (code/scripts/make_sample.py RIG, the same rig "
"as bevdet-p150's synthetic sample): 1600x900 pinhole cameras 1.55 m above the road, yaws 0 / "
"+-55 / +-110 / 180 deg, fx 1260 px (back camera 800 px), principal point (799.5, 480). "
"intrinsics = K of the image as sent; T_ref_from_camera = camera optical frame -> base_link "
"(the network's LIDAR_TOP is the default virtual one). Generated data, Apache-2.0."),
**cal}
(calib_dir / f"{NAME}.json").write_text(json.dumps(preset, indent=1) + "\n")
return calib_dir / f"{NAME}.json"
def finalise(a) -> int:
from PIL import Image
from tt_bevformer.host import postprocess as hpost
from tt_bevformer.host.inputs import load_frame
from tt_bevformer.reference.pipeline import BEVFormerReference
from tt_bevformer.tests import paths
cal = calibration()
base = renders(cal)
opt_imgs = np.clip(np.rint(np.load(a.images)), 0, 255).astype(np.uint8)
samples, calib_dir = PKG / "samples", PKG / "calib"
img_dir = samples / NAME
img_dir.mkdir(parents=True, exist_ok=True)
files = {}
for name in C.CAMERA_ORDER:
im = opt_imgs[OPT_CAMS.index(name)] if name in OPT_CAMS else base[name]
Image.fromarray(im).save(img_dir / f"{name}.png", format="PNG", optimize=True)
files[name] = f"{NAME}/{name}.png"
write_preset()
man = manifest_dict(files)
man_path = samples / f"{NAME}.json"
man_path.write_text(json.dumps(man, indent=1) + "\n")
import torch
torch.set_num_threads(a.threads)
wfile = Path(a.weights) if a.weights else paths.weights_file()
ref = BEVFormerReference(wfile)
frame = load_frame(man_path)
res = ref.step(frame, reset=True, port=False, workers=a.threads)
body = hpost.to_output(res.detections, model="bevformer-p150").to_dict()
body["timing_ms"] = {}
(samples / f"{NAME}.reference.json").write_text(json.dumps(body, indent=1) + "\n")
# margins: the published boxes and the strongest unpublished logit
cls = res.outputs_classes[-1, 0]
s = 1.0 / (1.0 + np.exp(-cls))
top = np.sort(s.reshape(-1))[::-1]
pub = body["detections"]
print(f"{NAME}: {len(pub)} detections >= 0.2; strongest scores {np.round(top[:len(pub) + 3], 3).tolist()}")
for d in pub:
print(" ", d["label"], d["label_id"], d["score"], d["center"], d["size"], d["yaw"])
sizes = {p.name: p.stat().st_size for p in sorted(img_dir.glob("*.png"))}
print("png bytes", sizes, sum(sizes.values()))
return 0
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
sub = ap.add_subparsers(dest="cmd", required=True)
o = sub.add_parser("optimise", help="gradient ascent on the front / back images -> <out>/images_f32.npy")
o.add_argument("--out", required=True)
o.add_argument("--iters", type=int, default=80)
o.add_argument("--lr", type=float, default=2.0, help="Adam step in uint8 pixel units")
o.add_argument("--outside", type=float, default=0.1, help="gradient scale outside the target masks")
o.add_argument("--tv", type=float, default=2.0, help="total-variation weight")
o.add_argument("--stop-score", type=float, default=0.75)
o.add_argument("--stop-bg", type=float, default=0.06)
o.add_argument("--resume", help="start from an images_f32.npy")
f = sub.add_parser("finalise", help="images -> sample PNGs, preset, manifest, CPU reference")
f.add_argument("--images", required=True)
for p in (o, f):
p.add_argument("--weights", help="weights file (default: tests/paths.weights_file())")
p.add_argument("--threads", type=int, default=4)
sub.add_parser("preset", help="write calib/synthetic_6cam.json only")
a = ap.parse_args()
if a.cmd == "preset":
print(write_preset())
return 0
return optimise(a) if a.cmd == "optimise" else finalise(a)
if __name__ == "__main__":
sys.exit(main())