git.sesse.net Git - ffmpeg/blob - libavcodec/rpzaenc.c

   1 /*
   2  * QuickTime RPZA Video Encoder
   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 rpzaenc.c
  23  * QT RPZA Video Encoder by Todd Kirby <doubleshot@pacbell.net> and David Adler
  24  */
  25
  26 #include "libavutil/avassert.h"
  27 #include "libavutil/common.h"
  28 #include "libavutil/opt.h"
  29
  30 #include "avcodec.h"
  31 #include "internal.h"
  32 #include "put_bits.h"
  33
  34 typedef struct RpzaContext {
  35     AVClass *avclass;
  36
  37     int skip_frame_thresh;
  38     int start_one_color_thresh;
  39     int continue_one_color_thresh;
  40     int sixteen_color_thresh;
  41
  42     AVFrame *prev_frame;    // buffer for previous source frame
  43     PutBitContext pb;       // buffer for encoded frame data.
  44
  45     int frame_width;        // width in pixels of source frame
  46     int frame_height;       // height in pixesl of source frame
  47
  48     int first_frame;        // flag set to one when the first frame is being processed
  49                             // so that comparisons with previous frame data in not attempted
  50 } RpzaContext;
  51
  52 typedef enum channel_offset {
  53     RED = 2,
  54     GREEN = 1,
  55     BLUE = 0,
  56 } channel_offset;
  57
  58 typedef struct rgb {
  59     uint8_t r;
  60     uint8_t g;
  61     uint8_t b;
  62 } rgb;
  63
  64 #define SQR(x) ((x) * (x))
  65
  66 /* 15 bit components */
  67 #define GET_CHAN(color, chan) (((color) >> ((chan) * 5) & 0x1F) * 8)
  68 #define R(color) GET_CHAN(color, RED)
  69 #define G(color) GET_CHAN(color, GREEN)
  70 #define B(color) GET_CHAN(color, BLUE)
  71
  72 typedef struct BlockInfo {
  73     int row;
  74     int col;
  75     int block_width;
  76     int block_height;
  77     int image_width;
  78     int image_height;
  79     int block_index;
  80     uint16_t start;
  81     int rowstride;
  82     int blocks_per_row;
  83     int total_blocks;
  84 } BlockInfo;
  85
  86 static void get_colors(uint8_t *min, uint8_t *max, uint8_t color4[4][3])
  87 {
  88     uint8_t step;
  89
  90     color4[0][0] = min[0];
  91     color4[0][1] = min[1];
  92     color4[0][2] = min[2];
  93
  94     color4[3][0] = max[0];
  95     color4[3][1] = max[1];
  96     color4[3][2] = max[2];
  97
  98     // red components
  99     step = (color4[3][0] - color4[0][0] + 1) / 3;
 100     color4[1][0] = color4[0][0] + step;
 101     color4[2][0] = color4[3][0] - step;
 102
 103     // green components
 104     step = (color4[3][1] - color4[0][1] + 1) / 3;
 105     color4[1][1] = color4[0][1] + step;
 106     color4[2][1] = color4[3][1] - step;
 107
 108     // blue components
 109     step = (color4[3][2] - color4[0][2] + 1) / 3;
 110     color4[1][2] = color4[0][2] + step;
 111     color4[2][2] = color4[3][2] - step;
 112 }
 113
 114 /* Fill BlockInfo struct with information about a 4x4 block of the image */
 115 static int get_block_info(BlockInfo *bi, int block)
 116 {
 117     bi->row = block / bi->blocks_per_row;
 118     bi->col = block % bi->blocks_per_row;
 119
 120     // test for right edge block
 121     if (bi->col == bi->blocks_per_row - 1 && (bi->image_width % 4) != 0) {
 122         bi->block_width = bi->image_width % 4;
 123     } else {
 124         bi->block_width = 4;
 125     }
 126
 127     // test for bottom edge block
 128     if (bi->row == (bi->image_height / 4) && (bi->image_height % 4) != 0) {
 129         bi->block_height = bi->image_height % 4;
 130     } else {
 131         bi->block_height = 4;
 132     }
 133
 134     return block ? (bi->col * 4) + (bi->row * bi->rowstride * 4) : 0;
 135 }
 136
 137 static uint16_t rgb24_to_rgb555(uint8_t *rgb24)
 138 {
 139     uint16_t rgb555 = 0;
 140     uint32_t r, g, b;
 141
 142     r = rgb24[0] >> 3;
 143     g = rgb24[1] >> 3;
 144     b = rgb24[2] >> 3;
 145
 146     rgb555 |= (r << 10);
 147     rgb555 |= (g << 5);
 148     rgb555 |= (b << 0);
 149
 150     return rgb555;
 151 }
 152
 153 /*
 154  * Returns the total difference between two 24 bit color values
 155  */
 156 static int diff_colors(uint8_t *colorA, uint8_t *colorB)
 157 {
 158     int tot;
 159
 160     tot  = SQR(colorA[0] - colorB[0]);
 161     tot += SQR(colorA[1] - colorB[1]);
 162     tot += SQR(colorA[2] - colorB[2]);
 163
 164     return tot;
 165 }
 166
 167 /*
 168  * Returns the maximum channel difference
 169  */
 170 static int max_component_diff(uint16_t *colorA, uint16_t *colorB)
 171 {
 172     int diff, max = 0;
 173
 174     diff = FFABS(R(colorA[0]) - R(colorB[0]));
 175     if (diff > max) {
 176         max = diff;
 177     }
 178     diff = FFABS(G(colorA[0]) - G(colorB[0]));
 179     if (diff > max) {
 180         max = diff;
 181     }
 182     diff = FFABS(B(colorA[0]) - B(colorB[0]));
 183     if (diff > max) {
 184         max = diff;
 185     }
 186     return max * 8;
 187 }
 188
 189 /*
 190  * Find the channel that has the largest difference between minimum and maximum
 191  * color values. Put the minimum value in min, maximum in max and the channel
 192  * in chan.
 193  */
 194 static void get_max_component_diff(BlockInfo *bi, uint16_t *block_ptr,
 195                                    uint8_t *min, uint8_t *max, channel_offset *chan)
 196 {
 197     int x, y;
 198     uint8_t min_r, max_r, min_g, max_g, min_b, max_b;
 199     uint8_t r, g, b;
 200
 201     // fix warning about uninitialized vars
 202     min_r = min_g = min_b = UINT8_MAX;
 203     max_r = max_g = max_b = 0;
 204
 205     // loop thru and compare pixels
 206     for (y = 0; y < bi->block_height; y++) {
 207         for (x = 0; x < bi->block_width; x++){
 208             // TODO:  optimize
 209             min_r = FFMIN(R(block_ptr[x]), min_r);
 210             min_g = FFMIN(G(block_ptr[x]), min_g);
 211             min_b = FFMIN(B(block_ptr[x]), min_b);
 212
 213             max_r = FFMAX(R(block_ptr[x]), max_r);
 214             max_g = FFMAX(G(block_ptr[x]), max_g);
 215             max_b = FFMAX(B(block_ptr[x]), max_b);
 216         }
 217         block_ptr += bi->rowstride;
 218     }
 219
 220     r = max_r - min_r;
 221     g = max_g - min_g;
 222     b = max_b - min_b;
 223
 224     if (r > g && r > b) {
 225         *max = max_r;
 226         *min = min_r;
 227         *chan = RED;
 228     } else if (g > b && g >= r) {
 229         *max = max_g;
 230         *min = min_g;
 231         *chan = GREEN;
 232     } else {
 233         *max = max_b;
 234         *min = min_b;
 235         *chan = BLUE;
 236     }
 237 }
 238
 239 /*
 240  * Compare two 4x4 blocks to determine if the total difference between the
 241  * blocks is greater than the thresh parameter. Returns -1 if difference
 242  * exceeds threshold or zero otherwise.
 243  */
 244 static int compare_blocks(uint16_t *block1, uint16_t *block2, BlockInfo *bi, int thresh)
 245 {
 246     int x, y, diff = 0;
 247     for (y = 0; y < bi->block_height; y++) {
 248         for (x = 0; x < bi->block_width; x++) {
 249             diff = max_component_diff(&block1[x], &block2[x]);
 250             if (diff >= thresh) {
 251                 return -1;
 252             }
 253         }
 254         block1 += bi->rowstride;
 255         block2 += bi->rowstride;
 256     }
 257     return 0;
 258 }
 259
 260 /*
 261  * Determine the fit of one channel to another within a 4x4 block. This
 262  * is used to determine the best palette choices for 4-color encoding.
 263  */
 264 static int leastsquares(uint16_t *block_ptr, BlockInfo *bi,
 265                         channel_offset xchannel, channel_offset ychannel,
 266                         double *slope, double *y_intercept, double *correlation_coef)
 267 {
 268     double sumx = 0, sumy = 0, sumx2 = 0, sumy2 = 0, sumxy = 0,
 269            sumx_sq = 0, sumy_sq = 0, tmp, tmp2;
 270     int i, j, count;
 271     uint8_t x, y;
 272
 273     count = bi->block_height * bi->block_width;
 274
 275     if (count < 2)
 276         return -1;
 277
 278     for (i = 0; i < bi->block_height; i++) {
 279         for (j = 0; j < bi->block_width; j++){
 280             x = GET_CHAN(block_ptr[j], xchannel);
 281             y = GET_CHAN(block_ptr[j], ychannel);
 282             sumx += x;
 283             sumy += y;
 284             sumx2 += x * x;
 285             sumy2 += y * y;
 286             sumxy += x * y;
 287         }
 288         block_ptr += bi->rowstride;
 289     }
 290
 291     sumx_sq = sumx * sumx;
 292     tmp = (count * sumx2 - sumx_sq);
 293
 294     // guard against div/0
 295     if (tmp == 0)
 296         return -2;
 297
 298     sumy_sq = sumy * sumy;
 299
 300     *slope = (sumx * sumy - sumxy) / tmp;
 301     *y_intercept = (sumy - (*slope) * sumx) / count;
 302
 303     tmp2 = count * sumy2 - sumy_sq;
 304     if (tmp2 == 0) {
 305         *correlation_coef = 0.0;
 306     } else {
 307         *correlation_coef = (count * sumxy - sumx * sumy) /
 308             sqrt(tmp * tmp2);
 309     }
 310
 311     return 0; // success
 312 }
 313
 314 /*
 315  * Determine the amount of error in the leastsquares fit.
 316  */
 317 static int calc_lsq_max_fit_error(uint16_t *block_ptr, BlockInfo *bi,
 318                                   int min, int max, int tmp_min, int tmp_max,
 319                                   channel_offset xchannel, channel_offset ychannel)
 320 {
 321     int i, j, x, y;
 322     int err;
 323     int max_err = 0;
 324
 325     for (i = 0; i < bi->block_height; i++) {
 326         for (j = 0; j < bi->block_width; j++){
 327             int x_inc, lin_y, lin_x;
 328             x = GET_CHAN(block_ptr[j], xchannel);
 329             y = GET_CHAN(block_ptr[j], ychannel);
 330
 331             /* calculate x_inc as the 4-color index (0..3) */
 332             x_inc = floor( (x - min) * 3.0 / (max - min) + 0.5);
 333             x_inc = FFMAX(FFMIN(3, x_inc), 0);
 334
 335             /* calculate lin_y corresponding to x_inc */
 336             lin_y = (int)(tmp_min + (tmp_max - tmp_min) * x_inc / 3.0 + 0.5);
 337
 338             err = FFABS(lin_y - y);
 339             if (err > max_err)
 340                 max_err = err;
 341
 342             /* calculate lin_x corresponding to x_inc */
 343             lin_x = (int)(min + (max - min) * x_inc / 3.0 + 0.5);
 344
 345             err = FFABS(lin_x - x);
 346             if (err > max_err)
 347                 max_err += err;
 348         }
 349         block_ptr += bi->rowstride;
 350     }
 351
 352     return max_err;
 353 }
 354
 355 /*
 356  * Find the closest match to a color within the 4-color palette
 357  */
 358 static int match_color(uint16_t *color, uint8_t colors[4][3])
 359 {
 360     int ret = 0;
 361     int smallest_variance = INT_MAX;
 362     uint8_t dithered_color[3];
 363
 364     for (int channel = 0; channel < 3; channel++) {
 365         dithered_color[channel] = GET_CHAN(color[0], channel);
 366     }
 367
 368     for (int palette_entry = 0; palette_entry < 4; palette_entry++) {
 369         int variance = diff_colors(dithered_color, colors[palette_entry]);
 370
 371         if (variance < smallest_variance) {
 372             smallest_variance = variance;
 373             ret = palette_entry;
 374         }
 375     }
 376
 377     return ret;
 378 }
 379
 380 /*
 381  * Encode a block using the 4-color opcode and palette. return number of
 382  * blocks encoded (until we implement multi-block 4 color runs this will
 383  * always be 1)
 384  */
 385 static int encode_four_color_block(uint8_t *min_color, uint8_t *max_color,
 386                                    PutBitContext *pb, uint16_t *block_ptr, BlockInfo *bi)
 387 {
 388     int x, y, idx;
 389     uint8_t color4[4][3];
 390     uint16_t rounded_max, rounded_min;
 391
 392     // round min and max wider
 393     rounded_min = rgb24_to_rgb555(min_color);
 394     rounded_max = rgb24_to_rgb555(max_color);
 395
 396     // put a and b colors
 397     // encode 4 colors = first 16 bit color with MSB zeroed and...
 398     put_bits(pb, 16, rounded_max & ~0x8000);
 399     // ...second 16 bit color with MSB on.
 400     put_bits(pb, 16, rounded_min | 0x8000);
 401
 402     get_colors(min_color, max_color, color4);
 403
 404     for (y = 0; y < 4; y++) {
 405         for (x = 0; x < 4; x++) {
 406             idx = match_color(&block_ptr[x], color4);
 407             put_bits(pb, 2, idx);
 408         }
 409         block_ptr += bi->rowstride;
 410     }
 411     return 1; // num blocks encoded
 412 }
 413
 414 /*
 415  * Copy a 4x4 block from the current frame buffer to the previous frame buffer.
 416  */
 417 static void update_block_in_prev_frame(const uint16_t *src_pixels,
 418                                        uint16_t *dest_pixels,
 419                                        const BlockInfo *bi, int block_counter)
 420 {
 421     for (int y = 0; y < 4; y++) {
 422         memcpy(dest_pixels, src_pixels, 8);
 423         dest_pixels += bi->rowstride;
 424         src_pixels += bi->rowstride;
 425     }
 426 }
 427
 428 /*
 429  * update statistics for the specified block. If first_block,
 430  * it initializes the statistics.  Otherwise it updates the statistics IF THIS
 431  * BLOCK IS SUITABLE TO CONTINUE A 1-COLOR RUN. That is, it checks whether
 432  * the range of colors (since the routine was called first_block != 0) are
 433  * all close enough intensities to be represented by a single color.
 434
 435  * The routine returns 0 if this block is too different to be part of
 436  * the same run of 1-color blocks. The routine returns 1 if this
 437  * block can be part of the same 1-color block run.
 438
 439  * If the routine returns 1, it also updates its arguments to include
 440  * the statistics of this block. Otherwise, the stats are unchanged
 441  * and don't include the current block.
 442  */
 443 static int update_block_stats(RpzaContext *s, BlockInfo *bi, uint16_t *block,
 444                               uint8_t min_color[3], uint8_t max_color[3],
 445                               int *total_rgb, int *total_pixels,
 446                               uint8_t avg_color[3], int first_block)
 447 {
 448     int x, y;
 449     int is_in_range;
 450     int total_pixels_blk;
 451     int threshold;
 452
 453     uint8_t min_color_blk[3], max_color_blk[3];
 454     int total_rgb_blk[3];
 455     uint8_t avg_color_blk[3];
 456
 457     if (first_block) {
 458         min_color[0] = UINT8_MAX;
 459         min_color[1] = UINT8_MAX;
 460         min_color[2] = UINT8_MAX;
 461         max_color[0] = 0;
 462         max_color[1] = 0;
 463         max_color[2] = 0;
 464         total_rgb[0] = 0;
 465         total_rgb[1] = 0;
 466         total_rgb[2] = 0;
 467         *total_pixels = 0;
 468         threshold = s->start_one_color_thresh;
 469     } else {
 470         threshold = s->continue_one_color_thresh;
 471     }
 472
 473     /*
 474        The *_blk variables will include the current block.
 475        Initialize them based on the blocks so far.
 476      */
 477     min_color_blk[0] = min_color[0];
 478     min_color_blk[1] = min_color[1];
 479     min_color_blk[2] = min_color[2];
 480     max_color_blk[0] = max_color[0];
 481     max_color_blk[1] = max_color[1];
 482     max_color_blk[2] = max_color[2];
 483     total_rgb_blk[0] = total_rgb[0];
 484     total_rgb_blk[1] = total_rgb[1];
 485     total_rgb_blk[2] = total_rgb[2];
 486     total_pixels_blk = *total_pixels + bi->block_height * bi->block_width;
 487
 488     /*
 489        Update stats for this block's pixels
 490      */
 491     for (y = 0; y < bi->block_height; y++) {
 492         for (x = 0; x < bi->block_width; x++) {
 493             total_rgb_blk[0] += R(block[x]);
 494             total_rgb_blk[1] += G(block[x]);
 495             total_rgb_blk[2] += B(block[x]);
 496
 497             min_color_blk[0] = FFMIN(R(block[x]), min_color_blk[0]);
 498             min_color_blk[1] = FFMIN(G(block[x]), min_color_blk[1]);
 499             min_color_blk[2] = FFMIN(B(block[x]), min_color_blk[2]);
 500
 501             max_color_blk[0] = FFMAX(R(block[x]), max_color_blk[0]);
 502             max_color_blk[1] = FFMAX(G(block[x]), max_color_blk[1]);
 503             max_color_blk[2] = FFMAX(B(block[x]), max_color_blk[2]);
 504         }
 505         block += bi->rowstride;
 506     }
 507
 508     /*
 509        Calculate average color including current block.
 510      */
 511     avg_color_blk[0] = total_rgb_blk[0] / total_pixels_blk;
 512     avg_color_blk[1] = total_rgb_blk[1] / total_pixels_blk;
 513     avg_color_blk[2] = total_rgb_blk[2] / total_pixels_blk;
 514
 515     /*
 516        Are all the pixels within threshold of the average color?
 517      */
 518     is_in_range = (max_color_blk[0] - avg_color_blk[0] <= threshold &&
 519                    max_color_blk[1] - avg_color_blk[1] <= threshold &&
 520                    max_color_blk[2] - avg_color_blk[2] <= threshold &&
 521                    avg_color_blk[0] - min_color_blk[0] <= threshold &&
 522                    avg_color_blk[1] - min_color_blk[1] <= threshold &&
 523                    avg_color_blk[2] - min_color_blk[2] <= threshold);
 524
 525     if (is_in_range) {
 526         /*
 527            Set the output variables to include this block.
 528          */
 529         min_color[0] = min_color_blk[0];
 530         min_color[1] = min_color_blk[1];
 531         min_color[2] = min_color_blk[2];
 532         max_color[0] = max_color_blk[0];
 533         max_color[1] = max_color_blk[1];
 534         max_color[2] = max_color_blk[2];
 535         total_rgb[0] = total_rgb_blk[0];
 536         total_rgb[1] = total_rgb_blk[1];
 537         total_rgb[2] = total_rgb_blk[2];
 538         *total_pixels = total_pixels_blk;
 539         avg_color[0] = avg_color_blk[0];
 540         avg_color[1] = avg_color_blk[1];
 541         avg_color[2] = avg_color_blk[2];
 542     }
 543
 544     return is_in_range;
 545 }
 546
 547 static void rpza_encode_stream(RpzaContext *s, const AVFrame *pict)
 548 {
 549     BlockInfo bi;
 550     int block_counter = 0;
 551     int n_blocks;
 552     int total_blocks;
 553     int prev_block_offset;
 554     int block_offset = 0;
 555     uint8_t min = 0, max = 0;
 556     channel_offset chan;
 557     int i;
 558     int tmp_min, tmp_max;
 559     int total_rgb[3];
 560     uint8_t avg_color[3];
 561     int pixel_count;
 562     uint8_t min_color[3], max_color[3];
 563     double slope, y_intercept, correlation_coef;
 564     uint16_t *src_pixels = (uint16_t *)pict->data[0];
 565     uint16_t *prev_pixels = (uint16_t *)s->prev_frame->data[0];
 566
 567     /* Number of 4x4 blocks in frame. */
 568     total_blocks = ((s->frame_width + 3) / 4) * ((s->frame_height + 3) / 4);
 569
 570     bi.image_width = s->frame_width;
 571     bi.image_height = s->frame_height;
 572     bi.rowstride = pict->linesize[0] / 2;
 573
 574     bi.blocks_per_row = (s->frame_width + 3) / 4;
 575
 576     while (block_counter < total_blocks) {
 577         // SKIP CHECK
 578         // make sure we have a valid previous frame and we're not writing
 579         // a key frame
 580         if (!s->first_frame) {
 581             n_blocks = 0;
 582             prev_block_offset = 0;
 583
 584             while (n_blocks < 32 && block_counter + n_blocks < total_blocks) {
 585
 586                 block_offset = get_block_info(&bi, block_counter + n_blocks);
 587
 588                 // multi-block opcodes cannot span multiple rows.
 589                 // If we're starting a new row, break out and write the opcode
 590                 /* TODO: Should eventually use bi.row here to determine when a
 591                    row break occurs, but that is currently breaking the
 592                    quicktime player. This is probably due to a bug in the
 593                    way I'm calculating the current row.
 594                  */
 595                 if (prev_block_offset && block_offset - prev_block_offset > 12) {
 596                     break;
 597                 }
 598
 599                 prev_block_offset = block_offset;
 600
 601                 if (compare_blocks(&prev_pixels[block_offset],
 602                                    &src_pixels[block_offset], &bi, s->skip_frame_thresh) != 0) {
 603                     // write out skipable blocks
 604                     if (n_blocks) {
 605
 606                         // write skip opcode
 607                         put_bits(&s->pb, 8, 0x80 | (n_blocks - 1));
 608                         block_counter += n_blocks;
 609
 610                         goto post_skip;
 611                     }
 612                     break;
 613                 }
 614
 615                 /*
 616                  * NOTE: we don't update skipped blocks in the previous frame buffer
 617                  * since skipped needs always to be compared against the first skipped
 618                  * block to avoid artifacts during gradual fade in/outs.
 619                  */
 620
 621                 // update_block_in_prev_frame(&src_pixels[block_offset],
 622                 //   &prev_pixels[block_offset], &bi, block_counter + n_blocks);
 623
 624                 n_blocks++;
 625             }
 626
 627             // we're either at the end of the frame or we've reached the maximum
 628             // of 32 blocks in a run. Write out the run.
 629             if (n_blocks) {
 630                 // write skip opcode
 631                 put_bits(&s->pb, 8, 0x80 | (n_blocks - 1));
 632                 block_counter += n_blocks;
 633
 634                 continue;
 635             }
 636
 637         } else {
 638             block_offset = get_block_info(&bi, block_counter);
 639         }
 640 post_skip :
 641
 642         // ONE COLOR CHECK
 643         if (update_block_stats(s, &bi, &src_pixels[block_offset],
 644                                min_color, max_color,
 645                                total_rgb, &pixel_count, avg_color, 1)) {
 646             prev_block_offset = block_offset;
 647
 648             n_blocks = 1;
 649
 650             /* update this block in the previous frame buffer */
 651             update_block_in_prev_frame(&src_pixels[block_offset],
 652                                        &prev_pixels[block_offset], &bi, block_counter + n_blocks);
 653
 654             // check for subsequent blocks with the same color
 655             while (n_blocks < 32 && block_counter + n_blocks < total_blocks) {
 656                 block_offset = get_block_info(&bi, block_counter + n_blocks);
 657
 658                 // multi-block opcodes cannot span multiple rows.
 659                 // If we've hit end of a row, break out and write the opcode
 660                 if (block_offset - prev_block_offset > 12) {
 661                     break;
 662                 }
 663
 664                 if (!update_block_stats(s, &bi, &src_pixels[block_offset],
 665                                         min_color, max_color,
 666                                         total_rgb, &pixel_count, avg_color, 0)) {
 667                     break;
 668                 }
 669
 670                 prev_block_offset = block_offset;
 671
 672                 /* update this block in the previous frame buffer */
 673                 update_block_in_prev_frame(&src_pixels[block_offset],
 674                                            &prev_pixels[block_offset], &bi, block_counter + n_blocks);
 675
 676                 n_blocks++;
 677             }
 678
 679             // write one color opcode.
 680             put_bits(&s->pb, 8, 0xa0 | (n_blocks - 1));
 681             // write color to encode.
 682             put_bits(&s->pb, 16, rgb24_to_rgb555(avg_color));
 683             // skip past the blocks we've just encoded.
 684             block_counter += n_blocks;
 685         } else { // FOUR COLOR CHECK
 686             int err = 0;
 687
 688             // get max component diff for block
 689             get_max_component_diff(&bi, &src_pixels[block_offset], &min, &max, &chan);
 690
 691             min_color[0] = 0;
 692             max_color[0] = 0;
 693             min_color[1] = 0;
 694             max_color[1] = 0;
 695             min_color[2] = 0;
 696             max_color[2] = 0;
 697
 698             // run least squares against other two components
 699             for (i = 0; i < 3; i++) {
 700                 if (i == chan) {
 701                     min_color[i] = min;
 702                     max_color[i] = max;
 703                     continue;
 704                 }
 705
 706                 slope = y_intercept = correlation_coef = 0;
 707
 708                 if (leastsquares(&src_pixels[block_offset], &bi, chan, i,
 709                                  &slope, &y_intercept, &correlation_coef)) {
 710                     min_color[i] = GET_CHAN(src_pixels[block_offset], i);
 711                     max_color[i] = GET_CHAN(src_pixels[block_offset], i);
 712                 } else {
 713                     tmp_min = (int)(0.5 + min * slope + y_intercept);
 714                     tmp_max = (int)(0.5 + max * slope + y_intercept);
 715
 716                     av_assert0(tmp_min <= tmp_max);
 717                     // clamp min and max color values
 718                     tmp_min = av_clip_uint8(tmp_min);
 719                     tmp_max = av_clip_uint8(tmp_max);
 720
 721                     err = FFMAX(calc_lsq_max_fit_error(&src_pixels[block_offset], &bi,
 722                                                        min, max, tmp_min, tmp_max, chan, i), err);
 723
 724                     min_color[i] = tmp_min;
 725                     max_color[i] = tmp_max;
 726                 }
 727             }
 728
 729             if (err > s->sixteen_color_thresh) { // DO SIXTEEN COLOR BLOCK
 730                 uint16_t *row_ptr;
 731                 int rgb555;
 732
 733                 block_offset = get_block_info(&bi, block_counter);
 734
 735                 row_ptr = &src_pixels[block_offset];
 736
 737                 for (int y = 0; y < 4; y++) {
 738                     for (int x = 0; x < 4; x++){
 739                         rgb555 = row_ptr[x] & ~0x8000;
 740
 741                         put_bits(&s->pb, 16, rgb555);
 742                     }
 743                     row_ptr += bi.rowstride;
 744                 }
 745
 746                 block_counter++;
 747             } else { // FOUR COLOR BLOCK
 748                 block_counter += encode_four_color_block(min_color, max_color,
 749                                                          &s->pb, &src_pixels[block_offset], &bi);
 750             }
 751
 752             /* update this block in the previous frame buffer */
 753             update_block_in_prev_frame(&src_pixels[block_offset],
 754                                        &prev_pixels[block_offset], &bi, block_counter);
 755         }
 756     }
 757 }
 758
 759 static int rpza_encode_init(AVCodecContext *avctx)
 760 {
 761     RpzaContext *s = avctx->priv_data;
 762
 763     s->frame_width = avctx->width;
 764     s->frame_height = avctx->height;
 765
 766     s->prev_frame = av_frame_alloc();
 767     if (!s->prev_frame)
 768         return AVERROR(ENOMEM);
 769
 770     return 0;
 771 }
 772
 773 static int rpza_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
 774                                 const AVFrame *frame, int *got_packet)
 775 {
 776     RpzaContext *s = avctx->priv_data;
 777     const AVFrame *pict = frame;
 778     uint8_t *buf;
 779     int ret;
 780
 781     if ((ret = ff_alloc_packet2(avctx, pkt, 6LL * avctx->height * avctx->width, 0)) < 0)
 782         return ret;
 783
 784     init_put_bits(&s->pb, pkt->data, pkt->size);
 785
 786     // skip 4 byte header, write it later once the size of the chunk is known
 787     put_bits32(&s->pb, 0x00);
 788
 789     if (!s->prev_frame->data[0]) {
 790         s->first_frame = 1;
 791         s->prev_frame->format = pict->format;
 792         s->prev_frame->width = pict->width;
 793         s->prev_frame->height = pict->height;
 794         ret = av_frame_get_buffer(s->prev_frame, 0);
 795         if (ret < 0)
 796             return ret;
 797     } else {
 798         s->first_frame = 0;
 799     }
 800
 801     rpza_encode_stream(s, pict);
 802
 803     flush_put_bits(&s->pb);
 804
 805     av_shrink_packet(pkt, put_bytes_output(&s->pb));
 806     buf = pkt->data;
 807
 808     // write header opcode
 809     buf[0] = 0xe1; // chunk opcode
 810
 811     // write chunk length
 812     AV_WB24(buf + 1, pkt->size);
 813
 814     *got_packet = 1;
 815
 816     return 0;
 817 }
 818
 819 static int rpza_encode_end(AVCodecContext *avctx)
 820 {
 821     RpzaContext *s = (RpzaContext *)avctx->priv_data;
 822
 823     av_frame_free(&s->prev_frame);
 824
 825     return 0;
 826 }
 827
 828 #define OFFSET(x) offsetof(RpzaContext, x)
 829 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
 830 static const AVOption options[] = {
 831     { "skip_frame_thresh", NULL, OFFSET(skip_frame_thresh), AV_OPT_TYPE_INT, {.i64=1}, 0, 24, VE},
 832     { "start_one_color_thresh", NULL, OFFSET(start_one_color_thresh), AV_OPT_TYPE_INT, {.i64=1}, 0, 24, VE},
 833     { "continue_one_color_thresh", NULL, OFFSET(continue_one_color_thresh), AV_OPT_TYPE_INT, {.i64=0}, 0, 24, VE},
 834     { "sixteen_color_thresh", NULL, OFFSET(sixteen_color_thresh), AV_OPT_TYPE_INT, {.i64=1}, 0, 24, VE},
 835     { NULL },
 836 };
 837
 838 static const AVClass rpza_class = {
 839     .class_name = "rpza",
 840     .item_name  = av_default_item_name,
 841     .option     = options,
 842     .version    = LIBAVUTIL_VERSION_INT,
 843 };
 844
 845 const AVCodec ff_rpza_encoder = {
 846     .name           = "rpza",
 847     .long_name      = NULL_IF_CONFIG_SMALL("QuickTime video (RPZA)"),
 848     .type           = AVMEDIA_TYPE_VIDEO,
 849     .id             = AV_CODEC_ID_RPZA,
 850     .priv_data_size = sizeof(RpzaContext),
 851     .priv_class     = &rpza_class,
 852     .init           = rpza_encode_init,
 853     .encode2        = rpza_encode_frame,
 854     .close          = rpza_encode_end,
 855     .caps_internal  = FF_CODEC_CAP_INIT_THREADSAFE,
 856     .pix_fmts       = (const enum AVPixelFormat[]) { AV_PIX_FMT_RGB555,
 857                                                      AV_PIX_FMT_NONE},
 858 };