Core ML kNearestNeighborsClassifier support (#193)

lutzroeder · Jul 24, 2019 · c9a07b7 · c9a07b7
1 parent 58f7df6
commit c9a07b7
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 0 deletions.
diff --git a/src/coreml.js b/src/coreml.js
@@ -400,6 +400,14 @@ coreml.Graph = class {
                 [ model.description.output[0].name ]);
             return 'Sound Analysis Preprocessing';
         }
+        else if (model.kNearestNeighborsClassifier) {
+            this._createNode(scope, group, 'kNearestNeighborsClassifier', null,
+                model.kNearestNeighborsClassifier,
+                [ model.description.input[0].name ],
+                [ model.description.output[0].name ]);
+            this._updateClassifierOutput(group, model.kNearestNeighborsClassifier);
+            return 'kNearestNeighborsClassifier';
+        }
         throw new coreml.Error("Unknown model type '" + JSON.stringify(Object.keys(model)) + "'.");
     }
 

diff --git a/test/models.json b/test/models.json
@@ -1769,6 +1769,20 @@
     "target": "TinyYOLO.mlmodel",
     "source": "https://raw.githubusercontent.com/hollance/YOLO-CoreML-MPSNNGraph/master/TinyYOLO-CoreML/TinyYOLO-CoreML/TinyYOLO.mlmodel"
   },
+  {
+    "type":   "coreml",
+    "target": "UpdatableKNN.mlmodel",
+    "source": "https://github.com/lutzroeder/netron/files/3424578/UpdatableKNN.zip[UpdatableKNN.mlmodel]",
+    "format": "Core ML v4",
+    "link":   "https://github.com/lutzroeder/netron/issues/193"
+  },
+  {
+    "type":   "coreml",
+    "target": "UpdatableMNISTDigitClassifier.mlmodel",
+    "source": "https://github.com/lutzroeder/netron/files/3424579/UpdatableMNISTDigitClassifier.zip[UpdatableMNISTDigitClassifier.mlmodel]",
+    "format": "Core ML v1",
+    "link":   "https://github.com/lutzroeder/netron/issues/193"
+  },
   {
     "type":   "coreml",
     "target": "VGG16.mlmodel",