1 #include <errno.h>
2 #include <stdlib.h>
3 #include <stdio.h>
4 #include <string.h>
5 #include <sys/stat.h>
6
7 #include <libvmaf/model.h>
8
9 #include "config.h"
10 #include "feature/feature_extractor.h"
11 #include "log.h"
12 #include "model.h"
13 #include "read_json_model.h"
14 #include "svm.h"
15 #include "unpickle.h"
16
17 typedef struct VmafBuiltInModel {
18 const char *version;
19 const char *data;
20 const int *data_len;
21 } VmafBuiltInModel;
22
23 #if VMAF_BUILT_IN_MODELS
24 #if VMAF_FLOAT_FEATURES
25 extern const char src_vmaf_float_v0_6_1neg_json;
26 extern const int src_vmaf_float_v0_6_1neg_json_len;
27 extern const char src_vmaf_float_v0_6_1_json;
28 extern const int src_vmaf_float_v0_6_1_json_len;
29 extern const char src_vmaf_float_b_v0_6_3_json;
30 extern const int src_vmaf_float_b_v0_6_3_json_len;
31 extern const char src_vmaf_float_4k_v0_6_1_json;
32 extern const int src_vmaf_float_4k_v0_6_1_json_len;
33 #endif
34 extern const char src_vmaf_v0_6_1_json;
35 extern const int src_vmaf_v0_6_1_json_len;
36 extern const char src_vmaf_b_v0_6_3_json;
37 extern const int src_vmaf_b_v0_6_3_json_len;
38 extern const char src_vmaf_v0_6_1neg_json;
39 extern const int src_vmaf_v0_6_1neg_json_len;
40 extern const char src_vmaf_4k_v0_6_1_json;
41 extern const int src_vmaf_4k_v0_6_1_json_len;
42 #endif
43
44 static const VmafBuiltInModel built_in_models[] = {
45 #if VMAF_BUILT_IN_MODELS
46 #if VMAF_FLOAT_FEATURES
47 {
48 .version = "vmaf_float_v0.6.1",
49 .data = &src_vmaf_float_v0_6_1_json,
50 .data_len = &src_vmaf_float_v0_6_1_json_len,
51 },
52 {
53 .version = "vmaf_float_b_v0.6.3",
54 .data = &src_vmaf_float_b_v0_6_3_json,
55 .data_len = &src_vmaf_float_b_v0_6_3_json_len,
56 },
57 {
58 .version = "vmaf_float_v0.6.1neg",
59 .data = &src_vmaf_float_v0_6_1neg_json,
60 .data_len = &src_vmaf_float_v0_6_1neg_json_len,
61 },
62 {
63 .version = "vmaf_float_4k_v0.6.1",
64 .data = &src_vmaf_float_4k_v0_6_1_json,
65 .data_len = &src_vmaf_float_4k_v0_6_1_json_len,
66 },
67 #endif
68 {
69 .version = "vmaf_v0.6.1",
70 .data = &src_vmaf_v0_6_1_json,
71 .data_len = &src_vmaf_v0_6_1_json_len,
72 },
73 {
74 .version = "vmaf_b_v0.6.3",
75 .data = &src_vmaf_b_v0_6_3_json,
76 .data_len = &src_vmaf_b_v0_6_3_json_len,
77 },
78 {
79 .version = "vmaf_v0.6.1neg",
80 .data = &src_vmaf_v0_6_1neg_json,
81 .data_len = &src_vmaf_v0_6_1neg_json_len,
82 },
83 {
84 .version = "vmaf_4k_v0.6.1",
85 .data = &src_vmaf_4k_v0_6_1_json,
86 .data_len = &src_vmaf_4k_v0_6_1_json_len,
87 },
88 #endif
89 { 0 }
90 };
91
92 #define BUILT_IN_MODEL_CNT \
93 ((sizeof(built_in_models)) / (sizeof(built_in_models[0]))) - 1
94
vmaf_model_load(VmafModel ** model,VmafModelConfig * cfg,const char * version)95 int vmaf_model_load(VmafModel **model, VmafModelConfig *cfg,
96 const char *version)
97 {
98 const VmafBuiltInModel *built_in_model = NULL;
99
100 for (unsigned i = 0; i < BUILT_IN_MODEL_CNT; i++) {
101 if (!strcmp(version, built_in_models[i].version)) {
102 built_in_model = &built_in_models[i];
103 break;
104 }
105 }
106
107 if (!built_in_model) {
108 vmaf_log(VMAF_LOG_LEVEL_WARNING,
109 "no such built-in model: \"%s\"\n", version);
110 return -EINVAL;
111 }
112
113 return vmaf_read_json_model_from_buffer(model, cfg, built_in_model->data,
114 *built_in_model->data_len);
115 }
116
vmaf_model_generate_name(VmafModelConfig * cfg)117 char *vmaf_model_generate_name(VmafModelConfig *cfg)
118 {
119 const char *default_name = "vmaf";
120 const size_t name_sz =
121 cfg->name ? strlen(cfg->name) + 1 : strlen(default_name) + 1;
122
123 char *name = malloc(name_sz);
124 if (!name) return NULL;
125 memset(name, 0, name_sz);
126
127 if (!cfg->name)
128 strncpy(name, default_name, name_sz);
129 else
130 strncpy(name, cfg->name, name_sz);
131
132 return name;
133 }
134
vmaf_model_load_from_path(VmafModel ** model,VmafModelConfig * cfg,const char * path)135 int vmaf_model_load_from_path(VmafModel **model, VmafModelConfig *cfg,
136 const char *path)
137 {
138 int err = vmaf_read_json_model_from_path(model, cfg, path);
139 if (err) {
140 vmaf_log(VMAF_LOG_LEVEL_ERROR,
141 "could not read model from path: \"%s\"\n", path);
142 char *ext = strrchr(path, '.');
143 if (ext && !strcmp(ext, ".pkl")) {
144 vmaf_log(
145 VMAF_LOG_LEVEL_ERROR,
146 "support for pkl model files has been removed, use json\n");
147 }
148 }
149 return err;
150 }
151
vmaf_model_feature_overload(VmafModel * model,const char * feature_name,VmafFeatureDictionary * opts_dict)152 int vmaf_model_feature_overload(VmafModel *model, const char *feature_name,
153 VmafFeatureDictionary *opts_dict)
154 {
155 if (!model) return -EINVAL;
156 if (!feature_name) return -EINVAL;
157 if (!opts_dict) return -EINVAL;
158
159 int err = 0;
160
161 for (unsigned i = 0; i < model->n_features; i++) {
162 VmafFeatureExtractor *fex =
163 vmaf_get_feature_extractor_by_feature_name(model->feature[i].name);
164 if (!fex) continue;
165 if (strcmp(feature_name, fex->name)) continue;
166 VmafDictionary *d =
167 vmaf_dictionary_merge(&model->feature[i].opts_dict,
168 &opts_dict, 0);
169 if (!d) return -ENOMEM;
170 err = vmaf_dictionary_free(&model->feature[i].opts_dict);
171 if (err) goto exit;
172 model->feature[i].opts_dict = d;
173 }
174
175 exit:
176 err |= vmaf_dictionary_free(&opts_dict);
177 return err;
178 }
179
vmaf_model_destroy(VmafModel * model)180 void vmaf_model_destroy(VmafModel *model)
181 {
182 if (!model) return;
183 free(model->path);
184 free(model->name);
185 svm_free_and_destroy_model(&(model->svm));
186 for (unsigned i = 0; i < model->n_features; i++) {
187 free(model->feature[i].name);
188 vmaf_dictionary_free(&model->feature[i].opts_dict);
189 }
190 free(model->feature);
191 free(model->score_transform.knots.list);
192 free(model);
193 }
194
vmaf_model_collection_append(VmafModelCollection ** model_collection,VmafModel * model)195 int vmaf_model_collection_append(VmafModelCollection **model_collection,
196 VmafModel *model)
197 {
198 if (!model_collection) return -EINVAL;
199 if (!model) return -EINVAL;
200
201 VmafModelCollection *mc = *model_collection;
202
203 if (!mc) {
204 mc = *model_collection = malloc(sizeof(*mc));
205 if (!mc) goto fail;
206 memset(mc, 0, sizeof(*mc));
207 const size_t initial_sz = 8 * sizeof(*mc->model);
208 mc->model = malloc(initial_sz);
209 if (!mc->model) goto fail_mc;
210 memset(mc->model, 0, initial_sz);
211 mc->size = 8;
212 mc->type = model->type;
213 const size_t name_sz = strlen(model->name) - 5 + 1;
214 mc->name = malloc(name_sz);
215 if (!mc->name) goto fail_model;
216 memset(mc->name, 0, name_sz);
217 strncpy(mc->name, model->name, name_sz - 1);
218 }
219
220 if (mc->type != model->type) return -EINVAL;
221
222 if (mc->cnt == mc->size) {
223 const size_t sz = mc->size * sizeof(*mc->model) * 2;
224 VmafModel **m = realloc(mc->model, sz);
225 if (!m) goto fail;
226 mc->model = m;
227 mc->size *= 2;
228 }
229
230 mc->model[mc->cnt++] = model;
231 return 0;
232
233 fail_model:
234 free(mc->model);
235 fail_mc:
236 free(mc);
237 fail:
238 *model_collection = NULL;
239 return -ENOMEM;
240 }
241
vmaf_model_collection_destroy(VmafModelCollection * model_collection)242 void vmaf_model_collection_destroy(VmafModelCollection *model_collection)
243 {
244 if (!model_collection) return;
245 for (unsigned i = 0; i < model_collection->cnt; i++)
246 vmaf_model_destroy(model_collection->model[i]);
247 free(model_collection->model);
248 free(model_collection->name);
249 free(model_collection);
250 }
251
vmaf_model_collection_load(VmafModel ** model,VmafModelCollection ** model_collection,VmafModelConfig * cfg,const char * version)252 int vmaf_model_collection_load(VmafModel **model,
253 VmafModelCollection **model_collection,
254 VmafModelConfig *cfg,
255 const char *version)
256 {
257 const VmafBuiltInModel *built_in_model = NULL;
258
259 for (unsigned i = 0; i < BUILT_IN_MODEL_CNT; i++) {
260 if (!strcmp(version, built_in_models[i].version)) {
261 built_in_model = &built_in_models[i];
262 break;
263 }
264 }
265
266 if (!built_in_model) {
267 vmaf_log(VMAF_LOG_LEVEL_WARNING,
268 "no such built-in model collection: \"%s\"\n", version);
269 return -EINVAL;
270 }
271
272 return vmaf_read_json_model_collection_from_buffer(model, model_collection,
273 cfg, built_in_model->data,
274 *built_in_model->data_len);
275 }
276
vmaf_model_collection_load_from_path(VmafModel ** model,VmafModelCollection ** model_collection,VmafModelConfig * cfg,const char * path)277 int vmaf_model_collection_load_from_path(VmafModel **model,
278 VmafModelCollection **model_collection,
279 VmafModelConfig *cfg,
280 const char *path)
281 {
282 int err =
283 vmaf_read_json_model_collection_from_path(model, model_collection,
284 cfg, path);
285 if (err) {
286 vmaf_log(VMAF_LOG_LEVEL_ERROR,
287 "could not read model collection from path: \"%s\"\n", path);
288 char *ext = strrchr(path, '.');
289 if (ext && !strcmp(ext, ".pkl")) {
290 vmaf_log(
291 VMAF_LOG_LEVEL_ERROR,
292 "support for pkl model files has been removed, use json\n");
293 }
294 }
295
296 return err;
297 }
298
vmaf_model_collection_feature_overload(VmafModel * model,VmafModelCollection ** model_collection,const char * feature_name,VmafFeatureDictionary * opts_dict)299 int vmaf_model_collection_feature_overload(VmafModel *model,
300 VmafModelCollection **model_collection,
301 const char *feature_name,
302 VmafFeatureDictionary *opts_dict)
303 {
304 if (!model_collection) return -EINVAL;
305 VmafModelCollection *mc = *model_collection;
306
307 int err = 0;
308 for (unsigned i = 0; i < mc->cnt; i++) {
309 VmafFeatureDictionary *d = NULL;
310 if (vmaf_dictionary_copy(&opts_dict, &d)) goto exit;
311 err |= vmaf_model_feature_overload(mc->model[i], feature_name, d);
312 }
313
314 exit:
315 err |= vmaf_model_feature_overload(model, feature_name, opts_dict);
316 return err;
317 }
318
319