Spaces:
Running
Running
| // The WhatDoWeSee page: pick a model, drop an image, show its tags. Runs entirely in the browser. | |
| import { MODELS, fileUrl, run } from "./tagger.js"; | |
| /* global ort */ | |
| ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/"; | |
| const $ = (id) => document.getElementById(id); | |
| const loaded = new Map(); // model id -> Promise<{session, config}> | |
| let current = MODELS[0]; | |
| let lastImage = null; // the File on screen, re-tagged when the model changes | |
| function load(model) { | |
| if (!loaded.has(model.id)) { | |
| loaded.set( | |
| model.id, | |
| (async () => { | |
| const config = await (await fetch(fileUrl(model, "config.json"))).json(); | |
| const session = await ort.InferenceSession.create(fileUrl(model, "model.onnx"), { | |
| executionProviders: ["wasm"], | |
| }); | |
| return { session, config }; | |
| })(), | |
| ); | |
| } | |
| return loaded.get(model.id); | |
| } | |
| /** The image as RGBA pixels squashed to size × size (no crop), like the training preprocessing. */ | |
| async function pixels(file, size) { | |
| let bitmap; | |
| try { | |
| bitmap = await createImageBitmap(file, { resizeWidth: size, resizeHeight: size, resizeQuality: "high" }); | |
| } catch { | |
| bitmap = await createImageBitmap(file); | |
| } | |
| const canvas = new OffscreenCanvas(size, size); | |
| const context = canvas.getContext("2d"); | |
| context.imageSmoothingEnabled = true; | |
| context.imageSmoothingQuality = "high"; | |
| context.drawImage(bitmap, 0, 0, size, size); | |
| bitmap.close(); | |
| return context.getImageData(0, 0, size, size).data; | |
| } | |
| function render(tags, ms) { | |
| const present = tags.filter((t) => t.present); | |
| const chips = $("present"); | |
| chips.replaceChildren( | |
| ...(present.length | |
| ? present.map((t) => { | |
| const li = document.createElement("li"); | |
| li.textContent = t.name; | |
| li.title = t.definition; | |
| return li; | |
| }) | |
| : [Object.assign(document.createElement("li"), { className: "none", textContent: "No attribute found" })]), | |
| ); | |
| $("rows").replaceChildren( | |
| ...tags.map((t) => { | |
| const tr = document.createElement("tr"); | |
| if (t.present) tr.className = "on"; | |
| const name = document.createElement("td"); | |
| name.textContent = t.name; | |
| name.title = t.definition; | |
| const barCell = document.createElement("td"); | |
| const track = document.createElement("div"); | |
| track.className = "track"; | |
| const fill = document.createElement("div"); | |
| fill.className = "fill"; | |
| fill.style.width = `${(t.p * 100).toFixed(1)}%`; | |
| const tick = document.createElement("div"); | |
| tick.className = "tick"; | |
| tick.style.left = `calc(${(t.threshold * 100).toFixed(1)}% - 1px)`; | |
| tick.title = `threshold ${(t.threshold * 100).toFixed(0)}%`; | |
| track.append(fill, tick); | |
| barCell.append(track); | |
| const pct = document.createElement("td"); | |
| pct.className = "num"; | |
| pct.textContent = `${Math.round(t.p * 100)}`; | |
| tr.append(name, barCell, pct); | |
| return tr; | |
| }), | |
| ); | |
| $("all").hidden = false; | |
| $("status").textContent = `${present.length} of ${tags.length} attributes · ${ms} ms in your browser · ${current.name}`; | |
| } | |
| async function tag(file) { | |
| if (!file || !file.type.startsWith("image/")) { | |
| $("status").textContent = "That is not an image."; | |
| return; | |
| } | |
| lastImage = file; | |
| const preview = $("preview"); | |
| if (preview.src) URL.revokeObjectURL(preview.src); | |
| preview.src = URL.createObjectURL(file); | |
| preview.hidden = false; | |
| $("drop").classList.add("has-image"); | |
| $("hint-title").textContent = "Drop, click or paste"; | |
| $("hint-sub").textContent = "another image to tag it"; | |
| $("status").textContent = "Tagging…"; | |
| try { | |
| const { session, config } = await load(current); | |
| const started = performance.now(); | |
| const tags = await run(ort, session, config, await pixels(file, config.input.shape[3])); | |
| render(tags, Math.round(performance.now() - started)); | |
| } catch (error) { | |
| $("status").textContent = `Could not tag this image: ${error.message ?? error}`; | |
| } | |
| } | |
| function setup() { | |
| const select = $("model"); | |
| for (const m of MODELS) { | |
| const option = document.createElement("option"); | |
| option.value = m.id; | |
| option.textContent = m.name; | |
| select.append(option); | |
| } | |
| select.addEventListener("change", () => { | |
| current = MODELS.find((m) => m.id === select.value); | |
| $("model-link").href = `https://huggingface.co/${current.repo}`; | |
| if (lastImage) tag(lastImage); | |
| else warm(); | |
| }); | |
| const drop = $("drop"); | |
| const input = $("file"); | |
| drop.addEventListener("click", () => input.click()); | |
| drop.addEventListener("keydown", (e) => { | |
| if (e.key === "Enter" || e.key === " ") { | |
| e.preventDefault(); | |
| input.click(); | |
| } | |
| }); | |
| input.addEventListener("change", () => tag(input.files[0])); | |
| for (const type of ["dragenter", "dragover"]) { | |
| drop.addEventListener(type, (e) => { | |
| e.preventDefault(); | |
| drop.classList.add("over"); | |
| }); | |
| } | |
| for (const type of ["dragleave", "drop"]) { | |
| drop.addEventListener(type, () => drop.classList.remove("over")); | |
| } | |
| drop.addEventListener("drop", (e) => { | |
| e.preventDefault(); | |
| tag(e.dataTransfer.files[0]); | |
| }); | |
| document.addEventListener("paste", (e) => { | |
| const item = [...e.clipboardData.items].find((i) => i.type.startsWith("image/")); | |
| if (item) tag(item.getAsFile()); | |
| }); | |
| warm(); | |
| } | |
| /** Load the model up front, so the first image is tagged right away. */ | |
| function warm() { | |
| $("status").textContent = "Loading the model (about 16 MB, once)…"; | |
| load(current).then( | |
| () => { | |
| if (!lastImage) $("status").textContent = "Ready. Drop an image to tag it."; | |
| }, | |
| (error) => { | |
| $("status").textContent = `Could not load the model: ${error.message ?? error}`; | |
| }, | |
| ); | |
| } | |
| setup(); | |