Files
NODEDC_MISSION_CORE/simulation/ai-polygon/ddrnet_server.py
T

84 lines
3.0 KiB
Python

"""Resident reference DDRNet adapter; run only inside the pinned module image.
Raw pinhole RGB is the only input. No device mask, scene labels or actor truth.
The container publishes this port on Windows loopback only and owns no weights.
"""
import hashlib
import importlib.util
import json
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
import numpy as np
from PIL import Image
def checked(path, expected):
value = Path(path)
if hashlib.sha256(value.read_bytes()).hexdigest() != expected:
raise RuntimeError("Pinned DDRNet asset changed: " + value.name)
return value
def main():
checkpoint = checked(
"/assets/ddrnet-checkpoint.pth",
"b99c2838051bcd7b092fd3970aa62a77d5c0bbb809c9b9afb2ff4b0ebdaa4ee6",
)
runner = checked(
"/assets/ddrnet-goose-runner.py",
"b18ad60f277eea69a240a28f290611b94627fb9707faf1bb3e6e22102dad67c1",
)
spec = importlib.util.spec_from_file_location("polygon_pinned_goose", runner)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
model, _, _ = module.load_model("ddrnet", checkpoint)
tensor, _ = module.preprocess(Image.fromarray(np.zeros((600, 800, 3), np.uint8)))
module.infer(model, tensor)
class Handler(BaseHTTPRequestHandler):
def setup(self):
super().setup()
self.connection.settimeout(10)
def reply(self, status, body, kind="application/octet-stream"):
self.send_response(status)
self.send_header("Content-Type", kind)
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def do_GET(self):
if self.path != "/ready":
self.reply(404, b"")
return
self.reply(200, json.dumps({"model": "ddrnet-goose-pytorch-reference"}).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")
rgb = np.frombuffer(raw, np.uint8).reshape(600, 800, 3)
tensor, _ = module.preprocess(Image.fromarray(rgb))
mask, _ = module.infer(model, tensor)
if mask.shape != (512, 512) or np.any(mask < 0) or np.any(mask >= 64):
raise ValueError("DDRNet output contract changed")
self.reply(200, mask.astype(np.uint8).tobytes())
except (TimeoutError, ValueError, RuntimeError):
self.reply(500, b"DDRNet inference failed")
def log_message(self, *_):
pass
# External exposure is restricted by the Compose loopback publication.
HTTPServer(("0.0.0.0", 8010), Handler).serve_forever()
if __name__ == "__main__":
main()