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

   1 /*
   2  * RoQ Video Encoder.
   3  *
   4  * Copyright (C) 2007 Vitor Sessak <vitor1001@gmail.com>
   5  * Copyright (C) 2004-2007 Eric Lasota
   6  *    Based on RoQ specs (C) 2001 Tim Ferguson
   7  *
   8  * This file is part of FFmpeg.
   9  *
  10  * FFmpeg is free software; you can redistribute it and/or
  11  * modify it under the terms of the GNU Lesser General Public
  12  * License as published by the Free Software Foundation; either
  13  * version 2.1 of the License, or (at your option) any later version.
  14  *
  15  * FFmpeg is distributed in the hope that it will be useful,
  16  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  17  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  18  * Lesser General Public License for more details.
  19  *
  20  * You should have received a copy of the GNU Lesser General Public
  21  * License along with FFmpeg; if not, write to the Free Software
  22  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  23  */
  24
  25 /**
  26  * @file
  27  * id RoQ encoder by Vitor. Based on the Switchblade3 library and the
  28  * Switchblade3 FFmpeg glue by Eric Lasota.
  29  */
  30
  31 /*
  32  * COSTS:
  33  * Level 1:
  34  *  SKIP - 2 bits
  35  *  MOTION - 2 + 8 bits
  36  *  CODEBOOK - 2 + 8 bits
  37  *  SUBDIVIDE - 2 + combined subcel cost
  38  *
  39  * Level 2:
  40  *  SKIP - 2 bits
  41  *  MOTION - 2 + 8 bits
  42  *  CODEBOOK - 2 + 8 bits
  43  *  SUBDIVIDE - 2 + 4*8 bits
  44  *
  45  * Maximum cost: 138 bits per cel
  46  *
  47  * Proper evaluation requires LCD fraction comparison, which requires
  48  * Squared Error (SE) loss * savings increase
  49  *
  50  * Maximum savings increase: 136 bits
  51  * Maximum SE loss without overflow: 31580641
  52  * Components in 8x8 supercel: 192
  53  * Maximum SE precision per component: 164482
  54  *    >65025, so no truncation is needed (phew)
  55  */
  56
  57 #include <string.h>
  58
  59 #include "libavutil/attributes.h"
  60 #include "libavutil/lfg.h"
  61 #include "libavutil/opt.h"
  62 #include "roqvideo.h"
  63 #include "bytestream.h"
  64 #include "elbg.h"
  65 #include "internal.h"
  66 #include "mathops.h"
  67
  68 #define CHROMA_BIAS 1
  69
  70 /**
  71  * Maximum number of generated 4x4 codebooks. Can't be 256 to workaround a
  72  * Quake 3 bug.
  73  */
  74 #define MAX_CBS_4x4 256
  75
  76 #define MAX_CBS_2x2 256 ///< Maximum number of 2x2 codebooks.
  77
  78 /* The cast is useful when multiplying it by INT_MAX */
  79 #define ROQ_LAMBDA_SCALE ((uint64_t) FF_LAMBDA_SCALE)
  80
  81 typedef struct SubcelEvaluation {
  82     int eval_dist[4];
  83     int best_bit_use;
  84     int best_coding;
  85
  86     int subCels[4];
  87     motion_vect motion;
  88     int cbEntry;
  89 } SubcelEvaluation;
  90
  91 typedef struct CelEvaluation {
  92     int eval_dist[4];
  93     int best_coding;
  94
  95     SubcelEvaluation subCels[4];
  96
  97     motion_vect motion;
  98     int cbEntry;
  99
 100     int sourceX, sourceY;
 101 } CelEvaluation;
 102
 103 typedef struct RoqEncContext {
 104     RoqContext common;
 105     AVLFG randctx;
 106     uint64_t lambda;
 107
 108     motion_vect *this_motion4;
 109     motion_vect *last_motion4;
 110
 111     motion_vect *this_motion8;
 112     motion_vect *last_motion8;
 113
 114     unsigned int framesSinceKeyframe;
 115
 116     const AVFrame *frame_to_enc;
 117     uint8_t *out_buf;
 118     struct RoqTempData *tmpData;
 119
 120     CelEvaluation *cel_evals;
 121     int *closest_cb;
 122     int *points;  // Allocated together with closest_cb
 123
 124     int first_frame;
 125     int quake3_compat; // Quake 3 compatibility option
 126 } RoqEncContext;
 127
 128 /* Macroblock support functions */
 129 static void unpack_roq_cell(roq_cell *cell, uint8_t u[4*3])
 130 {
 131     memcpy(u  , cell->y, 4);
 132     memset(u+4, cell->u, 4);
 133     memset(u+8, cell->v, 4);
 134 }
 135
 136 static void unpack_roq_qcell(uint8_t cb2[], roq_qcell *qcell, uint8_t u[4*4*3])
 137 {
 138     int i,cp;
 139     static const int offsets[4] = {0, 2, 8, 10};
 140
 141     for (cp=0; cp<3; cp++)
 142         for (i=0; i<4; i++) {
 143             u[4*4*cp + offsets[i]  ] = cb2[qcell->idx[i]*2*2*3 + 4*cp  ];
 144             u[4*4*cp + offsets[i]+1] = cb2[qcell->idx[i]*2*2*3 + 4*cp+1];
 145             u[4*4*cp + offsets[i]+4] = cb2[qcell->idx[i]*2*2*3 + 4*cp+2];
 146             u[4*4*cp + offsets[i]+5] = cb2[qcell->idx[i]*2*2*3 + 4*cp+3];
 147         }
 148 }
 149
 150
 151 static void enlarge_roq_mb4(uint8_t base[3*16], uint8_t u[3*64])
 152 {
 153     int x,y,cp;
 154
 155     for(cp=0; cp<3; cp++)
 156         for(y=0; y<8; y++)
 157             for(x=0; x<8; x++)
 158                 *u++ = base[(y/2)*4 + (x/2) + 16*cp];
 159 }
 160
 161 static inline int square(int x)
 162 {
 163     return x*x;
 164 }
 165
 166 static inline int eval_sse(const uint8_t *a, const uint8_t *b, int count)
 167 {
 168     int diff=0;
 169
 170     while(count--)
 171         diff += square(*b++ - *a++);
 172
 173     return diff;
 174 }
 175
 176 // FIXME Could use DSPContext.sse, but it is not so speed critical (used
 177 // just for motion estimation).
 178 static int block_sse(uint8_t * const *buf1, uint8_t * const *buf2, int x1, int y1,
 179                      int x2, int y2, const int *stride1, const int *stride2, int size)
 180 {
 181     int i, k;
 182     int sse=0;
 183
 184     for (k=0; k<3; k++) {
 185         int bias = (k ? CHROMA_BIAS : 4);
 186         for (i=0; i<size; i++)
 187             sse += bias*eval_sse(buf1[k] + (y1+i)*stride1[k] + x1,
 188                                  buf2[k] + (y2+i)*stride2[k] + x2, size);
 189     }
 190
 191     return sse;
 192 }
 193
 194 static int eval_motion_dist(RoqEncContext *enc, int x, int y, motion_vect vect,
 195                              int size)
 196 {
 197     RoqContext *const roq = &enc->common;
 198     int mx=vect.d[0];
 199     int my=vect.d[1];
 200
 201     if (mx < -7 || mx > 7)
 202         return INT_MAX;
 203
 204     if (my < -7 || my > 7)
 205         return INT_MAX;
 206
 207     mx += x;
 208     my += y;
 209
 210     if ((unsigned) mx > roq->width-size || (unsigned) my > roq->height-size)
 211         return INT_MAX;
 212
 213     return block_sse(enc->frame_to_enc->data, roq->last_frame->data, x, y,
 214                      mx, my,
 215                      enc->frame_to_enc->linesize, roq->last_frame->linesize,
 216                      size);
 217 }
 218
 219 /**
 220  * @return distortion between two macroblocks
 221  */
 222 static inline int squared_diff_macroblock(uint8_t a[], uint8_t b[], int size)
 223 {
 224     int cp, sdiff=0;
 225
 226     for(cp=0;cp<3;cp++) {
 227         int bias = (cp ? CHROMA_BIAS : 4);
 228         sdiff += bias*eval_sse(a, b, size*size);
 229         a += size*size;
 230         b += size*size;
 231     }
 232
 233     return sdiff;
 234 }
 235
 236 typedef struct RoqCodebooks {
 237     int numCB4;
 238     int numCB2;
 239     int usedCB2[MAX_CBS_2x2];
 240     int usedCB4[MAX_CBS_4x4];
 241     uint8_t unpacked_cb2[MAX_CBS_2x2*2*2*3];
 242     uint8_t unpacked_cb4[MAX_CBS_4x4*4*4*3];
 243     uint8_t unpacked_cb4_enlarged[MAX_CBS_4x4*8*8*3];
 244 } RoqCodebooks;
 245
 246 /**
 247  * Temporary vars
 248  */
 249 typedef struct RoqTempData
 250 {
 251     int f2i4[MAX_CBS_4x4];
 252     int i2f4[MAX_CBS_4x4];
 253     int f2i2[MAX_CBS_2x2];
 254     int i2f2[MAX_CBS_2x2];
 255
 256     int mainChunkSize;
 257
 258     int numCB4;
 259     int numCB2;
 260
 261     RoqCodebooks codebooks;
 262
 263     int used_option[4];
 264 } RoqTempdata;
 265
 266 /**
 267  * Initialize cel evaluators and set their source coordinates
 268  */
 269 static int create_cel_evals(RoqEncContext *enc)
 270 {
 271     RoqContext *const roq = &enc->common;
 272
 273     enc->cel_evals = av_malloc_array(roq->width * roq->height / 64, sizeof(CelEvaluation));
 274     if (!enc->cel_evals)
 275         return AVERROR(ENOMEM);
 276
 277     /* Map to the ROQ quadtree order */
 278     for (int y = 0, n = 0; y < roq->height; y += 16)
 279         for (int x = 0; x < roq->width; x += 16)
 280             for(int i = 0; i < 4; i++) {
 281                 enc->cel_evals[n  ].sourceX = x + (i&1)*8;
 282                 enc->cel_evals[n++].sourceY = y + (i&2)*4;
 283             }
 284
 285     return 0;
 286 }
 287
 288 /**
 289  * Get macroblocks from parts of the image
 290  */
 291 static void get_frame_mb(const AVFrame *frame, int x, int y, uint8_t mb[], int dim)
 292 {
 293     int i, j, cp;
 294
 295     for (cp=0; cp<3; cp++) {
 296         int stride = frame->linesize[cp];
 297         for (i=0; i<dim; i++)
 298             for (j=0; j<dim; j++)
 299                 *mb++ = frame->data[cp][(y+i)*stride + x + j];
 300     }
 301 }
 302
 303 /**
 304  * Find the codebook with the lowest distortion from an image
 305  */
 306 static int index_mb(uint8_t cluster[], uint8_t cb[], int numCB,
 307                     int *outIndex, int dim)
 308 {
 309     int i, lDiff = INT_MAX, pick=0;
 310
 311     /* Diff against the others */
 312     for (i=0; i<numCB; i++) {
 313         int diff = squared_diff_macroblock(cluster, cb + i*dim*dim*3, dim);
 314         if (diff < lDiff) {
 315             lDiff = diff;
 316             pick = i;
 317         }
 318     }
 319
 320     *outIndex = pick;
 321     return lDiff;
 322 }
 323
 324 #define EVAL_MOTION(MOTION) \
 325     do { \
 326         diff = eval_motion_dist(enc, j, i, MOTION, blocksize); \
 327             \
 328         if (diff < lowestdiff) { \
 329             lowestdiff = diff; \
 330             bestpick = MOTION; \
 331         } \
 332     } while(0)
 333
 334 static void motion_search(RoqEncContext *enc, int blocksize)
 335 {
 336     static const motion_vect offsets[8] = {
 337         {{ 0,-1}},
 338         {{ 0, 1}},
 339         {{-1, 0}},
 340         {{ 1, 0}},
 341         {{-1, 1}},
 342         {{ 1,-1}},
 343         {{-1,-1}},
 344         {{ 1, 1}},
 345     };
 346
 347     RoqContext *const roq = &enc->common;
 348     int diff, lowestdiff, oldbest;
 349     int off[3];
 350     motion_vect bestpick = {{0,0}};
 351     int i, j, k, offset;
 352
 353     motion_vect *last_motion;
 354     motion_vect *this_motion;
 355     motion_vect vect, vect2;
 356     const int max = (roq->width / blocksize) * roq->height / blocksize;
 357
 358     if (blocksize == 4) {
 359         last_motion = enc->last_motion4;
 360         this_motion = enc->this_motion4;
 361     } else {
 362         last_motion = enc->last_motion8;
 363         this_motion = enc->this_motion8;
 364     }
 365
 366     for (i = 0; i< roq->height; i += blocksize)
 367         for (j = 0; j < roq->width; j += blocksize) {
 368             lowestdiff = eval_motion_dist(enc, j, i, (motion_vect) {{0,0}},
 369                                           blocksize);
 370             bestpick.d[0] = 0;
 371             bestpick.d[1] = 0;
 372
 373             if (blocksize == 4)
 374                 EVAL_MOTION(enc->this_motion8[(i/8) * (roq->width/8) + j/8]);
 375
 376             offset = (i/blocksize) * roq->width / blocksize + j / blocksize;
 377             if (offset < max && offset >= 0)
 378                 EVAL_MOTION(last_motion[offset]);
 379
 380             offset++;
 381             if (offset < max && offset >= 0)
 382                 EVAL_MOTION(last_motion[offset]);
 383
 384             offset = (i/blocksize + 1) * roq->width / blocksize + j / blocksize;
 385             if (offset < max && offset >= 0)
 386                 EVAL_MOTION(last_motion[offset]);
 387
 388             off[0]= (i/blocksize) * roq->width / blocksize + j/blocksize - 1;
 389             off[1]= off[0] - roq->width / blocksize + 1;
 390             off[2]= off[1] + 1;
 391
 392             if (i) {
 393
 394                 for(k=0; k<2; k++)
 395                     vect.d[k]= mid_pred(this_motion[off[0]].d[k],
 396                                         this_motion[off[1]].d[k],
 397                                         this_motion[off[2]].d[k]);
 398
 399                 EVAL_MOTION(vect);
 400                 for(k=0; k<3; k++)
 401                     EVAL_MOTION(this_motion[off[k]]);
 402             } else if(j)
 403                 EVAL_MOTION(this_motion[off[0]]);
 404
 405             vect = bestpick;
 406
 407             oldbest = -1;
 408             while (oldbest != lowestdiff) {
 409                 oldbest = lowestdiff;
 410                 for (k=0; k<8; k++) {
 411                     vect2 = vect;
 412                     vect2.d[0] += offsets[k].d[0];
 413                     vect2.d[1] += offsets[k].d[1];
 414                     EVAL_MOTION(vect2);
 415                 }
 416                 vect = bestpick;
 417             }
 418             offset = (i/blocksize) * roq->width / blocksize + j/blocksize;
 419             this_motion[offset] = bestpick;
 420         }
 421 }
 422
 423 /**
 424  * Get distortion for all options available to a subcel
 425  */
 426 static void gather_data_for_subcel(SubcelEvaluation *subcel, int x,
 427                                    int y, RoqEncContext *enc, RoqTempdata *tempData)
 428 {
 429     RoqContext *const roq = &enc->common;
 430     uint8_t mb4[4*4*3];
 431     uint8_t mb2[2*2*3];
 432     int cluster_index;
 433     int i, best_dist;
 434
 435     static const int bitsUsed[4] = {2, 10, 10, 34};
 436
 437     if (enc->framesSinceKeyframe >= 1) {
 438         subcel->motion = enc->this_motion4[y * roq->width / 16 + x / 4];
 439
 440         subcel->eval_dist[RoQ_ID_FCC] =
 441             eval_motion_dist(enc, x, y,
 442                              enc->this_motion4[y * roq->width / 16 + x / 4], 4);
 443     } else
 444         subcel->eval_dist[RoQ_ID_FCC] = INT_MAX;
 445
 446     if (enc->framesSinceKeyframe >= 2)
 447         subcel->eval_dist[RoQ_ID_MOT] = block_sse(enc->frame_to_enc->data,
 448                                                   roq->current_frame->data, x,
 449                                                   y, x, y,
 450                                                   enc->frame_to_enc->linesize,
 451                                                   roq->current_frame->linesize,
 452                                                   4);
 453     else
 454         subcel->eval_dist[RoQ_ID_MOT] = INT_MAX;
 455
 456     cluster_index = y * roq->width / 16 + x / 4;
 457
 458     get_frame_mb(enc->frame_to_enc, x, y, mb4, 4);
 459
 460     subcel->eval_dist[RoQ_ID_SLD] = index_mb(mb4,
 461                                              tempData->codebooks.unpacked_cb4,
 462                                              tempData->codebooks.numCB4,
 463                                              &subcel->cbEntry, 4);
 464
 465     subcel->eval_dist[RoQ_ID_CCC] = 0;
 466
 467     for(i=0;i<4;i++) {
 468         subcel->subCels[i] = enc->closest_cb[cluster_index*4+i];
 469
 470         get_frame_mb(enc->frame_to_enc, x+2*(i&1),
 471                      y+(i&2), mb2, 2);
 472
 473         subcel->eval_dist[RoQ_ID_CCC] +=
 474             squared_diff_macroblock(tempData->codebooks.unpacked_cb2 + subcel->subCels[i]*2*2*3, mb2, 2);
 475     }
 476
 477     best_dist = INT_MAX;
 478     for (i=0; i<4; i++)
 479         if (ROQ_LAMBDA_SCALE*subcel->eval_dist[i] + enc->lambda*bitsUsed[i] <
 480             best_dist) {
 481             subcel->best_coding = i;
 482             subcel->best_bit_use = bitsUsed[i];
 483             best_dist = ROQ_LAMBDA_SCALE*subcel->eval_dist[i] +
 484                 enc->lambda*bitsUsed[i];
 485         }
 486 }
 487
 488 /**
 489  * Get distortion for all options available to a cel
 490  */
 491 static void gather_data_for_cel(CelEvaluation *cel, RoqEncContext *enc,
 492                                 RoqTempdata *tempData)
 493 {
 494     RoqContext *const roq = &enc->common;
 495     uint8_t mb8[8*8*3];
 496     int index = cel->sourceY * roq->width / 64 + cel->sourceX/8;
 497     int i, j, best_dist, divide_bit_use;
 498
 499     int bitsUsed[4] = {2, 10, 10, 0};
 500
 501     if (enc->framesSinceKeyframe >= 1) {
 502         cel->motion = enc->this_motion8[index];
 503
 504         cel->eval_dist[RoQ_ID_FCC] =
 505             eval_motion_dist(enc, cel->sourceX, cel->sourceY,
 506                              enc->this_motion8[index], 8);
 507     } else
 508         cel->eval_dist[RoQ_ID_FCC] = INT_MAX;
 509
 510     if (enc->framesSinceKeyframe >= 2)
 511         cel->eval_dist[RoQ_ID_MOT] = block_sse(enc->frame_to_enc->data,
 512                                                roq->current_frame->data,
 513                                                cel->sourceX, cel->sourceY,
 514                                                cel->sourceX, cel->sourceY,
 515                                                enc->frame_to_enc->linesize,
 516                                                roq->current_frame->linesize,8);
 517     else
 518         cel->eval_dist[RoQ_ID_MOT] = INT_MAX;
 519
 520     get_frame_mb(enc->frame_to_enc, cel->sourceX, cel->sourceY, mb8, 8);
 521
 522     cel->eval_dist[RoQ_ID_SLD] =
 523         index_mb(mb8, tempData->codebooks.unpacked_cb4_enlarged,
 524                  tempData->codebooks.numCB4, &cel->cbEntry, 8);
 525
 526     gather_data_for_subcel(cel->subCels + 0, cel->sourceX+0, cel->sourceY+0, enc, tempData);
 527     gather_data_for_subcel(cel->subCels + 1, cel->sourceX+4, cel->sourceY+0, enc, tempData);
 528     gather_data_for_subcel(cel->subCels + 2, cel->sourceX+0, cel->sourceY+4, enc, tempData);
 529     gather_data_for_subcel(cel->subCels + 3, cel->sourceX+4, cel->sourceY+4, enc, tempData);
 530
 531     cel->eval_dist[RoQ_ID_CCC] = 0;
 532     divide_bit_use = 0;
 533     for (i=0; i<4; i++) {
 534         cel->eval_dist[RoQ_ID_CCC] +=
 535             cel->subCels[i].eval_dist[cel->subCels[i].best_coding];
 536         divide_bit_use += cel->subCels[i].best_bit_use;
 537     }
 538
 539     best_dist = INT_MAX;
 540     bitsUsed[3] = 2 + divide_bit_use;
 541
 542     for (i=0; i<4; i++)
 543         if (ROQ_LAMBDA_SCALE*cel->eval_dist[i] + enc->lambda*bitsUsed[i] <
 544             best_dist) {
 545             cel->best_coding = i;
 546             best_dist = ROQ_LAMBDA_SCALE*cel->eval_dist[i] +
 547                 enc->lambda*bitsUsed[i];
 548         }
 549
 550     tempData->used_option[cel->best_coding]++;
 551     tempData->mainChunkSize += bitsUsed[cel->best_coding];
 552
 553     if (cel->best_coding == RoQ_ID_SLD)
 554         tempData->codebooks.usedCB4[cel->cbEntry]++;
 555
 556     if (cel->best_coding == RoQ_ID_CCC)
 557         for (i=0; i<4; i++) {
 558             if (cel->subCels[i].best_coding == RoQ_ID_SLD)
 559                 tempData->codebooks.usedCB4[cel->subCels[i].cbEntry]++;
 560             else if (cel->subCels[i].best_coding == RoQ_ID_CCC)
 561                 for (j=0; j<4; j++)
 562                     tempData->codebooks.usedCB2[cel->subCels[i].subCels[j]]++;
 563         }
 564 }
 565
 566 static void remap_codebooks(RoqEncContext *enc, RoqTempdata *tempData)
 567 {
 568     RoqContext *const roq = &enc->common;
 569     int i, j, idx=0;
 570
 571     /* Make remaps for the final codebook usage */
 572     for (i=0; i<(enc->quake3_compat ? MAX_CBS_4x4-1 : MAX_CBS_4x4); i++) {
 573         if (tempData->codebooks.usedCB4[i]) {
 574             tempData->i2f4[i] = idx;
 575             tempData->f2i4[idx] = i;
 576             for (j=0; j<4; j++)
 577                 tempData->codebooks.usedCB2[roq->cb4x4[i].idx[j]]++;
 578             idx++;
 579         }
 580     }
 581
 582     tempData->numCB4 = idx;
 583
 584     idx = 0;
 585     for (i=0; i<MAX_CBS_2x2; i++) {
 586         if (tempData->codebooks.usedCB2[i]) {
 587             tempData->i2f2[i] = idx;
 588             tempData->f2i2[idx] = i;
 589             idx++;
 590         }
 591     }
 592     tempData->numCB2 = idx;
 593
 594 }
 595
 596 /**
 597  * Write codebook chunk
 598  */
 599 static void write_codebooks(RoqEncContext *enc, RoqTempdata *tempData)
 600 {
 601     RoqContext *const roq = &enc->common;
 602     int i, j;
 603     uint8_t **outp= &enc->out_buf;
 604
 605     if (tempData->numCB2) {
 606         bytestream_put_le16(outp, RoQ_QUAD_CODEBOOK);
 607         bytestream_put_le32(outp, tempData->numCB2*6 + tempData->numCB4*4);
 608         bytestream_put_byte(outp, tempData->numCB4);
 609         bytestream_put_byte(outp, tempData->numCB2);
 610
 611         for (i=0; i<tempData->numCB2; i++) {
 612             bytestream_put_buffer(outp, roq->cb2x2[tempData->f2i2[i]].y, 4);
 613             bytestream_put_byte(outp, roq->cb2x2[tempData->f2i2[i]].u);
 614             bytestream_put_byte(outp, roq->cb2x2[tempData->f2i2[i]].v);
 615         }
 616
 617         for (i=0; i<tempData->numCB4; i++)
 618             for (j=0; j<4; j++)
 619                 bytestream_put_byte(outp, tempData->i2f2[roq->cb4x4[tempData->f2i4[i]].idx[j]]);
 620
 621     }
 622 }
 623
 624 static inline uint8_t motion_arg(motion_vect mot)
 625 {
 626     uint8_t ax = 8 - ((uint8_t) mot.d[0]);
 627     uint8_t ay = 8 - ((uint8_t) mot.d[1]);
 628     return ((ax&15)<<4) | (ay&15);
 629 }
 630
 631 typedef struct CodingSpool {
 632     int typeSpool;
 633     int typeSpoolLength;
 634     uint8_t argumentSpool[64];
 635     uint8_t *args;
 636     uint8_t **pout;
 637 } CodingSpool;
 638
 639 /* NOTE: Typecodes must be spooled AFTER arguments!! */
 640 static void write_typecode(CodingSpool *s, uint8_t type)
 641 {
 642     s->typeSpool |= (type & 3) << (14 - s->typeSpoolLength);
 643     s->typeSpoolLength += 2;
 644     if (s->typeSpoolLength == 16) {
 645         bytestream_put_le16(s->pout, s->typeSpool);
 646         bytestream_put_buffer(s->pout, s->argumentSpool,
 647                               s->args - s->argumentSpool);
 648         s->typeSpoolLength = 0;
 649         s->typeSpool = 0;
 650         s->args = s->argumentSpool;
 651     }
 652 }
 653
 654 static void reconstruct_and_encode_image(RoqEncContext *enc,
 655                                          RoqTempdata *tempData,
 656                                          int w, int h, int numBlocks)
 657 {
 658     RoqContext *const roq = &enc->common;
 659     int i, j, k;
 660     int x, y;
 661     int subX, subY;
 662     int dist=0;
 663
 664     roq_qcell *qcell;
 665     CelEvaluation *eval;
 666
 667     CodingSpool spool;
 668
 669     spool.typeSpool=0;
 670     spool.typeSpoolLength=0;
 671     spool.args = spool.argumentSpool;
 672     spool.pout = &enc->out_buf;
 673
 674     if (tempData->used_option[RoQ_ID_CCC]%2)
 675         tempData->mainChunkSize+=8; //FIXME
 676
 677     /* Write the video chunk header */
 678     bytestream_put_le16(&enc->out_buf, RoQ_QUAD_VQ);
 679     bytestream_put_le32(&enc->out_buf, tempData->mainChunkSize/8);
 680     bytestream_put_byte(&enc->out_buf, 0x0);
 681     bytestream_put_byte(&enc->out_buf, 0x0);
 682
 683     for (i=0; i<numBlocks; i++) {
 684         eval = enc->cel_evals + i;
 685
 686         x = eval->sourceX;
 687         y = eval->sourceY;
 688         dist += eval->eval_dist[eval->best_coding];
 689
 690         switch (eval->best_coding) {
 691         case RoQ_ID_MOT:
 692             write_typecode(&spool, RoQ_ID_MOT);
 693             break;
 694
 695         case RoQ_ID_FCC:
 696             bytestream_put_byte(&spool.args, motion_arg(eval->motion));
 697
 698             write_typecode(&spool, RoQ_ID_FCC);
 699             ff_apply_motion_8x8(roq, x, y,
 700                                 eval->motion.d[0], eval->motion.d[1]);
 701             break;
 702
 703         case RoQ_ID_SLD:
 704             bytestream_put_byte(&spool.args, tempData->i2f4[eval->cbEntry]);
 705             write_typecode(&spool, RoQ_ID_SLD);
 706
 707             qcell = roq->cb4x4 + eval->cbEntry;
 708             ff_apply_vector_4x4(roq, x  , y  , roq->cb2x2 + qcell->idx[0]);
 709             ff_apply_vector_4x4(roq, x+4, y  , roq->cb2x2 + qcell->idx[1]);
 710             ff_apply_vector_4x4(roq, x  , y+4, roq->cb2x2 + qcell->idx[2]);
 711             ff_apply_vector_4x4(roq, x+4, y+4, roq->cb2x2 + qcell->idx[3]);
 712             break;
 713
 714         case RoQ_ID_CCC:
 715             write_typecode(&spool, RoQ_ID_CCC);
 716
 717             for (j=0; j<4; j++) {
 718                 subX = x + 4*(j&1);
 719                 subY = y + 2*(j&2);
 720
 721                 switch(eval->subCels[j].best_coding) {
 722                 case RoQ_ID_MOT:
 723                     break;
 724
 725                 case RoQ_ID_FCC:
 726                     bytestream_put_byte(&spool.args,
 727                                         motion_arg(eval->subCels[j].motion));
 728
 729                     ff_apply_motion_4x4(roq, subX, subY,
 730                                         eval->subCels[j].motion.d[0],
 731                                         eval->subCels[j].motion.d[1]);
 732                     break;
 733
 734                 case RoQ_ID_SLD:
 735                     bytestream_put_byte(&spool.args,
 736                                         tempData->i2f4[eval->subCels[j].cbEntry]);
 737
 738                     qcell = roq->cb4x4 + eval->subCels[j].cbEntry;
 739
 740                     ff_apply_vector_2x2(roq, subX  , subY  ,
 741                                         roq->cb2x2 + qcell->idx[0]);
 742                     ff_apply_vector_2x2(roq, subX+2, subY  ,
 743                                         roq->cb2x2 + qcell->idx[1]);
 744                     ff_apply_vector_2x2(roq, subX  , subY+2,
 745                                         roq->cb2x2 + qcell->idx[2]);
 746                     ff_apply_vector_2x2(roq, subX+2, subY+2,
 747                                         roq->cb2x2 + qcell->idx[3]);
 748                     break;
 749
 750                 case RoQ_ID_CCC:
 751                     for (k=0; k<4; k++) {
 752                         int cb_idx = eval->subCels[j].subCels[k];
 753                         bytestream_put_byte(&spool.args,
 754                                             tempData->i2f2[cb_idx]);
 755
 756                         ff_apply_vector_2x2(roq, subX + 2*(k&1), subY + (k&2),
 757                                             roq->cb2x2 + cb_idx);
 758                     }
 759                     break;
 760                 }
 761                 write_typecode(&spool, eval->subCels[j].best_coding);
 762             }
 763             break;
 764         }
 765     }
 766
 767     /* Flush the remainder of the argument/type spool */
 768     while (spool.typeSpoolLength)
 769         write_typecode(&spool, 0x0);
 770 }
 771
 772
 773 /**
 774  * Create a single YUV cell from a 2x2 section of the image
 775  */
 776 static inline void frame_block_to_cell(uint8_t *block, uint8_t * const *data,
 777                                        int top, int left, const int *stride)
 778 {
 779     int i, j, u=0, v=0;
 780
 781     for (i=0; i<2; i++)
 782         for (j=0; j<2; j++) {
 783             int x = (top+i)*stride[0] + left + j;
 784             *block++ = data[0][x];
 785             x = (top+i)*stride[1] + left + j;
 786             u       += data[1][x];
 787             v       += data[2][x];
 788         }
 789
 790     *block++ = (u+2)/4;
 791     *block++ = (v+2)/4;
 792 }
 793
 794 /**
 795  * Create YUV clusters for the entire image
 796  */
 797 static void create_clusters(const AVFrame *frame, int w, int h, uint8_t *yuvClusters)
 798 {
 799     int i, j, k, l;
 800
 801     for (i=0; i<h; i+=4)
 802         for (j=0; j<w; j+=4) {
 803             for (k=0; k < 2; k++)
 804                 for (l=0; l < 2; l++)
 805                     frame_block_to_cell(yuvClusters + (l + 2*k)*6, frame->data,
 806                                         i+2*k, j+2*l, frame->linesize);
 807             yuvClusters += 24;
 808         }
 809 }
 810
 811 static int generate_codebook(RoqEncContext *enc,
 812                              int *points, int inputCount, roq_cell *results,
 813                              int size, int cbsize)
 814 {
 815     int i, j, k, ret = 0;
 816     int c_size = size*size/4;
 817     int *buf;
 818     int *codebook = av_malloc_array(6*c_size, cbsize*sizeof(int));
 819     int *closest_cb = enc->closest_cb;
 820
 821     if (!codebook)
 822         return AVERROR(ENOMEM);
 823
 824     ret = avpriv_init_elbg(points, 6 * c_size, inputCount, codebook,
 825                        cbsize, 1, closest_cb, &enc->randctx);
 826     if (ret < 0)
 827         goto out;
 828     ret = avpriv_do_elbg(points, 6 * c_size, inputCount, codebook,
 829                      cbsize, 1, closest_cb, &enc->randctx);
 830     if (ret < 0)
 831         goto out;
 832
 833     buf = codebook;
 834     for (i=0; i<cbsize; i++)
 835         for (k=0; k<c_size; k++) {
 836             for(j=0; j<4; j++)
 837                 results->y[j] = *buf++;
 838
 839             results->u =    (*buf++ + CHROMA_BIAS/2)/CHROMA_BIAS;
 840             results->v =    (*buf++ + CHROMA_BIAS/2)/CHROMA_BIAS;
 841             results++;
 842         }
 843 out:
 844     av_free(codebook);
 845     return ret;
 846 }
 847
 848 static int generate_new_codebooks(RoqEncContext *enc, RoqTempdata *tempData)
 849 {
 850     int i, j, ret = 0;
 851     RoqCodebooks *codebooks = &tempData->codebooks;
 852     RoqContext *const roq = &enc->common;
 853     int max = roq->width * roq->height / 16;
 854     uint8_t mb2[3*4];
 855     roq_cell *results4 = av_malloc(sizeof(roq_cell)*MAX_CBS_4x4*4);
 856     uint8_t *yuvClusters=av_malloc_array(max, sizeof(int)*6*4);
 857     int *points = enc->points;
 858     int bias;
 859
 860     if (!results4 || !yuvClusters) {
 861         ret = AVERROR(ENOMEM);
 862         goto out;
 863     }
 864
 865     /* Subsample YUV data */
 866     create_clusters(enc->frame_to_enc, roq->width, roq->height, yuvClusters);
 867
 868     /* Cast to integer and apply chroma bias */
 869     for (i=0; i<max*24; i++) {
 870         bias = ((i%6)<4) ? 1 : CHROMA_BIAS;
 871         points[i] = bias*yuvClusters[i];
 872     }
 873
 874     /* Create 4x4 codebooks */
 875     if ((ret = generate_codebook(enc, points, max,
 876                                  results4, 4, (enc->quake3_compat ? MAX_CBS_4x4-1 : MAX_CBS_4x4))) < 0)
 877         goto out;
 878
 879     codebooks->numCB4 = (enc->quake3_compat ? MAX_CBS_4x4-1 : MAX_CBS_4x4);
 880
 881     /* Create 2x2 codebooks */
 882     if ((ret = generate_codebook(enc, points, max * 4,
 883                                  roq->cb2x2, 2, MAX_CBS_2x2)) < 0)
 884         goto out;
 885
 886     codebooks->numCB2 = MAX_CBS_2x2;
 887
 888     /* Unpack 2x2 codebook clusters */
 889     for (i=0; i<codebooks->numCB2; i++)
 890         unpack_roq_cell(roq->cb2x2 + i, codebooks->unpacked_cb2 + i*2*2*3);
 891
 892     /* Index all 4x4 entries to the 2x2 entries, unpack, and enlarge */
 893     for (i=0; i<codebooks->numCB4; i++) {
 894         for (j=0; j<4; j++) {
 895             unpack_roq_cell(&results4[4*i + j], mb2);
 896             index_mb(mb2, codebooks->unpacked_cb2, codebooks->numCB2,
 897                      &roq->cb4x4[i].idx[j], 2);
 898         }
 899         unpack_roq_qcell(codebooks->unpacked_cb2, roq->cb4x4 + i,
 900                          codebooks->unpacked_cb4 + i*4*4*3);
 901         enlarge_roq_mb4(codebooks->unpacked_cb4 + i*4*4*3,
 902                         codebooks->unpacked_cb4_enlarged + i*8*8*3);
 903     }
 904 out:
 905     av_free(yuvClusters);
 906     av_free(results4);
 907     return ret;
 908 }
 909
 910 static int roq_encode_video(RoqEncContext *enc)
 911 {
 912     RoqTempdata *tempData = enc->tmpData;
 913     RoqContext *const roq = &enc->common;
 914     int ret;
 915
 916     memset(tempData, 0, sizeof(*tempData));
 917
 918     ret = generate_new_codebooks(enc, tempData);
 919     if (ret < 0)
 920         return ret;
 921
 922     if (enc->framesSinceKeyframe >= 1) {
 923         motion_search(enc, 8);
 924         motion_search(enc, 4);
 925     }
 926
 927  retry_encode:
 928     for (int i = 0; i < roq->width * roq->height / 64; i++)
 929         gather_data_for_cel(enc->cel_evals + i, enc, tempData);
 930
 931     /* Quake 3 can't handle chunks bigger than 65535 bytes */
 932     if (tempData->mainChunkSize/8 > 65535 && enc->quake3_compat) {
 933         if (enc->lambda > 100000) {
 934             av_log(roq->avctx, AV_LOG_ERROR, "Cannot encode video in Quake compatible form\n");
 935             return AVERROR(EINVAL);
 936         }
 937         av_log(roq->avctx, AV_LOG_ERROR,
 938                "Warning, generated a frame too big for Quake (%d > 65535), "
 939                "now switching to a bigger qscale value.\n",
 940                tempData->mainChunkSize/8);
 941         enc->lambda *= 1.5;
 942         tempData->mainChunkSize = 0;
 943         memset(tempData->used_option, 0, sizeof(tempData->used_option));
 944         memset(tempData->codebooks.usedCB4, 0,
 945                sizeof(tempData->codebooks.usedCB4));
 946         memset(tempData->codebooks.usedCB2, 0,
 947                sizeof(tempData->codebooks.usedCB2));
 948
 949         goto retry_encode;
 950     }
 951
 952     remap_codebooks(enc, tempData);
 953
 954     write_codebooks(enc, tempData);
 955
 956     reconstruct_and_encode_image(enc, tempData, roq->width, roq->height,
 957                                  roq->width * roq->height / 64);
 958
 959     /* Rotate frame history */
 960     FFSWAP(AVFrame *,    roq->current_frame,   roq->last_frame);
 961     FFSWAP(motion_vect *, enc->last_motion4, enc->this_motion4);
 962     FFSWAP(motion_vect *, enc->last_motion8, enc->this_motion8);
 963
 964     enc->framesSinceKeyframe++;
 965
 966     return 0;
 967 }
 968
 969 static av_cold int roq_encode_end(AVCodecContext *avctx)
 970 {
 971     RoqEncContext *const enc = avctx->priv_data;
 972
 973     av_frame_free(&enc->common.current_frame);
 974     av_frame_free(&enc->common.last_frame);
 975
 976     av_freep(&enc->tmpData);
 977     av_freep(&enc->cel_evals);
 978     av_freep(&enc->closest_cb);
 979     av_freep(&enc->this_motion4);
 980     av_freep(&enc->last_motion4);
 981     av_freep(&enc->this_motion8);
 982     av_freep(&enc->last_motion8);
 983
 984     return 0;
 985 }
 986
 987 static av_cold int roq_encode_init(AVCodecContext *avctx)
 988 {
 989     RoqEncContext *const enc = avctx->priv_data;
 990     RoqContext    *const roq = &enc->common;
 991
 992     av_lfg_init(&enc->randctx, 1);
 993
 994     roq->avctx = avctx;
 995
 996     enc->framesSinceKeyframe = 0;
 997     if ((avctx->width & 0xf) || (avctx->height & 0xf)) {
 998         av_log(avctx, AV_LOG_ERROR, "Dimensions must be divisible by 16\n");
 999         return AVERROR(EINVAL);
1000     }
1001
1002     if (avctx->width > 65535 || avctx->height > 65535) {
1003         av_log(avctx, AV_LOG_ERROR, "Dimensions are max %d\n", enc->quake3_compat ? 32768 : 65535);
1004         return AVERROR(EINVAL);
1005     }
1006
1007     if (((avctx->width)&(avctx->width-1))||((avctx->height)&(avctx->height-1)))
1008         av_log(avctx, AV_LOG_ERROR, "Warning: dimensions not power of two, this is not supported by quake\n");
1009
1010     roq->width  = avctx->width;
1011     roq->height = avctx->height;
1012
1013     enc->framesSinceKeyframe = 0;
1014     enc->first_frame = 1;
1015
1016     roq->last_frame    = av_frame_alloc();
1017     roq->current_frame = av_frame_alloc();
1018     if (!roq->last_frame || !roq->current_frame)
1019         return AVERROR(ENOMEM);
1020
1021     enc->tmpData      = av_malloc(sizeof(RoqTempdata));
1022
1023     enc->this_motion4 =
1024         av_mallocz_array(roq->width * roq->height / 16, sizeof(motion_vect));
1025
1026     enc->last_motion4 =
1027         av_malloc_array (roq->width * roq->height / 16, sizeof(motion_vect));
1028
1029     enc->this_motion8 =
1030         av_mallocz_array(roq->width * roq->height / 64, sizeof(motion_vect));
1031
1032     enc->last_motion8 =
1033         av_malloc_array (roq->width * roq->height / 64, sizeof(motion_vect));
1034
1035     /* 4x4 codebook needs 6 * 4 * 4 / 4 * width * height / 16 * sizeof(int);
1036      * and so does the points buffer. */
1037     enc->closest_cb   =
1038         av_malloc_array(roq->width * roq->height, 3 * sizeof(int));
1039
1040     if (!enc->tmpData || !enc->this_motion4 || !enc->last_motion4 ||
1041         !enc->this_motion8 || !enc->last_motion8 || !enc->closest_cb)
1042         return AVERROR(ENOMEM);
1043
1044     enc->points = enc->closest_cb + roq->width * roq->height * 3 / 2;
1045
1046     return create_cel_evals(enc);
1047 }
1048
1049 static void roq_write_video_info_chunk(RoqEncContext *enc)
1050 {
1051     /* ROQ info chunk */
1052     bytestream_put_le16(&enc->out_buf, RoQ_INFO);
1053
1054     /* Size: 8 bytes */
1055     bytestream_put_le32(&enc->out_buf, 8);
1056
1057     /* Unused argument */
1058     bytestream_put_byte(&enc->out_buf, 0x00);
1059     bytestream_put_byte(&enc->out_buf, 0x00);
1060
1061     /* Width */
1062     bytestream_put_le16(&enc->out_buf, enc->common.width);
1063
1064     /* Height */
1065     bytestream_put_le16(&enc->out_buf, enc->common.height);
1066
1067     /* Unused in Quake 3, mimics the output of the real encoder */
1068     bytestream_put_byte(&enc->out_buf, 0x08);
1069     bytestream_put_byte(&enc->out_buf, 0x00);
1070     bytestream_put_byte(&enc->out_buf, 0x04);
1071     bytestream_put_byte(&enc->out_buf, 0x00);
1072 }
1073
1074 static int roq_encode_frame(AVCodecContext *avctx, AVPacket *pkt,
1075                             const AVFrame *frame, int *got_packet)
1076 {
1077     RoqEncContext *const enc = avctx->priv_data;
1078     RoqContext    *const roq = &enc->common;
1079     int size, ret;
1080
1081     roq->avctx = avctx;
1082
1083     enc->frame_to_enc = frame;
1084
1085     if (frame->quality)
1086         enc->lambda = frame->quality - 1;
1087     else
1088         enc->lambda = 2*ROQ_LAMBDA_SCALE;
1089
1090     /* 138 bits max per 8x8 block +
1091      *     256 codebooks*(6 bytes 2x2 + 4 bytes 4x4) + 8 bytes frame header */
1092     size = ((roq->width * roq->height / 64) * 138 + 7) / 8 + 256 * (6 + 4) + 8;
1093     if ((ret = ff_alloc_packet2(avctx, pkt, size, 0)) < 0)
1094         return ret;
1095     enc->out_buf = pkt->data;
1096
1097     /* Check for I-frame */
1098     if (enc->framesSinceKeyframe == avctx->gop_size)
1099         enc->framesSinceKeyframe = 0;
1100
1101     if (enc->first_frame) {
1102         /* Alloc memory for the reconstruction data (we must know the stride
1103          for that) */
1104         if ((ret = ff_get_buffer(avctx, roq->current_frame, 0)) < 0 ||
1105             (ret = ff_get_buffer(avctx, roq->last_frame,    0)) < 0)
1106             return ret;
1107
1108         /* Before the first video frame, write a "video info" chunk */
1109         roq_write_video_info_chunk(enc);
1110
1111         enc->first_frame = 0;
1112     }
1113
1114     /* Encode the actual frame */
1115     ret = roq_encode_video(enc);
1116     if (ret < 0)
1117         return ret;
1118
1119     pkt->size   = enc->out_buf - pkt->data;
1120     if (enc->framesSinceKeyframe == 1)
1121         pkt->flags |= AV_PKT_FLAG_KEY;
1122     *got_packet = 1;
1123
1124     return 0;
1125 }
1126
1127 #define OFFSET(x) offsetof(RoqEncContext, x)
1128 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
1129 static const AVOption options[] = {
1130     { "quake3_compat", "Whether to respect known limitations in Quake 3 decoder", OFFSET(quake3_compat), AV_OPT_TYPE_BOOL, { .i64 = 1 }, 0, 1, VE },
1131     { NULL },
1132 };
1133
1134 static const AVClass roq_class = {
1135     .class_name = "RoQ",
1136     .item_name  = av_default_item_name,
1137     .option     = options,
1138     .version    = LIBAVUTIL_VERSION_INT,
1139 };
1140
1141 AVCodec ff_roq_encoder = {
1142     .name                 = "roqvideo",
1143     .long_name            = NULL_IF_CONFIG_SMALL("id RoQ video"),
1144     .type                 = AVMEDIA_TYPE_VIDEO,
1145     .id                   = AV_CODEC_ID_ROQ,
1146     .priv_data_size       = sizeof(RoqEncContext),
1147     .init                 = roq_encode_init,
1148     .encode2              = roq_encode_frame,
1149     .close                = roq_encode_end,
1150     .pix_fmts             = (const enum AVPixelFormat[]){ AV_PIX_FMT_YUVJ444P,
1151                                                         AV_PIX_FMT_NONE },
1152     .priv_class     = &roq_class,
1153     .caps_internal        = FF_CODEC_CAP_INIT_CLEANUP,
1154 };