git.sesse.net Git - ffmpeg/blob - libavfilter/vf_colorconstancy.c

   1 /*
   2  * Copyright (c) 2018 Mina Sami
   3  *
   4  * This file is part of FFmpeg.
   5  *
   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.
  10  *
  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.
  15  *
  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
  19  */
  20
  21 /**
  22  * @file
  23  * Color Constancy filter
  24  *
  25  * @see http://colorconstancy.com/
  26  *
  27  * @cite
  28  * J. van de Weijer, Th. Gevers, A. Gijsenij "Edge-Based Color Constancy".
  29  */
  30
  31 #include "libavutil/bprint.h"
  32 #include "libavutil/imgutils.h"
  33 #include "libavutil/opt.h"
  34 #include "libavutil/pixdesc.h"
  35
  36 #include "avfilter.h"
  37 #include "formats.h"
  38 #include "internal.h"
  39 #include "video.h"
  40
  41 #include <math.h>
  42
  43 #define GREY_EDGE "greyedge"
  44
  45 #define NUM_PLANES    3
  46 #define MAX_DIFF_ORD  2
  47 #define MAX_META_DATA 4
  48 #define MAX_DATA      4
  49
  50 #define INDEX_TEMP 0
  51 #define INDEX_DX   1
  52 #define INDEX_DY   2
  53 #define INDEX_DXY  3
  54 #define INDEX_NORM INDEX_DX
  55 #define INDEX_SRC  0
  56 #define INDEX_DST  1
  57 #define INDEX_ORD  2
  58 #define INDEX_DIR  3
  59 #define DIR_X 0
  60 #define DIR_Y 1
  61
  62 /**
  63  * Used for passing data between threads.
  64  */
  65 typedef struct ThreadData {
  66     AVFrame *in, *out;
  67     int meta_data[MAX_META_DATA];
  68     double  *data[MAX_DATA][NUM_PLANES];
  69 } ThreadData;
  70
  71 /**
  72  * Common struct for all algorithms contexts.
  73  */
  74 typedef struct ColorConstancyContext {
  75     const AVClass *class;
  76
  77     int difford;
  78     int minknorm; /**< @minknorm = 0 : getMax instead */
  79     double sigma;
  80
  81     int nb_threads;
  82     int planeheight[4];
  83     int planewidth[4];
  84
  85     int filtersize;
  86     double *gauss[MAX_DIFF_ORD+1];
  87
  88     double white[NUM_PLANES];
  89 } ColorConstancyContext;
  90
  91 #define OFFSET(x) offsetof(ColorConstancyContext, x)
  92 #define FLAGS AV_OPT_FLAG_FILTERING_PARAM|AV_OPT_FLAG_VIDEO_PARAM
  93
  94 #define GINDX(s, i) ( (i) - ((s) >> 2) )
  95
  96 /**
  97  * Sets gauss filters used for calculating gauss derivatives. Filter size
  98  * depends on sigma which is a user option hence we calculate these
  99  * filters each time. Also each higher order depends on lower ones. Sigma
 100  * can be zero only at difford = 0, then we only convert data to double
 101  * instead.
 102  *
 103  * @param ctx the filter context.
 104  *
 105  * @return 0 in case of success, a negative value corresponding to an
 106  * AVERROR code in case of failure.
 107  */
 108 static int set_gauss(AVFilterContext *ctx)
 109 {
 110     ColorConstancyContext *s = ctx->priv;
 111     int filtersize = s->filtersize;
 112     int difford    = s->difford;
 113     double sigma   = s->sigma;
 114     double sum1, sum2;
 115     int i;
 116
 117     for (i = 0; i <= difford; ++i) {
 118         s->gauss[i] = av_mallocz_array(filtersize, sizeof(*s->gauss[i]));
 119         if (!s->gauss[i]) {
 120             for (; i >= 0; --i) {
 121                 av_freep(&s->gauss[i]);
 122             }
 123             av_log(ctx, AV_LOG_ERROR, "Out of memory while allocating gauss buffers.\n");
 124             return AVERROR(ENOMEM);
 125         }
 126     }
 127
 128     // Order 0
 129     av_log(ctx, AV_LOG_TRACE, "Setting 0-d gauss with filtersize = %d.\n", filtersize);
 130     sum1 = 0.0;
 131     if (!sigma) {
 132         s->gauss[0][0] = 1; // Copying data to double instead of convolution
 133     } else {
 134         for (i = 0; i < filtersize; ++i) {
 135             s->gauss[0][i] = exp(- pow(GINDX(filtersize, i), 2.) / (2 * sigma * sigma)) / ( sqrt(2 * M_PI) * sigma );
 136             sum1 += s->gauss[0][i];
 137         }
 138         for (i = 0; i < filtersize; ++i) {
 139             s->gauss[0][i] /= sum1;
 140         }
 141     }
 142     // Order 1
 143     if (difford > 0) {
 144         av_log(ctx, AV_LOG_TRACE, "Setting 1-d gauss with filtersize = %d.\n", filtersize);
 145         sum1 = 0.0;
 146         for (i = 0; i < filtersize; ++i) {
 147             s->gauss[1][i] = - (GINDX(filtersize, i) / pow(sigma, 2)) * s->gauss[0][i];
 148             sum1 += s->gauss[1][i] *GINDX(filtersize, i);
 149         }
 150
 151         for (i = 0; i < filtersize; ++i) {
 152             s->gauss[1][i] /= sum1;
 153         }
 154
 155         // Order 2
 156         if (difford > 1) {
 157             av_log(ctx, AV_LOG_TRACE, "Setting 2-d gauss with filtersize = %d.\n", filtersize);
 158             sum1 = 0.0;
 159             for (i = 0; i < filtersize; ++i) {
 160                 s->gauss[2][i] = ( pow(GINDX(filtersize, i), 2) / pow(sigma, 4) - 1/pow(sigma, 2) )
 161                                  * s->gauss[0][i];
 162                 sum1 += s->gauss[2][i];
 163             }
 164
 165             sum2 = 0.0;
 166             for (i = 0; i < filtersize; ++i) {
 167                 s->gauss[2][i] -= sum1 / (filtersize);
 168                 sum2 += (0.5 * GINDX(filtersize, i) * GINDX(filtersize, i) * s->gauss[2][i]);
 169             }
 170             for (i = 0; i < filtersize ; ++i) {
 171                 s->gauss[2][i] /= sum2;
 172             }
 173         }
 174     }
 175     return 0;
 176 }
 177
 178 /**
 179  * Frees up buffers used by grey edge for storing derivatives final
 180  * and intermidiate results. Number of buffers and number of planes
 181  * for last buffer are given so it can be safely called at allocation
 182  * failure instances.
 183  *
 184  * @param td holds the buffers.
 185  * @param nb_buff number of buffers to be freed.
 186  * @param nb_planes number of planes for last buffer to be freed.
 187  */
 188 static void cleanup_derivative_buffers(ThreadData *td, int nb_buff, int nb_planes)
 189 {
 190     int b, p;
 191
 192     for (b = 0; b < nb_buff; ++b) {
 193         for (p = 0; p < NUM_PLANES; ++p) {
 194             av_freep(&td->data[b][p]);
 195         }
 196     }
 197     // Final buffer may not be fully allocated at fail cases
 198     for (p = 0; p < nb_planes; ++p) {
 199         av_freep(&td->data[b][p]);
 200     }
 201 }
 202
 203 /**
 204  * Allocates buffers used by grey edge for storing derivatives final
 205  * and intermidiate results.
 206  *
 207  * @param ctx the filter context.
 208  * @param td holds the buffers.
 209  *
 210  * @return 0 in case of success, a negative value corresponding to an
 211  * AVERROR code in case of failure.
 212  */
 213 static int setup_derivative_buffers(AVFilterContext* ctx, ThreadData *td)
 214 {
 215     ColorConstancyContext *s = ctx->priv;
 216     int nb_buff = s->difford + 1;
 217     int b, p;
 218
 219     av_log(ctx, AV_LOG_TRACE, "Allocating %d buffer(s) for grey edge.\n", nb_buff);
 220     for (b = 0; b <= nb_buff; ++b) { // We need difford + 1 buffers
 221         for (p = 0; p < NUM_PLANES; ++p) {
 222             td->data[b][p] = av_mallocz_array(s->planeheight[p] * s->planewidth[p], sizeof(*td->data[b][p]));
 223             if (!td->data[b][p]) {
 224                 cleanup_derivative_buffers(td, b + 1, p);
 225                 av_log(ctx, AV_LOG_ERROR, "Out of memory while allocating derivatives buffers.\n");
 226                 return AVERROR(ENOMEM);
 227             }
 228         }
 229     }
 230     return 0;
 231 }
 232
 233 #define CLAMP(x, mx) av_clip((x), 0, (mx-1))
 234 #define INDX2D(r, c, w) ( (r) * (w) + (c) )
 235 #define GAUSS(s, sr, sc, sls, sh, sw, g) ( (s)[ INDX2D(CLAMP((sr), (sh)), CLAMP((sc), (sw)), (sls)) ] * (g) )
 236
 237 /**
 238  * Slice calculation of gaussian derivatives. Applies 1-D gaussian derivative filter
 239  * either horizontally or vertically according to meta data given in thread data.
 240  * When convoluting horizontally source is always the in frame withing thread data
 241  * while when convoluting vertically source is a buffer.
 242  *
 243  * @param ctx the filter context.
 244  * @param arg data to be passed between threads.
 245  * @param jobnr current job nubmer.
 246  * @param nb_jobs total number of jobs.
 247  *
 248  * @return 0.
 249  */
 250 static int slice_get_derivative(AVFilterContext* ctx, void* arg, int jobnr, int nb_jobs)
 251 {
 252     ColorConstancyContext *s = ctx->priv;
 253     ThreadData *td = arg;
 254     AVFrame *in = td->in;
 255     const int ord = td->meta_data[INDEX_ORD];
 256     const int dir = td->meta_data[INDEX_DIR];
 257     const int src_index  = td->meta_data[INDEX_SRC];
 258     const int dst_index  = td->meta_data[INDEX_DST];
 259     const int filtersize = s->filtersize;
 260     const double *gauss  = s->gauss[ord];
 261     int plane;
 262
 263     for (plane = 0; plane < NUM_PLANES; ++plane) {
 264         const int height      = s->planeheight[plane];
 265         const int width       = s->planewidth[plane];
 266         const int in_linesize = in->linesize[plane];
 267         double *dst = td->data[dst_index][plane];
 268         int slice_start, slice_end;
 269         int r, c, g;
 270
 271         if (dir == DIR_X) {
 272             /** Applying gauss horizontally along each row */
 273             const uint8_t *src = in->data[plane];
 274             slice_start = (height * jobnr      ) / nb_jobs;
 275             slice_end   = (height * (jobnr + 1)) / nb_jobs;
 276
 277             for (r = slice_start; r < slice_end; ++r) {
 278                 for (c = 0; c < width; ++c) {
 279                     dst[INDX2D(r, c, width)] = 0;
 280                     for (g = 0; g < filtersize; ++g) {
 281                         dst[INDX2D(r, c, width)] += GAUSS(src, r,                        c + GINDX(filtersize, g),
 282                                                           in_linesize, height, width, gauss[GINDX(filtersize, g)]);
 283                     }
 284                 }
 285             }
 286         } else {
 287             /** Applying gauss vertically along each column */
 288             const double *src = td->data[src_index][plane];
 289             slice_start = (width * jobnr      ) / nb_jobs;
 290             slice_end   = (width * (jobnr + 1)) / nb_jobs;
 291
 292             for (c = slice_start; c < slice_end; ++c) {
 293                 for (r = 0; r < height; ++r) {
 294                     dst[INDX2D(r, c, width)] = 0;
 295                     for (g = 0; g < filtersize; ++g) {
 296                         dst[INDX2D(r, c, width)] += GAUSS(src, r + GINDX(filtersize, g), c,
 297                                                           width, height, width, gauss[GINDX(filtersize, g)]);
 298                     }
 299                 }
 300             }
 301         }
 302
 303     }
 304     return 0;
 305 }
 306
 307 /**
 308  * Slice Frobius normalization of gaussian derivatives. Only called for difford values of
 309  * 1 or 2.
 310  *
 311  * @param ctx the filter context.
 312  * @param arg data to be passed between threads.
 313  * @param jobnr current job nubmer.
 314  * @param nb_jobs total number of jobs.
 315  *
 316  * @return 0.
 317  */
 318 static int slice_normalize(AVFilterContext* ctx, void* arg, int jobnr, int nb_jobs)
 319 {
 320     ColorConstancyContext *s = ctx->priv;
 321     ThreadData *td = arg;
 322     const int difford = s->difford;
 323     int plane;
 324
 325     for (plane = 0; plane < NUM_PLANES; ++plane) {
 326         const int height = s->planeheight[plane];
 327         const int width  = s->planewidth[plane];
 328         const int64_t numpixels = width * (int64_t)height;
 329         const int slice_start   = (numpixels * jobnr    ) / nb_jobs;
 330         const int slice_end     = (numpixels * (jobnr+1)) / nb_jobs;
 331         const double *dx = td->data[INDEX_DX][plane];
 332         const double *dy = td->data[INDEX_DY][plane];
 333         double *norm = td->data[INDEX_NORM][plane];
 334         int i;
 335
 336         if (difford == 1) {
 337             for (i = slice_start; i < slice_end; ++i) {
 338                 norm[i] = sqrt( pow(dx[i], 2) + pow(dy[i], 2));
 339             }
 340         } else {
 341             const double *dxy = td->data[INDEX_DXY][plane];
 342             for (i = slice_start; i < slice_end; ++i) {
 343                 norm[i] = sqrt( pow(dx[i], 2) + 4 * pow(dxy[i], 2) + pow(dy[i], 2) );
 344             }
 345         }
 346     }
 347
 348     return 0;
 349 }
 350
 351 /**
 352  * Utility function for setting up differentiation data/metadata.
 353  *
 354  * @param ctx the filter context.
 355  * @param td to be used for passing data between threads.
 356  * @param ord ord of differentiation.
 357  * @param dir direction of differentiation.
 358  * @param src index of source used for differentiation.
 359  * @param dst index destination used for saving differentiation result.
 360  * @param dim maximum dimension in current direction.
 361  * @param nb_threads number of threads to use.
 362  */
 363 static void av_always_inline
 364 get_deriv(AVFilterContext *ctx, ThreadData *td, int ord, int dir,
 365           int src, int dst, int dim, int nb_threads) {
 366     td->meta_data[INDEX_ORD] = ord;
 367     td->meta_data[INDEX_DIR] = dir;
 368     td->meta_data[INDEX_SRC] = src;
 369     td->meta_data[INDEX_DST] = dst;
 370     ctx->internal->execute(ctx, slice_get_derivative, td, NULL, FFMIN(dim, nb_threads));
 371 }
 372
 373 /**
 374  * Main control function for calculating gaussian derivatives.
 375  *
 376  * @param ctx the filter context.
 377  * @param td holds the buffers used for storing results.
 378  *
 379  * @return 0 in case of success, a negative value corresponding to an
 380  * AVERROR code in case of failure.
 381  */
 382 static int get_derivative(AVFilterContext *ctx, ThreadData *td)
 383 {
 384     ColorConstancyContext *s = ctx->priv;
 385     int nb_threads = s->nb_threads;
 386     int height = s->planeheight[1];
 387     int width  = s->planewidth[1];
 388
 389     switch(s->difford) {
 390     case 0:
 391         if (!s->sigma) { // Only copy once
 392             get_deriv(ctx, td, 0, DIR_X, 0         , INDEX_NORM, height, nb_threads);
 393         } else {
 394             get_deriv(ctx, td, 0, DIR_X, 0,          INDEX_TEMP, height, nb_threads);
 395             get_deriv(ctx, td, 0, DIR_Y, INDEX_TEMP, INDEX_NORM, width , nb_threads);
 396             // save to INDEX_NORM because this will not be normalied and
 397             // end gry edge filter expects result to be found in INDEX_NORM
 398         }
 399         return 0;
 400
 401     case 1:
 402         get_deriv(ctx, td, 1, DIR_X, 0,          INDEX_TEMP, height, nb_threads);
 403         get_deriv(ctx, td, 0, DIR_Y, INDEX_TEMP, INDEX_DX,   width , nb_threads);
 404
 405         get_deriv(ctx, td, 0, DIR_X, 0,          INDEX_TEMP, height, nb_threads);
 406         get_deriv(ctx, td, 1, DIR_Y, INDEX_TEMP, INDEX_DY,   width , nb_threads);
 407         return 0;
 408
 409     case 2:
 410         get_deriv(ctx, td, 2, DIR_X, 0,          INDEX_TEMP, height, nb_threads);
 411         get_deriv(ctx, td, 0, DIR_Y, INDEX_TEMP, INDEX_DX,   width , nb_threads);
 412
 413         get_deriv(ctx, td, 0, DIR_X, 0,          INDEX_TEMP, height, nb_threads);
 414         get_deriv(ctx, td, 2, DIR_Y, INDEX_TEMP, INDEX_DY,   width , nb_threads);
 415
 416         get_deriv(ctx, td, 1, DIR_X, 0,          INDEX_TEMP, height, nb_threads);
 417         get_deriv(ctx, td, 1, DIR_Y, INDEX_TEMP, INDEX_DXY,  width , nb_threads);
 418         return 0;
 419
 420     default:
 421         av_log(ctx, AV_LOG_ERROR, "Unsupported difford value: %d.\n", s->difford);
 422         return AVERROR(EINVAL);
 423     }
 424
 425 }
 426
 427 /**
 428  * Slice function for grey edge algorithm that does partial summing/maximizing
 429  * of gaussian derivatives.
 430  *
 431  * @param ctx the filter context.
 432  * @param arg data to be passed between threads.
 433  * @param jobnr current job nubmer.
 434  * @param nb_jobs total number of jobs.
 435  *
 436  * @return 0.
 437  */
 438 static int filter_slice_grey_edge(AVFilterContext* ctx, void* arg, int jobnr, int nb_jobs)
 439 {
 440     ColorConstancyContext *s = ctx->priv;
 441     ThreadData *td = arg;
 442     AVFrame *in    = td->in;
 443     int minknorm   = s->minknorm;
 444     const uint8_t thresh = 255;
 445     int plane;
 446
 447     for (plane = 0; plane < NUM_PLANES; ++plane) {
 448         const int height        = s->planeheight[plane];
 449         const int width         = s->planewidth[plane];
 450         const int in_linesize   = in->linesize[plane];
 451         const int slice_start   = (height * jobnr) / nb_jobs;
 452         const int slice_end     = (height * (jobnr+1)) / nb_jobs;
 453         const uint8_t *img_data = in->data[plane];
 454         const double *src       = td->data[INDEX_NORM][plane];
 455         double *dst             = td->data[INDEX_DST][plane];
 456         int r, c;
 457
 458         dst[jobnr] = 0;
 459         if (!minknorm) {
 460             for (r = slice_start; r < slice_end; ++r) {
 461                 for (c = 0; c < width; ++c) {
 462                     dst[jobnr] = FFMAX( dst[jobnr], fabs(src[INDX2D(r, c, width)])
 463                                         * (img_data[INDX2D(r, c, in_linesize)] < thresh) );
 464                 }
 465             }
 466         } else {
 467             for (r = slice_start; r < slice_end; ++r) {
 468                 for (c = 0; c < width; ++c) {
 469                     dst[jobnr] += ( pow( fabs(src[INDX2D(r, c, width)] / 255.), minknorm)
 470                                     * (img_data[INDX2D(r, c, in_linesize)] < thresh) );
 471                 }
 472             }
 473         }
 474     }
 475     return 0;
 476 }
 477
 478 /**
 479  * Main control function for grey edge algorithm.
 480  *
 481  * @param ctx the filter context.
 482  * @param in frame to perfrom grey edge on.
 483  *
 484  * @return 0 in case of success, a negative value corresponding to an
 485  * AVERROR code in case of failure.
 486  */
 487 static int filter_grey_edge(AVFilterContext *ctx, AVFrame *in)
 488 {
 489     ColorConstancyContext *s = ctx->priv;
 490     ThreadData td;
 491     int minknorm  = s->minknorm;
 492     int difford   = s->difford;
 493     double *white = s->white;
 494     int nb_jobs   = FFMIN3(s->planeheight[1], s->planewidth[1], s->nb_threads);
 495     int plane, job, ret;
 496
 497     td.in = in;
 498     ret = setup_derivative_buffers(ctx, &td);
 499     if (ret) {
 500         return ret;
 501     }
 502     get_derivative(ctx, &td);
 503     if (difford > 0) {
 504         ctx->internal->execute(ctx, slice_normalize, &td, NULL, nb_jobs);
 505     }
 506
 507     ctx->internal->execute(ctx, filter_slice_grey_edge, &td, NULL, nb_jobs);
 508     if (!minknorm) {
 509         for (plane = 0; plane < NUM_PLANES; ++plane) {
 510             white[plane] = 0; // All values are absolute
 511             for (job = 0; job < nb_jobs; ++job) {
 512                 white[plane] = FFMAX(white[plane] , td.data[INDEX_DST][plane][job]);
 513             }
 514         }
 515     } else {
 516         for (plane = 0; plane < NUM_PLANES; ++plane) {
 517             white[plane] = 0;
 518             for (job = 0; job < nb_jobs; ++job) {
 519                 white[plane] += td.data[INDEX_DST][plane][job];
 520             }
 521             white[plane] = pow(white[plane], 1./minknorm);
 522         }
 523     }
 524
 525     cleanup_derivative_buffers(&td, difford + 1, NUM_PLANES);
 526     return 0;
 527 }
 528
 529 /**
 530  * Normalizes estimated illumination since only illumination vector
 531  * direction is required for color constancy.
 532  *
 533  * @param light the estimated illumination to be normalized in place
 534  */
 535 static void normalize_light(double *light)
 536 {
 537     double abs_val = pow( pow(light[0], 2.0) + pow(light[1], 2.0) + pow(light[2], 2.0), 0.5);
 538     int plane;
 539
 540     // TODO: check if setting to 1.0 when estimated = 0.0 is the best thing to do
 541
 542     if (!abs_val) {
 543         for (plane = 0; plane < NUM_PLANES; ++plane) {
 544             light[plane] = 1.0;
 545         }
 546     } else {
 547         for (plane = 0; plane < NUM_PLANES; ++plane) {
 548             light[plane] = (light[plane] / abs_val);
 549             if (!light[plane]) { // to avoid division by zero when correcting
 550                 light[plane] = 1.0;
 551             }
 552         }
 553     }
 554 }
 555
 556 /**
 557  * Redirects to corresponding algorithm estimation function and performs normalization
 558  * after estimation.
 559  *
 560  * @param ctx the filter context.
 561  * @param in frame to perfrom estimation on.
 562  *
 563  * @return 0 in case of success, a negative value corresponding to an
 564  * AVERROR code in case of failure.
 565  */
 566 static int illumination_estimation(AVFilterContext *ctx, AVFrame *in)
 567 {
 568     ColorConstancyContext *s = ctx->priv;
 569     int ret;
 570
 571     ret = filter_grey_edge(ctx, in);
 572
 573     av_log(ctx, AV_LOG_DEBUG, "Estimated illumination= %f %f %f\n",
 574            s->white[0], s->white[1], s->white[2]);
 575     normalize_light(s->white);
 576     av_log(ctx, AV_LOG_DEBUG, "Estimated illumination after normalization= %f %f %f\n",
 577            s->white[0], s->white[1], s->white[2]);
 578
 579     return ret;
 580 }
 581
 582 /**
 583  * Performs simple correction via diagonal transformation model.
 584  *
 585  * @param ctx the filter context.
 586  * @param arg data to be passed between threads.
 587  * @param jobnr current job nubmer.
 588  * @param nb_jobs total number of jobs.
 589  *
 590  * @return 0.
 591  */
 592 static int diagonal_transformation(AVFilterContext *ctx, void *arg, int jobnr, int nb_jobs)
 593 {
 594     ColorConstancyContext *s = ctx->priv;
 595     ThreadData *td = arg;
 596     AVFrame *in = td->in;
 597     AVFrame *out = td->out;
 598     double sqrt3 = pow(3.0, 0.5);
 599     int plane;
 600
 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];
 609         double temp;
 610         unsigned i;
 611
 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));
 615         }
 616     }
 617     return 0;
 618 }
 619
 620 /**
 621  * Main control function for correcting scene illumination based on
 622  * estimated illumination.
 623  *
 624  * @param ctx the filter context.
 625  * @param in holds frame to correct
 626  * @param out holds corrected frame
 627  */
 628 static void chromatic_adaptation(AVFilterContext *ctx, AVFrame *in, AVFrame *out)
 629 {
 630     ColorConstancyContext *s = ctx->priv;
 631     ThreadData td;
 632     int nb_jobs = FFMIN3(s->planeheight[1], s->planewidth[1], s->nb_threads);
 633
 634     td.in  = in;
 635     td.out = out;
 636     ctx->internal->execute(ctx, diagonal_transformation, &td, NULL, nb_jobs);
 637 }
 638
 639 static int query_formats(AVFilterContext *ctx)
 640 {
 641     static const enum AVPixelFormat pix_fmts[] = {
 642         // TODO: support more formats
 643         // FIXME: error when saving to .jpg
 644         AV_PIX_FMT_GBRP,
 645         AV_PIX_FMT_NONE
 646     };
 647
 648     return ff_set_common_formats(ctx, ff_make_format_list(pix_fmts));
 649 }
 650
 651 static int config_props(AVFilterLink *inlink)
 652 {
 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;
 658     int ret;
 659
 660     if (!sigma && s->difford) {
 661         av_log(ctx, AV_LOG_ERROR, "Sigma can't be set to 0 when difford > 0.\n");
 662         return AVERROR(EINVAL);
 663     }
 664
 665     s->filtersize = 2 * floor(break_off_sigma * s->sigma + 0.5) + 1;
 666     if (ret=set_gauss(ctx)) {
 667         return ret;
 668     }
 669
 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;
 675
 676     return 0;
 677 }
 678
 679 static int filter_frame(AVFilterLink *inlink, AVFrame *in)
 680 {
 681     AVFilterContext *ctx = inlink->dst;
 682     AVFilterLink *outlink = ctx->outputs[0];
 683     AVFrame *out;
 684     int ret;
 685
 686     ret = illumination_estimation(ctx, in);
 687     if (ret) {
 688         return ret;
 689     }
 690
 691     if (av_frame_is_writable(in)) {
 692         out = in;
 693     } else {
 694         out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
 695         if (!out) {
 696             av_log(ctx, AV_LOG_ERROR, "Out of memory while allocating output video buffer.\n");
 697             return AVERROR(ENOMEM);
 698         }
 699         av_frame_copy_props(out, in);
 700     }
 701     chromatic_adaptation(ctx, in, out);
 702
 703     return ff_filter_frame(outlink, out);
 704 }
 705
 706 static av_cold void uninit(AVFilterContext *ctx)
 707 {
 708     ColorConstancyContext *s = ctx->priv;
 709     int difford = s->difford;
 710     int i;
 711
 712     for (i = 0; i <= difford; ++i) {
 713         av_freep(&s->gauss[i]);
 714     }
 715 }
 716
 717 static const AVFilterPad colorconstancy_inputs[] = {
 718     {
 719         .name         = "default",
 720         .type         = AVMEDIA_TYPE_VIDEO,
 721         .config_props = config_props,
 722         .filter_frame = filter_frame,
 723     },
 724     { NULL }
 725 };
 726
 727 static const AVFilterPad colorconstancy_outputs[] = {
 728     {
 729         .name = "default",
 730         .type = AVMEDIA_TYPE_VIDEO,
 731     },
 732     { NULL }
 733 };
 734
 735 #if CONFIG_GREYEDGE_FILTER
 736
 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 },
 741     { NULL }
 742 };
 743
 744 AVFILTER_DEFINE_CLASS(greyedge);
 745
 746 AVFilter ff_vf_greyedge = {
 747     .name          = GREY_EDGE,
 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,
 752     .uninit        = uninit,
 753     .inputs        = colorconstancy_inputs,
 754     .outputs       = colorconstancy_outputs,
 755     .flags         = AVFILTER_FLAG_SUPPORT_TIMELINE_GENERIC | AVFILTER_FLAG_SLICE_THREADS,
 756 };
 757
 758 #endif /* CONFIG_GREY_EDGE_FILTER */