e1331a708ee881ac1b2312e9d2a9e96009e97b67
[ffmpeg.git] / libavcodec / cinepakenc.c
1 /*
2  * Cinepak encoder (c) 2011 Tomas Härdin
3  * http://titan.codemill.se/~tomhar/cinepakenc.patch
4
5 Permission is hereby granted, free of charge, to any person obtaining a
6 copy of this software and associated documentation files (the "Software"),
7 to deal in the Software without restriction, including without limitation
8 the rights to use, copy, modify, merge, publish, distribute, sublicense,
9 and/or sell copies of the Software, and to permit persons to whom the
10 Software is furnished to do so, subject to the following conditions:
11
12 The above copyright notice and this permission notice shall be included
13 in all copies or substantial portions of the Software.
14
15 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
18 THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
19 OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
20 ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
21 OTHER DEALINGS IN THE SOFTWARE.
22
23  */
24
25 #include "libavutil/intreadwrite.h"
26 #include "avcodec.h"
27 #include "libavutil/lfg.h"
28 #include "elbg.h"
29
30 #define CVID_HEADER_SIZE 10
31 #define STRIP_HEADER_SIZE 12
32 #define CHUNK_HEADER_SIZE 4
33
34 #define MB_SIZE 4           //4x4 MBs
35 #define MB_AREA (MB_SIZE*MB_SIZE)
36
37 #define VECTOR_MAX 6        //six or four entries per vector depending on format
38 #define CODEBOOK_MAX 256
39 #define CODEBOOK_NUM 5      //five potential codebooks (1, 4, 16, 64, 256) for V1 and V4
40
41 #define MAX_STRIPS  1       //Note: having fewer choices regarding the number of strip speeds up encoding (obviously)
42 #define MIN_STRIPS  1       //Note: having more strips speeds up encoding the frame (this is less obvious)
43
44 typedef enum {
45     MODE_V1_ONLY = 0,
46     MODE_V1_V4,
47     MODE_MC,
48
49     MODE_COUNT,
50 } CinepakMode;
51
52 typedef enum {
53     ENC_V1,
54     ENC_V4,
55     ENC_SKIP
56 } mb_encoding;
57
58 typedef struct {
59     int v1_vector;                  //index into v1 codebook
60     int v1_error;                   //error when using V1 encoding
61     int v4_vector[CODEBOOK_NUM][4]; //indices into v4 codebooks
62     int v4_error[CODEBOOK_NUM];     //error when using V4 encodings
63     int skip_error;                 //error when block is skipped (aka copied from last frame)
64     mb_encoding best_encoding;      //last result from calculate_mode_score()
65 } mb_info;
66
67 typedef struct {
68     int v1_codebook[CODEBOOK_MAX*VECTOR_MAX];
69     int *v4_codebook;
70 } strip_info;
71
72 typedef struct {
73     AVCodecContext *avctx;
74     unsigned char *pict_bufs[3], *strip_buf, *frame_buf;
75     AVFrame last_frame;
76     AVFrame best_frame;
77     AVFrame scratch_frame;
78     enum PixelFormat pix_fmt;
79     int w, h;
80     int curframe, keyint;
81     AVLFG randctx;
82     uint64_t lambda;
83     int *codebook_input;
84     int *codebook_closest;
85     mb_info *mb;                                //MB RD state
86 #ifdef CINEPAKENC_DEBUG
87     mb_info *best_mb;                           //TODO: remove. only used for printing stats
88 #endif
89     int num_v1_mode, num_v4_mode, num_mc_mode;
90     int num_v1_encs, num_v4_encs, num_skips;
91 } CinepakEncContext;
92
93 static av_cold int cinepak_encode_init(AVCodecContext *avctx)
94 {
95     CinepakEncContext *s = avctx->priv_data;
96     int x, mb_count, strip_buf_size, frame_buf_size;
97
98     if (avctx->width & 3 || avctx->height & 3) {
99         av_log(avctx, AV_LOG_ERROR, "width and height must be multiples of four (got %ix%i)\n",
100                 avctx->width, avctx->height);
101         return AVERROR(EINVAL);
102     }
103
104     if (!(s->codebook_input = av_malloc(sizeof(int) * (avctx->pix_fmt == PIX_FMT_YUV420P ? 6 : 4) * (avctx->width * avctx->height) >> 2)))
105         return AVERROR(ENOMEM);
106
107     if (!(s->codebook_closest = av_malloc(sizeof(int) * (avctx->width * avctx->height) >> 2)))
108         goto enomem;
109
110     for(x = 0; x < 3; x++)
111         if(!(s->pict_bufs[x] = av_malloc((avctx->pix_fmt == PIX_FMT_YUV420P ? 6 : 4) * (avctx->width * avctx->height) >> 2)))
112             goto enomem;
113
114     mb_count = avctx->width * avctx->height / MB_AREA;
115
116     //the largest possible chunk is 0x31 with all MBs encoded in V4 mode, which is 34 bits per MB
117     strip_buf_size = STRIP_HEADER_SIZE + 3 * CHUNK_HEADER_SIZE + 2 * VECTOR_MAX * CODEBOOK_MAX + 4 * (mb_count + (mb_count + 15) / 16);
118
119     frame_buf_size = CVID_HEADER_SIZE + MAX_STRIPS * strip_buf_size;
120
121     if (!(s->strip_buf = av_malloc(strip_buf_size)))
122         goto enomem;
123
124     if (!(s->frame_buf = av_malloc(frame_buf_size)))
125         goto enomem;
126
127     if (!(s->mb = av_malloc(mb_count*sizeof(mb_info))))
128         goto enomem;
129
130 #ifdef CINEPAKENC_DEBUG
131     if (!(s->best_mb = av_malloc(mb_count*sizeof(mb_info))))
132         goto enomem;
133 #endif
134
135     av_lfg_init(&s->randctx, 1);
136     s->avctx = avctx;
137     s->w = avctx->width;
138     s->h = avctx->height;
139     s->curframe = 0;
140     s->keyint = avctx->keyint_min;
141     s->pix_fmt = avctx->pix_fmt;
142
143     //set up AVFrames
144     s->last_frame.data[0]        = s->pict_bufs[0];
145     s->last_frame.linesize[0]    = s->w;
146     s->best_frame.data[0]        = s->pict_bufs[1];
147     s->best_frame.linesize[0]    = s->w;
148     s->scratch_frame.data[0]     = s->pict_bufs[2];
149     s->scratch_frame.linesize[0] = s->w;
150
151     if(s->pix_fmt == PIX_FMT_YUV420P) {
152         s->last_frame.data[1]        = s->last_frame.data[0] + s->w * s->h;
153         s->last_frame.data[2]        = s->last_frame.data[1] + ((s->w * s->h) >> 2);
154         s->last_frame.linesize[1]    = s->last_frame.linesize[2] = s->w >> 1;
155
156         s->best_frame.data[1]        = s->best_frame.data[0] + s->w * s->h;
157         s->best_frame.data[2]        = s->best_frame.data[1] + ((s->w * s->h) >> 2);
158         s->best_frame.linesize[1]    = s->best_frame.linesize[2] = s->w >> 1;
159
160         s->scratch_frame.data[1]     = s->scratch_frame.data[0] + s->w * s->h;
161         s->scratch_frame.data[2]     = s->scratch_frame.data[1] + ((s->w * s->h) >> 2);
162         s->scratch_frame.linesize[1] = s->scratch_frame.linesize[2] = s->w >> 1;
163     }
164
165     s->num_v1_mode = s->num_v4_mode = s->num_mc_mode = s->num_v1_encs = s->num_v4_encs = s->num_skips = 0;
166
167     return 0;
168
169 enomem:
170     av_free(s->codebook_input);
171     av_free(s->codebook_closest);
172     av_free(s->strip_buf);
173     av_free(s->frame_buf);
174     av_free(s->mb);
175 #ifdef CINEPAKENC_DEBUG
176     av_free(s->best_mb);
177 #endif
178
179     for(x = 0; x < 3; x++)
180         av_free(s->pict_bufs[x]);
181
182     return AVERROR(ENOMEM);
183 }
184
185 static int64_t calculate_mode_score(CinepakEncContext *s, CinepakMode mode, int h, int v1_size, int v4_size, int v4, strip_info *info)
186 {
187     //score = FF_LAMBDA_SCALE * error + lambda * bits
188     int x;
189     int entry_size = s->pix_fmt == PIX_FMT_YUV420P ? 6 : 4;
190     int mb_count = s->w * h / MB_AREA;
191     mb_info *mb;
192     int64_t score1, score2, score3;
193     int64_t ret = s->lambda * ((v1_size ? CHUNK_HEADER_SIZE + v1_size * entry_size : 0) +
194                    (v4_size ? CHUNK_HEADER_SIZE + v4_size * entry_size : 0) +
195                    CHUNK_HEADER_SIZE) << 3;
196
197     //av_log(s->avctx, AV_LOG_INFO, "sizes %3i %3i -> %9li score mb_count %i", v1_size, v4_size, ret, mb_count);
198
199     switch(mode) {
200     case MODE_V1_ONLY:
201         //one byte per MB
202         ret += s->lambda * 8 * mb_count;
203
204         for(x = 0; x < mb_count; x++) {
205             mb = &s->mb[x];
206             ret += FF_LAMBDA_SCALE * mb->v1_error;
207             mb->best_encoding = ENC_V1;
208         }
209
210         break;
211     case MODE_V1_V4:
212         //9 or 33 bits per MB
213         for(x = 0; x < mb_count; x++) {
214             mb = &s->mb[x];
215             score1 = s->lambda * 9  + FF_LAMBDA_SCALE * mb->v1_error;
216             score2 = s->lambda * 33 + FF_LAMBDA_SCALE * mb->v4_error[v4];
217
218             if(score1 <= score2) {
219                 ret += score1;
220                 mb->best_encoding = ENC_V1;
221             } else {
222                 ret += score2;
223                 mb->best_encoding = ENC_V4;
224             }
225         }
226
227         break;
228     case MODE_MC:
229         //1, 10 or 34 bits per MB
230         for(x = 0; x < mb_count; x++) {
231             mb = &s->mb[x];
232             score1 = s->lambda * 1  + FF_LAMBDA_SCALE * mb->skip_error;
233             score2 = s->lambda * 10 + FF_LAMBDA_SCALE * mb->v1_error;
234             score3 = s->lambda * 34 + FF_LAMBDA_SCALE * mb->v4_error[v4];
235
236
237             if(score1 <= score2 && score1 <= score3) {
238                 ret += score1;
239                 mb->best_encoding = ENC_SKIP;
240             } else if(score2 <= score1 && score2 <= score3) {
241                 ret += score2;
242                 mb->best_encoding = ENC_V1;
243             } else {
244                 ret += score3;
245                 mb->best_encoding = ENC_V4;
246             }
247         }
248
249         break;
250     }
251
252     return ret;
253 }
254
255 static int write_chunk_header(unsigned char *buf, int chunk_type, int chunk_size)
256 {
257     buf[0] = chunk_type;
258     AV_WB24(&buf[1], chunk_size + CHUNK_HEADER_SIZE);
259     return CHUNK_HEADER_SIZE;
260 }
261
262 static int encode_codebook(CinepakEncContext *s, int *codebook, int size, int chunk_type_yuv, int chunk_type_gray, unsigned char *buf)
263 {
264     int x, y, ret, entry_size = s->pix_fmt == PIX_FMT_YUV420P ? 6 : 4;
265
266     ret = write_chunk_header(buf, s->pix_fmt == PIX_FMT_YUV420P ? chunk_type_yuv : chunk_type_gray, entry_size * size);
267
268     for(x = 0; x < size; x++)
269         for(y = 0; y < entry_size; y++)
270             buf[ret++] = codebook[y + x*entry_size] ^ (y >= 4 ? 0x80 : 0);
271
272     return ret;
273 }
274
275 //sets out to the sub picture starting at (x,y) in in
276 static void get_sub_picture(CinepakEncContext *s, int x, int y, AVPicture *in, AVPicture *out)
277 {
278     out->data[0] = in->data[0] + x + y * in->linesize[0];
279     out->linesize[0] = in->linesize[0];
280
281     if(s->pix_fmt == PIX_FMT_YUV420P) {
282         out->data[1] = in->data[1] + (x >> 1) + (y >> 1) * in->linesize[1];
283         out->linesize[1] = in->linesize[1];
284
285         out->data[2] = in->data[2] + (x >> 1) + (y >> 1) * in->linesize[2];
286         out->linesize[2] = in->linesize[2];
287     }
288 }
289
290 //decodes the V1 vector in mb into the 4x4 MB pointed to by sub_pict
291 static void decode_v1_vector(CinepakEncContext *s, AVPicture *sub_pict, mb_info *mb, strip_info *info)
292 {
293     int entry_size = s->pix_fmt == PIX_FMT_YUV420P ? 6 : 4;
294
295     sub_pict->data[0][0] =
296             sub_pict->data[0][1] =
297             sub_pict->data[0][    sub_pict->linesize[0]] =
298             sub_pict->data[0][1+  sub_pict->linesize[0]] = info->v1_codebook[mb->v1_vector*entry_size];
299
300     sub_pict->data[0][2] =
301             sub_pict->data[0][3] =
302             sub_pict->data[0][2+  sub_pict->linesize[0]] =
303             sub_pict->data[0][3+  sub_pict->linesize[0]] = info->v1_codebook[mb->v1_vector*entry_size+1];
304
305     sub_pict->data[0][2*sub_pict->linesize[0]] =
306             sub_pict->data[0][1+2*sub_pict->linesize[0]] =
307             sub_pict->data[0][  3*sub_pict->linesize[0]] =
308             sub_pict->data[0][1+3*sub_pict->linesize[0]] = info->v1_codebook[mb->v1_vector*entry_size+2];
309
310     sub_pict->data[0][2+2*sub_pict->linesize[0]] =
311             sub_pict->data[0][3+2*sub_pict->linesize[0]] =
312             sub_pict->data[0][2+3*sub_pict->linesize[0]] =
313             sub_pict->data[0][3+3*sub_pict->linesize[0]] = info->v1_codebook[mb->v1_vector*entry_size+3];
314
315     if(s->pix_fmt == PIX_FMT_YUV420P) {
316         sub_pict->data[1][0] =
317             sub_pict->data[1][1] =
318             sub_pict->data[1][    sub_pict->linesize[1]] =
319             sub_pict->data[1][1+  sub_pict->linesize[1]] = info->v1_codebook[mb->v1_vector*entry_size+4];
320
321         sub_pict->data[2][0] =
322             sub_pict->data[2][1] =
323             sub_pict->data[2][    sub_pict->linesize[2]] =
324             sub_pict->data[2][1+  sub_pict->linesize[2]] = info->v1_codebook[mb->v1_vector*entry_size+5];
325     }
326 }
327
328 //decodes the V4 vectors in mb into the 4x4 MB pointed to by sub_pict
329 static void decode_v4_vector(CinepakEncContext *s, AVPicture *sub_pict, int *v4_vector, strip_info *info)
330 {
331     int i, x, y, entry_size = s->pix_fmt == PIX_FMT_YUV420P ? 6 : 4;
332
333     for(i = y = 0; y < 4; y += 2) {
334         for(x = 0; x < 4; x += 2, i++) {
335             sub_pict->data[0][x   +     y*sub_pict->linesize[0]] = info->v4_codebook[v4_vector[i]*entry_size];
336             sub_pict->data[0][x+1 +     y*sub_pict->linesize[0]] = info->v4_codebook[v4_vector[i]*entry_size+1];
337             sub_pict->data[0][x   + (y+1)*sub_pict->linesize[0]] = info->v4_codebook[v4_vector[i]*entry_size+2];
338             sub_pict->data[0][x+1 + (y+1)*sub_pict->linesize[0]] = info->v4_codebook[v4_vector[i]*entry_size+3];
339
340             if(s->pix_fmt == PIX_FMT_YUV420P) {
341                 sub_pict->data[1][(x>>1) + (y>>1)*sub_pict->linesize[1]] = info->v4_codebook[v4_vector[i]*entry_size+4];
342                 sub_pict->data[2][(x>>1) + (y>>1)*sub_pict->linesize[2]] = info->v4_codebook[v4_vector[i]*entry_size+5];
343             }
344         }
345     }
346 }
347
348 static int encode_mode(CinepakEncContext *s, CinepakMode mode, int h, int v1_size, int v4_size, int v4, AVPicture *scratch_pict, strip_info *info, unsigned char *buf)
349 {
350     int x, y, z, flags, bits, temp_size, header_ofs, ret = 0, mb_count = s->w * h / MB_AREA;
351     int needs_extra_bit, should_write_temp;
352     unsigned char temp[64]; //32/2 = 16 V4 blocks at 4 B each -> 64 B
353     mb_info *mb;
354     AVPicture sub_scratch;
355
356     //encode codebooks
357     if(v1_size)
358         ret += encode_codebook(s, info->v1_codebook, v1_size, 0x22, 0x26, buf + ret);
359
360     if(v4_size)
361         ret += encode_codebook(s, info->v4_codebook, v4_size, 0x20, 0x24, buf + ret);
362
363     //update scratch picture
364     for(z = y = 0; y < h; y += MB_SIZE) {
365         for(x = 0; x < s->w; x += MB_SIZE, z++) {
366             mb = &s->mb[z];
367
368             if(mode == MODE_MC && mb->best_encoding == ENC_SKIP)
369                 continue;
370
371             get_sub_picture(s, x, y, scratch_pict, &sub_scratch);
372
373             if(mode == MODE_V1_ONLY || mb->best_encoding == ENC_V1)
374                 decode_v1_vector(s, &sub_scratch, mb, info);
375             else if(mode != MODE_V1_ONLY && mb->best_encoding == ENC_V4)
376                 decode_v4_vector(s, &sub_scratch, mb->v4_vector[v4], info);
377         }
378     }
379
380     switch(mode) {
381     case MODE_V1_ONLY:
382         //av_log(s->avctx, AV_LOG_INFO, "mb_count = %i\n", mb_count);
383         ret += write_chunk_header(buf + ret, 0x32, mb_count);
384
385         for(x = 0; x < mb_count; x++)
386             buf[ret++] = s->mb[x].v1_vector;
387
388         break;
389     case MODE_V1_V4:
390         //remember header position
391         header_ofs = ret;
392         ret += CHUNK_HEADER_SIZE;
393
394         for(x = 0; x < mb_count; x += 32) {
395             flags = 0;
396             for(y = x; y < FFMIN(x+32, mb_count); y++)
397                 if(s->mb[y].best_encoding == ENC_V4)
398                     flags |= 1 << (31 - y + x);
399
400             AV_WB32(&buf[ret], flags);
401             ret += 4;
402
403             for(y = x; y < FFMIN(x+32, mb_count); y++) {
404                 mb = &s->mb[y];
405
406                 if(mb->best_encoding == ENC_V1)
407                     buf[ret++] = mb->v1_vector;
408                 else
409                     for(z = 0; z < 4; z++)
410                         buf[ret++] = mb->v4_vector[v4][z];
411             }
412         }
413
414         write_chunk_header(buf + header_ofs, 0x30, ret - header_ofs - CHUNK_HEADER_SIZE);
415
416         break;
417     case MODE_MC:
418         //remember header position
419         header_ofs = ret;
420         ret += CHUNK_HEADER_SIZE;
421         flags = bits = temp_size = 0;
422
423         for(x = 0; x < mb_count; x++) {
424             mb = &s->mb[x];
425             flags |= (mb->best_encoding != ENC_SKIP) << (31 - bits++);
426             needs_extra_bit = 0;
427             should_write_temp = 0;
428
429             if(mb->best_encoding != ENC_SKIP) {
430                 if(bits < 32)
431                     flags |= (mb->best_encoding == ENC_V4) << (31 - bits++);
432                 else
433                     needs_extra_bit = 1;
434             }
435
436             if(bits == 32) {
437                 AV_WB32(&buf[ret], flags);
438                 ret += 4;
439                 flags = bits = 0;
440
441                 if(mb->best_encoding == ENC_SKIP || needs_extra_bit) {
442                     memcpy(&buf[ret], temp, temp_size);
443                     ret += temp_size;
444                     temp_size = 0;
445                 } else
446                     should_write_temp = 1;
447             }
448
449             if(needs_extra_bit) {
450                 flags = (mb->best_encoding == ENC_V4) << 31;
451                 bits = 1;
452             }
453
454             if(mb->best_encoding == ENC_V1)
455                 temp[temp_size++] = mb->v1_vector;
456             else if(mb->best_encoding == ENC_V4)
457                 for(z = 0; z < 4; z++)
458                     temp[temp_size++] = mb->v4_vector[v4][z];
459
460             if(should_write_temp) {
461                 memcpy(&buf[ret], temp, temp_size);
462                 ret += temp_size;
463                 temp_size = 0;
464             }
465         }
466
467         if(bits > 0) {
468             AV_WB32(&buf[ret], flags);
469             ret += 4;
470             memcpy(&buf[ret], temp, temp_size);
471             ret += temp_size;
472         }
473
474         write_chunk_header(buf + header_ofs, 0x31, ret - header_ofs - CHUNK_HEADER_SIZE);
475
476         break;
477     }
478
479     return ret;
480 }
481
482 //computes distortion of 4x4 MB in b compared to a
483 static int compute_mb_distortion(CinepakEncContext *s, AVPicture *a, AVPicture *b)
484 {
485     int x, y, p, d, ret = 0;
486
487     for(y = 0; y < MB_SIZE; y++) {
488         for(x = 0; x < MB_SIZE; x++) {
489             d = a->data[0][x + y*a->linesize[0]] - b->data[0][x + y*b->linesize[0]];
490             ret += d*d;
491         }
492     }
493
494     if(s->pix_fmt == PIX_FMT_YUV420P) {
495         for(p = 1; p <= 2; p++) {
496             for(y = 0; y < MB_SIZE/2; y++) {
497                 for(x = 0; x < MB_SIZE/2; x++) {
498                     d = a->data[p][x + y*a->linesize[p]] - b->data[p][x + y*b->linesize[p]];
499                     ret += d*d;
500                 }
501             }
502         }
503     }
504
505     return ret;
506 }
507
508 static int quantize(CinepakEncContext *s, int h, AVPicture *pict, int v1mode, int size, int v4, strip_info *info)
509 {
510     int x, y, i, j, k, x2, y2, x3, y3, plane, shift;
511     int entry_size = s->pix_fmt == PIX_FMT_YUV420P ? 6 : 4;
512     int *codebook = v1mode ? info->v1_codebook : info->v4_codebook;
513     int64_t total_error = 0;
514     uint8_t vq_pict_buf[(MB_AREA*3)/2];
515     AVPicture sub_pict, vq_pict;
516
517     for(i = y = 0; y < h; y += MB_SIZE) {
518         for(x = 0; x < s->w; x += MB_SIZE, i += v1mode ? 1 : 4) {
519             int *base = s->codebook_input + i*entry_size;
520
521             if(v1mode) {
522                 //subsample
523                 for(j = y2 = 0; y2 < entry_size; y2 += 2) {
524                     for(x2 = 0; x2 < 4; x2 += 2, j++) {
525                         plane = y2 < 4 ? 0 : 1 + (x2 >> 1);
526                         shift = y2 < 4 ? 0 : 1;
527                         x3 = shift ? 0 : x2;
528                         y3 = shift ? 0 : y2;
529                         base[j] = (pict->data[plane][((x+x3) >> shift) +      ((y+y3) >> shift)      * pict->linesize[plane]] +
530                                    pict->data[plane][((x+x3) >> shift) + 1 +  ((y+y3) >> shift)      * pict->linesize[plane]] +
531                                    pict->data[plane][((x+x3) >> shift) +     (((y+y3) >> shift) + 1) * pict->linesize[plane]] +
532                                    pict->data[plane][((x+x3) >> shift) + 1 + (((y+y3) >> shift) + 1) * pict->linesize[plane]]) >> 2;
533                     }
534                 }
535             } else {
536                 //copy
537                 for(j = y2 = 0; y2 < MB_SIZE; y2 += 2) {
538                     for(x2 = 0; x2 < MB_SIZE; x2 += 2) {
539                         for(k = 0; k < entry_size; k++, j++) {
540                             plane = k >= 4 ? k - 3 : 0;
541
542                             if(k >= 4) {
543                                 x3 = (x+x2) >> 1;
544                                 y3 = (y+y2) >> 1;
545                             } else {
546                                 x3 = x + x2 + (k & 1);
547                                 y3 = y + y2 + (k >> 1);
548                             }
549
550                             base[j] = pict->data[plane][x3 + y3*pict->linesize[plane]];
551                         }
552                     }
553                 }
554             }
555         }
556     }
557
558     ff_init_elbg(s->codebook_input, entry_size, i, codebook, size, 1, s->codebook_closest, &s->randctx);
559     ff_do_elbg(s->codebook_input, entry_size, i, codebook, size, 1, s->codebook_closest, &s->randctx);
560
561     //setup vq_pict, which contains a single MB
562     vq_pict.data[0] = vq_pict_buf;
563     vq_pict.linesize[0] = MB_SIZE;
564     vq_pict.data[1] = &vq_pict_buf[MB_AREA];
565     vq_pict.data[2] = vq_pict.data[1] + (MB_AREA >> 2);
566     vq_pict.linesize[1] = vq_pict.linesize[2] = MB_SIZE >> 1;
567
568     //copy indices
569     for(i = j = y = 0; y < h; y += MB_SIZE) {
570         for(x = 0; x < s->w; x += MB_SIZE, j++, i += v1mode ? 1 : 4) {
571             mb_info *mb = &s->mb[j];
572
573             //point sub_pict to current MB
574             get_sub_picture(s, x, y, pict, &sub_pict);
575
576             if(v1mode) {
577                 mb->v1_vector = s->codebook_closest[i];
578
579                 //fill in vq_pict with V1 data
580                 decode_v1_vector(s, &vq_pict, mb, info);
581
582                 mb->v1_error = compute_mb_distortion(s, &sub_pict, &vq_pict);
583                 total_error += mb->v1_error;
584             } else {
585                 for(k = 0; k < 4; k++)
586                     mb->v4_vector[v4][k] = s->codebook_closest[i+k];
587
588                 //fill in vq_pict with V4 data
589                 decode_v4_vector(s, &vq_pict, mb->v4_vector[v4], info);
590
591                 mb->v4_error[v4] = compute_mb_distortion(s, &sub_pict, &vq_pict);
592                 total_error += mb->v4_error[v4];
593             }
594         }
595     }
596
597     //av_log(s->avctx, AV_LOG_INFO, "mode %i size %i i %i error %li\n", v1mode, size, i, total_error);
598
599     return 0;
600 }
601
602 static void calculate_skip_errors(CinepakEncContext *s, int h, AVPicture *last_pict, AVPicture *pict, strip_info *info)
603 {
604     int x, y, i;
605     AVPicture sub_last, sub_pict;
606
607     for(i = y = 0; y < h; y += MB_SIZE) {
608         for(x = 0; x < s->w; x += MB_SIZE, i++) {
609             get_sub_picture(s, x, y, last_pict, &sub_last);
610             get_sub_picture(s, x, y, pict,      &sub_pict);
611
612             s->mb[i].skip_error = compute_mb_distortion(s, &sub_last, &sub_pict);
613         }
614     }
615 }
616
617 static void write_strip_header(CinepakEncContext *s, int y, int h, int keyframe, unsigned char *buf, int strip_size)
618 {
619     buf[0] = keyframe ? 0x11: 0x10;
620     AV_WB24(&buf[1], strip_size + STRIP_HEADER_SIZE);
621     AV_WB16(&buf[4], y);
622     AV_WB16(&buf[6], 0);
623     AV_WB16(&buf[8], h);
624     AV_WB16(&buf[10], s->w);
625 }
626
627 static int rd_strip(CinepakEncContext *s, int y, int h, int keyframe, AVPicture *last_pict, AVPicture *pict, AVPicture *scratch_pict, unsigned char *buf, int64_t *best_score)
628 {
629     int64_t score = 0;
630     int best_size = 0, v1_size, v4_size, v4, mb_count = s->w * h / MB_AREA;
631     strip_info info;
632     CinepakMode best_mode;
633     int v4_codebooks[CODEBOOK_NUM][CODEBOOK_MAX*VECTOR_MAX];
634
635     if(!keyframe)
636         calculate_skip_errors(s, h, last_pict, pict, &info);
637
638     //precompute V4 codebooks
639     for(v4_size = 1, v4 = 0; v4_size <= 256; v4_size <<= 2, v4++) {
640         info.v4_codebook = v4_codebooks[v4];
641         quantize(s, h, pict, 0, v4_size, v4, &info);
642     }
643
644     //try all powers of 4 for the size of the codebooks
645     //constraint the v4 codebook to be no bigger than the v1 codebook
646     for(v1_size = 1; v1_size <= 256; v1_size <<= 2) {
647         //compute V1 codebook
648         quantize(s, h, pict, 1, v1_size, -1, &info);
649
650         for(v4_size = 0, v4 = -1; v4_size <= v1_size; v4_size = v4_size ? v4_size << 2 : v1_size >= 4 ? v1_size >> 2 : 1, v4++) {
651             //try all modes
652             for(CinepakMode mode = 0; mode < MODE_COUNT; mode++) {
653                 //don't allow MODE_MC in inter frames
654                 if(keyframe && mode == MODE_MC)
655                     continue;
656
657                 //only allow V1-only mode if v4 codebook is empty
658                 if(!v4_size && mode != MODE_V1_ONLY)
659                     continue;
660
661                 info.v4_codebook = v4 >= 0 ? v4_codebooks[v4] : NULL;
662                 score = calculate_mode_score(s, mode, h, v1_size, v4_size, v4, &info);
663
664                 //av_log(s->avctx, AV_LOG_INFO, "%3i %3i score = %li\n", v1_size, v4_size, score);
665
666                 if(best_size == 0 || score < *best_score) {
667                     *best_score = score;
668                     best_size = encode_mode(s, mode, h, v1_size, v4_size, v4, scratch_pict, &info, s->strip_buf + STRIP_HEADER_SIZE);
669                     best_mode = mode;
670
671                     av_log(s->avctx, AV_LOG_INFO, "mode %i, %3i, %3i: %18li %i B\n", mode, v1_size, v4_size, score, best_size);
672
673 #ifdef CINEPAKENC_DEBUG
674                     //save MB encoding choices
675                     memcpy(s->best_mb, s->mb, mb_count*sizeof(mb_info));
676 #endif
677
678                     //memcpy(strip_temp + STRIP_HEADER_SIZE, strip_temp, best_size);
679                     write_strip_header(s, y, h, keyframe, s->strip_buf, best_size);
680                 }
681             }
682         }
683     }
684
685 #ifdef CINEPAKENC_DEBUG
686     //gather stats. this will only work properly of MAX_STRIPS == 1
687     if(best_mode == MODE_V1_ONLY) {
688         s->num_v1_mode++;
689         s->num_v1_encs += s->w*h/MB_AREA;
690     } else {
691         if(best_mode == MODE_V1_V4)
692             s->num_v4_mode++;
693         else
694             s->num_mc_mode++;
695
696         int x;
697         for(x = 0; x < s->w*h/MB_AREA; x++)
698             if(s->best_mb[x].best_encoding == ENC_V1)
699                 s->num_v1_encs++;
700             else if(s->best_mb[x].best_encoding == ENC_V4)
701                 s->num_v4_encs++;
702             else
703                 s->num_skips++;
704     }
705 #endif
706
707     best_size += STRIP_HEADER_SIZE;
708     memcpy(buf, s->strip_buf, best_size);
709
710
711     return best_size;
712 }
713
714 static int write_cvid_header(CinepakEncContext *s, unsigned char *buf, int num_strips, int data_size)
715 {
716     buf[0] = 0;
717     AV_WB24(&buf[1], data_size + CVID_HEADER_SIZE);
718     AV_WB16(&buf[4], s->w);
719     AV_WB16(&buf[6], s->h);
720     AV_WB16(&buf[8], num_strips);
721
722     return CVID_HEADER_SIZE;
723 }
724
725 static int rd_frame(CinepakEncContext *s, AVFrame *frame, unsigned char *buf, int buf_size)
726 {
727     int num_strips, strip, h, i, y, size, temp_size, best_size;
728     AVPicture last_pict, pict, scratch_pict;
729     int64_t best_score = 0, score, score_temp;
730
731     //TODO: support encoding zero strips (meaning skip the whole frame)
732     for(num_strips = MIN_STRIPS; num_strips <= MAX_STRIPS && num_strips <= s->h / MB_SIZE; num_strips++) {
733         score = 0;
734         size = 0;
735         h = s->h / num_strips;
736         //make h into next multiple of 4
737         h += 4 - (h & 3);
738
739         for(strip = 0; strip < num_strips; strip++) {
740             y = strip*h;
741
742             get_sub_picture(s, 0, y, (AVPicture*)frame,            &pict);
743             get_sub_picture(s, 0, y, (AVPicture*)&s->last_frame,    &last_pict);
744             get_sub_picture(s, 0, y, (AVPicture*)&s->scratch_frame, &scratch_pict);
745
746             if((temp_size = rd_strip(s, y, FFMIN(h, s->h - y), frame->key_frame, &last_pict, &pict, &scratch_pict, s->frame_buf + CVID_HEADER_SIZE, &score_temp)) < 0)
747                 return temp_size;
748
749             score += score_temp;
750             size += temp_size;
751         }
752
753         if(best_score == 0 || score < best_score) {
754             best_score = score;
755             best_size = size + write_cvid_header(s, s->frame_buf, num_strips, size);
756             av_log(s->avctx, AV_LOG_INFO, "best number of strips so far: %2i, %12li, %i B\n", num_strips, score, best_size);
757
758             FFSWAP(AVFrame, s->best_frame, s->scratch_frame);
759         }
760     }
761
762     memcpy(buf, s->frame_buf, best_size);
763
764     return best_size;
765 }
766
767 static int cinepak_encode_frame(AVCodecContext *avctx, unsigned char *buf, int buf_size, void *data)
768 {
769     CinepakEncContext *s = avctx->priv_data;
770     AVFrame *frame = data;
771     int ret;
772
773     s->lambda = frame->quality ? frame->quality - 1 : 2 * FF_LAMBDA_SCALE;
774
775     frame->key_frame = s->curframe == 0;
776     frame->pict_type = frame->key_frame ? FF_I_TYPE : FF_P_TYPE;
777
778     ret = rd_frame(s, frame, buf, buf_size);
779
780     avctx->coded_frame = frame;
781
782     FFSWAP(AVFrame, s->last_frame, s->best_frame);
783
784     if (++s->curframe >= s->keyint)
785         s->curframe = 0;
786
787     return ret;
788 }
789
790 static av_cold int cinepak_encode_end(AVCodecContext *avctx)
791 {
792     CinepakEncContext *s = avctx->priv_data;
793     int x;
794
795     av_free(s->codebook_input);
796     av_free(s->codebook_closest);
797     av_free(s->strip_buf);
798     av_free(s->frame_buf);
799     av_free(s->mb);
800 #ifdef CINEPAKENC_DEBUG
801     av_free(s->best_mb);
802 #endif
803
804     for(x = 0; x < 3; x++)
805         av_free(s->pict_bufs[x]);
806
807     av_log(avctx, AV_LOG_INFO, "strip coding stats: %i V1 mode, %i V4 mode, %i MC mode (%i V1 encs, %i V4 encs, %i skips)\n",
808         s->num_v1_mode, s->num_v4_mode, s->num_mc_mode, s->num_v1_encs, s->num_v4_encs, s->num_skips);
809
810     return 0;
811 }
812
813 AVCodec ff_cinepak_encoder = {
814     "cinepak",
815     AVMEDIA_TYPE_VIDEO,
816     CODEC_ID_CINEPAK,
817     sizeof(CinepakEncContext),
818     cinepak_encode_init,
819     cinepak_encode_frame,
820     cinepak_encode_end,
821     .pix_fmts= (const enum PixelFormat[]){PIX_FMT_YUV420P, PIX_FMT_GRAY8, PIX_FMT_NONE},
822     .long_name= NULL_IF_CONFIG_SMALL("Cinepak / CVID"),
823 };