add imagenet class util function

demos/src/utils/index.js

@@ -1,7 +1,12 @@
 /* global ImageData */
+import unpack from 'ndarray-unpack'
 import sum from 'lodash/sum'
 import flatten from 'lodash/flatten'
-import unpack from 'ndarray-unpack'
+import isTypedArray from 'lodash/isTypedArray'
+import reverse from 'lodash/reverse'
+import sortBy from 'lodash/sortBy'
+import take from 'lodash/take'
+import { imagenetClasses } from './imagenet'
 
 /**
  * Find mindpoint of two points
@@ -174,3 +179,27 @@ export function unroll3Dtensor (tensor) {
     return new ImageData(imageData, shape[0], shape[1])
   })
 }
+
+/**
+ * Find top k imagenet classes
+ */
+export function imagenetClassesTopK (classProbabilities, k = 5) {
+  const probs = isTypedArray(classProbabilities)
+    ? Array.prototype.slice.call(classProbabilities)
+    : classProbabilities
+
+  const sorted = reverse(sortBy(
+    probs.map((prob, index) => [prob, index]),
+    probIndex => probIndex[0]
+  ))
+
+  const topK = take(sorted, k).map(probIndex => {
+    const iClass = imagenetClasses[probIndex[1]]
+    return {
+      id: iClass[0],
+      name: iClass[1].replace(/_/, ' '),
+      probability: probIndex[0]
+    }
+  })
+  return topK
+}