frontend/src/workers/fleetInference.worker.ts

import * as ort from "onnxruntime-web";
import type { FleetDetection, FleetHazardType } from "../lib/fleetTelemetry";

const HAZARD_CLASSES: FleetHazardType[] = [
  "pothole",
  "crack",
  "water_logging",
  "construction_zone",
  "traffic_block",
];

type WorkerRequest =
  | { type: "init"; modelUrl: string }
  | { type: "infer"; width: number; height: number; buffer: ArrayBuffer };

type WorkerResponse =
  | { type: "ready" }
  | { type: "result"; detections: FleetDetection[]; inferMs: number }
  | { type: "error"; message: string };

let session: ort.InferenceSession | null = null;

ort.env.wasm.numThreads = 1;

function post(message: WorkerResponse): void {
  self.postMessage(message);
}

async function init(modelUrl: string): Promise<void> {
  if (session) {
    post({ type: "ready" });
    return;
  }
  session = await ort.InferenceSession.create(modelUrl, {
    executionProviders: ["wasm"],
    graphOptimizationLevel: "all",
  });
  post({ type: "ready" });
}

function preprocess(buffer: ArrayBuffer, width: number, height: number): ort.Tensor {
  const pixels = new Uint8ClampedArray(buffer);
  const tensor = new Float32Array(3 * width * height);
  const plane = width * height;

  for (let i = 0; i < plane; i += 1) {
    const src = i * 4;
    tensor[i] = pixels[src] / 255;
    tensor[plane + i] = pixels[src + 1] / 255;
    tensor[plane * 2 + i] = pixels[src + 2] / 255;
  }
  return new ort.Tensor("float32", tensor, [1, 3, height, width]);
}

function clipBox(bbox: [number, number, number, number], size: number): [number, number, number, number] {
  const [x1, y1, x2, y2] = bbox;
  return [
    Math.max(0, Math.min(size, x1)),
    Math.max(0, Math.min(size, y1)),
    Math.max(0, Math.min(size, x2)),
    Math.max(0, Math.min(size, y2)),
  ];
}

function detectionFromRow(row: ArrayLike<number>, offset: number, stride: number, size: number): FleetDetection | null {
  const confidence = Number(row[offset + 4]);
  const classId = Math.round(Number(row[offset + 5]));
  const hazardType = HAZARD_CLASSES[classId];
  if (!hazardType || !Number.isFinite(confidence)) return null;

  const bbox: [number, number, number, number] = clipBox(
    [
      Number(row[offset]),
      Number(row[offset + 1]),
      Number(row[offset + 2]),
      Number(row[offset + 3]),
    ],
    size,
  );
  if (bbox[2] <= bbox[0] || bbox[3] <= bbox[1] || stride < 6) return null;
  return { hazard_type: hazardType, confidence, bbox };
}

function parseNmsOutput(output: ort.Tensor, size: number): FleetDetection[] | null {
  const data = output.data as Float32Array;
  const dims = output.dims;
  let rows = 0;
  let stride = 0;

  if (dims.length === 2 && dims[1] >= 6) {
    rows = dims[0];
    stride = dims[1];
  } else if (dims.length === 3 && dims[2] >= 6) {
    rows = dims[1];
    stride = dims[2];
  }
  if (!rows || !stride) return null;

  const detections: FleetDetection[] = [];
  for (let r = 0; r < rows; r += 1) {
    const det = detectionFromRow(data, r * stride, stride, size);
    if (det && det.confidence > 0.25) detections.push(det);
  }
  return detections;
}

function parseRawYoloOutput(output: ort.Tensor, size: number): FleetDetection[] {
  const data = output.data as Float32Array;
  const dims = output.dims;
  if (dims.length !== 3 || dims[1] < 9) return [];

  const featureCount = dims[1];
  const boxCount = dims[2];
  const detections: FleetDetection[] = [];
  for (let i = 0; i < boxCount; i += 1) {
    let bestClass = -1;
    let bestScore = 0;
    for (let c = 0; c < HAZARD_CLASSES.length; c += 1) {
      const score = data[(4 + c) * boxCount + i];
      if (score > bestScore) {
        bestScore = score;
        bestClass = c;
      }
    }
    const hazardType = HAZARD_CLASSES[bestClass];
    if (!hazardType || bestScore <= 0.25) continue;

    const cx = data[i];
    const cy = data[boxCount + i];
    const w = data[boxCount * 2 + i];
    const h = data[boxCount * 3 + i];
    if (!Number.isFinite(cx + cy + w + h) || featureCount < 9) continue;
    const bbox = clipBox([cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2], size);
    detections.push({ hazard_type: hazardType, confidence: bestScore, bbox });
  }
  return detections.sort((a, b) => b.confidence - a.confidence).slice(0, 12);
}

function parseDetections(results: Record<string, ort.Tensor>, size: number): FleetDetection[] {
  const output = Object.values(results)[0];
  if (!output) return [];
  return parseNmsOutput(output, size) ?? parseRawYoloOutput(output, size);
}

async function infer(width: number, height: number, buffer: ArrayBuffer): Promise<void> {
  if (!session) throw new Error("Inference worker is not initialized.");
  const started = performance.now();
  const tensor = preprocess(buffer, width, height);
  const inputName = session.inputNames[0];
  const results = (await session.run({ [inputName]: tensor })) as Record<string, ort.Tensor>;
  post({
    type: "result",
    detections: parseDetections(results, width),
    inferMs: Math.round(performance.now() - started),
  });
}

self.addEventListener("message", (event: MessageEvent<WorkerRequest>) => {
  const msg = event.data;
  const task = msg.type === "init" ? init(msg.modelUrl) : infer(msg.width, msg.height, msg.buffer);
  task.catch((err: unknown) => {
    post({ type: "error", message: err instanceof Error ? err.message : "Fleet worker failed." });
  });
});