OSDN Git Service

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