Blame for index.html

<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Selective Segmentation Demo</title>
<style>
  :root {
    --bg: #f7f7f5; --card: #ffffff; --text: #1c1d21; --muted: #5d6170; --line: #e3e3df;
    --accent: #2f5bea; --accent-soft: #e6ecfd; --bar: #dfe3ec;
  }
  @media (prefers-color-scheme: dark) {
    :root:not([data-theme="light"]) {
      --bg: #141518; --card: #1d1f24; --text: #ececee; --muted: #a0a4b0; --line: #2d3038;
      --accent: #7c9bff; --accent-soft: #263156; --bar: #2b2f38;
    }
  }
  :root[data-theme="dark"] {
    --bg: #141518; --card: #1d1f24; --text: #ececee; --muted: #a0a4b0; --line: #2d3038;
    --accent: #7c9bff; --accent-soft: #263156; --bar: #2b2f38;
  }
  * { box-sizing: border-box; }
  body { margin: 0; background: var(--bg); color: var(--text);
         font: 16px/1.5 system-ui, -apple-system, "Segoe UI", Roboto, sans-serif; }
  main { max-width: 1040px; margin: 0 auto; padding: 32px 16px 48px; }
  h1 { font-size: 1.7rem; margin: 0 0 4px; letter-spacing: -0.01em; }
  .sub { color: var(--muted); margin: 0 0 18px; max-width: 760px; }
  .stats { display: flex; flex-wrap: wrap; gap: 8px; margin-bottom: 22px; }
  .stat { background: var(--accent-soft); color: var(--accent); border-radius: 999px;
          padding: 4px 12px; font-size: 0.85rem; font-weight: 600; }
  .layout { display: grid; grid-template-columns: 300px 1fr; gap: 16px; align-items: start; }
  @media (max-width: 820px) { .layout { grid-template-columns: 1fr; } }
  .card { background: var(--card); border: 1px solid var(--line); border-radius: 14px; padding: 18px; }
  .card h2 { font-size: 0.8rem; text-transform: uppercase; letter-spacing: 0.08em; color: var(--muted); margin: 0 0 12px; }
  .examples { display: grid; grid-template-columns: repeat(4, 1fr); gap: 8px; margin-bottom: 12px; }
  .examples button { border: 2px solid transparent; border-radius: 8px; padding: 0; background: var(--bar);
                     cursor: pointer; aspect-ratio: 1; overflow: hidden; }
  .examples button.active, .examples button:focus-visible { border-color: var(--accent); }
  .examples img { width: 100%; height: 100%; object-fit: cover; display: block; }
  .upload { display: block; text-align: center; border: 1.5px dashed var(--line); border-radius: 10px;
            padding: 12px; color: var(--muted); cursor: pointer; margin-bottom: 16px; }
  .upload:hover { border-color: var(--accent); color: var(--accent); }
  .upload input { display: none; }
  label.field { display: block; font-size: 0.85rem; color: var(--muted); margin: 10px 0 4px; }
  select { width: 100%; padding: 8px 10px; border-radius: 8px; border: 1px solid var(--line);
           background: var(--card); color: var(--text); font: inherit; }
  .panels { display: grid; grid-template-columns: repeat(4, 1fr); gap: 10px; }
  @media (max-width: 820px) { .panels { grid-template-columns: repeat(2, 1fr); } }
  .panels figure { margin: 0; }
  .panels canvas { width: 100%; aspect-ratio: 1; border-radius: 8px; background: var(--bar); display: block; }
  .panels figcaption { font-size: 0.8rem; color: var(--muted); margin-top: 4px; }
  .legend { display: flex; flex-wrap: wrap; gap: 14px; font-size: 0.82rem; color: var(--muted); margin: 12px 0 4px; }
  .sw { display: inline-block; width: 11px; height: 11px; border-radius: 3px; margin-right: 5px; vertical-align: -1px; }
  .result { margin-top: 10px; font-size: 0.95rem; }
  .result p { margin: 4px 0; }
  .meta { color: var(--muted); font-size: 0.85rem; margin-top: 8px; min-height: 1.3em; }
  footer { color: var(--muted); font-size: 0.85rem; margin-top: 22px; max-width: 760px; }
