ai/pothole_server.py

#!/usr/bin/env python
"""Persistent YOLOv8 pothole classifier worker (real-time).

Loads the model ONCE at startup, then serves one request per line from stdin
and writes one JSON result per line to stdout. This avoids the multi-second
cold start of loading torch + the model on every request.

Protocol:
  startup -> emits {"ready": true}  (or {"error": "..."} and exits)
  stdin   -> one image file path per line ("__quit__" to stop)
  stdout  -> one JSON result per line:
             {"isPothole":bool,"confidence":float,"size":"Small|Medium|Large"|null,
              "count":int,"largestAreaFraction":float,"boxes":[...],"model":"..."}
             or {"error":"..."} for that request.
"""
import sys
import os
import json

# Size = area (as a fraction of the whole frame) of the largest *confident* pothole
# detection. A single 2D photo has no true scale, so box area is a heuristic proxy.
#   >= SIZE_LARGE  -> Large
#   >= SIZE_MEDIUM -> Medium
#   else           -> Small
# Thresholds are calibrated to the ground-truth pothole box-area distribution
# (3,625 annotated boxes): SIZE_MEDIUM ~= 40th percentile, SIZE_LARGE ~= 82nd
# percentile. That yields a sensible operational split of roughly 40% Small /
# 42% Medium / 18% Large instead of over-calling everything "Small".
SIZE_LARGE = 0.09
SIZE_MEDIUM = 0.013
# CONF_THRESHOLD: minimum confidence for a box to count as a detection at all.
CONF_THRESHOLD = 0.25
# SIZE_CONF_FLOOR: a box must clear this higher bar before its area is allowed to
# drive the size label. This stops low-confidence, sprawling false boxes (e.g. a
# shadow spanning the frame) from inflating a real, smaller pothole to "Large".
SIZE_CONF_FLOOR = 0.45
MODEL_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "pothole-yolov8.pt")


def emit(obj):
    sys.stdout.write(json.dumps(obj) + "\n")
    sys.stdout.flush()


def classify(model, image_path):
    if not os.path.exists(image_path):
        return {"error": "image not found: " + image_path}
    try:
        result = model.predict(image_path, conf=CONF_THRESHOLD, verbose=False)[0]
        h, w = result.orig_shape
        frame_area = float(h * w) if h and w else 1.0

        boxes = []
        max_frac = 0.0          # largest box area (any confidence) — reported for transparency
        max_conf = 0.0          # best detection confidence
        size_frac = 0.0         # area that actually drives the size label (confident boxes only)
        best_conf_frac = 0.0    # area of the single highest-confidence box (fallback)
        best_conf = -1.0
        for b in result.boxes:
            x1, y1, x2, y2 = (float(v) for v in b.xyxy[0].tolist())
            conf = float(b.conf[0])
            frac = ((x2 - x1) * (y2 - y1)) / frame_area
            boxes.append({
                "x1": round(x1), "y1": round(y1), "x2": round(x2), "y2": round(y2),
                "conf": round(conf, 3), "areaFraction": round(frac, 4),
            })
            max_frac = max(max_frac, frac)
            max_conf = max(max_conf, conf)
            if conf >= SIZE_CONF_FLOOR:
                size_frac = max(size_frac, frac)
            if conf > best_conf:
                best_conf = conf
                best_conf_frac = frac

        is_pothole = len(boxes) > 0
        # Prefer the largest confident pothole; if none clear the floor, fall back to
        # the area of the single most-confident detection (never a noisy wide box).
        sizing_frac = size_frac if size_frac > 0 else best_conf_frac
        if not is_pothole:
            size = None
        elif sizing_frac >= SIZE_LARGE:
            size = "Large"
        elif sizing_frac >= SIZE_MEDIUM:
            size = "Medium"
        else:
            size = "Small"

        return {
            "isPothole": is_pothole,
            "confidence": round(max_conf, 3),
            "size": size,
            "count": len(boxes),
            "largestAreaFraction": round(max_frac, 4),
            "sizeAreaFraction": round(sizing_frac, 4),
            "boxes": boxes,
            "model": "peterhdd/pothole-detection-yolov8",
        }
    except Exception as e:  # noqa: BLE001
        return {"error": str(e)}


def main():
    if not os.path.exists(MODEL_FILE):
        emit({"error": "model file missing: " + MODEL_FILE})
        return
    try:
        from ultralytics import YOLO
    except Exception as e:  # noqa: BLE001
        emit({"error": "ultralytics not installed: " + str(e)})
        return
    try:
        model = YOLO(MODEL_FILE)
    except Exception as e:  # noqa: BLE001
        emit({"error": "model load failed: " + str(e)})
        return

    emit({"ready": True})

    for line in sys.stdin:
        path = line.strip()
        if not path:
            continue
        if path == "__quit__":
            break
        emit(classify(model, path))


if __name__ == "__main__":
    main()