Lutz Roeder 8 лет назад
Родитель
Сommit
f9fb960ee1
2 измененных файлов с 23 добавлено и 2 удалено
  1. 11 1
      src/keras-model.js
  2. 12 1
      src/mxnet-model.js

+ 11 - 1
src/keras-model.js

@@ -4,7 +4,17 @@ class KerasModelFactory {
 
     match(buffer, identifier) {
         var extension = identifier.split('.').pop();
-        return (extension == 'keras' || extension == 'h5' || extension == 'json');
+        if (extension == 'keras' || extension == 'h5') {
+            return true;
+        }
+        if (extension == 'json') {
+            var decoder = new TextDecoder('utf-8');
+            var json = decoder.decode(buffer);
+            if (!json.includes('\"mxnet_version\":')) {
+                return true;
+            }
+        }
+        return false;
     }
 
     open(buffer, identifier, host, callback) {

+ 12 - 1
src/mxnet-model.js

@@ -5,7 +5,18 @@
 class MXNetModelFactory {
 
     match(buffer, identifier) {
-        return identifier.endsWith('-symbol.json');
+        if (identifier.endsWith('-symbol.json')) {
+            return true;
+        }
+        var extension = identifier.split('.').pop();
+        if (extension == 'json') {
+            var decoder = new TextDecoder('utf-8');
+            var json = decoder.decode(buffer);
+            if (json.includes('\"mxnet_version\":')) {
+                return true;
+            }
+        }
+        return false;
     }
 
     open(buffer, identifier, host, callback) {