</style>
</head>
<body>
<main>
  <h1>Selective segmentation with MC Dropout</h1>
  <p class="sub">A U-Net labels every pixel of a pet photo as pet, background or border (the outline). It runs 21 times in your browser: once normally and 20 times with random dropout. Where the 20 dropout runs disagree, the model is unsure and abstains (red). Lower coverage means it answers on fewer pixels but makes fewer mistakes.</p>
  <div class="stats">
    <span class="stat">Pixel error 8.13% → 2.18% at 80% coverage</span>
    <span class="stat">Selective mIoU 0.777 → 0.882</span>
    <span class="stat">Oxford-IIIT Pet, 3,669 test images</span>
  </div>

  <div class="layout">
    <section class="card">
      <h2>Input</h2>
      <div class="examples" id="examples"></div>
      <label class="upload">Upload a pet photo<input type="file" id="file" accept="image/*"></label>
      <label class="field" for="coverage">Coverage (share of pixels the model answers on)</label>
      <select id="coverage"></select>
      <label class="field" for="score">Uncertainty score</label>
      <select id="score">
        <option value="mc_entropy">MC Dropout entropy (20 passes)</option>
        <option value="mc_maxprob">MC Dropout max-probability (20 passes)</option>
        <option value="softmax_entropy">Softmax entropy (1 pass)</option>
      </select>
    </section>

    <section class="card">
      <h2>Output</h2>
      <div class="panels">
        <figure><canvas id="c_img"></canvas><figcaption>Input</figcaption></figure>
        <figure><canvas id="c_pred"></canvas><figcaption>Prediction</figcaption></figure>
        <figure><canvas id="c_unc"></canvas><figcaption>Uncertainty (brighter = less sure)</figcaption></figure>
        <figure><canvas id="c_kept"></canvas><figcaption>Kept answers (red = abstained)</figcaption></figure>
      </div>
      <div class="legend">
        <span><span class="sw" style="background:#ff8c00"></span>pet</span>
        <span><span class="sw" style="background:#285ac8"></span>background</span>
        <span><span class="sw" style="background:#fae63c"></span>border</span>
        <span><span class="sw" style="background:#dc1e3c"></span>abstained</span>
      </div>
      <div class="result" id="result"></div>
      <div class="meta" id="meta">Loading model (31 MB, once)…</div>
    </section>
  </div>

  <footer>Thresholds for each coverage and the temperature were chosen on validation data only. Images are resized to the model's 256×256 input. Everything runs locally with ONNX Runtime Web; no image leaves your browser.</footer>
</main>

<script src="https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/ort.webgpu.min.js"></script>
<script>
ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/";
ort.env.wasm.numThreads = self.crossOriginIsolated ? Math.min(4, navigator.hardwareConcurrency || 1) : 1;
const EXAMPLES = ["pet_0.jpg", "pet_900.jpg", "pet_1800.jpg", "pet_2700.jpg"];
const COLORS = [[255, 140, 0], [40, 90, 200], [250, 230, 60]];
const NAMES = ["pet", "background", "border"];
let CHUNK = 3;                // passes per run (keeps memory low); larger on the GPU
let DEVICE = "CPU";
const $ = id => document.getElementById(id);
let session, meta, S, cache = null, busy = false;

function rngMask(batch, c, p, first) {   // Dropout2d: drop whole channels, rescale the rest
  const m = new Float32Array(batch * c);
  for (let b = 0; b < batch; b++)
    for (let k = 0; k < c; k++)
      m[b * c + k] = (first && b === 0) ? 1 : (Math.random() >= p ? 1 / (1 - p) : 0);
  return new ort.Tensor("float32", m, [batch, c, 1, 1]);
}

