libavfilter/dnn: add header into native model file

Guo, Yejun committed 7 years ago

Commit 022f50d1 parent 3b3150c3 files · +70 −2Browse files

Unified · Split

3 files

libavfilter/dnn/dnn_backend_native.c +41 −2
tools/python/convert_from_tensorflow.py +3 −0
tools/python/convert_header.py +26 −0

libavfilter/dnn/dnn_backend_native.c+41−2

@@ -64,6 +64,10 @@ static DNNReturnType set_input_output_native(void model, DNNInputData input, c · 63 unmodified lines
64	DNNModel ff_dnn_load_model_native(const char model_filename)	64	DNNModel ff_dnn_load_model_native(const char model_filename)
65	{	65	{
66	DNNModel *model = NULL;	66	DNNModel *model = NULL;
		67	char header_expected[] = "FFMPEGDNNNATIVE";
		68	char *buf;
		69	size_t size;
		70	int version, header_size, major_version_expected = 0;
67	ConvolutionalNetwork *network = NULL;	71	ConvolutionalNetwork *network = NULL;
68	AVIOContext *model_file_context;	72	AVIOContext *model_file_context;
69	int file_size, dnn_size, kernel_size, i;	73	int file_size, dnn_size, kernel_size, i;
@@ -84,6 +88,41 @@ DNNModel ff_dnn_load_model_native(const char model_filename) · 14 unmodified lines
84	}	88	}
85	file_size = avio_size(model_file_context);	89	file_size = avio_size(model_file_context);
86		90
		91	/**
		92	* check file header with string and version
		93	*/
		94	size = sizeof(header_expected);
		95	buf = av_malloc(size);
		96	if (!buf) {
		97	avio_closep(&model_file_context);
		98	av_freep(&model);
		99	return NULL;
		100	}
		101
		102	// size - 1 to skip the ending '\0' which is not saved in file
		103	avio_get_str(model_file_context, size - 1, buf, size);
		104	dnn_size = size - 1;
		105	if (strncmp(buf, header_expected, size) != 0) {
		106	av_freep(&buf);
		107	avio_closep(&model_file_context);
		108	av_freep(&model);
		109	return NULL;
		110	}
		111	av_freep(&buf);
		112
		113	version = (int32_t)avio_rl32(model_file_context);
		114	dnn_size += 4;
		115	if (version != major_version_expected) {
		116	avio_closep(&model_file_context);
		117	av_freep(&model);
		118	return NULL;
		119	}
		120
		121	// currently no need to check minor version
		122	version = (int32_t)avio_rl32(model_file_context);
		123	dnn_size += 4;
		124	header_size = dnn_size;
		125
87	network = av_mallocz(sizeof(ConvolutionalNetwork));	126	network = av_mallocz(sizeof(ConvolutionalNetwork));
88	if (!network){	127	if (!network){
89	avio_closep(&model_file_context);	128	avio_closep(&model_file_context);
@@ -95,8 +134,8 @@ DNNModel ff_dnn_load_model_native(const char model_filename) · 5 unmodified lines
95	avio_seek(model_file_context, file_size - 8, SEEK_SET);	134	avio_seek(model_file_context, file_size - 8, SEEK_SET);
96	network->layers_num = (int32_t)avio_rl32(model_file_context);	135	network->layers_num = (int32_t)avio_rl32(model_file_context);
97	network->operands_num = (int32_t)avio_rl32(model_file_context);	136	network->operands_num = (int32_t)avio_rl32(model_file_context);
98	dnn_size = 8;	137	dnn_size += 8;
99	avio_seek(model_file_context, 0, SEEK_SET);	138	avio_seek(model_file_context, header_size, SEEK_SET);
100		139
101	network->layers = av_mallocz(network->layers_num * sizeof(Layer));	140	network->layers = av_mallocz(network->layers_num * sizeof(Layer));
102	if (!network->layers){	141	if (!network->layers){

tools/python/convert_from_tensorflow.py+3−0

 import tensorflow as tf
 import numpy as np
 import sys, struct
+import convert_header as header
 __all__ = ['convert_from_tensorflow']
     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)

tools/python/convert_header.py+26−0

@@ -0,0 +1,26 @@
	1	# Copyright (c) 2019
	2	#
	3	# This file is part of FFmpeg.
	4	#
	5	# FFmpeg is free software; you can redistribute it and/or
	6	# modify it under the terms of the GNU Lesser General Public
	7	# License as published by the Free Software Foundation; either
	8	# version 2.1 of the License, or (at your option) any later version.
	9	#
	10	# FFmpeg is distributed in the hope that it will be useful,
	11	# but WITHOUT ANY WARRANTY; without even the implied warranty of
	12	# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
	13	# Lesser General Public License for more details.
	14	#
	15	# You should have received a copy of the GNU Lesser General Public
	16	# License along with FFmpeg; if not, write to the Free Software
	17	# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
	18	# ==============================================================================
	19
	20	str = 'FFMPEGDNNNATIVE'
	21
	22	# increase major and reset minor when we have to re-convert the model file
	23	major = 0
	24
	25	# increase minor when we don't have to re-convert the model file
	26	minor = 1

libavfilter/dnn: add header into native model file

Keyboard shortcuts

General

Issues and pull requests

Issues, pull requests and Inbox

Repository code