File size: 5,007 Bytes
787389e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 |
import {
AutoTokenizer,
CLIPTextModelWithProjection,
AutoProcessor,
CLIPVisionModelWithProjection,
RawImage,
dot,
softmax,
} from "https://cdn.jsdelivr.net/npm/@huggingface/[email protected]";
// Reference the elements that we will need
const status = document.getElementById("status");
const container = document.getElementById("container");
const video = document.getElementById("video");
const labelsInput = document.getElementById("labels");
const templateInput = document.getElementById("template");
const overlay = document.getElementById("overlay");
status.textContent = "Loading model (88MB)...";
const model_id = "Xenova/mobileclip_s0";
let tokenizer, text_model, processor, vision_model;
try {
// Load tokenizer and text model
tokenizer = await AutoTokenizer.from_pretrained(model_id);
text_model = await CLIPTextModelWithProjection.from_pretrained(model_id, {
device: "wasm",
dtype: "q8",
});
// Load processor and vision model
processor = await AutoProcessor.from_pretrained(model_id);
vision_model = await CLIPVisionModelWithProjection.from_pretrained(model_id, {
device: "webgpu",
dtype: "fp32",
});
} catch (err) {
console.error(err);
status.textContent = err.message;
alert(err.message);
throw err;
}
labelsInput.disabled = false;
templateInput.disabled = false;
status.textContent = "Ready";
// See `model.logit_scale` parameter of original model
const exp_logit_scale = Math.exp(4.6052);
const IMAGE_SIZE = 224;
const canvas = document.createElement("canvas");
canvas.width = canvas.height = IMAGE_SIZE;
const context = canvas.getContext("2d", { willReadFrequently: true });
let isProcessing = false;
let previousTime;
let textEmbeddings;
let prevTextInputs;
let prevTemplate;
let labels;
function onFrameUpdate() {
if (!isProcessing) {
isProcessing = true;
(async function () {
// If text inputs have changed, update the embeddings
if (
prevTextInputs !== labelsInput.value ||
prevTemplate !== templateInput.value
) {
textEmbeddings = null;
prevTextInputs = labelsInput.value;
prevTemplate = templateInput.value;
labels = prevTextInputs.split(/\s*,\s*/).filter((x) => x);
if (labels.length > 0) {
const texts = labels.map((x) =>
templateInput.value.replaceAll("{}", x),
);
const text_inputs = tokenizer(texts, {
padding: "max_length", // NB: the model requires max_length padding
truncation: true,
});
// Compute embeddings
const { text_embeds } = await text_model(text_inputs);
textEmbeddings = text_embeds.normalize().tolist();
} else {
overlay.innerHTML = "";
}
}
if (textEmbeddings) {
// Read the current frame from the video
context.drawImage(video, 0, 0, IMAGE_SIZE, IMAGE_SIZE);
const pixelData = context.getImageData(
0,
0,
IMAGE_SIZE,
IMAGE_SIZE,
).data;
const image = new RawImage(pixelData, IMAGE_SIZE, IMAGE_SIZE, 4);
const image_inputs = await processor(image);
// Compute embeddings
const { image_embeds } = await vision_model(image_inputs);
const imageEmbedding = image_embeds.normalize().tolist()[0];
// Compute similarity
const similarities = textEmbeddings.map(
(x) => dot(x, imageEmbedding) * exp_logit_scale,
);
const sortedIndices = softmax(similarities)
.map((x, i) => [x, i])
.sort((a, b) => b[0] - a[0]);
// Update UI
overlay.innerHTML = "";
for (const [score, index] of sortedIndices) {
overlay.appendChild(
document.createTextNode(`${labels[index]}: ${score.toFixed(2)}`),
);
overlay.appendChild(document.createElement("br"));
}
}
if (previousTime !== undefined) {
const fps = 1000 / (performance.now() - previousTime);
status.textContent = `FPS: ${fps.toFixed(2)}`;
}
previousTime = performance.now();
isProcessing = false;
})();
}
window.requestAnimationFrame(onFrameUpdate);
}
// Start the video stream
navigator.mediaDevices
.getUserMedia(
{ video: true }, // Ask for video
)
.then((stream) => {
// Set up the video and canvas elements.
video.srcObject = stream;
video.play();
const videoTrack = stream.getVideoTracks()[0];
const { width, height } = videoTrack.getSettings();
video.width = width;
video.height = height;
// Set container width and height depending on the image aspect ratio
const ar = width / height;
const [cw, ch] = ar > 720 / 405 ? [720, 720 / ar] : [405 * ar, 405];
container.style.width = `${cw}px`;
container.style.height = `${ch}px`;
// Start the animation loop
window.requestAnimationFrame(onFrameUpdate);
})
.catch((error) => {
alert(error);
});
|