dnn/dnn_backend_native: Avoid allocation for checking file magic

Reviewed-by: Guo, Yejun <yejun.guo@intel.com>
Signed-off-by: Andreas Rheinhardt <andreas.rheinhardt@gmail.com>
This commit is contained in:
Andreas Rheinhardt 2021-03-10 02:34:14 +01:00
parent b7bf631f91
commit 2e2ed39dac
1 changed files with 6 additions and 16 deletions

View File

@ -114,10 +114,10 @@ static DNNReturnType get_output_native(void *model, const char *input_name, int
// For DEPTH_TO_SPACE layer: block_size
DNNModel *ff_dnn_load_model_native(const char *model_filename, DNNFunctionType func_type, const char *options, AVFilterContext *filter_ctx)
{
#define DNN_NATIVE_MAGIC "FFMPEGDNNNATIVE"
DNNModel *model = NULL;
char header_expected[] = "FFMPEGDNNNATIVE";
char *buf;
size_t size;
// sizeof - 1 to skip the terminating '\0' which is not written in the file
char buf[sizeof(DNN_NATIVE_MAGIC) - 1];
int version, header_size, major_version_expected = 1;
NativeModel *native_model = NULL;
AVIOContext *model_file_context;
@ -138,20 +138,10 @@ DNNModel *ff_dnn_load_model_native(const char *model_filename, DNNFunctionType f
/**
* check file header with string and version
*/
size = sizeof(header_expected);
buf = av_malloc(size);
if (!buf) {
if (avio_read(model_file_context, buf, sizeof(buf)) != sizeof(buf) ||
memcmp(buf, DNN_NATIVE_MAGIC, sizeof(buf)))
goto fail;
}
// 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);
goto fail;
}
av_freep(&buf);
dnn_size = sizeof(buf);
version = (int32_t)avio_rl32(model_file_context);
dnn_size += 4;