blob: cf925bd211a2c110d2ad27fa1ddbb75178c23b1b
1 | /* |
2 | * Copyright (c) 2003-2013 Loren Merritt |
3 | * Copyright (c) 2015 Paul B Mahol |
4 | * |
5 | * This file is part of FFmpeg. |
6 | * |
7 | * FFmpeg is free software; you can redistribute it and/or |
8 | * modify it under the terms of the GNU Lesser General Public |
9 | * License as published by the Free Software Foundation; either |
10 | * version 2.1 of the License, or (at your option) any later version. |
11 | * |
12 | * FFmpeg is distributed in the hope that it will be useful, |
13 | * but WITHOUT ANY WARRANTY; without even the implied warranty of |
14 | * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU |
15 | * Lesser General Public License for more details. |
16 | * |
17 | * You should have received a copy of the GNU Lesser General Public |
18 | * License along with FFmpeg; if not, write to the Free Software |
19 | * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA |
20 | */ |
21 | |
22 | /* Computes the Structural Similarity Metric between two video streams. |
23 | * original algorithm: |
24 | * Z. Wang, A. C. Bovik, H. R. Sheikh and E. P. Simoncelli, |
25 | * "Image quality assessment: From error visibility to structural similarity," |
26 | * IEEE Transactions on Image Processing, vol. 13, no. 4, pp. 600-612, Apr. 2004. |
27 | * |
28 | * To improve speed, this implementation uses the standard approximation of |
29 | * overlapped 8x8 block sums, rather than the original gaussian weights. |
30 | */ |
31 | |
32 | /* |
33 | * @file |
34 | * Caculate the SSIM between two input videos. |
35 | */ |
36 | |
37 | #include "libavutil/avstring.h" |
38 | #include "libavutil/opt.h" |
39 | #include "libavutil/pixdesc.h" |
40 | #include "avfilter.h" |
41 | #include "dualinput.h" |
42 | #include "drawutils.h" |
43 | #include "formats.h" |
44 | #include "internal.h" |
45 | #include "ssim.h" |
46 | #include "video.h" |
47 | |
48 | typedef struct SSIMContext { |
49 | const AVClass *class; |
50 | FFDualInputContext dinput; |
51 | FILE *stats_file; |
52 | char *stats_file_str; |
53 | int nb_components; |
54 | int max; |
55 | uint64_t nb_frames; |
56 | double ssim[4], ssim_total; |
57 | char comps[4]; |
58 | float coefs[4]; |
59 | uint8_t rgba_map[4]; |
60 | int planewidth[4]; |
61 | int planeheight[4]; |
62 | int *temp; |
63 | int is_rgb; |
64 | float (*ssim_plane)(SSIMDSPContext *dsp, |
65 | uint8_t *main, int main_stride, |
66 | uint8_t *ref, int ref_stride, |
67 | int width, int height, void *temp, |
68 | int max); |
69 | SSIMDSPContext dsp; |
70 | } SSIMContext; |
71 | |
72 | #define OFFSET(x) offsetof(SSIMContext, x) |
73 | #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM |
74 | |
75 | static const AVOption ssim_options[] = { |
76 | {"stats_file", "Set file where to store per-frame difference information", OFFSET(stats_file_str), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 0, FLAGS }, |
77 | {"f", "Set file where to store per-frame difference information", OFFSET(stats_file_str), AV_OPT_TYPE_STRING, {.str=NULL}, 0, 0, FLAGS }, |
78 | { NULL } |
79 | }; |
80 | |
81 | AVFILTER_DEFINE_CLASS(ssim); |
82 | |
83 | static void set_meta(AVDictionary **metadata, const char *key, char comp, float d) |
84 | { |
85 | char value[128]; |
86 | snprintf(value, sizeof(value), "%0.2f", d); |
87 | if (comp) { |
88 | char key2[128]; |
89 | snprintf(key2, sizeof(key2), "%s%c", key, comp); |
90 | av_dict_set(metadata, key2, value, 0); |
91 | } else { |
92 | av_dict_set(metadata, key, value, 0); |
93 | } |
94 | } |
95 | |
96 | static void ssim_4x4xn_16bit(const uint8_t *main8, ptrdiff_t main_stride, |
97 | const uint8_t *ref8, ptrdiff_t ref_stride, |
98 | int64_t (*sums)[4], int width) |
99 | { |
100 | const uint16_t *main16 = (const uint16_t *)main8; |
101 | const uint16_t *ref16 = (const uint16_t *)ref8; |
102 | int x, y, z; |
103 | |
104 | main_stride >>= 1; |
105 | ref_stride >>= 1; |
106 | |
107 | for (z = 0; z < width; z++) { |
108 | uint64_t s1 = 0, s2 = 0, ss = 0, s12 = 0; |
109 | |
110 | for (y = 0; y < 4; y++) { |
111 | for (x = 0; x < 4; x++) { |
112 | int a = main16[x + y * main_stride]; |
113 | int b = ref16[x + y * ref_stride]; |
114 | |
115 | s1 += a; |
116 | s2 += b; |
117 | ss += a*a; |
118 | ss += b*b; |
119 | s12 += a*b; |
120 | } |
121 | } |
122 | |
123 | sums[z][0] = s1; |
124 | sums[z][1] = s2; |
125 | sums[z][2] = ss; |
126 | sums[z][3] = s12; |
127 | main16 += 4; |
128 | ref16 += 4; |
129 | } |
130 | } |
131 | |
132 | static void ssim_4x4xn_8bit(const uint8_t *main, ptrdiff_t main_stride, |
133 | const uint8_t *ref, ptrdiff_t ref_stride, |
134 | int (*sums)[4], int width) |
135 | { |
136 | int x, y, z; |
137 | |
138 | for (z = 0; z < width; z++) { |
139 | uint32_t s1 = 0, s2 = 0, ss = 0, s12 = 0; |
140 | |
141 | for (y = 0; y < 4; y++) { |
142 | for (x = 0; x < 4; x++) { |
143 | int a = main[x + y * main_stride]; |
144 | int b = ref[x + y * ref_stride]; |
145 | |
146 | s1 += a; |
147 | s2 += b; |
148 | ss += a*a; |
149 | ss += b*b; |
150 | s12 += a*b; |
151 | } |
152 | } |
153 | |
154 | sums[z][0] = s1; |
155 | sums[z][1] = s2; |
156 | sums[z][2] = ss; |
157 | sums[z][3] = s12; |
158 | main += 4; |
159 | ref += 4; |
160 | } |
161 | } |
162 | |
163 | static float ssim_end1x(int64_t s1, int64_t s2, int64_t ss, int64_t s12, int max) |
164 | { |
165 | int64_t ssim_c1 = (int64_t)(.01*.01*max*max*64 + .5); |
166 | int64_t ssim_c2 = (int64_t)(.03*.03*max*max*64*63 + .5); |
167 | |
168 | int64_t fs1 = s1; |
169 | int64_t fs2 = s2; |
170 | int64_t fss = ss; |
171 | int64_t fs12 = s12; |
172 | int64_t vars = fss * 64 - fs1 * fs1 - fs2 * fs2; |
173 | int64_t covar = fs12 * 64 - fs1 * fs2; |
174 | |
175 | return (float)(2 * fs1 * fs2 + ssim_c1) * (float)(2 * covar + ssim_c2) |
176 | / ((float)(fs1 * fs1 + fs2 * fs2 + ssim_c1) * (float)(vars + ssim_c2)); |
177 | } |
178 | |
179 | static float ssim_end1(int s1, int s2, int ss, int s12) |
180 | { |
181 | static const int ssim_c1 = (int)(.01*.01*255*255*64 + .5); |
182 | static const int ssim_c2 = (int)(.03*.03*255*255*64*63 + .5); |
183 | |
184 | int fs1 = s1; |
185 | int fs2 = s2; |
186 | int fss = ss; |
187 | int fs12 = s12; |
188 | int vars = fss * 64 - fs1 * fs1 - fs2 * fs2; |
189 | int covar = fs12 * 64 - fs1 * fs2; |
190 | |
191 | return (float)(2 * fs1 * fs2 + ssim_c1) * (float)(2 * covar + ssim_c2) |
192 | / ((float)(fs1 * fs1 + fs2 * fs2 + ssim_c1) * (float)(vars + ssim_c2)); |
193 | } |
194 | |
195 | static float ssim_endn_16bit(const int64_t (*sum0)[4], const int64_t (*sum1)[4], int width, int max) |
196 | { |
197 | float ssim = 0.0; |
198 | int i; |
199 | |
200 | for (i = 0; i < width; i++) |
201 | ssim += ssim_end1x(sum0[i][0] + sum0[i + 1][0] + sum1[i][0] + sum1[i + 1][0], |
202 | sum0[i][1] + sum0[i + 1][1] + sum1[i][1] + sum1[i + 1][1], |
203 | sum0[i][2] + sum0[i + 1][2] + sum1[i][2] + sum1[i + 1][2], |
204 | sum0[i][3] + sum0[i + 1][3] + sum1[i][3] + sum1[i + 1][3], |
205 | max); |
206 | return ssim; |
207 | } |
208 | |
209 | static float ssim_endn_8bit(const int (*sum0)[4], const int (*sum1)[4], int width) |
210 | { |
211 | float ssim = 0.0; |
212 | int i; |
213 | |
214 | for (i = 0; i < width; i++) |
215 | ssim += ssim_end1(sum0[i][0] + sum0[i + 1][0] + sum1[i][0] + sum1[i + 1][0], |
216 | sum0[i][1] + sum0[i + 1][1] + sum1[i][1] + sum1[i + 1][1], |
217 | sum0[i][2] + sum0[i + 1][2] + sum1[i][2] + sum1[i + 1][2], |
218 | sum0[i][3] + sum0[i + 1][3] + sum1[i][3] + sum1[i + 1][3]); |
219 | return ssim; |
220 | } |
221 | |
222 | static float ssim_plane_16bit(SSIMDSPContext *dsp, |
223 | uint8_t *main, int main_stride, |
224 | uint8_t *ref, int ref_stride, |
225 | int width, int height, void *temp, |
226 | int max) |
227 | { |
228 | int z = 0, y; |
229 | float ssim = 0.0; |
230 | int64_t (*sum0)[4] = temp; |
231 | int64_t (*sum1)[4] = sum0 + (width >> 2) + 3; |
232 | |
233 | width >>= 2; |
234 | height >>= 2; |
235 | |
236 | for (y = 1; y < height; y++) { |
237 | for (; z <= y; z++) { |
238 | FFSWAP(void*, sum0, sum1); |
239 | ssim_4x4xn_16bit(&main[4 * z * main_stride], main_stride, |
240 | &ref[4 * z * ref_stride], ref_stride, |
241 | sum0, width); |
242 | } |
243 | |
244 | ssim += ssim_endn_16bit((const int64_t (*)[4])sum0, (const int64_t (*)[4])sum1, width - 1, max); |
245 | } |
246 | |
247 | return ssim / ((height - 1) * (width - 1)); |
248 | } |
249 | |
250 | static float ssim_plane(SSIMDSPContext *dsp, |
251 | uint8_t *main, int main_stride, |
252 | uint8_t *ref, int ref_stride, |
253 | int width, int height, void *temp, |
254 | int max) |
255 | { |
256 | int z = 0, y; |
257 | float ssim = 0.0; |
258 | int (*sum0)[4] = temp; |
259 | int (*sum1)[4] = sum0 + (width >> 2) + 3; |
260 | |
261 | width >>= 2; |
262 | height >>= 2; |
263 | |
264 | for (y = 1; y < height; y++) { |
265 | for (; z <= y; z++) { |
266 | FFSWAP(void*, sum0, sum1); |
267 | dsp->ssim_4x4_line(&main[4 * z * main_stride], main_stride, |
268 | &ref[4 * z * ref_stride], ref_stride, |
269 | sum0, width); |
270 | } |
271 | |
272 | ssim += dsp->ssim_end_line((const int (*)[4])sum0, (const int (*)[4])sum1, width - 1); |
273 | } |
274 | |
275 | return ssim / ((height - 1) * (width - 1)); |
276 | } |
277 | |
278 | static double ssim_db(double ssim, double weight) |
279 | { |
280 | return 10 * log10(weight / (weight - ssim)); |
281 | } |
282 | |
283 | static AVFrame *do_ssim(AVFilterContext *ctx, AVFrame *main, |
284 | const AVFrame *ref) |
285 | { |
286 | AVDictionary **metadata = avpriv_frame_get_metadatap(main); |
287 | SSIMContext *s = ctx->priv; |
288 | float c[4], ssimv = 0.0; |
289 | int i; |
290 | |
291 | s->nb_frames++; |
292 | |
293 | for (i = 0; i < s->nb_components; i++) { |
294 | c[i] = s->ssim_plane(&s->dsp, main->data[i], main->linesize[i], |
295 | ref->data[i], ref->linesize[i], |
296 | s->planewidth[i], s->planeheight[i], s->temp, |
297 | s->max); |
298 | ssimv += s->coefs[i] * c[i]; |
299 | s->ssim[i] += c[i]; |
300 | } |
301 | for (i = 0; i < s->nb_components; i++) { |
302 | int cidx = s->is_rgb ? s->rgba_map[i] : i; |
303 | set_meta(metadata, "lavfi.ssim.", s->comps[i], c[cidx]); |
304 | } |
305 | s->ssim_total += ssimv; |
306 | |
307 | set_meta(metadata, "lavfi.ssim.All", 0, ssimv); |
308 | set_meta(metadata, "lavfi.ssim.dB", 0, ssim_db(ssimv, 1.0)); |
309 | |
310 | if (s->stats_file) { |
311 | fprintf(s->stats_file, "n:%"PRId64" ", s->nb_frames); |
312 | |
313 | for (i = 0; i < s->nb_components; i++) { |
314 | int cidx = s->is_rgb ? s->rgba_map[i] : i; |
315 | fprintf(s->stats_file, "%c:%f ", s->comps[i], c[cidx]); |
316 | } |
317 | |
318 | fprintf(s->stats_file, "All:%f (%f)\n", ssimv, ssim_db(ssimv, 1.0)); |
319 | } |
320 | |
321 | return main; |
322 | } |
323 | |
324 | static av_cold int init(AVFilterContext *ctx) |
325 | { |
326 | SSIMContext *s = ctx->priv; |
327 | |
328 | if (s->stats_file_str) { |
329 | if (!strcmp(s->stats_file_str, "-")) { |
330 | s->stats_file = stdout; |
331 | } else { |
332 | s->stats_file = fopen(s->stats_file_str, "w"); |
333 | if (!s->stats_file) { |
334 | int err = AVERROR(errno); |
335 | char buf[128]; |
336 | av_strerror(err, buf, sizeof(buf)); |
337 | av_log(ctx, AV_LOG_ERROR, "Could not open stats file %s: %s\n", |
338 | s->stats_file_str, buf); |
339 | return err; |
340 | } |
341 | } |
342 | } |
343 | |
344 | s->dinput.process = do_ssim; |
345 | s->dinput.shortest = 1; |
346 | s->dinput.repeatlast = 0; |
347 | return 0; |
348 | } |
349 | |
350 | static int query_formats(AVFilterContext *ctx) |
351 | { |
352 | static const enum AVPixelFormat pix_fmts[] = { |
353 | AV_PIX_FMT_GRAY8, AV_PIX_FMT_GRAY10, |
354 | AV_PIX_FMT_GRAY12, AV_PIX_FMT_GRAY16, |
355 | AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV422P, AV_PIX_FMT_YUV444P, |
356 | AV_PIX_FMT_YUV440P, AV_PIX_FMT_YUV411P, AV_PIX_FMT_YUV410P, |
357 | AV_PIX_FMT_YUVJ411P, AV_PIX_FMT_YUVJ420P, AV_PIX_FMT_YUVJ422P, |
358 | AV_PIX_FMT_YUVJ440P, AV_PIX_FMT_YUVJ444P, |
359 | AV_PIX_FMT_GBRP, |
360 | #define PF(suf) AV_PIX_FMT_YUV420##suf, AV_PIX_FMT_YUV422##suf, AV_PIX_FMT_YUV444##suf, AV_PIX_FMT_GBR##suf |
361 | PF(P9), PF(P10), PF(P12), PF(P14), PF(P16), |
362 | AV_PIX_FMT_NONE |
363 | }; |
364 | |
365 | AVFilterFormats *fmts_list = ff_make_format_list(pix_fmts); |
366 | if (!fmts_list) |
367 | return AVERROR(ENOMEM); |
368 | return ff_set_common_formats(ctx, fmts_list); |
369 | } |
370 | |
371 | static int config_input_ref(AVFilterLink *inlink) |
372 | { |
373 | const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format); |
374 | AVFilterContext *ctx = inlink->dst; |
375 | SSIMContext *s = ctx->priv; |
376 | int sum = 0, i; |
377 | |
378 | s->nb_components = desc->nb_components; |
379 | |
380 | if (ctx->inputs[0]->w != ctx->inputs[1]->w || |
381 | ctx->inputs[0]->h != ctx->inputs[1]->h) { |
382 | av_log(ctx, AV_LOG_ERROR, "Width and height of input videos must be same.\n"); |
383 | return AVERROR(EINVAL); |
384 | } |
385 | if (ctx->inputs[0]->format != ctx->inputs[1]->format) { |
386 | av_log(ctx, AV_LOG_ERROR, "Inputs must be of same pixel format.\n"); |
387 | return AVERROR(EINVAL); |
388 | } |
389 | |
390 | s->is_rgb = ff_fill_rgba_map(s->rgba_map, inlink->format) >= 0; |
391 | s->comps[0] = s->is_rgb ? 'R' : 'Y'; |
392 | s->comps[1] = s->is_rgb ? 'G' : 'U'; |
393 | s->comps[2] = s->is_rgb ? 'B' : 'V'; |
394 | s->comps[3] = 'A'; |
395 | |
396 | s->planeheight[1] = s->planeheight[2] = AV_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h); |
397 | s->planeheight[0] = s->planeheight[3] = inlink->h; |
398 | s->planewidth[1] = s->planewidth[2] = AV_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w); |
399 | s->planewidth[0] = s->planewidth[3] = inlink->w; |
400 | for (i = 0; i < s->nb_components; i++) |
401 | sum += s->planeheight[i] * s->planewidth[i]; |
402 | for (i = 0; i < s->nb_components; i++) |
403 | s->coefs[i] = (double) s->planeheight[i] * s->planewidth[i] / sum; |
404 | |
405 | s->temp = av_malloc_array((2 * inlink->w + 12), sizeof(*s->temp) * (1 + (desc->comp[0].depth > 8))); |
406 | if (!s->temp) |
407 | return AVERROR(ENOMEM); |
408 | s->max = (1 << desc->comp[0].depth) - 1; |
409 | |
410 | s->ssim_plane = desc->comp[0].depth > 8 ? ssim_plane_16bit : ssim_plane; |
411 | s->dsp.ssim_4x4_line = ssim_4x4xn_8bit; |
412 | s->dsp.ssim_end_line = ssim_endn_8bit; |
413 | if (ARCH_X86) |
414 | ff_ssim_init_x86(&s->dsp); |
415 | |
416 | return 0; |
417 | } |
418 | |
419 | static int config_output(AVFilterLink *outlink) |
420 | { |
421 | AVFilterContext *ctx = outlink->src; |
422 | SSIMContext *s = ctx->priv; |
423 | AVFilterLink *mainlink = ctx->inputs[0]; |
424 | int ret; |
425 | |
426 | outlink->w = mainlink->w; |
427 | outlink->h = mainlink->h; |
428 | outlink->time_base = mainlink->time_base; |
429 | outlink->sample_aspect_ratio = mainlink->sample_aspect_ratio; |
430 | outlink->frame_rate = mainlink->frame_rate; |
431 | |
432 | if ((ret = ff_dualinput_init(ctx, &s->dinput)) < 0) |
433 | return ret; |
434 | |
435 | return 0; |
436 | } |
437 | |
438 | static int filter_frame(AVFilterLink *inlink, AVFrame *buf) |
439 | { |
440 | SSIMContext *s = inlink->dst->priv; |
441 | return ff_dualinput_filter_frame(&s->dinput, inlink, buf); |
442 | } |
443 | |
444 | static int request_frame(AVFilterLink *outlink) |
445 | { |
446 | SSIMContext *s = outlink->src->priv; |
447 | return ff_dualinput_request_frame(&s->dinput, outlink); |
448 | } |
449 | |
450 | static av_cold void uninit(AVFilterContext *ctx) |
451 | { |
452 | SSIMContext *s = ctx->priv; |
453 | |
454 | if (s->nb_frames > 0) { |
455 | char buf[256]; |
456 | int i; |
457 | buf[0] = 0; |
458 | for (i = 0; i < s->nb_components; i++) { |
459 | int c = s->is_rgb ? s->rgba_map[i] : i; |
460 | av_strlcatf(buf, sizeof(buf), " %c:%f (%f)", s->comps[i], s->ssim[c] / s->nb_frames, |
461 | ssim_db(s->ssim[c], s->nb_frames)); |
462 | } |
463 | av_log(ctx, AV_LOG_INFO, "SSIM%s All:%f (%f)\n", buf, |
464 | s->ssim_total / s->nb_frames, ssim_db(s->ssim_total, s->nb_frames)); |
465 | } |
466 | |
467 | ff_dualinput_uninit(&s->dinput); |
468 | |
469 | if (s->stats_file && s->stats_file != stdout) |
470 | fclose(s->stats_file); |
471 | |
472 | av_freep(&s->temp); |
473 | } |
474 | |
475 | static const AVFilterPad ssim_inputs[] = { |
476 | { |
477 | .name = "main", |
478 | .type = AVMEDIA_TYPE_VIDEO, |
479 | .filter_frame = filter_frame, |
480 | },{ |
481 | .name = "reference", |
482 | .type = AVMEDIA_TYPE_VIDEO, |
483 | .filter_frame = filter_frame, |
484 | .config_props = config_input_ref, |
485 | }, |
486 | { NULL } |
487 | }; |
488 | |
489 | static const AVFilterPad ssim_outputs[] = { |
490 | { |
491 | .name = "default", |
492 | .type = AVMEDIA_TYPE_VIDEO, |
493 | .config_props = config_output, |
494 | .request_frame = request_frame, |
495 | }, |
496 | { NULL } |
497 | }; |
498 | |
499 | AVFilter ff_vf_ssim = { |
500 | .name = "ssim", |
501 | .description = NULL_IF_CONFIG_SMALL("Calculate the SSIM between two video streams."), |
502 | .init = init, |
503 | .uninit = uninit, |
504 | .query_formats = query_formats, |
505 | .priv_size = sizeof(SSIMContext), |
506 | .priv_class = &ssim_class, |
507 | .inputs = ssim_inputs, |
508 | .outputs = ssim_outputs, |
509 | }; |
510 |