diff --git a/libavfilter/dnn/dnn_backend_native.c b/libavfilter/dnn/dnn_backend_native.c index daa4f5008c..5d393537a8 100644 --- a/libavfilter/dnn/dnn_backend_native.c +++ b/libavfilter/dnn/dnn_backend_native.c @@ -93,8 +93,10 @@ DNNModel *ff_dnn_load_model_native(const char *model_filename) } model->model = (void *)network; + avio_seek(model_file_context, file_size - 4, SEEK_SET); network->layers_num = (int32_t)avio_rl32(model_file_context); dnn_size = 4; + avio_seek(model_file_context, 0, 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 34454b86de..cbc76a9f03 100644 --- a/tools/python/convert_from_tensorflow.py +++ b/tools/python/convert_from_tensorflow.py @@ -129,15 +129,6 @@ class TFConverter: self.converted_nodes.add(node.name) - def generate_layer_number(self): - # in current hard code implementation, the layer number is the first data written to the native model file - # it is not easy to know it at the beginning time in the general converter, so first do a dry run for compatibility - # will be refined later. - with open('/tmp/tmp.model', 'wb') as f: - self.dump_layers_to_file(f) - self.converted_nodes.clear() - - def dump_layers_to_file(self, f): for node in self.nodes: if node.name in self.converted_nodes: @@ -157,10 +148,9 @@ class TFConverter: def dump_to_file(self): - self.generate_layer_number() with open(self.outfile, 'wb') as f: - np.array([self.layer_number], dtype=np.uint32).tofile(f) self.dump_layers_to_file(f) + np.array([self.layer_number], dtype=np.uint32).tofile(f) def generate_name_node_dict(self):