example.py

"""Run a checkpoint on one image:  python example.py frame.jpg [drive-v5.pt]"""
import sys

import numpy as np
from PIL import Image

from driver import Driver

image = sys.argv[1] if len(sys.argv) > 1 else None
weights = sys.argv[2] if len(sys.argv) > 2 else "drive-v5.pt"
driver = Driver(weights)
frame = (
    np.asarray(Image.open(image).convert("RGB").resize((480, 270)))
    if image
    else np.full((270, 480, 3), 110, dtype=np.uint8)
)
mask, line, speed = driver.predict_drive(frame)
print(f"{weights}: pavement share {float((mask > 0.5).mean()):.2f}")
for k, ((x, y), v) in enumerate(zip(line, speed, strict=True)):
    print(f"  waypoint {k:2d}: {x:5.1f} m ahead, {y:+5.2f} m left, wants {v:4.1f} m/s")