2 * Copyright (c) 2018 Mina Sami
4 * This file is part of FFmpeg.
6 * FFmpeg is free software; you can redistribute it and/or
7 * modify it under the terms of the GNU Lesser General Public
8 * License as published by the Free Software Foundation; either
9 * version 2.1 of the License, or (at your option) any later version.
11 * FFmpeg is distributed in the hope that it will be useful,
12 * but WITHOUT ANY WARRANTY; without even the implied warranty of
13 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
14 * Lesser General Public License for more details.
16 * You should have received a copy of the GNU Lesser General Public
17 * License along with FFmpeg; if not, write to the Free Software
18 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
23 * Color Constancy filter
25 * @see http://colorconstancy.com/
28 * J. van de Weijer, Th. Gevers, A. Gijsenij "Edge-Based Color Constancy".
31 #include "libavutil/imgutils.h"
32 #include "libavutil/opt.h"
33 #include "libavutil/pixdesc.h"
42 #define GREY_EDGE "greyedge"
44 #define SQRT3 1.73205080757
47 #define MAX_DIFF_ORD 2
48 #define MAX_META_DATA 4
55 #define INDEX_NORM INDEX_DX
64 * Used for passing data between threads.
66 typedef struct ThreadData {
68 int meta_data[MAX_META_DATA];
69 double *data[MAX_DATA][NUM_PLANES];
73 * Common struct for all algorithms contexts.
75 typedef struct ColorConstancyContext {
79 int minknorm; /**< @minknorm = 0 : getMax instead */
87 double *gauss[MAX_DIFF_ORD+1];
89 double white[NUM_PLANES];
90 } ColorConstancyContext;
92 #define OFFSET(x) offsetof(ColorConstancyContext, x)
93 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
95 #define GINDX(s, i) ( (i) - ((s) >> 2) )
98 * Sets gauss filters used for calculating gauss derivatives. Filter size
99 * depends on sigma which is a user option hence we calculate these
100 * filters each time. Also each higher order depends on lower ones. Sigma
101 * can be zero only at difford = 0, then we only convert data to double
104 * @param ctx the filter context.
106 * @return 0 in case of success, a negative value corresponding to an
107 * AVERROR code in case of failure.
109 static int set_gauss(AVFilterContext *ctx)
111 ColorConstancyContext *s = ctx->priv;
112 int filtersize = s->filtersize;
113 int difford = s->difford;
114 double sigma = s->sigma;
118 for (i = 0; i <= difford; ++i) {
119 s->gauss[i] = av_mallocz_array(filtersize, sizeof(*s->gauss[i]));
121 for (; i >= 0; --i) {
122 av_freep(&s->gauss[i]);
124 av_log(ctx, AV_LOG_ERROR, "Out of memory while allocating gauss buffers.\n");
125 return AVERROR(ENOMEM);
130 av_log(ctx, AV_LOG_TRACE, "Setting 0-d gauss with filtersize = %d.\n", filtersize);
133 s->gauss[0][0] = 1; // Copying data to double instead of convolution
135 for (i = 0; i < filtersize; ++i) {
136 s->gauss[0][i] = exp(- pow(GINDX(filtersize, i), 2.) / (2 * sigma * sigma)) / ( sqrt(2 * M_PI) * sigma );
137 sum1 += s->gauss[0][i];
139 for (i = 0; i < filtersize; ++i) {
140 s->gauss[0][i] /= sum1;
145 av_log(ctx, AV_LOG_TRACE, "Setting 1-d gauss with filtersize = %d.\n", filtersize);
147 for (i = 0; i < filtersize; ++i) {
148 s->gauss[1][i] = - (GINDX(filtersize, i) / pow(sigma, 2)) * s->gauss[0][i];
149 sum1 += s->gauss[1][i] * GINDX(filtersize, i);
152 for (i = 0; i < filtersize; ++i) {
153 s->gauss[1][i] /= sum1;
158 av_log(ctx, AV_LOG_TRACE, "Setting 2-d gauss with filtersize = %d.\n", filtersize);
160 for (i = 0; i < filtersize; ++i) {
161 s->gauss[2][i] = ( pow(GINDX(filtersize, i), 2) / pow(sigma, 4) - 1/pow(sigma, 2) )
163 sum1 += s->gauss[2][i];
167 for (i = 0; i < filtersize; ++i) {
168 s->gauss[2][i] -= sum1 / (filtersize);
169 sum2 += (0.5 * GINDX(filtersize, i) * GINDX(filtersize, i) * s->gauss[2][i]);
171 for (i = 0; i < filtersize ; ++i) {
172 s->gauss[2][i] /= sum2;
180 * Frees up buffers used by grey edge for storing derivatives final
181 * and intermidiate results. Number of buffers and number of planes
182 * for last buffer are given so it can be safely called at allocation
185 * @param td holds the buffers.
186 * @param nb_buff number of buffers to be freed.
187 * @param nb_planes number of planes for last buffer to be freed.
189 static void cleanup_derivative_buffers(ThreadData *td, int nb_buff, int nb_planes)
193 for (b = 0; b < nb_buff; ++b) {
194 for (p = 0; p < NUM_PLANES; ++p) {
195 av_freep(&td->data[b][p]);
198 // Final buffer may not be fully allocated at fail cases
199 for (p = 0; p < nb_planes; ++p) {
200 av_freep(&td->data[b][p]);
205 * Allocates buffers used by grey edge for storing derivatives final
206 * and intermidiate results.
208 * @param ctx the filter context.
209 * @param td holds the buffers.
211 * @return 0 in case of success, a negative value corresponding to an
212 * AVERROR code in case of failure.
214 static int setup_derivative_buffers(AVFilterContext* ctx, ThreadData *td)
216 ColorConstancyContext *s = ctx->priv;
217 int nb_buff = s->difford + 1;
220 av_log(ctx, AV_LOG_TRACE, "Allocating %d buffer(s) for grey edge.\n", nb_buff);
221 for (b = 0; b <= nb_buff; ++b) { // We need difford + 1 buffers
222 for (p = 0; p < NUM_PLANES; ++p) {
223 td->data[b][p] = av_mallocz_array(s->planeheight[p] * s->planewidth[p], sizeof(*td->data[b][p]));
224 if (!td->data[b][p]) {
225 cleanup_derivative_buffers(td, b + 1, p);
226 av_log(ctx, AV_LOG_ERROR, "Out of memory while allocating derivatives buffers.\n");
227 return AVERROR(ENOMEM);
234 #define CLAMP(x, mx) av_clip((x), 0, (mx-1))
235 #define INDX2D(r, c, w) ( (r) * (w) + (c) )
236 #define GAUSS(s, sr, sc, sls, sh, sw, g) ( (s)[ INDX2D(CLAMP((sr), (sh)), CLAMP((sc), (sw)), (sls)) ] * (g) )
239 * Slice calculation of gaussian derivatives. Applies 1-D gaussian derivative filter
240 * either horizontally or vertically according to meta data given in thread data.
241 * When convoluting horizontally source is always the in frame withing thread data
242 * while when convoluting vertically source is a buffer.
244 * @param ctx the filter context.
245 * @param arg data to be passed between threads.
246 * @param jobnr current job nubmer.
247 * @param nb_jobs total number of jobs.
251 static int slice_get_derivative(AVFilterContext* ctx, void* arg, int jobnr, int nb_jobs)
253 ColorConstancyContext *s = ctx->priv;
254 ThreadData *td = arg;
255 AVFrame *in = td->in;
256 const int ord = td->meta_data[INDEX_ORD];
257 const int dir = td->meta_data[INDEX_DIR];
258 const int src_index = td->meta_data[INDEX_SRC];
259 const int dst_index = td->meta_data[INDEX_DST];
260 const int filtersize = s->filtersize;
261 const double *gauss = s->gauss[ord];
264 for (plane = 0; plane < NUM_PLANES; ++plane) {
265 const int height = s->planeheight[plane];
266 const int width = s->planewidth[plane];
267 const int in_linesize = in->linesize[plane];
268 double *dst = td->data[dst_index][plane];
269 int slice_start, slice_end;
273 /** Applying gauss horizontally along each row */
274 const uint8_t *src = in->data[plane];
275 slice_start = (height * jobnr ) / nb_jobs;
276 slice_end = (height * (jobnr + 1)) / nb_jobs;
278 for (r = slice_start; r < slice_end; ++r) {
279 for (c = 0; c < width; ++c) {
280 dst[INDX2D(r, c, width)] = 0;
281 for (g = 0; g < filtersize; ++g) {
282 dst[INDX2D(r, c, width)] += GAUSS(src, r, c + GINDX(filtersize, g),
283 in_linesize, height, width, gauss[GINDX(filtersize, g)]);
288 /** Applying gauss vertically along each column */
289 const double *src = td->data[src_index][plane];
290 slice_start = (width * jobnr ) / nb_jobs;
291 slice_end = (width * (jobnr + 1)) / nb_jobs;
293 for (c = slice_start; c < slice_end; ++c) {
294 for (r = 0; r < height; ++r) {
295 dst[INDX2D(r, c, width)] = 0;
296 for (g = 0; g < filtersize; ++g) {
297 dst[INDX2D(r, c, width)] += GAUSS(src, r + GINDX(filtersize, g), c,
298 width, height, width, gauss[GINDX(filtersize, g)]);
309 * Slice Frobius normalization of gaussian derivatives. Only called for difford values of
312 * @param ctx the filter context.
313 * @param arg data to be passed between threads.
314 * @param jobnr current job nubmer.
315 * @param nb_jobs total number of jobs.
319 static int slice_normalize(AVFilterContext* ctx, void* arg, int jobnr, int nb_jobs)
321 ColorConstancyContext *s = ctx->priv;
322 ThreadData *td = arg;
323 const int difford = s->difford;
326 for (plane = 0; plane < NUM_PLANES; ++plane) {
327 const int height = s->planeheight[plane];
328 const int width = s->planewidth[plane];
329 const int64_t numpixels = width * (int64_t)height;
330 const int slice_start = (numpixels * jobnr ) / nb_jobs;
331 const int slice_end = (numpixels * (jobnr+1)) / nb_jobs;
332 const double *dx = td->data[INDEX_DX][plane];
333 const double *dy = td->data[INDEX_DY][plane];
334 double *norm = td->data[INDEX_NORM][plane];
338 for (i = slice_start; i < slice_end; ++i) {
339 norm[i] = sqrt( pow(dx[i], 2) + pow(dy[i], 2));
342 const double *dxy = td->data[INDEX_DXY][plane];
343 for (i = slice_start; i < slice_end; ++i) {
344 norm[i] = sqrt( pow(dx[i], 2) + 4 * pow(dxy[i], 2) + pow(dy[i], 2) );
353 * Utility function for setting up differentiation data/metadata.
355 * @param ctx the filter context.
356 * @param td to be used for passing data between threads.
357 * @param ord ord of differentiation.
358 * @param dir direction of differentiation.
359 * @param src index of source used for differentiation.
360 * @param dst index destination used for saving differentiation result.
361 * @param dim maximum dimension in current direction.
362 * @param nb_threads number of threads to use.
364 static void av_always_inline
365 get_deriv(AVFilterContext *ctx, ThreadData *td, int ord, int dir,
366 int src, int dst, int dim, int nb_threads) {
367 td->meta_data[INDEX_ORD] = ord;
368 td->meta_data[INDEX_DIR] = dir;
369 td->meta_data[INDEX_SRC] = src;
370 td->meta_data[INDEX_DST] = dst;
371 ctx->internal->execute(ctx, slice_get_derivative, td, NULL, FFMIN(dim, nb_threads));
375 * Main control function for calculating gaussian derivatives.
377 * @param ctx the filter context.
378 * @param td holds the buffers used for storing results.
380 * @return 0 in case of success, a negative value corresponding to an
381 * AVERROR code in case of failure.
383 static int get_derivative(AVFilterContext *ctx, ThreadData *td)
385 ColorConstancyContext *s = ctx->priv;
386 int nb_threads = s->nb_threads;
387 int height = s->planeheight[1];
388 int width = s->planewidth[1];
392 if (!s->sigma) { // Only copy once
393 get_deriv(ctx, td, 0, DIR_X, 0 , INDEX_NORM, height, nb_threads);
395 get_deriv(ctx, td, 0, DIR_X, 0, INDEX_TEMP, height, nb_threads);
396 get_deriv(ctx, td, 0, DIR_Y, INDEX_TEMP, INDEX_NORM, width , nb_threads);
397 // save to INDEX_NORM because this will not be normalied and
398 // end gry edge filter expects result to be found in INDEX_NORM
403 get_deriv(ctx, td, 1, DIR_X, 0, INDEX_TEMP, height, nb_threads);
404 get_deriv(ctx, td, 0, DIR_Y, INDEX_TEMP, INDEX_DX, width , nb_threads);
406 get_deriv(ctx, td, 0, DIR_X, 0, INDEX_TEMP, height, nb_threads);
407 get_deriv(ctx, td, 1, DIR_Y, INDEX_TEMP, INDEX_DY, width , nb_threads);
411 get_deriv(ctx, td, 2, DIR_X, 0, INDEX_TEMP, height, nb_threads);
412 get_deriv(ctx, td, 0, DIR_Y, INDEX_TEMP, INDEX_DX, width , nb_threads);
414 get_deriv(ctx, td, 0, DIR_X, 0, INDEX_TEMP, height, nb_threads);
415 get_deriv(ctx, td, 2, DIR_Y, INDEX_TEMP, INDEX_DY, width , nb_threads);
417 get_deriv(ctx, td, 1, DIR_X, 0, INDEX_TEMP, height, nb_threads);
418 get_deriv(ctx, td, 1, DIR_Y, INDEX_TEMP, INDEX_DXY, width , nb_threads);
422 av_log(ctx, AV_LOG_ERROR, "Unsupported difford value: %d.\n", s->difford);
423 return AVERROR(EINVAL);
429 * Slice function for grey edge algorithm that does partial summing/maximizing
430 * of gaussian derivatives.
432 * @param ctx the filter context.
433 * @param arg data to be passed between threads.
434 * @param jobnr current job nubmer.
435 * @param nb_jobs total number of jobs.
439 static int filter_slice_grey_edge(AVFilterContext* ctx, void* arg, int jobnr, int nb_jobs)
441 ColorConstancyContext *s = ctx->priv;
442 ThreadData *td = arg;
443 AVFrame *in = td->in;
444 int minknorm = s->minknorm;
445 const uint8_t thresh = 255;
448 for (plane = 0; plane < NUM_PLANES; ++plane) {
449 const int height = s->planeheight[plane];
450 const int width = s->planewidth[plane];
451 const int in_linesize = in->linesize[plane];
452 const int slice_start = (height * jobnr) / nb_jobs;
453 const int slice_end = (height * (jobnr+1)) / nb_jobs;
454 const uint8_t *img_data = in->data[plane];
455 const double *src = td->data[INDEX_NORM][plane];
456 double *dst = td->data[INDEX_DST][plane];
461 for (r = slice_start; r < slice_end; ++r) {
462 for (c = 0; c < width; ++c) {
463 dst[jobnr] = FFMAX( dst[jobnr], fabs(src[INDX2D(r, c, width)])
464 * (img_data[INDX2D(r, c, in_linesize)] < thresh) );
468 for (r = slice_start; r < slice_end; ++r) {
469 for (c = 0; c < width; ++c) {
470 dst[jobnr] += ( pow( fabs(src[INDX2D(r, c, width)] / 255.), minknorm)
471 * (img_data[INDX2D(r, c, in_linesize)] < thresh) );
480 * Main control function for grey edge algorithm.
482 * @param ctx the filter context.
483 * @param in frame to perfrom grey edge on.
485 * @return 0 in case of success, a negative value corresponding to an
486 * AVERROR code in case of failure.
488 static int filter_grey_edge(AVFilterContext *ctx, AVFrame *in)
490 ColorConstancyContext *s = ctx->priv;
492 int minknorm = s->minknorm;
493 int difford = s->difford;
494 double *white = s->white;
495 int nb_jobs = FFMIN3(s->planeheight[1], s->planewidth[1], s->nb_threads);
499 ret = setup_derivative_buffers(ctx, &td);
503 get_derivative(ctx, &td);
505 ctx->internal->execute(ctx, slice_normalize, &td, NULL, nb_jobs);
508 ctx->internal->execute(ctx, filter_slice_grey_edge, &td, NULL, nb_jobs);
510 for (plane = 0; plane < NUM_PLANES; ++plane) {
511 white[plane] = 0; // All values are absolute
512 for (job = 0; job < nb_jobs; ++job) {
513 white[plane] = FFMAX(white[plane] , td.data[INDEX_DST][plane][job]);
517 for (plane = 0; plane < NUM_PLANES; ++plane) {
519 for (job = 0; job < nb_jobs; ++job) {
520 white[plane] += td.data[INDEX_DST][plane][job];
522 white[plane] = pow(white[plane], 1./minknorm);
526 cleanup_derivative_buffers(&td, difford + 1, NUM_PLANES);
531 * Normalizes estimated illumination since only illumination vector
532 * direction is required for color constancy.
534 * @param light the estimated illumination to be normalized in place
536 static void normalize_light(double *light)
538 double abs_val = pow( pow(light[0], 2.0) + pow(light[1], 2.0) + pow(light[2], 2.0), 0.5);
541 // TODO: check if setting to 1.0 when estimated = 0.0 is the best thing to do
544 for (plane = 0; plane < NUM_PLANES; ++plane) {
548 for (plane = 0; plane < NUM_PLANES; ++plane) {
549 light[plane] = (light[plane] / abs_val);
550 if (!light[plane]) { // to avoid division by zero when correcting
558 * Redirects to corresponding algorithm estimation function and performs normalization
561 * @param ctx the filter context.
562 * @param in frame to perfrom estimation on.
564 * @return 0 in case of success, a negative value corresponding to an
565 * AVERROR code in case of failure.
567 static int illumination_estimation(AVFilterContext *ctx, AVFrame *in)
569 ColorConstancyContext *s = ctx->priv;
572 ret = filter_grey_edge(ctx, in);
574 av_log(ctx, AV_LOG_DEBUG, "Estimated illumination= %f %f %f\n",
575 s->white[0], s->white[1], s->white[2]);
576 normalize_light(s->white);
577 av_log(ctx, AV_LOG_DEBUG, "Estimated illumination after normalization= %f %f %f\n",
578 s->white[0], s->white[1], s->white[2]);
584 * Performs simple correction via diagonal transformation model.
586 * @param ctx the filter context.
587 * @param arg data to be passed between threads.
588 * @param jobnr current job nubmer.
589 * @param nb_jobs total number of jobs.
593 static int diagonal_transformation(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
595 ColorConstancyContext *s = ctx->priv;
596 ThreadData *td = arg;
597 AVFrame *in = td->in;
598 AVFrame *out = td->out;
601 for (plane = 0; plane < NUM_PLANES; ++plane) {
602 const int height = s->planeheight[plane];
603 const int width = s->planewidth[plane];
604 const int64_t numpixels = width * (int64_t)height;
605 const int slice_start = (numpixels * jobnr) / nb_jobs;
606 const int slice_end = (numpixels * (jobnr+1)) / nb_jobs;
607 const uint8_t *src = in->data[plane];
608 uint8_t *dst = out->data[plane];
612 for (i = slice_start; i < slice_end; ++i) {
613 temp = src[i] / (s->white[plane] * SQRT3);
614 dst[i] = av_clip_uint8((int)(temp + 0.5));
621 * Main control function for correcting scene illumination based on
622 * estimated illumination.
624 * @param ctx the filter context.
625 * @param in holds frame to correct
626 * @param out holds corrected frame
628 static void chromatic_adaptation(AVFilterContext *ctx, AVFrame *in, AVFrame *out)
630 ColorConstancyContext *s = ctx->priv;
632 int nb_jobs = FFMIN3(s->planeheight[1], s->planewidth[1], s->nb_threads);
636 ctx->internal->execute(ctx, diagonal_transformation, &td, NULL, nb_jobs);
639 static int query_formats(AVFilterContext *ctx)
641 static const enum AVPixelFormat pix_fmts[] = {
642 // TODO: support more formats
643 // FIXME: error when saving to .jpg
648 return ff_set_common_formats(ctx, ff_make_format_list(pix_fmts));
651 static int config_props(AVFilterLink *inlink)
653 AVFilterContext *ctx = inlink->dst;
654 ColorConstancyContext *s = ctx->priv;
655 const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(inlink->format);
656 const double break_off_sigma = 3.0;
657 double sigma = s->sigma;
660 if (!floor(break_off_sigma * sigma + 0.5) && s->difford) {
661 av_log(ctx, AV_LOG_ERROR, "floor(%f * sigma) must be > 0 when difford > 0.\n", break_off_sigma);
662 return AVERROR(EINVAL);
665 s->filtersize = 2 * floor(break_off_sigma * sigma + 0.5) + 1;
666 if (ret=set_gauss(ctx)) {
670 s->nb_threads = ff_filter_get_nb_threads(ctx);
671 s->planewidth[1] = s->planewidth[2] = AV_CEIL_RSHIFT(inlink->w, desc->log2_chroma_w);
672 s->planewidth[0] = s->planewidth[3] = inlink->w;
673 s->planeheight[1] = s->planeheight[2] = AV_CEIL_RSHIFT(inlink->h, desc->log2_chroma_h);
674 s->planeheight[0] = s->planeheight[3] = inlink->h;
679 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
681 AVFilterContext *ctx = inlink->dst;
682 AVFilterLink *outlink = ctx->outputs[0];
686 ret = illumination_estimation(ctx, in);
691 if (av_frame_is_writable(in)) {
694 out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
696 av_log(ctx, AV_LOG_ERROR, "Out of memory while allocating output video buffer.\n");
697 return AVERROR(ENOMEM);
699 av_frame_copy_props(out, in);
701 chromatic_adaptation(ctx, in, out);
703 return ff_filter_frame(outlink, out);
706 static av_cold void uninit(AVFilterContext *ctx)
708 ColorConstancyContext *s = ctx->priv;
709 int difford = s->difford;
712 for (i = 0; i <= difford; ++i) {
713 av_freep(&s->gauss[i]);
717 static const AVFilterPad colorconstancy_inputs[] = {
720 .type = AVMEDIA_TYPE_VIDEO,
721 .config_props = config_props,
722 .filter_frame = filter_frame,
727 static const AVFilterPad colorconstancy_outputs[] = {
730 .type = AVMEDIA_TYPE_VIDEO,
735 #if CONFIG_GREYEDGE_FILTER
737 static const AVOption greyedge_options[] = {
738 { "difford", "set differentiation order", OFFSET(difford), AV_OPT_TYPE_INT, {.i64=1}, 0, 2, FLAGS },
739 { "minknorm", "set Minkowski norm", OFFSET(minknorm), AV_OPT_TYPE_INT, {.i64=1}, 0, 20, FLAGS },
740 { "sigma", "set sigma", OFFSET(sigma), AV_OPT_TYPE_DOUBLE, {.dbl=1}, 0.0, 1024.0, FLAGS },
744 AVFILTER_DEFINE_CLASS(greyedge);
746 AVFilter ff_vf_greyedge = {
748 .description = NULL_IF_CONFIG_SMALL("Estimates scene illumination by grey edge assumption."),
749 .priv_size = sizeof(ColorConstancyContext),
750 .priv_class = &greyedge_class,
751 .query_formats = query_formats,
753 .inputs = colorconstancy_inputs,
754 .outputs = colorconstancy_outputs,
755 .flags = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS,
758 #endif /* CONFIG_GREY_EDGE_FILTER */