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