Update index.js

index.js

@@ -1,4 +1,4 @@
-import { env, AutoModel, AutoProcessor, RawImage } from 'https://cdn.jsdelivr.net/npm/@xenova/transformers@2.17.1';
+import { env, AutoProcessor, AutoModel, RawImage } from 'https://cdn.jsdelivr.net/npm/@xenova/transformers@2.15.1';
 
 // Since we will download the model from the Hugging Face Hub, we can skip the local model check
 env.allowLocalModels = false;
@@ -10,15 +10,13 @@ const imageContainer = document.getElementById('container');
 const example = document.getElementById('example');
 
 const EXAMPLE_URL = 'https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/city-streets.jpg';
-const IMAGE_SIZE = 640;
+const THRESHOLD = 0.25;
 
 // Create a new object detection pipeline
 status.textContent = 'Loading model...';
-const model_id = 'onnx-community/yolov10n';
-const model = await AutoModel.from_pretrained(model_id, {
-    quantized: false,    // (Optional) Use unquantized version.
-});
+const model_id = 'onnx-community/yolov10s';
 const processor = await AutoProcessor.from_pretrained(model_id);
+const model = await AutoModel.from_pretrained(model_id);
 status.textContent = 'Ready';
 
 example.addEventListener('click', (e) => {
@@ -42,25 +40,38 @@ fileUpload.addEventListener('change', function (e) {
 
 
 // Detect objects in the image
-async function detect(img) {
+async function detect(url) {
+    // Update UI
     imageContainer.innerHTML = '';
-    imageContainer.style.backgroundImage = `url(${img})`;
+
+    // Read image
+    const image = await RawImage.fromURL(url);
+
+    // Set container width and height depending on the image aspect ratio
+    const ar = image.width / image.height;
+    const [cw, ch] = (ar > 1) ? [640, 640 / ar] : [640 * ar, 640];
+    imageContainer.style.width = `${cw}px`;
+    imageContainer.style.height = `${ch}px`;
+    imageContainer.style.backgroundImage = `url(${url})`;
 
     status.textContent = 'Analysing...';
-    const image = await RawImage.read(img);
-    const { pixel_values } = await processor(image);
-    const { output0 } = await model({ images: pixel_values });
-    const predictions = output0.tolist()[0];
-    const threshold = 0.3;
-    for (const [xmin, ymin, xmax, ymax, score, id] of predictions) {
-        if (score < threshold) continue;
-        renderBox(xmin, ymin, xmax, ymax, score, model.config.id2label[id]);
-    }
+
+    // Preprocess image
+    const inputs = await processor(image);
+
+    // Predict bounding boxes
+    const { output0 } = await model(inputs);
+
     status.textContent = '';
+
+    const sizes = inputs.reshaped_input_sizes[0].reverse();
+    outputs.tolist()[0].forEach(x => renderBox(x, sizes));
 }
 
 // Render a bounding box and label on the image
-function renderBox(xmin, ymin, xmax, ymax, score, label) {
+function renderBox([xmin, ymin, xmax, ymax, score, id], [w, h]) {
+    if (score < THRESHOLD) return; // Skip boxes with low confidence
+
     // Generate a random color for the box
     const color = '#' + Math.floor(Math.random() * 0xFFFFFF).toString(16).padStart(6, 0);
 
@@ -69,15 +80,15 @@ function renderBox(xmin, ymin, xmax, ymax, score, label) {
     boxElement.className = 'bounding-box';
     Object.assign(boxElement.style, {
         borderColor: color,
-        left: 100 * xmin / IMAGE_SIZE + '%',
-        top: 100 * ymin / IMAGE_SIZE + '%',
-        width: 100 * (xmax - xmin) / IMAGE_SIZE + '%',
-        height: 100 * (ymax - ymin) / IMAGE_SIZE + '%',
+        left: 100 * xmin / w + '%',
+        top: 100 * ymin / h + '%',
+        width: 100 * (xmax - xmin) / w + '%',
+        height: 100 * (ymax - ymin) / h + '%',
     })
 
     // Draw label
     const labelElement = document.createElement('span');
-    labelElement.textContent = `${label} (${(score * 100).toFixed(2)}%)`;
+    labelElement.textContent = model.config.id2label[id];
     labelElement.className = 'bounding-box-label';
     labelElement.style.backgroundColor = color;