test/encode/mpeg2enc.c

   1 /*
   2  * Copyright (c) 2012 Intel Corporation. All Rights Reserved.
   3  *
   4  * Permission is hereby granted, free of charge, to any person obtaining a
   5  * copy of this software and associated documentation files (the
   6  * "Software"), to deal in the Software without restriction, including
   7  * without limitation the rights to use, copy, modify, merge, publish,
   8  * distribute, sub license, and/or sell copies of the Software, and to
   9  * permit persons to whom the Software is furnished to do so, subject to
  10  * the following conditions:
  11  *
  12  * The above copyright notice and this permission notice (including the
  13  * next paragraph) shall be included in all copies or substantial portions
  14  * of the Software.
  15  *
  16  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
  17  * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
  18  * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT.
  19  * IN NO EVENT SHALL PRECISION INSIGHT AND/OR ITS SUPPLIERS BE LIABLE FOR
  20  * ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
  21  * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
  22  * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  23  */
  24 /*
  25  * Simple MPEG-2 encoder based on libVA.
  26  *
  27  */
  28
  29 #include "sysdeps.h"
  30
  31 #include <getopt.h>
  32 #include <unistd.h>
  33
  34 #include <sys/time.h>
  35 #include <sys/types.h>
  36 #include <fcntl.h>
  37 #include <time.h>
  38 #include <pthread.h>
  39
  40 #include <va/va.h>
  41 #include <va/va_enc_mpeg2.h>
  42
  43 #include "va_display.h"
  44
  45 #define START_CODE_PICUTRE      0x00000100
  46 #define START_CODE_SLICE        0x00000101
  47 #define START_CODE_USER         0x000001B2
  48 #define START_CODE_SEQ          0x000001B3
  49 #define START_CODE_EXT          0x000001B5
  50 #define START_CODE_GOP          0x000001B8
  51
  52 #define CHROMA_FORMAT_RESERVED  0
  53 #define CHROMA_FORMAT_420       1
  54 #define CHROMA_FORMAT_422       2
  55 #define CHROMA_FORMAT_444       3
  56
  57 #define MAX_SLICES              128
  58
  59 #define CHECK_VASTATUS(va_status, func)                                 \
  60     if (va_status != VA_STATUS_SUCCESS) {                               \
  61         fprintf(stderr, "%s:%s (%d) failed, exit\n", __func__, func, __LINE__); \
  62         exit(1);                                                        \
  63     }
  64
  65 static int const picture_type_pattern[][2] = {{VAEncPictureTypeIntra, 1},
  66                                              {VAEncPictureTypePredictive, 3}, {VAEncPictureTypePredictive, 3},{VAEncPictureTypePredictive, 3},
  67                                              {VAEncPictureTypePredictive, 3}, {VAEncPictureTypePredictive, 3},{VAEncPictureTypePredictive, 3},
  68                                              {VAEncPictureTypePredictive, 3}, {VAEncPictureTypePredictive, 3},{VAEncPictureTypePredictive, 3},
  69                                              {VAEncPictureTypePredictive, 2}};
  70
  71 struct mpeg2enc_context {
  72     /* args */
  73     int rate_control_mode;
  74     int fps;
  75     int mode; /* 0:I, 1:I/P, 2:I/P/B */
  76     int width;
  77     int height;
  78     int frame_size;
  79     int num_pictures;
  80     int qp;
  81     FILE *ifp;
  82     FILE *ofp;
  83     unsigned char *frame_data_buffer;
  84     int intra_period;
  85     int bit_rate; /* in kbps */
  86
  87     /* VA resource */
  88     VADisplay va_dpy;
  89     VAEncSequenceParameterBufferMPEG2 seq_param;
  90     VAEncPictureParameterBufferMPEG2 pic_param;
  91     VAEncSliceParameterBufferMPEG2 slice_param[MAX_SLICES];
  92     VAContextID context_id;
  93     VAConfigID config_id;
  94     VABufferID seq_param_buf_id;                /* Sequence level parameter */
  95     VABufferID pic_param_buf_id;                /* Picture level parameter */
  96     VABufferID slice_param_buf_id[MAX_SLICES];  /* Slice level parameter, multil slices */
  97     VABufferID codedbuf_buf_id;                 /* Output buffer, compressed data */
  98     VABufferID packed_seq_header_param_buf_id;
  99     VABufferID packed_seq_buf_id;
 100     VABufferID packed_pic_header_param_buf_id;
 101     VABufferID packed_pic_buf_id;
 102     int num_slice_groups;
 103     int codedbuf_i_size;
 104     int codedbuf_pb_size;
 105
 106     /* thread */
 107     pthread_t upload_thread_id;
 108     int upload_thread_value;
 109     int current_input_surface;
 110     int current_upload_surface;
 111 };
 112
 113 /*
 114  * mpeg2enc helpers
 115  */
 116 #define BITSTREAM_ALLOCATE_STEPPING     4096
 117
 118 struct __bitstream {
 119     unsigned int *buffer;
 120     int bit_offset;
 121     int max_size_in_dword;
 122 };
 123
 124 typedef struct __bitstream bitstream;
 125
 126 static unsigned int
 127 swap32(unsigned int val)
 128 {
 129     unsigned char *pval = (unsigned char *)&val;
 130
 131     return ((pval[0] << 24)     |
 132             (pval[1] << 16)     |
 133             (pval[2] << 8)      |
 134             (pval[3] << 0));
 135 }
 136
 137 static void
 138 bitstream_start(bitstream *bs)
 139 {
 140     bs->max_size_in_dword = BITSTREAM_ALLOCATE_STEPPING;
 141     bs->buffer = calloc(bs->max_size_in_dword * sizeof(int), 1);
 142     bs->bit_offset = 0;
 143 }
 144
 145 static void
 146 bitstream_end(bitstream *bs)
 147 {
 148     int pos = (bs->bit_offset >> 5);
 149     int bit_offset = (bs->bit_offset & 0x1f);
 150     int bit_left = 32 - bit_offset;
 151
 152     if (bit_offset) {
 153         bs->buffer[pos] = swap32((bs->buffer[pos] << bit_left));
 154     }
 155 }
 156
 157 static void
 158 bitstream_put_ui(bitstream *bs, unsigned int val, int size_in_bits)
 159 {
 160     int pos = (bs->bit_offset >> 5);
 161     int bit_offset = (bs->bit_offset & 0x1f);
 162     int bit_left = 32 - bit_offset;
 163
 164     if (!size_in_bits)
 165         return;
 166
 167     if (size_in_bits < 32)
 168         val &= ((1 << size_in_bits) - 1);
 169
 170     bs->bit_offset += size_in_bits;
 171
 172     if (bit_left > size_in_bits) {
 173         bs->buffer[pos] = (bs->buffer[pos] << size_in_bits | val);
 174     } else {
 175         size_in_bits -= bit_left;
 176         bs->buffer[pos] = (bs->buffer[pos] << bit_left) | (val >> size_in_bits);
 177         bs->buffer[pos] = swap32(bs->buffer[pos]);
 178
 179         if (pos + 1 == bs->max_size_in_dword) {
 180             bs->max_size_in_dword += BITSTREAM_ALLOCATE_STEPPING;
 181             bs->buffer = realloc(bs->buffer, bs->max_size_in_dword * sizeof(unsigned int));
 182         }
 183
 184         bs->buffer[pos + 1] = val;
 185     }
 186 }
 187
 188 static void
 189 bitstream_byte_aligning(bitstream *bs, int bit)
 190 {
 191     int bit_offset = (bs->bit_offset & 0x7);
 192     int bit_left = 8 - bit_offset;
 193     int new_val;
 194
 195     if (!bit_offset)
 196         return;
 197
 198     assert(bit == 0 || bit == 1);
 199
 200     if (bit)
 201         new_val = (1 << bit_left) - 1;
 202     else
 203         new_val = 0;
 204
 205     bitstream_put_ui(bs, new_val, bit_left);
 206 }
 207
 208 static struct mpeg2_frame_rate {
 209     int code;
 210     float value;
 211 } frame_rate_tab[] = {
 212     {1, 23.976},
 213     {2, 24.0},
 214     {3, 25.0},
 215     {4, 29.97},
 216     {5, 30},
 217     {6, 50},
 218     {7, 59.94},
 219     {8, 60}
 220 };
 221
 222 static int
 223 find_frame_rate_code(const VAEncSequenceParameterBufferMPEG2 *seq_param)
 224 {
 225     unsigned int delta = -1;
 226     int code = 1, i;
 227     float frame_rate_value = seq_param->frame_rate *
 228         (seq_param->sequence_extension.bits.frame_rate_extension_d + 1) /
 229         (seq_param->sequence_extension.bits.frame_rate_extension_n + 1);
 230
 231     for (i = 0; i < sizeof(frame_rate_tab) / sizeof(frame_rate_tab[0]); i++) {
 232
 233         if (abs(1000 * frame_rate_tab[i].value - 1000 * frame_rate_value) < delta) {
 234             code = frame_rate_tab[i].code;
 235             delta = abs(1000 * frame_rate_tab[i].value - 1000 * frame_rate_value);
 236         }
 237     }
 238
 239     return code;
 240 }
 241
 242 static void
 243 sps_rbsp(const VAEncSequenceParameterBufferMPEG2 *seq_param,
 244          bitstream *bs)
 245 {
 246     int frame_rate_code = find_frame_rate_code(seq_param);
 247
 248     bitstream_put_ui(bs, START_CODE_SEQ, 32);
 249     bitstream_put_ui(bs, seq_param->picture_width, 12);
 250     bitstream_put_ui(bs, seq_param->picture_height, 12);
 251     bitstream_put_ui(bs, seq_param->aspect_ratio_information, 4);
 252     bitstream_put_ui(bs, frame_rate_code, 4); /* frame_rate_code */
 253     bitstream_put_ui(bs, (seq_param->bits_per_second + 399) / 400, 18); /* the low 18 bits of bit_rate */
 254     bitstream_put_ui(bs, 1, 1); /* marker_bit */
 255     bitstream_put_ui(bs, seq_param->vbv_buffer_size, 10);
 256     bitstream_put_ui(bs, 0, 1); /* constraint_parameter_flag, always 0 for MPEG-2 */
 257     bitstream_put_ui(bs, 0, 1); /* load_intra_quantiser_matrix */
 258     bitstream_put_ui(bs, 0, 1); /* load_non_intra_quantiser_matrix */
 259
 260     bitstream_byte_aligning(bs, 0);
 261
 262     bitstream_put_ui(bs, START_CODE_EXT, 32);
 263     bitstream_put_ui(bs, 1, 4); /* sequence_extension id */
 264     bitstream_put_ui(bs, seq_param->sequence_extension.bits.profile_and_level_indication, 8);
 265     bitstream_put_ui(bs, seq_param->sequence_extension.bits.progressive_sequence, 1);
 266     bitstream_put_ui(bs, seq_param->sequence_extension.bits.chroma_format, 2);
 267     bitstream_put_ui(bs, seq_param->picture_width >> 12, 2);
 268     bitstream_put_ui(bs, seq_param->picture_height >> 12, 2);
 269     bitstream_put_ui(bs, ((seq_param->bits_per_second + 399) / 400) >> 18, 12); /* bit_rate_extension */
 270     bitstream_put_ui(bs, 1, 1); /* marker_bit */
 271     bitstream_put_ui(bs, seq_param->vbv_buffer_size >> 10, 8);
 272     bitstream_put_ui(bs, seq_param->sequence_extension.bits.low_delay, 1);
 273     bitstream_put_ui(bs, seq_param->sequence_extension.bits.frame_rate_extension_n, 2);
 274     bitstream_put_ui(bs, seq_param->sequence_extension.bits.frame_rate_extension_d, 5);
 275
 276     bitstream_byte_aligning(bs, 0);
 277 }
 278
 279 static void
 280 pps_rbsp(const VAEncSequenceParameterBufferMPEG2 *seq_param,
 281          const VAEncPictureParameterBufferMPEG2 *pic_param,
 282          bitstream *bs)
 283 {
 284     int i;
 285     int chroma_420_type;
 286
 287     if (pic_param->temporal_reference == 0) {
 288         bitstream_put_ui(bs, START_CODE_GOP, 32);
 289         bitstream_put_ui(bs, seq_param->gop_header.bits.time_code, 25);
 290         bitstream_put_ui(bs, seq_param->gop_header.bits.closed_gop, 1);
 291         bitstream_put_ui(bs, seq_param->gop_header.bits.broken_link, 1);
 292
 293         bitstream_byte_aligning(bs, 0);
 294     }
 295
 296     if (seq_param->sequence_extension.bits.chroma_format == CHROMA_FORMAT_420)
 297         chroma_420_type = pic_param->picture_coding_extension.bits.progressive_frame;
 298     else
 299         chroma_420_type = 0;
 300
 301     bitstream_put_ui(bs, START_CODE_PICUTRE, 32);
 302     bitstream_put_ui(bs, pic_param->temporal_reference, 10);
 303     bitstream_put_ui(bs,
 304                      pic_param->picture_type == VAEncPictureTypeIntra ? 1 :
 305                      pic_param->picture_type == VAEncPictureTypePredictive ? 2 : 3,
 306                      3);
 307     bitstream_put_ui(bs, 0xFFFF, 16); /* vbv_delay, always 0xFFFF */
 308
 309     if (pic_param->picture_type == VAEncPictureTypePredictive) {
 310         bitstream_put_ui(bs, 0, 1); /* full_pel_forward_vector, always 0 for MPEG-2 */
 311         bitstream_put_ui(bs, 7, 3); /* forward_f_code, always 7 for MPEG-2 */
 312     }
 313
 314     if (pic_param->picture_type == VAEncPictureTypeBidirectional) {
 315         bitstream_put_ui(bs, 0, 1); /* full_pel_backward_vector, always 0 for MPEG-2 */
 316         bitstream_put_ui(bs, 7, 3); /* backward_f_code, always 7 for MPEG-2 */
 317     }
 318
 319     bitstream_put_ui(bs, 0, 1); /* extra_bit_picture, 0 */
 320
 321     bitstream_byte_aligning(bs, 0);
 322
 323     bitstream_put_ui(bs, START_CODE_EXT, 32);
 324     bitstream_put_ui(bs, 8, 4); /* Picture Coding Extension ID: 8 */
 325     bitstream_put_ui(bs, pic_param->f_code[0][0], 4);
 326     bitstream_put_ui(bs, pic_param->f_code[0][1], 4);
 327     bitstream_put_ui(bs, pic_param->f_code[1][0], 4);
 328     bitstream_put_ui(bs, pic_param->f_code[1][1], 4);
 329
 330     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.intra_dc_precision, 2);
 331     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.picture_structure, 2);
 332     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.top_field_first, 1);
 333     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.frame_pred_frame_dct, 1);
 334     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.concealment_motion_vectors, 1);
 335     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.q_scale_type, 1);
 336     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.intra_vlc_format, 1);
 337     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.alternate_scan, 1);
 338     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.repeat_first_field, 1);
 339     bitstream_put_ui(bs, chroma_420_type, 1);
 340     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.progressive_frame, 1);
 341     bitstream_put_ui(bs, pic_param->picture_coding_extension.bits.composite_display_flag, 1);
 342
 343     bitstream_byte_aligning(bs, 0);
 344
 345     if (pic_param->user_data_length) {
 346         bitstream_put_ui(bs, START_CODE_USER, 32);
 347
 348         for (i = 0; i < pic_param->user_data_length; i++) {
 349             bitstream_put_ui(bs, pic_param->user_data[i], 8);
 350         }
 351
 352         bitstream_byte_aligning(bs, 0);
 353     }
 354 }
 355
 356 static int
 357 build_packed_pic_buffer(const VAEncSequenceParameterBufferMPEG2 *seq_param,
 358                         const VAEncPictureParameterBufferMPEG2 *pic_param,
 359                         unsigned char **header_buffer)
 360 {
 361     bitstream bs;
 362
 363     bitstream_start(&bs);
 364     pps_rbsp(seq_param, pic_param, &bs);
 365     bitstream_end(&bs);
 366
 367     *header_buffer = (unsigned char *)bs.buffer;
 368     return bs.bit_offset;
 369 }
 370
 371 static int
 372 build_packed_seq_buffer(const VAEncSequenceParameterBufferMPEG2 *seq_param,
 373                         unsigned char **header_buffer)
 374 {
 375     bitstream bs;
 376
 377     bitstream_start(&bs);
 378     sps_rbsp(seq_param, &bs);
 379     bitstream_end(&bs);
 380
 381     *header_buffer = (unsigned char *)bs.buffer;
 382     return bs.bit_offset;
 383 }
 384
 385 /*
 386  * mpeg2enc
 387  */
 388 #define SID_INPUT_PICTURE_0                     0
 389 #define SID_INPUT_PICTURE_1                     1
 390 #define SID_REFERENCE_PICTURE_L0                2
 391 #define SID_REFERENCE_PICTURE_L1                3
 392 #define SID_RECON_PICTURE                       4
 393 #define SID_NUMBER                              SID_RECON_PICTURE + 1
 394
 395 static VASurfaceID surface_ids[SID_NUMBER];
 396
 397 /*
 398  * upload thread function
 399  */
 400 static void *
 401 upload_yuv_to_surface(void *data)
 402 {
 403     struct mpeg2enc_context *ctx = data;
 404     VAImage surface_image;
 405     VAStatus va_status;
 406     void *surface_p = NULL;
 407     unsigned char *y_src, *u_src, *v_src;
 408     unsigned char *y_dst, *u_dst, *v_dst;
 409     int y_size = ctx->width * ctx->height;
 410     int u_size = (ctx->width >> 1) * (ctx->height >> 1);
 411     int row, col;
 412     size_t n_items;
 413
 414     do {
 415         n_items = fread(ctx->frame_data_buffer, ctx->frame_size, 1, ctx->ifp);
 416     } while (n_items != 1);
 417
 418     va_status = vaDeriveImage(ctx->va_dpy, surface_ids[ctx->current_upload_surface], &surface_image);
 419     CHECK_VASTATUS(va_status,"vaDeriveImage");
 420
 421     vaMapBuffer(ctx->va_dpy, surface_image.buf, &surface_p);
 422     assert(VA_STATUS_SUCCESS == va_status);
 423
 424     y_src = ctx->frame_data_buffer;
 425     u_src = ctx->frame_data_buffer + y_size; /* UV offset for NV12 */
 426     v_src = ctx->frame_data_buffer + y_size + u_size;
 427
 428     y_dst = surface_p + surface_image.offsets[0];
 429     u_dst = surface_p + surface_image.offsets[1]; /* UV offset for NV12 */
 430     v_dst = surface_p + surface_image.offsets[2];
 431
 432     /* Y plane */
 433     for (row = 0; row < surface_image.height; row++) {
 434         memcpy(y_dst, y_src, surface_image.width);
 435         y_dst += surface_image.pitches[0];
 436         y_src += ctx->width;
 437     }
 438
 439     if (surface_image.format.fourcc == VA_FOURCC_NV12) { /* UV plane */
 440         for (row = 0; row < surface_image.height / 2; row++) {
 441             for (col = 0; col < surface_image.width / 2; col++) {
 442                 u_dst[col * 2] = u_src[col];
 443                 u_dst[col * 2 + 1] = v_src[col];
 444             }
 445
 446             u_dst += surface_image.pitches[1];
 447             u_src += (ctx->width / 2);
 448             v_src += (ctx->width / 2);
 449         }
 450     } else {
 451         for (row = 0; row < surface_image.height / 2; row++) {
 452             for (col = 0; col < surface_image.width / 2; col++) {
 453                 u_dst[col] = u_src[col];
 454                 v_dst[col] = v_src[col];
 455             }
 456
 457             u_dst += surface_image.pitches[1];
 458             v_dst += surface_image.pitches[2];
 459             u_src += (ctx->width / 2);
 460             v_src += (ctx->width / 2);
 461         }
 462     }
 463
 464     vaUnmapBuffer(ctx->va_dpy, surface_image.buf);
 465     vaDestroyImage(ctx->va_dpy, surface_image.image_id);
 466
 467     return NULL;
 468 }
 469
 470 static void
 471 mpeg2enc_exit(struct mpeg2enc_context *ctx, int exit_code)
 472 {
 473     if (ctx->frame_data_buffer) {
 474         free(ctx->frame_data_buffer);
 475         ctx->frame_data_buffer = NULL;
 476     }
 477
 478     if (ctx->ifp) {
 479         fclose(ctx->ifp);
 480         ctx->ifp = NULL;
 481     }
 482
 483     if (ctx->ofp) {
 484         fclose(ctx->ofp);
 485         ctx->ofp = NULL;
 486     }
 487
 488     exit(exit_code);
 489 }
 490
 491 static void
 492 usage(char *program)
 493 {
 494     fprintf(stderr, "Usage: %s --help\n", program);
 495     fprintf(stderr, "\t--help   print this message\n");
 496     fprintf(stderr, "Usage: %s <width> <height> <ifile> <ofile> [options]\n", program);
 497     fprintf(stderr, "\t<width>  specifies the frame width\n");
 498     fprintf(stderr, "\t<height> specifies the frame height\n");
 499     fprintf(stderr, "\t<ifile>  specifies the I420/IYUV YUV file\n");
 500     fprintf(stderr, "\t<ofile>  specifies the encoded MPEG-2 file\n");
 501     fprintf(stderr, "where options include:\n");
 502     fprintf(stderr, "\t--cqp <QP>       const qp mode with specified <QP>\n");
 503     fprintf(stderr, "\t--fps <FPS>      specify the frame rate\n");
 504     fprintf(stderr, "\t--mode <MODE>    specify the mode 0 (I), 1 (I/P) and 2 (I/P/B)\n");
 505 }
 506
 507 static void
 508 parse_args(struct mpeg2enc_context *ctx, int argc, char **argv)
 509 {
 510     int c, tmp, i = 0;
 511     int option_index = 0;
 512     long file_size;
 513     static struct option long_options[] = {
 514         {"help",        no_argument,            0,      'h'},
 515         {"cqp",         required_argument,      0,      'c'},
 516         {"fps",         required_argument,      0,      'f'},
 517         {"mode",        required_argument,      0,      'm'},
 518         { NULL,         0,                      NULL,   0 }
 519     };
 520
 521     if ((argc == 2 && strcmp(argv[1], "--help") == 0) ||
 522         (argc < 5))
 523         goto print_usage;
 524
 525     ctx->width = atoi(argv[1]);
 526     ctx->height = atoi(argv[2]);
 527
 528     if (ctx->width <= 0 || ctx->height <= 0) {
 529         fprintf(stderr, "<width> and <height> must be greater than 0\n");
 530         goto err_exit;
 531     }
 532
 533     ctx->ifp = fopen(argv[3], "rb");
 534
 535     if (ctx->ifp == NULL) {
 536         fprintf(stderr, "Can't open the input file\n");
 537         goto err_exit;
 538     }
 539
 540     fseek(ctx->ifp, 0l, SEEK_END);
 541     file_size = ftell(ctx->ifp);
 542     ctx->frame_size = ctx->width * ctx->height * 3 / 2;
 543
 544     if ((file_size < ctx->frame_size) ||
 545         (file_size % ctx->frame_size)) {
 546         fprintf(stderr, "The input file size %ld isn't a multiple of the frame size %d\n", file_size, ctx->frame_size);
 547         goto err_exit;
 548     }
 549
 550     ctx->num_pictures = file_size / ctx->frame_size;
 551     fseek(ctx->ifp, 0l, SEEK_SET);
 552
 553     ctx->ofp = fopen(argv[4], "wb");
 554
 555     if (ctx->ofp == NULL) {
 556         fprintf(stderr, "Can't create the output file\n");
 557         goto err_exit;
 558     }
 559
 560     opterr = 0;
 561     ctx->fps = 30;
 562     ctx->qp = 8;
 563     ctx->rate_control_mode = VA_RC_CQP;
 564     ctx->mode = 1;
 565
 566     optind = 5;
 567
 568     while((c = getopt_long(argc, argv,
 569                            "",
 570                            long_options,
 571                            &option_index)) != -1) {
 572         switch(c) {
 573         case 'c':
 574             tmp = atoi(optarg);
 575
 576             /* only support q_scale_type = 0 */
 577             if (tmp > 62 || tmp < 2) {
 578                 fprintf(stderr, "Warning: QP must be in [2, 62]\n");
 579
 580                 if (tmp > 62)
 581                     tmp = 62;
 582
 583                 if (tmp < 2)
 584                     tmp = 2;
 585             }
 586
 587             ctx->qp = tmp & 0xFE;
 588             ctx->rate_control_mode = VA_RC_CQP;
 589
 590             break;
 591
 592         case 'f':
 593             tmp = atoi(optarg);
 594
 595             if (tmp <= 0)
 596                 fprintf(stderr, "Warning: FPS must be greater than 0\n");
 597             else
 598                 ctx->fps = tmp;
 599
 600             ctx->rate_control_mode = VA_RC_CBR;
 601
 602             break;
 603
 604         case 'm':
 605             tmp = atoi(optarg);
 606
 607             if (tmp < 0 || tmp > 2)
 608                 fprintf(stderr, "Waning: MODE must be 0, 1, or 2\n");
 609             else
 610                 ctx->mode = tmp;
 611
 612             break;
 613
 614         case '?':
 615             fprintf(stderr, "Error: unkown command options\n");
 616
 617         case 'h':
 618             goto print_usage;
 619         }
 620     }
 621
 622     return;
 623
 624 print_usage:
 625     usage(argv[0]);
 626 err_exit:
 627     mpeg2enc_exit(ctx, 1);
 628 }
 629
 630 /*
 631  * init
 632  */
 633 static void
 634 mpeg2enc_init_sequence_parameter(struct mpeg2enc_context *ctx,
 635                                 VAEncSequenceParameterBufferMPEG2 *seq_param)
 636 {
 637     seq_param->intra_period = ctx->intra_period;
 638     seq_param->ip_period = 0;   /* FIXME: ??? */
 639     seq_param->picture_width = ctx->width;
 640     seq_param->picture_height = ctx->height;
 641
 642     if (ctx->bit_rate > 0)
 643         seq_param->bits_per_second = 1024 * ctx->bit_rate; /* use kbps as input */
 644     else
 645         seq_param->bits_per_second = 0x3FFFF * 400;
 646
 647     seq_param->frame_rate = ctx->fps;
 648     seq_param->aspect_ratio_information = 1;
 649     seq_param->vbv_buffer_size = 3; /* B = 16 * 1024 * vbv_buffer_size */
 650
 651     seq_param->sequence_extension.bits.profile_and_level_indication = 4 << 4 | 0x8; /* FIXME: Main/Main ??? */
 652     seq_param->sequence_extension.bits.progressive_sequence = 1; /* progressive frame-pictures */
 653     seq_param->sequence_extension.bits.chroma_format = CHROMA_FORMAT_420; /* 4:2:0 */
 654     seq_param->sequence_extension.bits.low_delay = 0; /* FIXME */
 655     seq_param->sequence_extension.bits.frame_rate_extension_n = 0;
 656     seq_param->sequence_extension.bits.frame_rate_extension_d = 0;
 657
 658     seq_param->gop_header.bits.time_code = (1 << 12); /* bit12: marker_bit */
 659     seq_param->gop_header.bits.closed_gop = 0;
 660     seq_param->gop_header.bits.broken_link = 0;
 661 }
 662
 663 static void
 664 mpeg2enc_init_picture_parameter(struct mpeg2enc_context *ctx,
 665                                VAEncPictureParameterBufferMPEG2 *pic_param)
 666 {
 667     pic_param->forward_reference_picture = VA_INVALID_ID;
 668     pic_param->backward_reference_picture = VA_INVALID_ID;
 669     pic_param->reconstructed_picture = VA_INVALID_ID;
 670     pic_param->coded_buf = VA_INVALID_ID;
 671     pic_param->picture_type = VAEncPictureTypeIntra;
 672
 673     pic_param->temporal_reference = 0;
 674     pic_param->f_code[0][0] = 0xf;
 675     pic_param->f_code[0][1] = 0xf;
 676     pic_param->f_code[1][0] = 0xf;
 677     pic_param->f_code[1][1] = 0xf;
 678
 679     pic_param->picture_coding_extension.bits.intra_dc_precision = 0; /* 8bits */
 680     pic_param->picture_coding_extension.bits.picture_structure = 3; /* frame picture */
 681     pic_param->picture_coding_extension.bits.top_field_first = 0;
 682     pic_param->picture_coding_extension.bits.frame_pred_frame_dct = 1; /* FIXME */
 683     pic_param->picture_coding_extension.bits.concealment_motion_vectors = 0;
 684     pic_param->picture_coding_extension.bits.q_scale_type = 0;
 685     pic_param->picture_coding_extension.bits.intra_vlc_format = 0;
 686     pic_param->picture_coding_extension.bits.alternate_scan = 0;
 687     pic_param->picture_coding_extension.bits.repeat_first_field = 0;
 688     pic_param->picture_coding_extension.bits.progressive_frame = 1;
 689     pic_param->picture_coding_extension.bits.composite_display_flag = 0;
 690
 691     pic_param->user_data_length = 0;
 692 }
 693
 694 static void
 695 mpeg2enc_alloc_va_resources(struct mpeg2enc_context *ctx)
 696 {
 697     VAEntrypoint *entrypoint_list;
 698     VAConfigAttrib attrib_list[2];
 699     VAStatus va_status;
 700     int max_entrypoints, num_entrypoints, entrypoint;
 701     int major_ver, minor_ver;
 702
 703     ctx->va_dpy = va_open_display();
 704     va_status = vaInitialize(ctx->va_dpy,
 705                              &major_ver,
 706                              &minor_ver);
 707     CHECK_VASTATUS(va_status, "vaInitialize");
 708
 709     max_entrypoints = vaMaxNumEntrypoints(ctx->va_dpy);
 710     entrypoint_list = malloc(max_entrypoints * sizeof(VAEntrypoint));
 711     vaQueryConfigEntrypoints(ctx->va_dpy,
 712                              VAProfileMPEG2Simple,
 713                              entrypoint_list,
 714                              &num_entrypoints);
 715
 716     for (entrypoint = 0; entrypoint < num_entrypoints; entrypoint++) {
 717         if (entrypoint_list[entrypoint] == VAEntrypointEncSlice)
 718             break;
 719     }
 720
 721     free(entrypoint_list);
 722
 723     if (entrypoint == num_entrypoints) {
 724         /* not find Slice entry point */
 725         assert(0);
 726     }
 727
 728     /* find out the format for the render target, and rate control mode */
 729     attrib_list[0].type = VAConfigAttribRTFormat;
 730     attrib_list[1].type = VAConfigAttribRateControl;
 731     vaGetConfigAttributes(ctx->va_dpy,
 732                           VAProfileMPEG2Simple,
 733                           VAEntrypointEncSlice,
 734                           &attrib_list[0],
 735                           2);
 736
 737     if ((attrib_list[0].value & VA_RT_FORMAT_YUV420) == 0) {
 738         /* not find desired YUV420 RT format */
 739         assert(0);
 740     }
 741
 742     if ((attrib_list[1].value & ctx->rate_control_mode) == 0) {
 743         /* Can't find matched RC mode */
 744         fprintf(stderr, "RC mode %d isn't found, exit\n", ctx->rate_control_mode);
 745         assert(0);
 746     }
 747
 748     attrib_list[0].value = VA_RT_FORMAT_YUV420; /* set to desired RT format */
 749     attrib_list[1].value = ctx->rate_control_mode; /* set to desired RC mode */
 750
 751     va_status = vaCreateConfig(ctx->va_dpy,
 752                                VAProfileMPEG2Simple,
 753                                VAEntrypointEncSlice,
 754                                attrib_list,
 755                                2,
 756                                &ctx->config_id);
 757     CHECK_VASTATUS(va_status, "vaCreateConfig");
 758
 759     /* Create a context for this decode pipe */
 760     va_status = vaCreateContext(ctx->va_dpy,
 761                                 ctx->config_id,
 762                                 ctx->width,
 763                                 ctx->height,
 764                                 VA_PROGRESSIVE,
 765                                 0,
 766                                 0,
 767                                 &ctx->context_id);
 768     CHECK_VASTATUS(va_status, "vaCreateContext");
 769
 770     va_status = vaCreateSurfaces(ctx->va_dpy,
 771                                  VA_RT_FORMAT_YUV420,
 772                                  ctx->width,
 773                                  ctx->height,
 774                                  surface_ids,
 775                                  SID_NUMBER,
 776                                  NULL,
 777                                  0);
 778     CHECK_VASTATUS(va_status, "vaCreateSurfaces");
 779 }
 780
 781 static void
 782 mpeg2enc_init(struct mpeg2enc_context *ctx)
 783 {
 784     int i;
 785
 786     ctx->frame_data_buffer = (unsigned char *)malloc(ctx->frame_size);
 787     ctx->seq_param_buf_id = VA_INVALID_ID;
 788     ctx->pic_param_buf_id = VA_INVALID_ID;
 789     ctx->packed_seq_header_param_buf_id = VA_INVALID_ID;
 790     ctx->packed_seq_buf_id = VA_INVALID_ID;
 791     ctx->packed_pic_header_param_buf_id = VA_INVALID_ID;
 792     ctx->packed_pic_buf_id = VA_INVALID_ID;
 793     ctx->codedbuf_buf_id = VA_INVALID_ID;
 794     ctx->codedbuf_i_size = ctx->frame_size;
 795     ctx->codedbuf_pb_size = 0;
 796
 797     if (ctx->mode == 0)
 798         ctx->intra_period = 1;
 799     else
 800         ctx->intra_period = 16;
 801
 802     ctx->bit_rate = -1;
 803
 804     for (i = 0; i < MAX_SLICES; i++) {
 805         ctx->slice_param_buf_id[i] = VA_INVALID_ID;
 806     }
 807
 808     mpeg2enc_init_sequence_parameter(ctx, &ctx->seq_param);
 809     mpeg2enc_init_picture_parameter(ctx, &ctx->pic_param);
 810     mpeg2enc_alloc_va_resources(ctx);
 811
 812     /* thread */
 813     ctx->current_input_surface = SID_INPUT_PICTURE_0;
 814     ctx->current_upload_surface = SID_INPUT_PICTURE_1;
 815     ctx->upload_thread_value = pthread_create(&ctx->upload_thread_id,
 816                                               NULL,
 817                                               upload_yuv_to_surface,
 818                                               ctx);
 819 }
 820
 821 static int
 822 mpeg2enc_time_code(VAEncSequenceParameterBufferMPEG2 *seq_param,
 823                    int num_frames)
 824 {
 825     int fps = (int)(seq_param->frame_rate + 0.5);
 826     int time_code = 0;
 827     int time_code_pictures, time_code_seconds, time_code_minutes, time_code_hours;
 828     int drop_frame_flag = 0;
 829
 830     assert(fps <= 60);
 831
 832     time_code_seconds = num_frames / fps;
 833     time_code_pictures = num_frames % fps;
 834     time_code |= time_code_pictures;
 835
 836     time_code_minutes = time_code_minutes / 60;
 837     time_code_seconds = time_code_minutes % 60;
 838     time_code |= (time_code_seconds << 6);
 839
 840     time_code_hours = time_code_minutes / 60;
 841     time_code_minutes = time_code_minutes % 60;
 842
 843     time_code |= (1 << 12);     /* marker_bit */
 844     time_code |= (time_code_minutes << 13);
 845
 846     time_code_hours = time_code_hours % 24;
 847     time_code |= (time_code_hours << 19);
 848
 849     time_code |= (drop_frame_flag << 24);
 850
 851     return time_code;
 852 }
 853
 854 /*
 855  * run
 856  */
 857 static void
 858 mpeg2enc_update_sequence_parameter(struct mpeg2enc_context *ctx,
 859                                    VAEncPictureType picture_type,
 860                                    int coded_order,
 861                                    int display_order)
 862 {
 863     VAEncSequenceParameterBufferMPEG2 *seq_param = &ctx->seq_param;
 864
 865     /* update the GOP info for the new GOP */
 866     if (display_order % ctx->intra_period == 0) {
 867         seq_param->gop_header.bits.time_code = mpeg2enc_time_code(seq_param, display_order);
 868     }
 869 }
 870
 871 static void
 872 mpeg2enc_update_picture_parameter(struct mpeg2enc_context *ctx,
 873                                   VAEncPictureType picture_type,
 874                                   int coded_order,
 875                                   int display_order)
 876 {
 877     VAEncPictureParameterBufferMPEG2 *pic_param = &ctx->pic_param;
 878
 879     pic_param->picture_type = picture_type;
 880     pic_param->temporal_reference = display_order % ctx->intra_period;
 881     pic_param->reconstructed_picture = surface_ids[SID_RECON_PICTURE];
 882     pic_param->forward_reference_picture = surface_ids[SID_REFERENCE_PICTURE_L0];
 883     pic_param->backward_reference_picture = surface_ids[SID_REFERENCE_PICTURE_L1];
 884 }
 885
 886 static void
 887 mpeg2enc_update_picture_parameter_buffer(struct mpeg2enc_context *ctx,
 888                                          VAEncPictureType picture_type,
 889                                          int coded_order,
 890                                          int display_order)
 891 {
 892     VAEncPictureParameterBufferMPEG2 *pic_param = &ctx->pic_param;
 893     VAStatus va_status;
 894
 895     /* update the coded buffer id */
 896     pic_param->coded_buf = ctx->codedbuf_buf_id;
 897     va_status = vaCreateBuffer(ctx->va_dpy,
 898                                ctx->context_id,
 899                                VAEncPictureParameterBufferType,
 900                                sizeof(*pic_param),
 901                                1,
 902                                pic_param,
 903                                &ctx->pic_param_buf_id);
 904     CHECK_VASTATUS(va_status, "vaCreateBuffer");
 905 }
 906
 907 static void
 908 mpeg2enc_update_slice_parameter(struct mpeg2enc_context *ctx, VAEncPictureType picture_type)
 909 {
 910     VAEncSequenceParameterBufferMPEG2 *seq_param;
 911     VAEncPictureParameterBufferMPEG2 *pic_param;
 912     VAEncSliceParameterBufferMPEG2 *slice_param;
 913     VAStatus va_status;
 914     int i, width_in_mbs, height_in_mbs;
 915
 916     pic_param = &ctx->pic_param;
 917     assert(pic_param->picture_coding_extension.bits.q_scale_type == 0);
 918
 919     seq_param = &ctx->seq_param;
 920     width_in_mbs = (seq_param->picture_width + 15) / 16;
 921     height_in_mbs = (seq_param->picture_height + 15) / 16;
 922     ctx->num_slice_groups = 1;
 923
 924     for (i = 0; i < height_in_mbs; i++) {
 925         slice_param = &ctx->slice_param[i];
 926         slice_param->macroblock_address = i * width_in_mbs;
 927         slice_param->num_macroblocks = width_in_mbs;
 928         slice_param->is_intra_slice = (picture_type == VAEncPictureTypeIntra);
 929         slice_param->quantiser_scale_code = ctx->qp / 2;
 930     }
 931
 932     va_status = vaCreateBuffer(ctx->va_dpy,
 933                                ctx->context_id,
 934                                VAEncSliceParameterBufferType,
 935                                sizeof(*slice_param),
 936                                height_in_mbs,
 937                                ctx->slice_param,
 938                                ctx->slice_param_buf_id);
 939     CHECK_VASTATUS(va_status, "vaCreateBuffer");;
 940 }
 941
 942 static int
 943 begin_picture(struct mpeg2enc_context *ctx,
 944               int coded_order,
 945               int display_order,
 946               VAEncPictureType picture_type)
 947 {
 948     VAStatus va_status;
 949     int tmp;
 950     VAEncPackedHeaderParameterBuffer packed_header_param_buffer;
 951     unsigned int length_in_bits;
 952     unsigned char *packed_seq_buffer = NULL, *packed_pic_buffer = NULL;
 953
 954     if (ctx->upload_thread_value != 0) {
 955         fprintf(stderr, "FATAL error!!!\n");
 956         exit(1);
 957     }
 958
 959     pthread_join(ctx->upload_thread_id, NULL);
 960
 961     ctx->upload_thread_value = -1;
 962     tmp = ctx->current_input_surface;
 963     ctx->current_input_surface = ctx->current_upload_surface;
 964     ctx->current_upload_surface = tmp;
 965
 966     mpeg2enc_update_sequence_parameter(ctx, picture_type, coded_order, display_order);
 967     mpeg2enc_update_picture_parameter(ctx, picture_type, coded_order, display_order);
 968
 969     if (coded_order == 0) {
 970         assert(picture_type == VAEncPictureTypeIntra);
 971         length_in_bits = build_packed_seq_buffer(&ctx->seq_param, &packed_seq_buffer);
 972         packed_header_param_buffer.type = VAEncPackedHeaderMPEG2_SPS;
 973         packed_header_param_buffer.has_emulation_bytes = 0;
 974         packed_header_param_buffer.bit_length = length_in_bits;
 975         va_status = vaCreateBuffer(ctx->va_dpy,
 976                                    ctx->context_id,
 977                                    VAEncPackedHeaderParameterBufferType,
 978                                    sizeof(packed_header_param_buffer), 1, &packed_header_param_buffer,
 979                                    &ctx->packed_seq_header_param_buf_id);
 980         CHECK_VASTATUS(va_status,"vaCreateBuffer");
 981
 982         va_status = vaCreateBuffer(ctx->va_dpy,
 983                                    ctx->context_id,
 984                                    VAEncPackedHeaderDataBufferType,
 985                                    (length_in_bits + 7) / 8, 1, packed_seq_buffer,
 986                                    &ctx->packed_seq_buf_id);
 987         CHECK_VASTATUS(va_status,"vaCreateBuffer");
 988
 989         free(packed_seq_buffer);
 990     }
 991
 992     length_in_bits = build_packed_pic_buffer(&ctx->seq_param, &ctx->pic_param, &packed_pic_buffer);
 993     packed_header_param_buffer.type = VAEncPackedHeaderMPEG2_PPS;
 994     packed_header_param_buffer.has_emulation_bytes = 0;
 995     packed_header_param_buffer.bit_length = length_in_bits;
 996
 997     va_status = vaCreateBuffer(ctx->va_dpy,
 998                                ctx->context_id,
 999                                VAEncPackedHeaderParameterBufferType,
1000                                sizeof(packed_header_param_buffer), 1, &packed_header_param_buffer,
1001                                &ctx->packed_pic_header_param_buf_id);
1002     CHECK_VASTATUS(va_status,"vaCreateBuffer");
1003
1004     va_status = vaCreateBuffer(ctx->va_dpy,
1005                                ctx->context_id,
1006                                VAEncPackedHeaderDataBufferType,
1007                                (length_in_bits + 7) / 8, 1, packed_pic_buffer,
1008                                &ctx->packed_pic_buf_id);
1009     CHECK_VASTATUS(va_status,"vaCreateBuffer");
1010
1011     free(packed_pic_buffer);
1012
1013     /* sequence parameter set */
1014     VAEncSequenceParameterBufferMPEG2 *seq_param = &ctx->seq_param;
1015     va_status = vaCreateBuffer(ctx->va_dpy,
1016                                ctx->context_id,
1017                                VAEncSequenceParameterBufferType,
1018                                sizeof(*seq_param),
1019                                1,
1020                                seq_param,
1021                                &ctx->seq_param_buf_id);
1022     CHECK_VASTATUS(va_status,"vaCreateBuffer");;
1023
1024     /* slice parameter */
1025     mpeg2enc_update_slice_parameter(ctx, picture_type);
1026
1027     return 0;
1028 }
1029
1030 static int
1031 mpeg2enc_render_picture(struct mpeg2enc_context *ctx)
1032 {
1033     VAStatus va_status;
1034     VABufferID va_buffers[16];
1035     unsigned int num_va_buffers = 0;
1036
1037     va_buffers[num_va_buffers++] = ctx->seq_param_buf_id;
1038     va_buffers[num_va_buffers++] = ctx->pic_param_buf_id;
1039
1040     if (ctx->packed_seq_header_param_buf_id != VA_INVALID_ID)
1041         va_buffers[num_va_buffers++] = ctx->packed_seq_header_param_buf_id;
1042
1043     if (ctx->packed_seq_buf_id != VA_INVALID_ID)
1044         va_buffers[num_va_buffers++] = ctx->packed_seq_buf_id;
1045
1046     if (ctx->packed_pic_header_param_buf_id != VA_INVALID_ID)
1047         va_buffers[num_va_buffers++] = ctx->packed_pic_header_param_buf_id;
1048
1049     if (ctx->packed_pic_buf_id != VA_INVALID_ID)
1050         va_buffers[num_va_buffers++] = ctx->packed_pic_buf_id;
1051
1052     va_status = vaBeginPicture(ctx->va_dpy,
1053                                ctx->context_id,
1054                                surface_ids[ctx->current_input_surface]);
1055     CHECK_VASTATUS(va_status,"vaBeginPicture");
1056
1057     va_status = vaRenderPicture(ctx->va_dpy,
1058                                 ctx->context_id,
1059                                 va_buffers,
1060                                 num_va_buffers);
1061     CHECK_VASTATUS(va_status,"vaRenderPicture");
1062
1063     va_status = vaRenderPicture(ctx->va_dpy,
1064                                 ctx->context_id,
1065                                 &ctx->slice_param_buf_id[0],
1066                                 ctx->num_slice_groups);
1067     CHECK_VASTATUS(va_status,"vaRenderPicture");
1068
1069     va_status = vaEndPicture(ctx->va_dpy, ctx->context_id);
1070     CHECK_VASTATUS(va_status,"vaEndPicture");
1071
1072     return 0;
1073 }
1074
1075 static int
1076 mpeg2enc_destroy_buffers(struct mpeg2enc_context *ctx, VABufferID *va_buffers, unsigned int num_va_buffers)
1077 {
1078     VAStatus va_status;
1079     unsigned int i;
1080
1081     for (i = 0; i < num_va_buffers; i++) {
1082         if (va_buffers[i] != VA_INVALID_ID) {
1083             va_status = vaDestroyBuffer(ctx->va_dpy, va_buffers[i]);
1084             CHECK_VASTATUS(va_status,"vaDestroyBuffer");
1085             va_buffers[i] = VA_INVALID_ID;
1086         }
1087     }
1088
1089     return 0;
1090 }
1091
1092 static void
1093 end_picture(struct mpeg2enc_context *ctx, VAEncPictureType picture_type, int next_is_bpic)
1094 {
1095     VABufferID tempID;
1096
1097     /* Prepare for next picture */
1098     tempID = surface_ids[SID_RECON_PICTURE];
1099
1100     if (picture_type != VAEncPictureTypeBidirectional) {
1101         if (next_is_bpic) {
1102             surface_ids[SID_RECON_PICTURE] = surface_ids[SID_REFERENCE_PICTURE_L1];
1103             surface_ids[SID_REFERENCE_PICTURE_L1] = tempID;
1104         } else {
1105             surface_ids[SID_RECON_PICTURE] = surface_ids[SID_REFERENCE_PICTURE_L0];
1106             surface_ids[SID_REFERENCE_PICTURE_L0] = tempID;
1107         }
1108     } else {
1109         if (!next_is_bpic) {
1110             surface_ids[SID_RECON_PICTURE] = surface_ids[SID_REFERENCE_PICTURE_L0];
1111             surface_ids[SID_REFERENCE_PICTURE_L0] = surface_ids[SID_REFERENCE_PICTURE_L1];
1112             surface_ids[SID_REFERENCE_PICTURE_L1] = tempID;
1113         }
1114     }
1115
1116     mpeg2enc_destroy_buffers(ctx, &ctx->seq_param_buf_id, 1);
1117     mpeg2enc_destroy_buffers(ctx, &ctx->pic_param_buf_id, 1);
1118     mpeg2enc_destroy_buffers(ctx, &ctx->packed_seq_header_param_buf_id, 1);
1119     mpeg2enc_destroy_buffers(ctx, &ctx->packed_seq_buf_id, 1);
1120     mpeg2enc_destroy_buffers(ctx, &ctx->packed_pic_header_param_buf_id, 1);
1121     mpeg2enc_destroy_buffers(ctx, &ctx->packed_pic_buf_id, 1);
1122     mpeg2enc_destroy_buffers(ctx, &ctx->slice_param_buf_id[0], ctx->num_slice_groups);
1123     mpeg2enc_destroy_buffers(ctx, &ctx->codedbuf_buf_id, 1);
1124     memset(ctx->slice_param, 0, sizeof(ctx->slice_param));
1125     ctx->num_slice_groups = 0;
1126 }
1127
1128 static int
1129 store_coded_buffer(struct mpeg2enc_context *ctx, VAEncPictureType picture_type)
1130 {
1131     VACodedBufferSegment *coded_buffer_segment;
1132     unsigned char *coded_mem;
1133     int slice_data_length;
1134     VAStatus va_status;
1135     VASurfaceStatus surface_status;
1136     size_t w_items;
1137
1138     va_status = vaSyncSurface(ctx->va_dpy, surface_ids[ctx->current_input_surface]);
1139     CHECK_VASTATUS(va_status,"vaSyncSurface");
1140
1141     surface_status = 0;
1142     va_status = vaQuerySurfaceStatus(ctx->va_dpy, surface_ids[ctx->current_input_surface], &surface_status);
1143     CHECK_VASTATUS(va_status,"vaQuerySurfaceStatus");
1144
1145     va_status = vaMapBuffer(ctx->va_dpy, ctx->codedbuf_buf_id, (void **)(&coded_buffer_segment));
1146     CHECK_VASTATUS(va_status,"vaMapBuffer");
1147     coded_mem = coded_buffer_segment->buf;
1148
1149     if (coded_buffer_segment->status & VA_CODED_BUF_STATUS_SLICE_OVERFLOW_MASK) {
1150         if (picture_type == VAEncPictureTypeIntra)
1151             ctx->codedbuf_i_size *= 2;
1152         else
1153             ctx->codedbuf_pb_size *= 2;
1154
1155         vaUnmapBuffer(ctx->va_dpy, ctx->codedbuf_buf_id);
1156         return -1;
1157     }
1158
1159     slice_data_length = coded_buffer_segment->size;
1160
1161     do {
1162         w_items = fwrite(coded_mem, slice_data_length, 1, ctx->ofp);
1163     } while (w_items != 1);
1164
1165     if (picture_type == VAEncPictureTypeIntra) {
1166         if (ctx->codedbuf_i_size > slice_data_length * 3 / 2) {
1167             ctx->codedbuf_i_size = slice_data_length * 3 / 2;
1168         }
1169
1170         if (ctx->codedbuf_pb_size < slice_data_length) {
1171             ctx->codedbuf_pb_size = slice_data_length;
1172         }
1173     } else {
1174         if (ctx->codedbuf_pb_size > slice_data_length * 3 / 2) {
1175             ctx->codedbuf_pb_size = slice_data_length * 3 / 2;
1176         }
1177     }
1178
1179     vaUnmapBuffer(ctx->va_dpy, ctx->codedbuf_buf_id);
1180
1181     return 0;
1182 }
1183
1184 static void
1185 encode_picture(struct mpeg2enc_context *ctx,
1186                int coded_order,
1187                int display_order,
1188                VAEncPictureType picture_type,
1189                int next_is_bpic,
1190                int next_display_order)
1191 {
1192     VAStatus va_status;
1193     int ret = 0, codedbuf_size;
1194
1195     begin_picture(ctx, coded_order, display_order, picture_type);
1196
1197     if (1) {
1198         /* upload YUV data to VA surface for next frame */
1199         if (next_display_order >= ctx->num_pictures)
1200             next_display_order = ctx->num_pictures - 1;
1201
1202         fseek(ctx->ifp, ctx->frame_size * next_display_order, SEEK_SET);
1203         ctx->upload_thread_value = pthread_create(&ctx->upload_thread_id,
1204                                                   NULL,
1205                                                   upload_yuv_to_surface,
1206                                                   ctx);
1207     }
1208
1209     do {
1210         mpeg2enc_destroy_buffers(ctx, &ctx->codedbuf_buf_id, 1);
1211         mpeg2enc_destroy_buffers(ctx, &ctx->pic_param_buf_id, 1);
1212
1213
1214         if (VAEncPictureTypeIntra == picture_type) {
1215             codedbuf_size = ctx->codedbuf_i_size;
1216         } else {
1217             codedbuf_size = ctx->codedbuf_pb_size;
1218         }
1219
1220         /* coded buffer */
1221         va_status = vaCreateBuffer(ctx->va_dpy,
1222                                    ctx->context_id,
1223                                    VAEncCodedBufferType,
1224                                    codedbuf_size, 1, NULL,
1225                                    &ctx->codedbuf_buf_id);
1226         CHECK_VASTATUS(va_status,"vaCreateBuffer");
1227
1228         /* picture parameter set */
1229         mpeg2enc_update_picture_parameter_buffer(ctx, picture_type, coded_order, display_order);
1230
1231         mpeg2enc_render_picture(ctx);
1232
1233         ret = store_coded_buffer(ctx, picture_type);
1234     } while (ret);
1235
1236     end_picture(ctx, picture_type, next_is_bpic);
1237 }
1238
1239 static void
1240 encode_pb_pictures(struct mpeg2enc_context *ctx,
1241                    int coded_order,
1242                    int f,
1243                    int nbframes,
1244                    int next_f)
1245 {
1246     int i;
1247
1248     encode_picture(ctx,
1249                    coded_order,
1250                    f + nbframes,
1251                    VAEncPictureTypePredictive,
1252                    1,
1253                    f);
1254
1255     for( i = 0; i < nbframes - 1; i++) {
1256         encode_picture(ctx,
1257                        coded_order + 1,
1258                        f + i,
1259                        VAEncPictureTypeBidirectional,
1260                        1,
1261                        f + i + 1);
1262     }
1263
1264     encode_picture(ctx,
1265                    coded_order + 1,
1266                    f + nbframes - 1,
1267                    VAEncPictureTypeBidirectional,
1268                    0,
1269                    next_f);
1270 }
1271
1272 static void
1273 mpeg2enc_run(struct mpeg2enc_context *ctx)
1274 {
1275     int display_order = 0, coded_order = 0;
1276
1277     for (display_order = 0; display_order < ctx->num_pictures;) {
1278         if (ctx->mode == 0) {
1279             encode_picture(ctx,
1280                            coded_order,
1281                            display_order,
1282                            VAEncPictureTypeIntra,
1283                            0,
1284                            display_order + 1);
1285             display_order++;
1286             coded_order++;
1287         } else if (ctx->mode == 1) {
1288             assert(0);
1289             if ((display_order % ctx->intra_period) == 0) {
1290                 encode_picture(ctx,
1291                                coded_order,
1292                                display_order,
1293                                VAEncPictureTypeIntra,
1294                                0,
1295                                display_order + 1);
1296                 display_order++;
1297                 coded_order++;
1298             } else {
1299                 encode_picture(ctx,
1300                                coded_order,
1301                                display_order,
1302                                VAEncPictureTypePredictive,
1303                                0,
1304                                display_order + 1);
1305                 display_order++;
1306                 coded_order++;
1307             }
1308         } else { // follow the i,p,b pattern
1309             static int fcurrent = 0;
1310             int fnext;
1311
1312             assert(0);
1313
1314             fcurrent = fcurrent % (sizeof(picture_type_pattern)/sizeof(int[2]));
1315             fnext = (fcurrent+1) % (sizeof(picture_type_pattern)/sizeof(int[2]));
1316
1317             if ( picture_type_pattern[fcurrent][0] == VAEncPictureTypeIntra ) {
1318                 encode_picture(ctx,
1319                                coded_order,
1320                                display_order,
1321                                VAEncPictureTypeIntra,
1322                                0,
1323                                display_order + picture_type_pattern[fnext][1]);
1324                 display_order++;
1325                 coded_order++;
1326             } else {
1327                 encode_pb_pictures(ctx,
1328                                    coded_order,
1329                                    display_order,
1330                                    picture_type_pattern[fcurrent][1] - 1,
1331                                    display_order + picture_type_pattern[fcurrent][1] + picture_type_pattern[fnext][1] - 1);
1332                 display_order += picture_type_pattern[fcurrent][1];
1333                 coded_order++;
1334             }
1335
1336             fcurrent++;
1337         }
1338
1339         fprintf(stderr, "\r %d/%d ...", display_order + 1, ctx->num_pictures);
1340         fflush(stdout);
1341     }
1342 }
1343
1344 /*
1345  * end
1346  */
1347 static void
1348 mpeg2enc_release_va_resources(struct mpeg2enc_context *ctx)
1349 {
1350     vaDestroySurfaces(ctx->va_dpy, surface_ids, SID_NUMBER);
1351     vaDestroyContext(ctx->va_dpy, ctx->context_id);
1352     vaDestroyConfig(ctx->va_dpy, ctx->config_id);
1353     vaTerminate(ctx->va_dpy);
1354     va_close_display(ctx->va_dpy);
1355 }
1356
1357 static void
1358 mpeg2enc_end(struct mpeg2enc_context *ctx)
1359 {
1360     pthread_join(ctx->upload_thread_id, NULL);
1361     mpeg2enc_release_va_resources(ctx);
1362 }
1363
1364 int
1365 main(int argc, char *argv[])
1366 {
1367     struct mpeg2enc_context ctx;
1368     struct timeval tpstart, tpend;
1369     float timeuse;
1370
1371     gettimeofday(&tpstart, NULL);
1372
1373     memset(&ctx, 0, sizeof(ctx));
1374     parse_args(&ctx, argc, argv);
1375     mpeg2enc_init(&ctx);
1376     mpeg2enc_run(&ctx);
1377     mpeg2enc_end(&ctx);
1378
1379     gettimeofday(&tpend, NULL);
1380     timeuse = 1000000 * (tpend.tv_sec - tpstart.tv_sec) + tpend.tv_usec - tpstart.tv_usec;
1381     timeuse /= 1000000;
1382     fprintf(stderr, "\ndone!\n");
1383     fprintf(stderr, "encode %d frames in %f secondes, FPS is %.1f\n", ctx.num_pictures, timeuse, ctx.num_pictures / timeuse);
1384
1385     mpeg2enc_exit(&ctx, 0);
1386
1387     return 0;
1388 }