script.js

const MODEL_PATH = "https://huggingface.co/fique5/watermeter/resolve/main/best.onnx";
const INPUT_SIZE = 640;
const CLASS_NAMES = [
  "meter",
  "window",
  "0",
  "1",
  "2",
  "3",
  "4",
  "5",
  "6",
  "7",
  "8",
  "9",
  "u"
];
const SCORE_THRESHOLD = 0.16;
const IOU_THRESHOLD = 0.45;
const MAX_BOXES = 200;

let session = null;
let selectedImage = null;

const canvas = document.getElementById("canvas");
const ctx = canvas.getContext("2d");
const imageInput = document.getElementById("imageInput");
const browseBtn = document.getElementById("browseBtn");
const dropZone = document.getElementById("dropZone");
const predictBtn = document.getElementById("predictBtn");
const modelStatus = document.getElementById("modelStatus");
const meterReading = document.getElementById("meterReading");
const confidence = document.getElementById("confidence");
const detectionsList = document.getElementById("detectionsList");
const imageInfo = document.getElementById("imageInfo");

async function loadModel() {
  modelStatus.textContent = "Loading model...";
  modelStatus.className = "badge loading";

  try {
    session = await ort.InferenceSession.create(MODEL_PATH, {
      executionProviders: ["wasm"]
    });

    modelStatus.textContent = "Model ready";
    modelStatus.className = "badge ready";
  } catch (error) {
    console.error("Model load error:", error);
    modelStatus.textContent = "Model failed";
    modelStatus.className = "badge error";
    detectionsList.innerHTML =
      "<p class=\"small-text\">Unable to load the ONNX model. Make sure best.onnx is public and accessible.</p>";
  }
}

function setDropZoneState(active) {
  dropZone.classList.toggle("drag-over", active);
}

browseBtn.addEventListener("click", () => imageInput.click());
imageInput.addEventListener("change", (event) => {
  const file = event.target.files[0];
  if (file) {
    loadImage(file);
  }
});

dropZone.addEventListener("dragover", (event) => {
  event.preventDefault();
  setDropZoneState(true);
});

dropZone.addEventListener("dragleave", () => setDropZoneState(false));

dropZone.addEventListener("drop", (event) => {
  event.preventDefault();
  setDropZoneState(false);
  const file = event.dataTransfer.files[0];
  if (file) {
    loadImage(file);
  }
});

window.addEventListener("dragover", (event) => event.preventDefault());
window.addEventListener("drop", (event) => event.preventDefault());

function loadImage(file) {
  const img = new Image();
  img.onload = () => {
    selectedImage = img;
    canvas.width = img.width;
    canvas.height = img.height;
    ctx.clearRect(0, 0, canvas.width, canvas.height);
    ctx.drawImage(img, 0, 0);
    imageInfo.textContent = `${img.width}px × ${img.height}px`;
    meterReading.textContent = "--";
    confidence.textContent = "--";
    detectionsList.innerHTML =
      "<p class=\"small-text\">Ready to analyze. Press Analyze Image.</p>";
    URL.revokeObjectURL(img.src);
  };
  img.src = URL.createObjectURL(file);
}

predictBtn.addEventListener("click", runPrediction);

async function runPrediction() {
  if (!selectedImage) {
    alert("Please upload an image first.");
    return;
  }
  if (!session) {
    alert("Model is not ready yet. Wait until the model finishes loading.");
    return;
  }

  modelStatus.textContent = "Running inference...";
  modelStatus.className = "badge loading";

  try {
    const prediction = await predictImage(selectedImage);
    drawDetectionResults(prediction.detections);
    updatePredictionUI(prediction);
    if (!prediction.detections.length) {
      detectionsList.innerHTML =
        "<p class=\"small-text\">No detections found. Try a different photo or use a clearer meter image.</p>";
    }
    modelStatus.textContent = "Ready";
    modelStatus.className = "badge ready";
  } catch (error) {
    console.error("Inference error:", error);
    modelStatus.textContent = "Inference failed";
    modelStatus.className = "badge error";
    detectionsList.innerHTML =
      `<p class="small-text">Inference failed: ${error.message}</p>`;
  }
}

async function predictImage(image) {
  const {tensor, ratio, pad, originalWidth, originalHeight} = prepareInput(image);
  const inputName = session.inputNames[0];
  const feeds = {
    [inputName]: new ort.Tensor("float32", tensor, [1, 3, INPUT_SIZE, INPUT_SIZE])
  };

  const results = await session.run(feeds);
  const outputName = session.outputNames[0];
  const rawOutput = results[outputName];

  if (!rawOutput) {
    throw new Error("Model did not return an output tensor.");
  }

  console.log("Model output names:", session.outputNames);
  console.log("Raw model dims:", rawOutput.dims);

  const normalized = normalizeOutput(rawOutput);
  const detections = decodeOutput(normalized, ratio, pad, originalWidth, originalHeight);

  return {
    detections,
    reading: extractMeterReading(detections),
    averageConfidence: computeAverageConfidence(detections)
  };
}

