LGTM Pushed, thanks!
Em seg, 2 de set de 2019 às 01:40, Guo, Yejun <yejun....@intel.com> escreveu: > > 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". _______________________________________________ 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".