1 #include <stdlib.h>
2 #include <string.h>
3 #include <time.h>
4 #include <unistd.h>
5
6 #include "cli_parse.h"
7 #include "spinner.h"
8 #include "vidinput.h"
9
10 #include "libvmaf/picture.h"
11 #include "libvmaf/libvmaf.h"
12
pix_fmt_map(int pf)13 static enum VmafPixelFormat pix_fmt_map(int pf)
14 {
15 switch (pf) {
16 case PF_420:
17 return VMAF_PIX_FMT_YUV420P;
18 case PF_422:
19 return VMAF_PIX_FMT_YUV422P;
20 case PF_444:
21 return VMAF_PIX_FMT_YUV444P;
22 default:
23 return VMAF_PIX_FMT_UNKNOWN;
24 }
25 }
26
validate_videos(video_input * vid1,video_input * vid2)27 static int validate_videos(video_input *vid1, video_input *vid2)
28 {
29 int err_cnt = 0;
30
31 video_input_info info1, info2;
32 video_input_get_info(vid1, &info1);
33 video_input_get_info(vid2, &info2);
34
35 if ((info1.frame_w != info2.frame_w) || (info1.frame_h != info2.frame_h)) {
36 fprintf(stderr, "dimensions do not match: %dx%d, %dx%d\n",
37 info1.frame_w, info1.frame_h, info2.frame_w, info2.frame_h);
38 err_cnt++;
39 }
40
41 if (info1.pixel_fmt != info2.pixel_fmt) {
42 fprintf(stderr, "pixel formats do not match: %d, %d\n",
43 info1.pixel_fmt, info2.pixel_fmt);
44 err_cnt++;
45 }
46
47 if (!pix_fmt_map(info1.pixel_fmt) || !pix_fmt_map(info2.pixel_fmt)) {
48 fprintf(stderr, "unsupported pixel format: %d\n", info1.pixel_fmt);
49 err_cnt++;
50 }
51
52 if (info1.depth != info2.depth) {
53 fprintf(stderr, "bitdepths do not match: %d, %d\n",
54 info1.depth, info2.depth);
55 err_cnt++;
56 }
57
58 if (info1.depth < 8 || info1.depth > 16) {
59 fprintf(stderr, "unsupported bitdepth: %d\n", info1.depth);
60 err_cnt++;
61 }
62
63 //TODO: more validations are possible.
64
65 return err_cnt;
66 }
67
fetch_picture(video_input * vid,VmafPicture * pic)68 static int fetch_picture(video_input *vid, VmafPicture *pic)
69 {
70 int ret;
71 video_input_ycbcr ycbcr;
72 video_input_info info;
73
74 ret = video_input_fetch_frame(vid, ycbcr, NULL);
75 if (ret < 1) return !ret;
76
77 video_input_get_info(vid, &info);
78 ret = vmaf_picture_alloc(pic, pix_fmt_map(info.pixel_fmt), info.depth,
79 info.pic_w, info.pic_h);
80 if (ret) {
81 fprintf(stderr, "problem allocating picture.\n");
82 return -1;
83 }
84
85 if (info.depth == 8) {
86 for (unsigned i = 0; i < 3; i++) {
87 int xdec = i&&!(info.pixel_fmt&1);
88 int ydec = i&&!(info.pixel_fmt&2);
89 int xstride = info.depth > 8 ? 2 : 1;
90 uint8_t *ycbcr_data = ycbcr[i].data +
91 (info.pic_y >> ydec) * ycbcr[i].stride +
92 (info.pic_x * xstride >> xdec);
93 // ^ gross, but this is how the daala y4m API works. FIXME.
94 uint8_t *pic_data = pic->data[i];
95
96 for (unsigned j = 0; j < pic->h[i]; j++) {
97 memcpy(pic_data, ycbcr_data, sizeof(*pic_data) * pic->w[i]);
98 pic_data += pic->stride[i];
99 ycbcr_data += ycbcr[i].stride;
100 }
101 }
102 } else {
103 for (unsigned i = 0; i < 3; i++) {
104 int xdec = i&&!(info.pixel_fmt&1);
105 int ydec = i&&!(info.pixel_fmt&2);
106 int xstride = info.depth > 8 ? 2 : 1;
107 uint16_t *ycbcr_data = (uint16_t*) ycbcr[i].data +
108 (info.pic_y >> ydec) * (ycbcr[i].stride / 2) +
109 (info.pic_x * xstride >> xdec);
110 // ^ gross, but this is how the daala y4m API works. FIXME.
111 uint16_t *pic_data = pic->data[i];
112
113 for (unsigned j = 0; j < pic->h[i]; j++) {
114 memcpy(pic_data, ycbcr_data, sizeof(*pic_data) * pic->w[i]);
115 pic_data += pic->stride[i] / 2;
116 ycbcr_data += ycbcr[i].stride / 2;
117 }
118 }
119 }
120
121 return 0;
122 }
123
main(int argc,char * argv[])124 int main(int argc, char *argv[])
125 {
126 int err = 0;
127 const int istty = isatty(fileno(stderr));
128
129 CLISettings c;
130 cli_parse(argc, argv, &c);
131
132 if (istty && !c.quiet) {
133 fprintf(stderr, "VMAF version %s\n", vmaf_version());
134 }
135
136 FILE *file_ref = fopen(c.path_ref, "rb");
137 if (!file_ref) {
138 fprintf(stderr, "could not open file: %s\n", c.path_ref);
139 return -1;
140 }
141
142 FILE *file_dist = fopen(c.path_dist, "rb");
143 if (!file_dist) {
144 fprintf(stderr, "could not open file: %s\n", c.path_dist);
145 return -1;
146 }
147
148 video_input vid_ref;
149 if (c.use_yuv) {
150 err = raw_input_open(&vid_ref, file_ref,
151 c.width, c.height, c.pix_fmt, c.bitdepth);
152 } else {
153 err = video_input_open(&vid_ref, file_ref);
154 }
155 if (err) {
156 fprintf(stderr, "problem with reference file: %s\n", c.path_ref);
157 return -1;
158 }
159
160 video_input vid_dist;
161 if (c.use_yuv) {
162 err = raw_input_open(&vid_dist, file_dist,
163 c.width, c.height, c.pix_fmt, c.bitdepth);
164 } else {
165 err = video_input_open(&vid_dist, file_dist);
166 }
167 if (err) {
168 fprintf(stderr, "problem with distorted file: %s\n", c.path_dist);
169 return -1;
170 }
171
172 err = validate_videos(&vid_ref, &vid_dist);
173 if (err) {
174 fprintf(stderr, "videos are incompatible, %d %s.\n",
175 err, err == 1 ? "problem" : "problems");
176 return -1;
177 }
178
179 VmafConfiguration cfg = {
180 .log_level = VMAF_LOG_LEVEL_INFO,
181 .n_threads = c.thread_cnt,
182 .n_subsample = c.subsample,
183 .cpumask = c.cpumask,
184 };
185
186 VmafContext *vmaf;
187 err = vmaf_init(&vmaf, cfg);
188 if (err) {
189 fprintf(stderr, "problem initializing VMAF context\n");
190 return -1;
191 }
192
193 VmafModel **model;
194 const size_t model_sz = sizeof(*model) * c.model_cnt;
195 model = malloc(model_sz);
196 memset(model, 0, model_sz);
197
198 VmafModelCollection **model_collection;
199 const size_t model_collection_sz =
200 sizeof(*model_collection) * c.model_cnt;
201 model_collection = malloc(model_sz);
202 memset(model_collection, 0, model_collection_sz);
203
204 const char *model_collection_label[c.model_cnt];
205 unsigned model_collection_cnt = 0;
206
207 for (unsigned i = 0; i < c.model_cnt; i++) {
208 if (c.model_config[i].version) {
209 err = vmaf_model_load(&model[i], &c.model_config[i].cfg,
210 c.model_config[i].version);
211 } else {
212 err = vmaf_model_load_from_path(&model[i], &c.model_config[i].cfg,
213 c.model_config[i].path);
214 }
215
216 if (err) {
217 // check for model_collection before failing
218 // this is implicit because the `--model` option could take either
219 // a model or model_collection
220 if (c.model_config[i].version) {
221 err = vmaf_model_collection_load(&model[i],
222 &model_collection[model_collection_cnt],
223 &c.model_config[i].cfg,
224 c.model_config[i].version);
225 } else {
226 err = vmaf_model_collection_load_from_path(&model[i],
227 &model_collection[model_collection_cnt],
228 &c.model_config[i].cfg,
229 c.model_config[i].path);
230 }
231
232 if (err) {
233 fprintf(stderr, "problem loading model: %s\n",
234 c.model_config[i].version ?
235 c.model_config[i].version : c.model_config[i].path);
236 return -1;
237 }
238
239 model_collection_label[model_collection_cnt] =
240 c.model_config[i].version ?
241 c.model_config[i].version : c.model_config[i].path;
242
243 for (unsigned j = 0; j < c.model_config[i].overload_cnt; j++) {
244 err = vmaf_model_collection_feature_overload(
245 model[i],
246 &model_collection[model_collection_cnt],
247 c.model_config[i].feature_overload[j].name,
248 c.model_config[i].feature_overload[j].opts_dict);
249 if (err) {
250 fprintf(stderr,
251 "problem overloading feature extractors from "
252 "model collection: %s\n",
253 c.model_config[i].version ?
254 c.model_config[i].version : c.model_config[i].path);
255 return -1;
256 }
257 }
258
259 err = vmaf_use_features_from_model_collection(vmaf,
260 model_collection[model_collection_cnt]);
261 if (err) {
262 fprintf(stderr,
263 "problem loading feature extractors from "
264 "model collection: %s\n",
265 c.model_config[i].version ?
266 c.model_config[i].version : c.model_config[i].path);
267 return -1;
268 }
269
270 model_collection_cnt++;
271 continue;
272 }
273
274 for (unsigned j = 0; j < c.model_config[i].overload_cnt; j++) {
275 err = vmaf_model_feature_overload(model[i],
276 c.model_config[i].feature_overload[j].name,
277 c.model_config[i].feature_overload[j].opts_dict);
278 if (err) {
279 fprintf(stderr,
280 "problem overloading feature extractors from "
281 "model: %s\n",
282 c.model_config[i].version ?
283 c.model_config[i].version : c.model_config[i].path);
284 return -1;
285
286 }
287 }
288
289 err = vmaf_use_features_from_model(vmaf, model[i]);
290 if (err) {
291 fprintf(stderr,
292 "problem loading feature extractors from model: %s\n",
293 c.model_config[i].version ?
294 c.model_config[i].version : c.model_config[i].path);
295 return -1;
296 }
297 }
298
299 for (unsigned i = 0; i < c.feature_cnt; i++) {
300 err = vmaf_use_feature(vmaf, c.feature_cfg[i].name,
301 c.feature_cfg[i].opts_dict);
302 if (err) {
303 fprintf(stderr, "problem loading feature extractor: %s\n",
304 c.feature_cfg[i].name);
305 return -1;
306 }
307 }
308
309 float fps = 0.;
310 const time_t t0 = clock();
311 unsigned picture_index;
312 for (picture_index = 0 ;; picture_index++) {
313 VmafPicture pic_ref, pic_dist;
314 int ret1 = fetch_picture(&vid_ref, &pic_ref);
315 int ret2 = fetch_picture(&vid_dist, &pic_dist);
316
317 if (ret1 && ret2) {
318 break;
319 } else if (ret1 < 0 || ret2 < 0) {
320 fprintf(stderr, "\nproblem while reading pictures\n");
321 break;
322 } else if (ret1) {
323 fprintf(stderr, "\n\"%s\" ended before \"%s\".\n",
324 c.path_ref, c.path_dist);
325 break;
326 } else if (ret2) {
327 fprintf(stderr, "\n\"%s\" ended before \"%s\".\n",
328 c.path_dist, c.path_ref);
329 break;
330 }
331
332 if (istty && !c.quiet) {
333 if (picture_index > 0 && !(picture_index % 10)) {
334 fps = (picture_index + 1) /
335 (((float)clock() - t0) / CLOCKS_PER_SEC);
336 }
337
338 fprintf(stderr, "\r%d frame%s %s %.2f FPS\033[K",
339 picture_index + 1, picture_index ? "s" : " ",
340 spinner[picture_index % spinner_length], fps);
341 fflush(stderr);
342 }
343
344 err = vmaf_read_pictures(vmaf, &pic_ref, &pic_dist, picture_index);
345 if (err) {
346 fprintf(stderr, "\nproblem reading pictures\n");
347 break;
348 }
349 }
350 if (istty && !c.quiet)
351 fprintf(stderr, "\n");
352
353 err |= vmaf_read_pictures(vmaf, NULL, NULL, 0);
354 if (err) {
355 fprintf(stderr, "problem flushing context\n");
356 return err;
357 }
358
359 if (!c.no_prediction) {
360 for (unsigned i = 0; i < c.model_cnt; i++) {
361 double vmaf_score;
362 err = vmaf_score_pooled(vmaf, model[i], VMAF_POOL_METHOD_MEAN,
363 &vmaf_score, 0, picture_index - 1);
364 if (err) {
365 fprintf(stderr, "problem generating pooled VMAF score\n");
366 return -1;
367 }
368
369 if (istty && (!c.quiet || !c.output_path)) {
370 fprintf(stderr, "%s: %f\n",
371 c.model_config[i].version ?
372 c.model_config[i].version : c.model_config[i].path,
373 vmaf_score);
374 }
375 }
376
377 for (unsigned i = 0; i < model_collection_cnt; i++) {
378 VmafModelCollectionScore score = { 0 };
379 err = vmaf_score_pooled_model_collection(vmaf, model_collection[i],
380 VMAF_POOL_METHOD_MEAN, &score,
381 0, picture_index - 1);
382 if (err) {
383 fprintf(stderr, "problem generating pooled VMAF score\n");
384 return -1;
385 }
386
387 switch (score.type) {
388 case VMAF_MODEL_COLLECTION_SCORE_BOOTSTRAP:
389 if (istty && (!c.quiet || !c.output_path)) {
390 fprintf(stderr, "%s: %f, ci.p95: [%f, %f], stddev: %f\n",
391 model_collection_label[i],
392 score.bootstrap.bagging_score, score.bootstrap.ci.p95.lo,
393 score.bootstrap.ci.p95.hi,
394 score.bootstrap.stddev);
395 }
396 break;
397 default:
398 break;
399 }
400 }
401 }
402
403 if (c.output_path)
404 vmaf_write_output(vmaf, c.output_path, c.output_fmt);
405
406 for (unsigned i = 0; i < c.model_cnt; i++)
407 vmaf_model_destroy(model[i]);
408 free(model);
409
410 for (unsigned i = 0; i < model_collection_cnt; i++)
411 vmaf_model_collection_destroy(model_collection[i]);
412 free(model_collection);
413
414 video_input_close(&vid_ref);
415 video_input_close(&vid_dist);
416 vmaf_close(vmaf);
417 cli_free(&c);
418 return err;
419 }
420