function normalizeOutput(output) {
  const expectedAttrs = 4 + CLASS_NAMES.length;
  let dims = Array.from(output.dims);

  while (dims.length > 3 && dims.some((d) => d === 1)) {
    const idx = dims.findIndex((d) => d === 1);
    dims.splice(idx, 1);
  }

  if (dims.length === 2) {
    dims = [1, dims[0], dims[1]];
  }

  if (dims.length !== 3) {
    throw new Error(`Unsupported output tensor shape: ${output.dims.join("x")}`);
  }

  if (dims[0] === 1 && dims[1] === expectedAttrs) {
    return {data: output.data, dims, layout: "chw", boxes: dims[2]};
  }

  if (dims[0] === 1 && dims[2] === expectedAttrs) {
    return {data: output.data, dims: [dims[0], dims[2], dims[1]], layout: "hwc", boxes: dims[1]};
  }

  throw new Error(`Unsupported ONNX output layout. Expected attrs=${expectedAttrs}, got: ${dims.join("x")}`);
}

function prepareInput(image) {
  const letterbox = letterboxImage(image, INPUT_SIZE);
  const imageData = letterbox.imageData;
  const floatArray = new Float32Array(1 * 3 * INPUT_SIZE * INPUT_SIZE);

  for (let y = 0; y < INPUT_SIZE; y++) {
    for (let x = 0; x < INPUT_SIZE; x++) {
      const idx = (y * INPUT_SIZE + x) * 4;
      floatArray[y * INPUT_SIZE + x] = imageData.data[idx] / 255;
      floatArray[INPUT_SIZE * INPUT_SIZE + y * INPUT_SIZE + x] = imageData.data[idx + 1] / 255;
      floatArray[2 * INPUT_SIZE * INPUT_SIZE + y * INPUT_SIZE + x] = imageData.data[idx + 2] / 255;
    }
  }

  return {
    tensor: floatArray,
    ratio: letterbox.ratio,
    pad: letterbox.pad,
    originalWidth: image.width,
    originalHeight: image.height
  };
}

function letterboxImage(image, size) {
  const offscreen = document.createElement("canvas");
  offscreen.width = size;
  offscreen.height = size;
  const ctxOff = offscreen.getContext("2d");
  ctxOff.fillStyle = "#000";
  ctxOff.fillRect(0, 0, size, size);

  const ratio = Math.min(size / image.width, size / image.height);
  const newWidth = Math.round(image.width * ratio);
  const newHeight = Math.round(image.height * ratio);
  const padX = Math.round((size - newWidth) / 2);
  const padY = Math.round((size - newHeight) / 2);

  ctxOff.drawImage(image, 0, 0, image.width, image.height, padX, padY, newWidth, newHeight);

  return {
    imageData: ctxOff.getImageData(0, 0, size, size),
    ratio,
    pad: { x: padX, y: padY }
  };
}

function decodeOutput(normalized, ratio, pad, originalWidth, originalHeight) {
  const {data, dims, layout, boxes} = normalized;
  const attributes = dims[1];
  const classCount = attributes - 4;
  const detections = [];

  for (let i = 0; i < boxes; i++) {
    const x = layout === "chw" ? data[0 * boxes + i] : data[i * attributes + 0];
    const y = layout === "chw" ? data[1 * boxes + i] : data[i * attributes + 1];
    const w = layout === "chw" ? data[2 * boxes + i] : data[i * attributes + 2];
    const h = layout === "chw" ? data[3 * boxes + i] : data[i * attributes + 3];

    let bestClass = -1;
    let bestScore = 0;

    for (let c = 0; c < classCount; c++) {
      const classScore = layout === "chw" ? data[(4 + c) * boxes + i] : data[i * attributes + 4 + c];
      if (classScore > bestScore) {
        bestScore = classScore;
        bestClass = c;
      }
    }

    if (bestScore < SCORE_THRESHOLD || bestClass < 0) {
      continue;
    }

    const x1 = (x - w / 2 - pad.x) / ratio;
    const y1 = (y - h / 2 - pad.y) / ratio;
    const x2 = (x + w / 2 - pad.x) / ratio;
    const y2 = (y + h / 2 - pad.y) / ratio;
    const label = bestClass < CLASS_NAMES.length ? CLASS_NAMES[bestClass] : `class_${bestClass}`;

    detections.push({
      classIndex: bestClass,
      label,
      score: bestScore,
      x1: clamp(x1, 0, originalWidth),
      y1: clamp(y1, 0, originalHeight),
      x2: clamp(x2, 0, originalWidth),
      y2: clamp(y2, 0, originalHeight)
    });
  }

  return nonMaxSuppression(detections, IOU_THRESHOLD, MAX_BOXES);
}