function imageTensor(img) {
  const cv = $("c_img"); cv.width = cv.height = S;
  const ctx = cv.getContext("2d", { willReadFrequently: true });
  ctx.imageSmoothingQuality = "high";
  ctx.drawImage(img, 0, 0, S, S);               // same as training: full image resized to S x S
  const px = ctx.getImageData(0, 0, S, S).data, n = S * S, d = new Float32Array(3 * n);
  for (let i = 0; i < n; i++) { d[i] = px[4*i] / 255; d[n+i] = px[4*i+1] / 255; d[2*n+i] = px[4*i+2] / 255; }
  return { data: d, rgb: px };
}

async function infer(img) {
  const { data, rgb } = imageTensor(img);
  const n = S * S, T = meta.mc_samples, total = T + 1;
  const p1 = new Float32Array(3 * n), pm = new Float32Array(3 * n);
  const t0 = performance.now();
  for (let start = 0; start < total; start += CHUNK) {
    const b = Math.min(CHUNK, total - start);
    const x = new Float32Array(b * 3 * n);
    for (let k = 0; k < b; k++) x.set(data, k * 3 * n);
    const feeds = { image: new ort.Tensor("float32", x, [b, 3, S, S]) };
    meta.mask_inputs.forEach((name, i) => { feeds[name] = rngMask(b, meta.mask_channels[i], meta.dropout_p, start === 0); });
    const out = (await session.run(feeds)).probs.data;
    for (let k = 0; k < b; k++) {
      const off = k * 3 * n;
      if (start + k === 0) p1.set(out.subarray(off, off + 3 * n));
      else for (let j = 0; j < 3 * n; j++) pm[j] += out[off + j] / T;
    }
    $("meta").textContent = `Running pass ${Math.min(start + b, total)} of ${total}…`;
    await new Promise(r => setTimeout(r, 0));
  }
  const ms = performance.now() - t0;
  const ent = (p, i) => { let h = 0; for (let c = 0; c < 3; c++) { const v = Math.max(p[c*n+i], 1e-8); h -= v * Math.log(v); } return h; };
  const argmax = (p, i) => { let best = 0; for (let c = 1; c < 3; c++) if (p[c*n+i] > p[best*n+i]) best = c; return best; };
  const pred1 = new Uint8Array(n), predm = new Uint8Array(n);
  const u = { softmax_entropy: new Float32Array(n), mc_entropy: new Float32Array(n), mc_maxprob: new Float32Array(n) };
  for (let i = 0; i < n; i++) {
    pred1[i] = argmax(p1, i); predm[i] = argmax(pm, i);
    u.softmax_entropy[i] = ent(p1, i); u.mc_entropy[i] = ent(pm, i); u.mc_maxprob[i] = 1 - pm[predm[i]*n+i];
  }
  cache = { rgb, pred1, predm, u, ms };
  render();
}

function heat(v) {   // 0..1 -> dark blue .. red .. yellow
  const lo = [26, 26, 89], mid = [217, 64, 89], hi = [255, 242, 89];
  const [a, b, t] = v < 0.5 ? [lo, mid, v / 0.5] : [mid, hi, (v - 0.5) / 0.5];
  return a.map((x, k) => x + (b[k] - x) * t);
}

function paint(id, fn) {
  const cv = $(id); cv.width = cv.height = S;
  const ctx = cv.getContext("2d"), im = ctx.createImageData(S, S);
  for (let i = 0; i < S * S; i++) { const c = fn(i); im.data[4*i] = c[0]; im.data[4*i+1] = c[1]; im.data[4*i+2] = c[2]; im.data[4*i+3] = 255; }
  ctx.putImageData(im, 0, 0);
}

