Fix crash in --me esa/tesa introduced in r1058
[x262.git] / muxers.c
1 /*****************************************************************************
2  * muxers.c: h264 file i/o plugins
3  *****************************************************************************
4  * Copyright (C) 2003-2008 x264 project
5  *
6  * Authors: Laurent Aimar <fenrir@via.ecp.fr>
7  *          Loren Merritt <lorenm@u.washington.edu>
8  *
9  * This program is free software; you can redistribute it and/or modify
10  * it under the terms of the GNU General Public License as published by
11  * the Free Software Foundation; either version 2 of the License, or
12  * (at your option) any later version.
13  *
14  * This program is distributed in the hope that it will be useful,
15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
17  * GNU General Public License for more details.
18  *
19  * You should have received a copy of the GNU General Public License
20  * along with this program; if not, write to the Free Software
21  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02111, USA.
22  *****************************************************************************/
23
24 #include "common/common.h"
25 #include "x264.h"
26 #include "matroska.h"
27 #include "muxers.h"
28
29 #ifndef _MSC_VER
30 #include "config.h"
31 #endif
32
33 #include <sys/types.h>
34
35 #ifdef AVIS_INPUT
36 #include <windows.h>
37 #include <vfw.h>
38 #endif
39
40 #ifdef MP4_OUTPUT
41 #include <gpac/isomedia.h>
42 #endif
43
44 static int64_t gcd( int64_t a, int64_t b )
45 {
46     while (1)
47     {
48         int64_t c = a % b;
49         if( !c )
50             return b;
51         a = b;
52         b = c;
53     }
54 }
55
56 typedef struct {
57     FILE *fh;
58     int width, height;
59     int next_frame;
60 } yuv_input_t;
61
62 /* raw 420 yuv file operation */
63 int open_file_yuv( char *psz_filename, hnd_t *p_handle, x264_param_t *p_param )
64 {
65     yuv_input_t *h = malloc(sizeof(yuv_input_t));
66     h->width = p_param->i_width;
67     h->height = p_param->i_height;
68     h->next_frame = 0;
69
70     if( !strcmp(psz_filename, "-") )
71         h->fh = stdin;
72     else
73         h->fh = fopen(psz_filename, "rb");
74     if( h->fh == NULL )
75         return -1;
76
77     *p_handle = (hnd_t)h;
78     return 0;
79 }
80
81 int get_frame_total_yuv( hnd_t handle )
82 {
83     yuv_input_t *h = handle;
84     int i_frame_total = 0;
85
86     if( !fseek( h->fh, 0, SEEK_END ) )
87     {
88         uint64_t i_size = ftell( h->fh );
89         fseek( h->fh, 0, SEEK_SET );
90         i_frame_total = (int)(i_size / ( h->width * h->height * 3 / 2 ));
91     }
92
93     return i_frame_total;
94 }
95
96 int read_frame_yuv( x264_picture_t *p_pic, hnd_t handle, int i_frame )
97 {
98     yuv_input_t *h = handle;
99
100     if( i_frame != h->next_frame )
101         if( fseek( h->fh, (uint64_t)i_frame * h->width * h->height * 3 / 2, SEEK_SET ) )
102             return -1;
103
104     if( fread( p_pic->img.plane[0], 1, h->width * h->height, h->fh ) <= 0
105             || fread( p_pic->img.plane[1], 1, h->width * h->height / 4, h->fh ) <= 0
106             || fread( p_pic->img.plane[2], 1, h->width * h->height / 4, h->fh ) <= 0 )
107         return -1;
108
109     h->next_frame = i_frame+1;
110
111     return 0;
112 }
113
114 int close_file_yuv(hnd_t handle)
115 {
116     yuv_input_t *h = handle;
117     if( !h || !h->fh )
118         return 0;
119     fclose( h->fh );
120     free( h );
121     return 0;
122 }
123
124 /* YUV4MPEG2 raw 420 yuv file operation */
125 typedef struct {
126     FILE *fh;
127     int width, height;
128     int next_frame;
129     int seq_header_len, frame_header_len;
130     int frame_size;
131 } y4m_input_t;
132
133 #define Y4M_MAGIC "YUV4MPEG2"
134 #define MAX_YUV4_HEADER 80
135 #define Y4M_FRAME_MAGIC "FRAME"
136 #define MAX_FRAME_HEADER 80
137
138 int open_file_y4m( char *psz_filename, hnd_t *p_handle, x264_param_t *p_param )
139 {
140     int  i, n, d;
141     int  interlaced;
142     char header[MAX_YUV4_HEADER+10];
143     char *tokstart, *tokend, *header_end;
144     y4m_input_t *h = malloc(sizeof(y4m_input_t));
145
146     h->next_frame = 0;
147
148     if( !strcmp(psz_filename, "-") )
149         h->fh = stdin;
150     else
151         h->fh = fopen(psz_filename, "rb");
152     if( h->fh == NULL )
153         return -1;
154
155     h->frame_header_len = strlen(Y4M_FRAME_MAGIC)+1;
156
157     /* Read header */
158     for( i=0; i<MAX_YUV4_HEADER; i++ )
159     {
160         header[i] = fgetc(h->fh);
161         if( header[i] == '\n' )
162         {
163             /* Add a space after last option. Makes parsing "444" vs
164                "444alpha" easier. */
165             header[i+1] = 0x20;
166             header[i+2] = 0;
167             break;
168         }
169     }
170     if( i == MAX_YUV4_HEADER || strncmp(header, Y4M_MAGIC, strlen(Y4M_MAGIC)) )
171         return -1;
172
173     /* Scan properties */
174     header_end = &header[i+1]; /* Include space */
175     h->seq_header_len = i+1;
176     for( tokstart = &header[strlen(Y4M_MAGIC)+1]; tokstart < header_end; tokstart++ )
177     {
178         if(*tokstart==0x20) continue;
179         switch(*tokstart++)
180         {
181         case 'W': /* Width. Required. */
182             h->width = p_param->i_width = strtol(tokstart, &tokend, 10);
183             tokstart=tokend;
184             break;
185         case 'H': /* Height. Required. */
186             h->height = p_param->i_height = strtol(tokstart, &tokend, 10);
187             tokstart=tokend;
188             break;
189         case 'C': /* Color space */
190             if( strncmp("420", tokstart, 3) )
191             {
192                 fprintf(stderr, "Colorspace unhandled\n");
193                 return -1;
194             }
195             tokstart = strchr(tokstart, 0x20);
196             break;
197         case 'I': /* Interlace type */
198             switch(*tokstart++)
199             {
200             case 'p': interlaced = 0; break;
201             case '?':
202             case 't':
203             case 'b':
204             case 'm':
205             default: interlaced = 1;
206                 fprintf(stderr, "Warning, this sequence might be interlaced\n");
207             }
208             break;
209         case 'F': /* Frame rate - 0:0 if unknown */
210             if( sscanf(tokstart, "%d:%d", &n, &d) == 2 && n && d )
211             {
212                 x264_reduce_fraction( &n, &d );
213                 p_param->i_fps_num = n;
214                 p_param->i_fps_den = d;
215             }
216             tokstart = strchr(tokstart, 0x20);
217             break;
218         case 'A': /* Pixel aspect - 0:0 if unknown */
219             /* Don't override the aspect ratio if sar has been explicitly set on the commandline. */
220             if( sscanf(tokstart, "%d:%d", &n, &d) == 2 && n && d && !p_param->vui.i_sar_width && !p_param->vui.i_sar_height )
221             {
222                 x264_reduce_fraction( &n, &d );
223                 p_param->vui.i_sar_width = n;
224                 p_param->vui.i_sar_height = d;
225             }
226             tokstart = strchr(tokstart, 0x20);
227             break;
228         case 'X': /* Vendor extensions */
229             if( !strncmp("YSCSS=",tokstart,6) )
230             {
231                 /* Older nonstandard pixel format representation */
232                 tokstart += 6;
233                 if( strncmp("420JPEG",tokstart,7) &&
234                     strncmp("420MPEG2",tokstart,8) &&
235                     strncmp("420PALDV",tokstart,8) )
236                 {
237                     fprintf(stderr, "Unsupported extended colorspace\n");
238                     return -1;
239                 }
240             }
241             tokstart = strchr(tokstart, 0x20);
242             break;
243         }
244     }
245
246     fprintf(stderr, "yuv4mpeg: %ix%i@%i/%ifps, %i:%i\n",
247             h->width, h->height, p_param->i_fps_num, p_param->i_fps_den,
248             p_param->vui.i_sar_width, p_param->vui.i_sar_height);
249
250     *p_handle = (hnd_t)h;
251     return 0;
252 }
253
254 /* Most common case: frame_header = "FRAME" */
255 int get_frame_total_y4m( hnd_t handle )
256 {
257     y4m_input_t *h             = handle;
258     int          i_frame_total = 0;
259     uint64_t     init_pos      = ftell(h->fh);
260
261     if( !fseek( h->fh, 0, SEEK_END ) )
262     {
263         uint64_t i_size = ftell( h->fh );
264         fseek( h->fh, init_pos, SEEK_SET );
265         i_frame_total = (int)((i_size - h->seq_header_len) /
266                               (3*(h->width*h->height)/2+h->frame_header_len));
267     }
268
269     return i_frame_total;
270 }
271
272 int read_frame_y4m( x264_picture_t *p_pic, hnd_t handle, int i_frame )
273 {
274     int          slen = strlen(Y4M_FRAME_MAGIC);
275     int          i    = 0;
276     char         header[16];
277     y4m_input_t *h    = handle;
278
279     if( i_frame != h->next_frame )
280     {
281         if (fseek(h->fh, (uint64_t)i_frame*(3*(h->width*h->height)/2+h->frame_header_len)
282                   + h->seq_header_len, SEEK_SET))
283             return -1;
284     }
285
286     /* Read frame header - without terminating '\n' */
287     if (fread(header, 1, slen, h->fh) != slen)
288         return -1;
289
290     header[slen] = 0;
291     if (strncmp(header, Y4M_FRAME_MAGIC, slen))
292     {
293         fprintf(stderr, "Bad header magic (%"PRIx32" <=> %s)\n",
294                 *((uint32_t*)header), header);
295         return -1;
296     }
297
298     /* Skip most of it */
299     while (i<MAX_FRAME_HEADER && fgetc(h->fh) != '\n')
300         i++;
301     if (i == MAX_FRAME_HEADER)
302     {
303         fprintf(stderr, "Bad frame header!\n");
304         return -1;
305     }
306     h->frame_header_len = i+slen+1;
307
308     if( fread(p_pic->img.plane[0], 1, h->width*h->height, h->fh) <= 0
309         || fread(p_pic->img.plane[1], 1, h->width * h->height / 4, h->fh) <= 0
310         || fread(p_pic->img.plane[2], 1, h->width * h->height / 4, h->fh) <= 0)
311         return -1;
312
313     h->next_frame = i_frame+1;
314
315     return 0;
316 }
317
318 int close_file_y4m(hnd_t handle)
319 {
320     y4m_input_t *h = handle;
321     if( !h || !h->fh )
322         return 0;
323     fclose( h->fh );
324     free( h );
325     return 0;
326 }
327
328 /* avs/avi input file support under cygwin */
329
330 #ifdef AVIS_INPUT
331 typedef struct {
332     PAVISTREAM p_avi;
333     int width, height;
334 } avis_input_t;
335
336 int open_file_avis( char *psz_filename, hnd_t *p_handle, x264_param_t *p_param )
337 {
338     avis_input_t *h = malloc(sizeof(avis_input_t));
339     AVISTREAMINFO info;
340     int i;
341
342     *p_handle = (hnd_t)h;
343
344     AVIFileInit();
345     if( AVIStreamOpenFromFile( &h->p_avi, psz_filename, streamtypeVIDEO, 0, OF_READ, NULL ) )
346     {
347         AVIFileExit();
348         return -1;
349     }
350
351     if( AVIStreamInfo(h->p_avi, &info, sizeof(AVISTREAMINFO)) )
352     {
353         AVIStreamRelease(h->p_avi);
354         AVIFileExit();
355         return -1;
356     }
357
358     // check input format
359     if (info.fccHandler != MAKEFOURCC('Y', 'V', '1', '2'))
360     {
361         fprintf( stderr, "avis [error]: unsupported input format (%c%c%c%c)\n",
362             (char)(info.fccHandler & 0xff), (char)((info.fccHandler >> 8) & 0xff),
363             (char)((info.fccHandler >> 16) & 0xff), (char)((info.fccHandler >> 24)) );
364
365         AVIStreamRelease(h->p_avi);
366         AVIFileExit();
367
368         return -1;
369     }
370
371     h->width =
372     p_param->i_width = info.rcFrame.right - info.rcFrame.left;
373     h->height =
374     p_param->i_height = info.rcFrame.bottom - info.rcFrame.top;
375     i = gcd(info.dwRate, info.dwScale);
376     p_param->i_fps_den = info.dwScale / i;
377     p_param->i_fps_num = info.dwRate / i;
378
379     fprintf( stderr, "avis [info]: %dx%d @ %.2f fps (%d frames)\n",
380         p_param->i_width, p_param->i_height,
381         (double)p_param->i_fps_num / (double)p_param->i_fps_den,
382         (int)info.dwLength );
383
384     return 0;
385 }
386
387 int get_frame_total_avis( hnd_t handle )
388 {
389     avis_input_t *h = handle;
390     AVISTREAMINFO info;
391
392     if( AVIStreamInfo(h->p_avi, &info, sizeof(AVISTREAMINFO)) )
393         return -1;
394
395     return info.dwLength;
396 }
397
398 int read_frame_avis( x264_picture_t *p_pic, hnd_t handle, int i_frame )
399 {
400     avis_input_t *h = handle;
401
402     p_pic->img.i_csp = X264_CSP_YV12;
403
404     if( AVIStreamRead(h->p_avi, i_frame, 1, p_pic->img.plane[0], h->width * h->height * 3 / 2, NULL, NULL ) )
405         return -1;
406
407     return 0;
408 }
409
410 int close_file_avis( hnd_t handle )
411 {
412     avis_input_t *h = handle;
413     AVIStreamRelease(h->p_avi);
414     AVIFileExit();
415     free(h);
416     return 0;
417 }
418 #endif
419
420
421 #ifdef HAVE_PTHREAD
422 typedef struct {
423     int (*p_read_frame)( x264_picture_t *p_pic, hnd_t handle, int i_frame );
424     int (*p_close_infile)( hnd_t handle );
425     hnd_t p_handle;
426     x264_picture_t pic;
427     x264_pthread_t tid;
428     int next_frame;
429     int frame_total;
430     int in_progress;
431     struct thread_input_arg_t *next_args;
432 } thread_input_t;
433
434 typedef struct thread_input_arg_t {
435     thread_input_t *h;
436     x264_picture_t *pic;
437     int i_frame;
438     int status;
439 } thread_input_arg_t;
440
441 int open_file_thread( char *psz_filename, hnd_t *p_handle, x264_param_t *p_param )
442 {
443     thread_input_t *h = malloc(sizeof(thread_input_t));
444     x264_picture_alloc( &h->pic, X264_CSP_I420, p_param->i_width, p_param->i_height );
445     h->p_read_frame = p_read_frame;
446     h->p_close_infile = p_close_infile;
447     h->p_handle = *p_handle;
448     h->in_progress = 0;
449     h->next_frame = -1;
450     h->next_args = malloc(sizeof(thread_input_arg_t));
451     h->next_args->h = h;
452     h->next_args->status = 0;
453     h->frame_total = p_get_frame_total( h->p_handle );
454
455     *p_handle = (hnd_t)h;
456     return 0;
457 }
458
459 int get_frame_total_thread( hnd_t handle )
460 {
461     thread_input_t *h = handle;
462     return h->frame_total;
463 }
464
465 static void read_frame_thread_int( thread_input_arg_t *i )
466 {
467     i->status = i->h->p_read_frame( i->pic, i->h->p_handle, i->i_frame );
468 }
469
470 int read_frame_thread( x264_picture_t *p_pic, hnd_t handle, int i_frame )
471 {
472     thread_input_t *h = handle;
473     UNUSED void *stuff;
474     int ret = 0;
475
476     if( h->next_frame >= 0 )
477     {
478         x264_pthread_join( h->tid, &stuff );
479         ret |= h->next_args->status;
480         h->in_progress = 0;
481     }
482
483     if( h->next_frame == i_frame )
484     {
485         XCHG( x264_picture_t, *p_pic, h->pic );
486     }
487     else
488     {
489         ret |= h->p_read_frame( p_pic, h->p_handle, i_frame );
490     }
491
492     if( !h->frame_total || i_frame+1 < h->frame_total )
493     {
494         h->next_frame =
495         h->next_args->i_frame = i_frame+1;
496         h->next_args->pic = &h->pic;
497         x264_pthread_create( &h->tid, NULL, (void*)read_frame_thread_int, h->next_args );
498         h->in_progress = 1;
499     }
500     else
501         h->next_frame = -1;
502
503     return ret;
504 }
505
506 int close_file_thread( hnd_t handle )
507 {
508     thread_input_t *h = handle;
509     h->p_close_infile( h->p_handle );
510     x264_picture_clean( &h->pic );
511     if( h->in_progress )
512         x264_pthread_join( h->tid, NULL );
513     free( h->next_args );
514     free( h );
515     return 0;
516 }
517 #endif
518
519
520 int open_file_bsf( char *psz_filename, hnd_t *p_handle )
521 {
522     if ((*p_handle = fopen(psz_filename, "w+b")) == NULL)
523         return -1;
524
525     return 0;
526 }
527
528 int set_param_bsf( hnd_t handle, x264_param_t *p_param )
529 {
530     return 0;
531 }
532
533 int write_nalu_bsf( hnd_t handle, uint8_t *p_nalu, int i_size )
534 {
535     if (fwrite(p_nalu, i_size, 1, (FILE *)handle) > 0)
536         return i_size;
537     return -1;
538 }
539
540 int set_eop_bsf( hnd_t handle,  x264_picture_t *p_picture )
541 {
542     return 0;
543 }
544
545 int close_file_bsf( hnd_t handle )
546 {
547     if ((handle == NULL) || (handle == stdout))
548         return 0;
549
550     return fclose((FILE *)handle);
551 }
552
553 /* -- mp4 muxing support ------------------------------------------------- */
554 #ifdef MP4_OUTPUT
555
556 typedef struct
557 {
558     GF_ISOFile *p_file;
559     GF_AVCConfig *p_config;
560     GF_ISOSample *p_sample;
561     int i_track;
562     uint32_t i_descidx;
563     int i_time_inc;
564     int i_time_res;
565     int i_numframe;
566     int i_init_delay;
567     uint8_t b_sps;
568     uint8_t b_pps;
569 } mp4_t;
570
571
572 static void recompute_bitrate_mp4(GF_ISOFile *p_file, int i_track)
573 {
574     u32 i, count, di, timescale, time_wnd, rate;
575     u64 offset;
576     Double br;
577     GF_ESD *esd;
578
579     esd = gf_isom_get_esd(p_file, i_track, 1);
580     if (!esd) return;
581
582     esd->decoderConfig->avgBitrate = 0;
583     esd->decoderConfig->maxBitrate = 0;
584     rate = time_wnd = 0;
585
586     timescale = gf_isom_get_media_timescale(p_file, i_track);
587     count = gf_isom_get_sample_count(p_file, i_track);
588     for (i=0; i<count; i++) {
589         GF_ISOSample *samp = gf_isom_get_sample_info(p_file, i_track, i+1, &di, &offset);
590
591         if (samp->dataLength>esd->decoderConfig->bufferSizeDB) esd->decoderConfig->bufferSizeDB = samp->dataLength;
592
593         if (esd->decoderConfig->bufferSizeDB < samp->dataLength) esd->decoderConfig->bufferSizeDB = samp->dataLength;
594         esd->decoderConfig->avgBitrate += samp->dataLength;
595         rate += samp->dataLength;
596         if (samp->DTS > time_wnd + timescale) {
597             if (rate > esd->decoderConfig->maxBitrate) esd->decoderConfig->maxBitrate = rate;
598             time_wnd = samp->DTS;
599             rate = 0;
600         }
601
602         gf_isom_sample_del(&samp);
603     }
604
605     br = (Double) (s64) gf_isom_get_media_duration(p_file, i_track);
606     br /= timescale;
607     esd->decoderConfig->avgBitrate = (u32) (esd->decoderConfig->avgBitrate / br);
608     /*move to bps*/
609     esd->decoderConfig->avgBitrate *= 8;
610     esd->decoderConfig->maxBitrate *= 8;
611
612     gf_isom_change_mpeg4_description(p_file, i_track, 1, esd);
613     gf_odf_desc_del((GF_Descriptor *) esd);
614 }
615
616
617 int close_file_mp4( hnd_t handle )
618 {
619     mp4_t *p_mp4 = (mp4_t *)handle;
620
621     if (p_mp4 == NULL)
622         return 0;
623
624     if (p_mp4->p_config)
625         gf_odf_avc_cfg_del(p_mp4->p_config);
626
627     if (p_mp4->p_sample)
628     {
629         if (p_mp4->p_sample->data)
630             free(p_mp4->p_sample->data);
631
632         gf_isom_sample_del(&p_mp4->p_sample);
633     }
634
635     if (p_mp4->p_file)
636     {
637         recompute_bitrate_mp4(p_mp4->p_file, p_mp4->i_track);
638         gf_isom_set_pl_indication(p_mp4->p_file, GF_ISOM_PL_VISUAL, 0x15);
639         gf_isom_set_storage_mode(p_mp4->p_file, GF_ISOM_STORE_FLAT);
640         gf_isom_close(p_mp4->p_file);
641     }
642
643     free(p_mp4);
644
645     return 0;
646 }
647
648 int open_file_mp4( char *psz_filename, hnd_t *p_handle )
649 {
650     mp4_t *p_mp4;
651
652     *p_handle = NULL;
653
654     if ((p_mp4 = (mp4_t *)malloc(sizeof(mp4_t))) == NULL)
655         return -1;
656
657     memset(p_mp4, 0, sizeof(mp4_t));
658     p_mp4->p_file = gf_isom_open(psz_filename, GF_ISOM_OPEN_WRITE, NULL);
659
660     if ((p_mp4->p_sample = gf_isom_sample_new()) == NULL)
661     {
662         close_file_mp4( p_mp4 );
663         return -1;
664     }
665
666     gf_isom_set_brand_info(p_mp4->p_file, GF_ISOM_BRAND_AVC1, 0);
667
668     *p_handle = p_mp4;
669
670     return 0;
671 }
672
673
674 int set_param_mp4( hnd_t handle, x264_param_t *p_param )
675 {
676     mp4_t *p_mp4 = (mp4_t *)handle;
677
678     p_mp4->i_track = gf_isom_new_track(p_mp4->p_file, 0, GF_ISOM_MEDIA_VISUAL,
679         p_param->i_fps_num);
680
681     p_mp4->p_config = gf_odf_avc_cfg_new();
682     gf_isom_avc_config_new(p_mp4->p_file, p_mp4->i_track, p_mp4->p_config,
683         NULL, NULL, &p_mp4->i_descidx);
684
685     gf_isom_set_track_enabled(p_mp4->p_file, p_mp4->i_track, 1);
686
687     gf_isom_set_visual_info(p_mp4->p_file, p_mp4->i_track, p_mp4->i_descidx,
688         p_param->i_width, p_param->i_height);
689
690     if( p_param->vui.i_sar_width && p_param->vui.i_sar_height )
691     {
692         uint64_t dw = p_param->i_width << 16;
693         uint64_t dh = p_param->i_height << 16;
694         double sar = (double)p_param->vui.i_sar_width / p_param->vui.i_sar_height;
695         if( sar > 1.0 )
696             dw *= sar ;
697         else
698             dh /= sar;
699         gf_isom_set_track_layout_info( p_mp4->p_file, p_mp4->i_track, dw, dh, 0, 0, 0 );
700     }
701
702     p_mp4->p_sample->data = (char *)malloc(p_param->i_width * p_param->i_height * 3 / 2);
703     if (p_mp4->p_sample->data == NULL)
704         return -1;
705
706     p_mp4->i_time_res = p_param->i_fps_num;
707     p_mp4->i_time_inc = p_param->i_fps_den;
708     p_mp4->i_init_delay = p_param->i_bframe ? (p_param->b_bframe_pyramid ? 2 : 1) : 0;
709     p_mp4->i_init_delay *= p_mp4->i_time_inc;
710     fprintf(stderr, "mp4 [info]: initial delay %d (scale %d)\n",
711         p_mp4->i_init_delay, p_mp4->i_time_res);
712
713     return 0;
714 }
715
716
717 int write_nalu_mp4( hnd_t handle, uint8_t *p_nalu, int i_size )
718 {
719     mp4_t *p_mp4 = (mp4_t *)handle;
720     GF_AVCConfigSlot *p_slot;
721     uint8_t type = p_nalu[4] & 0x1f;
722     int psize;
723
724     switch(type)
725     {
726     // sps
727     case 0x07:
728         if (!p_mp4->b_sps)
729         {
730             p_mp4->p_config->configurationVersion = 1;
731             p_mp4->p_config->AVCProfileIndication = p_nalu[5];
732             p_mp4->p_config->profile_compatibility = p_nalu[6];
733             p_mp4->p_config->AVCLevelIndication = p_nalu[7];
734             p_slot = (GF_AVCConfigSlot *)malloc(sizeof(GF_AVCConfigSlot));
735             p_slot->size = i_size - 4;
736             p_slot->data = (char *)malloc(p_slot->size);
737             memcpy(p_slot->data, p_nalu + 4, i_size - 4);
738             gf_list_add(p_mp4->p_config->sequenceParameterSets, p_slot);
739             p_slot = NULL;
740             p_mp4->b_sps = 1;
741         }
742         break;
743
744     // pps
745     case 0x08:
746         if (!p_mp4->b_pps)
747         {
748             p_slot = (GF_AVCConfigSlot *)malloc(sizeof(GF_AVCConfigSlot));
749             p_slot->size = i_size - 4;
750             p_slot->data = (char *)malloc(p_slot->size);
751             memcpy(p_slot->data, p_nalu + 4, i_size - 4);
752             gf_list_add(p_mp4->p_config->pictureParameterSets, p_slot);
753             p_slot = NULL;
754             p_mp4->b_pps = 1;
755             if (p_mp4->b_sps)
756                 gf_isom_avc_config_update(p_mp4->p_file, p_mp4->i_track, 1, p_mp4->p_config);
757         }
758         break;
759
760     // slice, sei
761     case 0x1:
762     case 0x5:
763     case 0x6:
764         psize = i_size - 4 ;
765         memcpy(p_mp4->p_sample->data + p_mp4->p_sample->dataLength, p_nalu, i_size);
766         p_mp4->p_sample->data[p_mp4->p_sample->dataLength + 0] = (psize >> 24) & 0xff;
767         p_mp4->p_sample->data[p_mp4->p_sample->dataLength + 1] = (psize >> 16) & 0xff;
768         p_mp4->p_sample->data[p_mp4->p_sample->dataLength + 2] = (psize >> 8) & 0xff;
769         p_mp4->p_sample->data[p_mp4->p_sample->dataLength + 3] = (psize >> 0) & 0xff;
770         p_mp4->p_sample->dataLength += i_size;
771         break;
772     }
773
774     return i_size;
775 }
776
777 int set_eop_mp4( hnd_t handle, x264_picture_t *p_picture )
778 {
779     mp4_t *p_mp4 = (mp4_t *)handle;
780     uint64_t dts = (uint64_t)p_mp4->i_numframe * p_mp4->i_time_inc;
781     uint64_t pts = (uint64_t)p_picture->i_pts;
782     int32_t offset = p_mp4->i_init_delay + pts - dts;
783
784     p_mp4->p_sample->IsRAP = p_picture->i_type == X264_TYPE_IDR ? 1 : 0;
785     p_mp4->p_sample->DTS = dts;
786     p_mp4->p_sample->CTS_Offset = offset;
787     gf_isom_add_sample(p_mp4->p_file, p_mp4->i_track, p_mp4->i_descidx, p_mp4->p_sample);
788
789     p_mp4->p_sample->dataLength = 0;
790     p_mp4->i_numframe++;
791
792     return 0;
793 }
794
795 #endif
796
797
798 /* -- mkv muxing support ------------------------------------------------- */
799 typedef struct
800 {
801     mk_Writer *w;
802
803     uint8_t   *sps, *pps;
804     int       sps_len, pps_len;
805
806     int       width, height, d_width, d_height;
807
808     int64_t   frame_duration;
809     int       fps_num;
810
811     int       b_header_written;
812     char      b_writing_frame;
813 } mkv_t;
814
815 static int write_header_mkv( mkv_t *p_mkv )
816 {
817     int       ret;
818     uint8_t   *avcC;
819     int  avcC_len;
820
821     if( p_mkv->sps == NULL || p_mkv->pps == NULL ||
822         p_mkv->width == 0 || p_mkv->height == 0 ||
823         p_mkv->d_width == 0 || p_mkv->d_height == 0)
824         return -1;
825
826     avcC_len = 5 + 1 + 2 + p_mkv->sps_len + 1 + 2 + p_mkv->pps_len;
827     avcC = malloc(avcC_len);
828     if (avcC == NULL)
829         return -1;
830
831     avcC[0] = 1;
832     avcC[1] = p_mkv->sps[1];
833     avcC[2] = p_mkv->sps[2];
834     avcC[3] = p_mkv->sps[3];
835     avcC[4] = 0xff; // nalu size length is four bytes
836     avcC[5] = 0xe1; // one sps
837
838     avcC[6] = p_mkv->sps_len >> 8;
839     avcC[7] = p_mkv->sps_len;
840
841     memcpy(avcC+8, p_mkv->sps, p_mkv->sps_len);
842
843     avcC[8+p_mkv->sps_len] = 1; // one pps
844     avcC[9+p_mkv->sps_len] = p_mkv->pps_len >> 8;
845     avcC[10+p_mkv->sps_len] = p_mkv->pps_len;
846
847     memcpy( avcC+11+p_mkv->sps_len, p_mkv->pps, p_mkv->pps_len );
848
849     ret = mk_writeHeader( p_mkv->w, "x264", "V_MPEG4/ISO/AVC",
850                           avcC, avcC_len, p_mkv->frame_duration, 50000,
851                           p_mkv->width, p_mkv->height,
852                           p_mkv->d_width, p_mkv->d_height );
853
854     free( avcC );
855
856     p_mkv->b_header_written = 1;
857
858     return ret;
859 }
860
861 int open_file_mkv( char *psz_filename, hnd_t *p_handle )
862 {
863     mkv_t *p_mkv;
864
865     *p_handle = NULL;
866
867     p_mkv  = malloc(sizeof(*p_mkv));
868     if (p_mkv == NULL)
869         return -1;
870
871     memset(p_mkv, 0, sizeof(*p_mkv));
872
873     p_mkv->w = mk_createWriter(psz_filename);
874     if (p_mkv->w == NULL)
875     {
876         free(p_mkv);
877         return -1;
878     }
879
880     *p_handle = p_mkv;
881
882     return 0;
883 }
884
885 int set_param_mkv( hnd_t handle, x264_param_t *p_param )
886 {
887     mkv_t   *p_mkv = handle;
888     int64_t dw, dh;
889
890     if( p_param->i_fps_num > 0 )
891     {
892         p_mkv->frame_duration = (int64_t)p_param->i_fps_den *
893                                 (int64_t)1000000000 / p_param->i_fps_num;
894         p_mkv->fps_num = p_param->i_fps_num;
895     }
896     else
897     {
898         p_mkv->frame_duration = 0;
899         p_mkv->fps_num = 1;
900     }
901
902     p_mkv->width = p_param->i_width;
903     p_mkv->height = p_param->i_height;
904
905     if( p_param->vui.i_sar_width && p_param->vui.i_sar_height )
906     {
907         dw = (int64_t)p_param->i_width  * p_param->vui.i_sar_width;
908         dh = (int64_t)p_param->i_height * p_param->vui.i_sar_height;
909     }
910     else
911     {
912         dw = p_param->i_width;
913         dh = p_param->i_height;
914     }
915
916     if( dw > 0 && dh > 0 )
917     {
918         int64_t x = gcd( dw, dh );
919         dw /= x;
920         dh /= x;
921     }
922
923     p_mkv->d_width = (int)dw;
924     p_mkv->d_height = (int)dh;
925
926     return 0;
927 }
928
929 int write_nalu_mkv( hnd_t handle, uint8_t *p_nalu, int i_size )
930 {
931     mkv_t *p_mkv = handle;
932     uint8_t type = p_nalu[4] & 0x1f;
933     uint8_t dsize[4];
934     int psize;
935
936     switch( type )
937     {
938     // sps
939     case 0x07:
940         if( !p_mkv->sps )
941         {
942             p_mkv->sps = malloc(i_size - 4);
943             if (p_mkv->sps == NULL)
944                 return -1;
945             p_mkv->sps_len = i_size - 4;
946             memcpy(p_mkv->sps, p_nalu + 4, i_size - 4);
947         }
948         break;
949
950     // pps
951     case 0x08:
952         if( !p_mkv->pps )
953         {
954             p_mkv->pps = malloc(i_size - 4);
955             if (p_mkv->pps == NULL)
956                 return -1;
957             p_mkv->pps_len = i_size - 4;
958             memcpy(p_mkv->pps, p_nalu + 4, i_size - 4);
959         }
960         break;
961
962     // slice, sei
963     case 0x1:
964     case 0x5:
965     case 0x6:
966         if( !p_mkv->b_writing_frame )
967         {
968             if( mk_startFrame(p_mkv->w) < 0 )
969                 return -1;
970             p_mkv->b_writing_frame = 1;
971         }
972         psize = i_size - 4 ;
973         dsize[0] = psize >> 24;
974         dsize[1] = psize >> 16;
975         dsize[2] = psize >> 8;
976         dsize[3] = psize;
977         if( mk_addFrameData(p_mkv->w, dsize, 4) < 0 ||
978             mk_addFrameData(p_mkv->w, p_nalu + 4, i_size - 4) < 0 )
979             return -1;
980         break;
981
982     default:
983         break;
984     }
985
986     if( !p_mkv->b_header_written && p_mkv->pps && p_mkv->sps &&
987         write_header_mkv(p_mkv) < 0 )
988         return -1;
989
990     return i_size;
991 }
992
993 int set_eop_mkv( hnd_t handle, x264_picture_t *p_picture )
994 {
995     mkv_t *p_mkv = handle;
996     int64_t i_stamp = (int64_t)(p_picture->i_pts * 1e9 / p_mkv->fps_num);
997
998     p_mkv->b_writing_frame = 0;
999
1000     return mk_setFrameFlags( p_mkv->w, i_stamp,
1001                              p_picture->i_type == X264_TYPE_IDR );
1002 }
1003
1004 int close_file_mkv( hnd_t handle )
1005 {
1006     mkv_t *p_mkv = handle;
1007     int   ret;
1008
1009     if( p_mkv->sps )
1010         free( p_mkv->sps );
1011     if( p_mkv->pps )
1012         free( p_mkv->pps );
1013
1014     ret = mk_close(p_mkv->w);
1015
1016     free( p_mkv );
1017
1018     return ret;
1019 }
1020