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