Normalize input and output

Author: pearmini

examples/imageClassifier-transformer-single-image/sketch.js

@@ -15,16 +15,15 @@ let img;
 // Variables for displaying the results on the canvas
 let label = "";
 let confidence = "";
-let url = "images/bird.jpg";
 
 function preload() {
   classifier = ml5.imageClassifier("transformer");
-  img = loadImage(url);
+  img = loadImage("images/bird.jpg");
 }
 
 function setup() {
   createCanvas(400, 400);
-  classifier.classify(url, gotResult);
+  classifier.classify(img, gotResult);
   image(img, 0, 0, width, height);
 }
 
@@ -39,7 +38,7 @@ function gotResult(results) {
   stroke(0);
   textSize(18);
   label = "Label: " + results[0].label;
-  confidence = "Confidence: " + nf(results[0].score, 0, 2);
+  confidence = "Confidence: " + nf(results[0].confidence, 0, 2);
   text(label, 10, 360);
   text(confidence, 10, 380);
 }

src/ImageClassifier/transformer.js

@@ -1,5 +1,9 @@
 import { pipeline } from "@huggingface/transformers";
+import handleArguments from "../utils/handleArguments";
 
+/**
+ * @reference https://huggingface.co/docs/transformers.js/en/api/pipelines#module_pipelines.ImageClassificationPipeline
+ */
 export class ImageClassifierTransformer {
   constructor(options, callback) {
     this.classifier = null;
@@ -16,18 +20,30 @@ export class ImageClassifierTransformer {
     });
   }
 
-  async classify(input, callback) {
-    if (this.isClassifying) return;
-    if (!this.classifier) return;
+  async classify(inputNumOrCallback, numOrCallback, cb) {
+    if (this.isClassifying || !this.classifier) return;
     this.isClassifying = true;
-    const results = await this.classifier(input);
-    callback(results);
+    const { image, number, callback } = handleArguments(
+      inputNumOrCallback,
+      numOrCallback,
+      cb
+    ).require(
+      "image",
+      "No input image provided. If you want to classify a video, use classifyStart."
+    );
+    const options = number !== undefined ? { top_k: number } : {};
+    const results = await this.classifier(image, options);
+    const normalized = results.map((result) => ({
+      label: result.label,
+      confidence: result.score,
+    }));
+    callback(normalized);
     this.isClassifying = false;
+    return normalized;
   }
 
   async classifyStart(input, callback) {
-    if (this.isClassifying) return;
-    if (!this.classifier) return;
+    if (this.isClassifying || !this.classifier) return;
     this.needToStop = false;
     const next = (...args) => {
       if (this.needToStop) return;