Signed-off-by: Guo, Yejun <yejun....@intel.com>
---
libavfilter/dnn/dnn_backend_native.c | 43 +++++++++++++++++++++++++++++++--
tools/python/convert_from_tensorflow.py | 3 +++
tools/python/convert_header.py | 26 ++++++++++++++++++++
3 files changed, 70 insertions(+), 2 deletions(-)
create mode 100644 tools/python/convert_header.py
diff --git a/libavfilter/dnn/dnn_backend_native.c
b/libavfilter/dnn/dnn_backend_native.c
index 8b05bec..f56cd81 100644
--- a/libavfilter/dnn/dnn_backend_native.c
+++ b/libavfilter/dnn/dnn_backend_native.c
@@ -64,6 +64,10 @@ static DNNReturnType set_input_output_native(void *model,
DNNInputData *input, c
DNNModel *ff_dnn_load_model_native(const char *model_filename)
{
DNNModel *model = NULL;
+ char header_expected[] = "FFMPEGDNNNATIVE";
+ char *buf;
+ size_t size;
+ int version, header_size, major_version_expected = 0;
ConvolutionalNetwork *network = NULL;
AVIOContext *model_file_context;
int file_size, dnn_size, kernel_size, i;
@@ -84,6 +88,41 @@ DNNModel *ff_dnn_load_model_native(const char
*model_filename)
}
file_size = avio_size(model_file_context);
+ /**
+ * check file header with string and version
+ */
+ size = sizeof(header_expected);
+ buf = av_malloc(size);
+ if (!buf) {
+ avio_closep(&model_file_context);
+ av_freep(&model);
+ return NULL;
+ }
+
+ // size - 1 to skip the ending '\0' which is not saved in file
+ avio_get_str(model_file_context, size - 1, buf, size);
+ dnn_size = size - 1;
+ if (strncmp(buf, header_expected, size) != 0) {
+ av_freep(&buf);
+ avio_closep(&model_file_context);
+ av_freep(&model);
+ return NULL;
+ }
+ av_freep(&buf);
+
+ version = (int32_t)avio_rl32(model_file_context);
+ dnn_size += 4;
+ if (version != major_version_expected) {
+ avio_closep(&model_file_context);
+ av_freep(&model);
+ return NULL;
+ }
+
+ // currently no need to check minor version
+ version = (int32_t)avio_rl32(model_file_context);
+ dnn_size += 4;
+ header_size = dnn_size;
+
network = av_mallocz(sizeof(ConvolutionalNetwork));
if (!network){
avio_closep(&model_file_context);
@@ -95,8 +134,8 @@ DNNModel *ff_dnn_load_model_native(const char
*model_filename)
avio_seek(model_file_context, file_size - 8, SEEK_SET);
network->layers_num = (int32_t)avio_rl32(model_file_context);
network->operands_num = (int32_t)avio_rl32(model_file_context);
- dnn_size = 8;
- avio_seek(model_file_context, 0, SEEK_SET);
+ dnn_size += 8;
+ avio_seek(model_file_context, header_size, SEEK_SET);
network->layers = av_mallocz(network->layers_num * sizeof(Layer));
if (!network->layers){
diff --git a/tools/python/convert_from_tensorflow.py
b/tools/python/convert_from_tensorflow.py
index bab11a5..1437ad3 100644
--- a/tools/python/convert_from_tensorflow.py
+++ b/tools/python/convert_from_tensorflow.py
@@ -20,6 +20,7 @@
import tensorflow as tf
import numpy as np
import sys, struct
+import convert_header as header
__all__ = ['convert_from_tensorflow']
@@ -229,6 +230,8 @@ class TFConverter:
def dump_to_file(self):
with open(self.outfile, 'wb') as f:
+ f.write(header.str.encode('utf-8'))
+ np.array([header.major, header.minor], dtype=np.uint32).tofile(f)
self.dump_layers_to_file(f)
self.dump_operands_to_file(f)
np.array([self.layer_number, len(self.name_operand_dict)],
dtype=np.uint32).tofile(f)
diff --git a/tools/python/convert_header.py b/tools/python/convert_header.py
new file mode 100644
index 0000000..6a7e4af
--- /dev/null
+++ b/tools/python/convert_header.py
@@ -0,0 +1,26 @@
+# Copyright (c) 2019
+#
+# This file is part of FFmpeg.
+#
+# FFmpeg is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License as published by the Free Software Foundation; either
+# version 2.1 of the License, or (at your option) any later version.
+#
+# FFmpeg is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public
+# License along with FFmpeg; if not, write to the Free Software
+# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+#
==============================================================================
+
+str = 'FFMPEGDNNNATIVE'
+
+# increase major and reset minor when we have to re-convert the model file
+major = 0
+
+# increase minor when we don't have to re-convert the model file
+minor = 1
--
2.7.4
_______________________________________________
ffmpeg-devel mailing list
ffmpeg-devel@ffmpeg.org
https://ffmpeg.org/mailman/listinfo/ffmpeg-devel
To unsubscribe, visit link above, or email
ffmpeg-devel-requ...@ffmpeg.org with subject "unsubscribe".