function render() {
  if (!cache) return;
  const score = $("score").value, cov = $("coverage").value, n = S * S;
  const pred = score === "softmax_entropy" ? cache.pred1 : cache.predm;
  const u = cache.u[score], thr = meta.thresholds[score][cov], rgb = cache.rgb;
  let umax = 1e-6; for (let i = 0; i < n; i++) umax = Math.max(umax, u[i]);
  const mix = (i, c, a) => [0, 1, 2].map(k => (1 - a) * rgb[4*i+k] + a * c[k]);
  paint("c_pred", i => mix(i, COLORS[pred[i]], 0.55));
  paint("c_unc", i => heat(u[i] / umax));
  let abst = 0, borderAbst = 0; const kept = [0, 0, 0];
  paint("c_kept", i => {
    if (u[i] >= thr) { abst++; if (pred[i] === 2) borderAbst++; return mix(i, [220, 30, 60], 0.65); }
    kept[pred[i]]++; return mix(i, COLORS[pred[i]], 0.55);
  });
  const nk = Math.max(1, n - abst), pct = v => (100 * v).toFixed(0) + "%";
  $("result").innerHTML =
    `<p><strong>Abstained on ${(100 * abst / n).toFixed(1)}% of pixels</strong> at the ${pct(+cov)}-coverage threshold.</p>` +
    `<p>Of the abstained pixels, <strong>${pct(abst ? borderAbst / abst : 0)}</strong> were predicted as the pet's outline (border).</p>` +
    `<p>Kept pixels: ${NAMES.map((nm, c) => `${nm} ${pct(kept[c] / nk)}`).join(" · ")}</p>`;
  $("meta").textContent = `21 forward passes in ${(cache.ms / 1000).toFixed(1)} s on your ${DEVICE}` +
    (DEVICE === "CPU" ? " (a browser with WebGPU, e.g. recent Chrome, Edge or Safari, is much faster)" : "");
}

async function run(src) {
  if (busy || !session) return;
  busy = true;
  try { const img = new Image(); img.src = src; await img.decode(); await infer(img); }
  catch (e) { $("meta").textContent = "Error: " + e.message; console.error(e); }
  busy = false;
}

(async () => {
  try {
    meta = await (await fetch("model_meta.json")).json();
    S = meta.img_size;
    const covs = Object.keys(meta.thresholds.mc_entropy).sort((a, b) => b - a);
    $("coverage").innerHTML = covs.map(c => `<option value="${c}"${+c === 0.8 ? " selected" : ""}>${Math.round(c * 100)}%</option>`).join("");
    if (navigator.gpu) {
      try {
        session = await ort.InferenceSession.create("unet_mc.onnx", { executionProviders: ["webgpu"] });
        DEVICE = "GPU (WebGPU)"; CHUNK = 7;
      } catch (e) { console.warn("WebGPU unavailable, using CPU:", e.message); }
    }
    if (!session) session = await ort.InferenceSession.create("unet_mc.onnx", { executionProviders: ["wasm"] });
    const box = $("examples");
    EXAMPLES.forEach(f => {
      const b = document.createElement("button");
      b.title = "Example from the Oxford-IIIT Pet test set";
      b.innerHTML = `<img src="examples/${f}" alt="Example pet photo">`;
      b.onclick = () => { box.querySelectorAll("button").forEach(x => x.classList.remove("active")); b.classList.add("active"); run(`examples/${f}`); };
      box.appendChild(b);
    });
    $("meta").textContent = `Model loaded on your ${DEVICE}. Pick an example or upload a photo.`;
    box.querySelector("button").click();
  } catch (e) { $("meta").textContent = "Could not load the model: " + e.message; console.error(e); }
})();

$("coverage").onchange = render;
$("score").onchange = render;
$("file").onchange = e => {
  const f = e.target.files[0]; if (!f) return;
  $("examples").querySelectorAll("button").forEach(x => x.classList.remove("active"));
  const r = new FileReader(); r.onload = () => run(r.result); r.readAsDataURL(f);
};
</script>
</body>
</html>