function nonMaxSuppression(detections, iouThreshold, maxBoxes) {
  const results = [];
  const sorted = detections.sort((a, b) => b.score - a.score);

  while (sorted.length && results.length < maxBoxes) {
    const current = sorted.shift();
    results.push(current);

    for (let i = sorted.length - 1; i >= 0; i--) {
      if (current.classIndex !== sorted[i].classIndex) {
        continue;
      }
      if (intersectionOverUnion(current, sorted[i]) > iouThreshold) {
        sorted.splice(i, 1);
      }
    }
  }

  return results;
}

function intersectionOverUnion(a, b) {
  const x1 = Math.max(a.x1, b.x1);
  const y1 = Math.max(a.y1, b.y1);
  const x2 = Math.min(a.x2, b.x2);
  const y2 = Math.min(a.y2, b.y2);

  const width = Math.max(0, x2 - x1);
  const height = Math.max(0, y2 - y1);
  const intersection = width * height;
  const union =
    (a.x2 - a.x1) * (a.y2 - a.y1) +
    (b.x2 - b.x1) * (b.y2 - b.y1) -
    intersection;

  return union === 0 ? 0 : intersection / union;
}

function drawDetectionResults(detections) {
  if (!selectedImage) {
    return;
  }

  canvas.width = selectedImage.width;
  canvas.height = selectedImage.height;
  ctx.clearRect(0, 0, canvas.width, canvas.height);
  ctx.drawImage(selectedImage, 0, 0);

  detections.forEach((detection) => {
    const width = detection.x2 - detection.x1;
    const height = detection.y2 - detection.y1;

    ctx.strokeStyle = detection.classIndex === 0 ? "#00d4ff" : "#ffb703";
    ctx.lineWidth = Math.max(2, Math.round(canvas.width / 360));
    ctx.strokeRect(detection.x1, detection.y1, width, height);

    const label = `${detection.label} ${(detection.score * 100).toFixed(1)}%`;
    ctx.font = `${Math.max(12, Math.round(canvas.width / 60))}px Inter`;
    ctx.textBaseline = "top";
    ctx.fillStyle = "rgba(0, 0, 0, 0.65)";
    const textWidth = ctx.measureText(label).width + 16;
    const textHeight = parseInt(ctx.font, 10) + 10;
    const textX = detection.x1;
    const textY = Math.max(0, detection.y1 - textHeight - 4);

    ctx.fillRect(textX, textY, textWidth, textHeight);
    ctx.fillStyle = "#ffffff";
    ctx.fillText(label, textX + 8, textY + 5);
  });
}

function extractMeterReading(detections) {
  const digits = detections.filter(
    (item) => item.classIndex >= 2 && item.classIndex <= 11
  );
  const unknown = detections.some((item) => item.classIndex === 12);

  if (!digits.length) {
    if (unknown) {
      return "Unreadable";
    }
    return "No digits detected";
  }

  const ordered = digits.sort((a, b) => a.x1 - b.x1);
  return ordered.map((item) => item.label).join("");
}

function computeAverageConfidence(detections) {
  const digits = detections.filter(
    (item) => item.classIndex >= 2 && item.classIndex <= 11
  );

  if (!digits.length) {
    return 0;
  }

  const sum = digits.reduce((acc, item) => acc + item.score, 0);
  return sum / digits.length;
}

function updatePredictionUI(prediction) {
  const { detections, reading, averageConfidence } = prediction;
  meterReading.textContent = reading;
  confidence.textContent = averageConfidence
    ? `${(averageConfidence * 100).toFixed(1)}%`
    : "--";

  if (!detections.length) {
    detectionsList.innerHTML =
      "<p class=\"small-text\">No objects detected in this image.</p>";
    return;
  }

  detectionsList.innerHTML = detections
    .slice(0, 20)
    .map(
      (item) =>
        `<div class="detection-card"><strong>${item.label}</strong><span>Score: ${(item.score * 100).toFixed(1)}%</span></div>`
    )
    .join("");
}

function clamp(value, min, max) {
  return Math.max(min, Math.min(value, max));
}

loadModel();