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