diff --git a/libvmaf/include/libvmaf/model.h b/libvmaf/include/libvmaf/model.h index 98211466f..c5455a1a2 100644 --- a/libvmaf/include/libvmaf/model.h +++ b/libvmaf/include/libvmaf/model.h @@ -70,6 +70,20 @@ typedef struct VmafModelCollectionScore { } bootstrap; } VmafModelCollectionScore; +typedef struct VmafModelDescriptor{ + const char *version; +} VmafModelDescriptor; + +/** + * Iterate through all built in vmaf models + * + * @param prev previous model. NULL to get the first model + * + * @return next model or NULL after the last model + * +*/ +const VmafModelDescriptor *vmaf_model_descriptor_next(const VmafModelDescriptor *prev); + int vmaf_model_collection_load(VmafModel **model, VmafModelCollection **model_collection, VmafModelConfig *cfg, diff --git a/libvmaf/src/model.c b/libvmaf/src/model.c index 4661ec148..09e7002f4 100644 --- a/libvmaf/src/model.c +++ b/libvmaf/src/model.c @@ -13,12 +13,6 @@ #include "read_json_model.h" #include "svm.h" -typedef struct VmafBuiltInModel { - const char *version; - const char *data; - const int *data_len; -} VmafBuiltInModel; - #if VMAF_BUILT_IN_MODELS #if VMAF_FLOAT_FEATURES extern const char src_vmaf_float_v0_6_1neg_json; @@ -42,6 +36,17 @@ extern const char src_vmaf_4k_v0_6_1neg_json; extern const int src_vmaf_4k_v0_6_1neg_json_len; #endif +/** + * For the purposes of working with vmaf_model_descriptor_next, + * the version property of VmafBuiltInModel must be the first property + * declared on this struct +*/ +typedef struct VmafBuiltInModel { + const char *version; + const char *data; + const int *data_len; +} VmafBuiltInModel; + static const VmafBuiltInModel built_in_models[] = { #if VMAF_BUILT_IN_MODELS #if VMAF_FLOAT_FEATURES @@ -322,3 +327,15 @@ int vmaf_model_collection_feature_overload(VmafModel *model, return err; } +const VmafModelDescriptor *vmaf_model_descriptor_next(const VmafModelDescriptor *prev){ + // Cast *VmafModelDescriptor into *VmafBuiltInModel + VmafBuiltInModel *builtin = (VmafBuiltInModel*)prev; + + if(!builtin) + return (VmafModelDescriptor*)&built_in_models[0]; + + if(builtin - built_in_models < BUILT_IN_MODEL_CNT) + return (VmafModelDescriptor*)(builtin + 1); + + return NULL; +} \ No newline at end of file diff --git a/libvmaf/src/model.h b/libvmaf/src/model.h index e0a41953b..fd780469a 100644 --- a/libvmaf/src/model.h +++ b/libvmaf/src/model.h @@ -83,6 +83,16 @@ typedef struct VmafModelCollection { const char *name; } VmafModelCollection; +/** + * Iterate through all built in vmaf models + * + * @param prev previous model. NULL to get the first model + * + * @return next model or NULL after the last model + * +*/ +const VmafModelDescriptor *vmaf_model_descriptor_next(const VmafModelDescriptor *prev); + char *vmaf_model_generate_name(VmafModelConfig *cfg); int vmaf_model_collection_append(VmafModelCollection **model_collection, diff --git a/libvmaf/test/test_model.c b/libvmaf/test/test_model.c index ff965cd38..e5e8342bc 100644 --- a/libvmaf/test/test_model.c +++ b/libvmaf/test/test_model.c @@ -371,6 +371,29 @@ static char *test_model_set_flags() return NULL; } +/** + * This test may fail if VmafBuiltInModel's memory layout has changed + * (version MUST be the first defined property in the struct) + */ +char *test_model_descriptor_next() +{ + VmafModelDescriptor *next = NULL; + VmafModelConfig config = {.name = NULL, .flags = VMAF_MODEL_FLAGS_DEFAULT}; + while (next = vmaf_model_descriptor_next(next)) + { + VmafBuiltInModel *builtin = next; + mu_assert("Version field should match on both builtin and next", next->version == builtin->version); + VmafModel *model; + printf(next->version); + printf("\n"); + int err = vmaf_model_load(&model, &config, next->version); + + mu_assert("Model load should not return an error",err==0); + vmaf_model_destroy(model); + } + return NULL; +} + char *run_tests() { mu_run_test(test_json_model); @@ -382,5 +405,6 @@ char *run_tests() mu_run_test(test_model_check_default_behavior_set_flags); mu_run_test(test_model_set_flags); mu_run_test(test_model_feature); + mu_run_test(test_model_descriptor_next); return NULL; }