88 lines
3.3 KiB
Python
88 lines
3.3 KiB
Python
"""Resident, pinned ADE20K surface provider. RGB only; no scene truth input."""
|
|
|
|
import hashlib
|
|
import json
|
|
from http.server import BaseHTTPRequestHandler, HTTPServer
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
|
|
|
|
# These are surface candidates, not permission to drive. The metric terrain
|
|
# planner still checks step height, slope and the full rover footprint.
|
|
SURFACES = (3, 6, 9, 11, 13, 29, 34, 46, 52, 91)
|
|
FILES = {
|
|
"config.json": "ee7400840fdb1e5045f0b2eba78bf053df8e33a309c4acec31705a48c8cc5c00",
|
|
"preprocessor_config.json": "8039d1d210abaa7117ad78e58cdfd6141a2ec72c03dae891b3cd76737e422c6c",
|
|
"pytorch_model.bin": "187ca07bea003a5717c63d04ea90b07f33cd033c0ebf44b4b89fce5070d6c8f3",
|
|
}
|
|
|
|
|
|
def main():
|
|
for name, expected in FILES.items():
|
|
if hashlib.sha256((Path("/assets") / name).read_bytes()).hexdigest() != expected:
|
|
raise RuntimeError("SegFormer asset identity changed: " + name)
|
|
processor = SegformerImageProcessor.from_pretrained("/assets", local_files_only=True)
|
|
model = (
|
|
SegformerForSemanticSegmentation.from_pretrained("/assets", local_files_only=True)
|
|
.cuda()
|
|
.eval()
|
|
)
|
|
torch.set_num_threads(2)
|
|
|
|
def infer(rgb):
|
|
image = Image.fromarray(rgb).crop((100, 0, 700, 600))
|
|
tensor = processor(images=image, return_tensors="pt")["pixel_values"].cuda()
|
|
with torch.inference_mode():
|
|
logits = model(tensor).logits
|
|
logits = torch.nn.functional.interpolate(
|
|
logits, size=(512, 512), mode="bilinear", align_corners=False
|
|
)
|
|
confidence, labels = logits.softmax(1).max(1)
|
|
labels = labels[0].cpu().numpy().astype(np.uint8)
|
|
candidate = np.isin(labels, SURFACES) & (confidence[0].cpu().numpy() >= 0.55)
|
|
return labels.tobytes() + candidate.astype(np.uint8).tobytes()
|
|
|
|
infer(np.zeros((600, 800, 3), dtype=np.uint8))
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
def setup(self):
|
|
super().setup()
|
|
self.connection.settimeout(10)
|
|
|
|
def reply(self, status, body):
|
|
self.send_response(status)
|
|
self.send_header("Content-Type", "application/octet-stream")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def do_GET(self):
|
|
self.reply(
|
|
200 if self.path == "/ready" else 404,
|
|
json.dumps({"model": "segformer-b2-ade150"}).encode(),
|
|
)
|
|
|
|
def do_POST(self):
|
|
if self.path != "/infer" or self.headers.get("Content-Length") != "1440000":
|
|
self.reply(400, b"Expected 800x600 raw RGB uint8")
|
|
return
|
|
try:
|
|
raw = self.rfile.read(1440000)
|
|
if len(raw) != 1440000:
|
|
raise ValueError("Incomplete camera frame")
|
|
self.reply(200, infer(np.frombuffer(raw, np.uint8).reshape(600, 800, 3)))
|
|
except (TimeoutError, ValueError, RuntimeError):
|
|
self.reply(500, b"Surface inference failed")
|
|
|
|
def log_message(self, *_):
|
|
pass
|
|
|
|
HTTPServer(("0.0.0.0", 8010), Handler).serve_forever()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|