Update index.js

index.js

@@ -11,16 +11,16 @@ const example = document.getElementById('example');
 
 const EXAMPLE_URL = 'https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/city-streets.jpg';
 
-const IMAGE_SIZE = 256;
+const THRESHOLD = 0.25;
 
 // Create a new object detection pipeline
 status.textContent = 'Loading model...';
-const processor = await AutoProcessor.from_pretrained('Xenova/yolov9-c');
+const processor = await AutoProcessor.from_pretrained('Xenova/yolov9-c_all');
 
-// For this demo, we resize the image to IMAGE_SIZE x IMAGE_SIZE
-processor.feature_extractor.size = { width: IMAGE_SIZE, height: IMAGE_SIZE }
+// For this demo, we resize the image so that its shortest edge is 256px
+processor.feature_extractor.size = { shortest_edge: 256 }
 
-const model = await AutoModel.from_pretrained('Xenova/yolov9-c');
+const model = await AutoModel.from_pretrained('Xenova/yolov9-c_all');
 status.textContent = 'Ready';
 
 example.addEventListener('click', (e) => {
@@ -48,7 +48,7 @@ async function detect(url) {
     // Update UI
     imageContainer.innerHTML = '';
     imageContainer.style.backgroundImage = `url(${url})`;
-  
+
     // Read image
     const image = await RawImage.fromURL(url);
 
@@ -61,17 +61,21 @@ async function detect(url) {
     status.textContent = 'Analysing...';
 
     // Preprocess image
-    const { pixel_values } = await processor(image);
+    const inputs = await processor(image);
 
     // Predict bounding boxes
-    const { outputs } = await model({ images: pixel_values });
-  
+    const { outputs } = await model(inputs);
+
     status.textContent = '';
-    outputs.tolist().forEach(renderBox);
+
+    const sizes = inputs.reshaped_input_sizes[0].reverse();
+    outputs.tolist().forEach(x => renderBox(x, sizes));
 }
 
 // Render a bounding box and label on the image
-function renderBox([xmin, ymin, xmax, ymax, score, id]) {
+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);
 
@@ -80,10 +84,10 @@ function renderBox([xmin, ymin, xmax, ymax, score, id]) {
     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