Lutz Roeder 4 лет назад
Родитель
Сommit
114cb74d07
2 измененных файлов с 42 добавлено и 0 удалено
  1. 28 0
      source/view.js
  2. 14 0
      test/models.json

+ 28 - 0
source/view.js

@@ -8,6 +8,7 @@ var gzip = gzip || require('./gzip');
 var tar = tar || require('./tar');
 var json = json || require('./json');
 var protobuf = protobuf || require('./protobuf');
+var flatbuffers = flatbuffers || require('./flatbuffers');
 var python = python || require('./python');
 var sidebar = sidebar || require('./view-sidebar');
 var grapher = grapher || require('./view-grapher');
@@ -1312,6 +1313,17 @@ view.ModelContext = class {
                                 tags = reader.decode();
                                 break;
                             }
+                            case 'flatbuffers': {
+                                if (stream.length >= 8) {
+                                    const buffer = stream.peek(Math.min(32, stream.length));
+                                    const reader = flatbuffers.BinaryReader.open(buffer);
+                                    const identifier = reader.identifier;
+                                    if (identifier.length > 0) {
+                                        tags.set('file_identifier', identifier);
+                                    }
+                                }
+                                break;
+                            }
                         }
                     }
                     catch (error) {
@@ -1634,6 +1646,21 @@ view.ModelFactoryService = class {
                 throw new view.Error("Unsupported Protocol Buffers content '" + (content.length > 64 ? content.substring(0, 100) + '...' : content) + "' for extension '." + extension + "' in '" + identifier + "'.", !skip());
             }
         };
+        const flatbuffers = () => {
+            const tags = context.tags('flatbuffers');
+            if (tags.has('file_identifier')) {
+                const file_identifier = tags.get('file_identifier');
+                const formats = [
+                    { name: 'onnxruntime.experimental.fbs.InferenceSession data', identifier: 'ORTM' }
+                ];
+                for (const format of formats) {
+                    if (file_identifier === format.identifier) {
+                        throw new view.Error('Invalid file content. File contains ' + format.name + '.', true);
+                    }
+                }
+                throw new view.Error("Unsupported FlatBuffers content '" + file_identifier + "' in '" + context.identifier + "'.", !skip());
+            }
+        };
         const unknown = () => {
             stream.seek(0);
             const buffer = stream.peek(Math.min(16, stream.length));
@@ -1644,6 +1671,7 @@ view.ModelFactoryService = class {
         json();
         pbtxt();
         pb();
+        flatbuffers();
         unknown();
     }
 

+ 14 - 0
test/models.json

@@ -3224,6 +3224,20 @@
     "format": "ONNX v3",
     "link":   "https://github.com/Microsoft/onnxruntime/tree/master/onnxruntime/test/testdata"
   },
+  {
+    "type":   "onnx",
+    "target": "mobilenet_v2_float.ort",
+    "source": "https://github.com/lutzroeder/netron/files/6805664/mobilenet_v2_float.ort.zip[mobilenet_v2_float.ort]",
+    "error":  "Invalid file content. File contains onnxruntime.experimental.fbs.InferenceSession data.",
+    "link":   "https://github.com/lutzroeder/netron/issues/767"
+  },
+  {
+    "type":   "onnx",
+    "target": "mobilenet_v2_float.ort",
+    "source": "https://github.com/lutzroeder/netron/files/6805665/mobilenet_v2_uint8.ort.zip[mobilenet_v2_uint8.ort]",
+    "error":  "Invalid file content. File contains onnxruntime.experimental.fbs.InferenceSession data.",
+    "link":   "https://github.com/lutzroeder/netron/issues/767"
+  },
   {
     "type":   "onnx",
     "target": "netron_issue_119.onnx",