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