[FFmpeg-devel] [PATCH] libavfilter/dnn: add header into native model file

[Pedro Arthur]

LGTM

Pushed, thanks!

Em seg, 2 de set de 2019 às 01:40, Guo, Yejun <yejun.guo at intel.com> escreveu:
>
> Signed-off-by: Guo, Yejun <yejun.guo at 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 at ffmpeg.org
> https://ffmpeg.org/mailman/listinfo/ffmpeg-devel
>
> To unsubscribe, visit link above, or email
> ffmpeg-devel-request at ffmpeg.org with subject "unsubscribe".