output example: use OutputStream for audio streams as well
[ffmpeg.git] / doc / examples / output.c
1 /*
2  * Copyright (c) 2003 Fabrice Bellard
3  *
4  * Permission is hereby granted, free of charge, to any person obtaining a copy
5  * of this software and associated documentation files (the "Software"), to deal
6  * in the Software without restriction, including without limitation the rights
7  * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
8  * copies of the Software, and to permit persons to whom the Software is
9  * furnished to do so, subject to the following conditions:
10  *
11  * The above copyright notice and this permission notice shall be included in
12  * all copies or substantial portions of the Software.
13  *
14  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
15  * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
16  * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
17  * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
18  * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
19  * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
20  * THE SOFTWARE.
21  */
22
23 /**
24  * @file
25  * libavformat API example.
26  *
27  * @example output.c
28  * Output a media file in any supported libavformat format. The default
29  * codecs are used.
30  */
31
32 #include <stdlib.h>
33 #include <stdio.h>
34 #include <string.h>
35 #include <math.h>
36
37 #include "libavutil/channel_layout.h"
38 #include "libavutil/mathematics.h"
39 #include "libavformat/avformat.h"
40 #include "libswscale/swscale.h"
41
42 /* 5 seconds stream duration */
43 #define STREAM_DURATION   5.0
44 #define STREAM_FRAME_RATE 25 /* 25 images/s */
45 #define STREAM_NB_FRAMES  ((int)(STREAM_DURATION * STREAM_FRAME_RATE))
46 #define STREAM_PIX_FMT    AV_PIX_FMT_YUV420P /* default pix_fmt */
47
48 static int sws_flags = SWS_BICUBIC;
49
50 // a wrapper around a single output AVStream
51 typedef struct OutputStream {
52     AVStream *st;
53
54     AVFrame *frame;
55     AVFrame *tmp_frame;
56
57     float t, tincr, tincr2;
58     int audio_input_frame_size;
59 } OutputStream;
60
61 /**************************************************************/
62 /* audio output */
63
64 /*
65  * add an audio output stream
66  */
67 static void add_audio_stream(OutputStream *ost, AVFormatContext *oc,
68                              enum AVCodecID codec_id)
69 {
70     AVCodecContext *c;
71     AVCodec *codec;
72
73     /* find the audio encoder */
74     codec = avcodec_find_encoder(codec_id);
75     if (!codec) {
76         fprintf(stderr, "codec not found\n");
77         exit(1);
78     }
79
80     ost->st = avformat_new_stream(oc, codec);
81     if (!ost->st) {
82         fprintf(stderr, "Could not alloc stream\n");
83         exit(1);
84     }
85
86     c = ost->st->codec;
87
88     /* put sample parameters */
89     c->sample_fmt  = AV_SAMPLE_FMT_S16;
90     c->bit_rate    = 64000;
91     c->sample_rate = 44100;
92     c->channels    = 2;
93     c->channel_layout = AV_CH_LAYOUT_STEREO;
94
95     // some formats want stream headers to be separate
96     if (oc->oformat->flags & AVFMT_GLOBALHEADER)
97         c->flags |= CODEC_FLAG_GLOBAL_HEADER;
98 }
99
100 static void open_audio(AVFormatContext *oc, OutputStream *ost)
101 {
102     AVCodecContext *c;
103
104     c = ost->st->codec;
105
106     /* open it */
107     if (avcodec_open2(c, NULL, NULL) < 0) {
108         fprintf(stderr, "could not open codec\n");
109         exit(1);
110     }
111
112     /* init signal generator */
113     ost->t     = 0;
114     ost->tincr = 2 * M_PI * 110.0 / c->sample_rate;
115     /* increment frequency by 110 Hz per second */
116     ost->tincr2 = 2 * M_PI * 110.0 / c->sample_rate / c->sample_rate;
117
118     if (c->codec->capabilities & CODEC_CAP_VARIABLE_FRAME_SIZE)
119         ost->audio_input_frame_size = 10000;
120     else
121         ost->audio_input_frame_size = c->frame_size;
122 }
123
124 /* Prepare a 16 bit dummy audio frame of 'frame_size' samples and
125  * 'nb_channels' channels. */
126 static void get_audio_frame(OutputStream *ost, AVFrame *frame, int nb_channels)
127 {
128     int j, i, v, ret;
129     int16_t *q = (int16_t*)frame->data[0];
130
131     /* when we pass a frame to the encoder, it may keep a reference to it
132      * internally;
133      * make sure we do not overwrite it here
134      */
135     ret = av_frame_make_writable(frame);
136     if (ret < 0)
137         exit(1);
138
139     for (j = 0; j < frame->nb_samples; j++) {
140         v = (int)(sin(ost->t) * 10000);
141         for (i = 0; i < nb_channels; i++)
142             *q++ = v;
143         ost->t     += ost->tincr;
144         ost->tincr += ost->tincr2;
145     }
146 }
147
148 static void write_audio_frame(AVFormatContext *oc, OutputStream *ost)
149 {
150     AVCodecContext *c;
151     AVPacket pkt = { 0 }; // data and size must be 0;
152     AVFrame *frame = av_frame_alloc();
153     int got_packet, ret;
154
155     av_init_packet(&pkt);
156     c = ost->st->codec;
157
158     frame->sample_rate    = c->sample_rate;
159     frame->nb_samples     = ost->audio_input_frame_size;
160     frame->format         = AV_SAMPLE_FMT_S16;
161     frame->channel_layout = c->channel_layout;
162     ret = av_frame_get_buffer(frame, 0);
163     if (ret < 0) {
164         fprintf(stderr, "Could not allocate an audio frame.\n");
165         exit(1);
166     }
167
168     get_audio_frame(ost, frame, c->channels);
169
170     avcodec_encode_audio2(c, &pkt, frame, &got_packet);
171     if (!got_packet)
172         return;
173
174     pkt.stream_index = ost->st->index;
175
176     /* Write the compressed frame to the media file. */
177     if (av_interleaved_write_frame(oc, &pkt) != 0) {
178         fprintf(stderr, "Error while writing audio frame\n");
179         exit(1);
180     }
181     av_frame_free(&frame);
182 }
183
184 static void close_audio(AVFormatContext *oc, OutputStream *ost)
185 {
186     avcodec_close(ost->st->codec);
187 }
188
189 /**************************************************************/
190 /* video output */
191
192 static int frame_count;
193
194 /* Add a video output stream. */
195 static void add_video_stream(OutputStream *ost, AVFormatContext *oc,
196                              enum AVCodecID codec_id)
197 {
198     AVCodecContext *c;
199     AVCodec *codec;
200
201     /* find the video encoder */
202     codec = avcodec_find_encoder(codec_id);
203     if (!codec) {
204         fprintf(stderr, "codec not found\n");
205         exit(1);
206     }
207
208     ost->st = avformat_new_stream(oc, codec);
209     if (!ost->st) {
210         fprintf(stderr, "Could not alloc stream\n");
211         exit(1);
212     }
213
214     c = ost->st->codec;
215
216     /* Put sample parameters. */
217     c->bit_rate = 400000;
218     /* Resolution must be a multiple of two. */
219     c->width    = 352;
220     c->height   = 288;
221     /* timebase: This is the fundamental unit of time (in seconds) in terms
222      * of which frame timestamps are represented. For fixed-fps content,
223      * timebase should be 1/framerate and timestamp increments should be
224      * identical to 1. */
225     c->time_base.den = STREAM_FRAME_RATE;
226     c->time_base.num = 1;
227     c->gop_size      = 12; /* emit one intra frame every twelve frames at most */
228     c->pix_fmt       = STREAM_PIX_FMT;
229     if (c->codec_id == AV_CODEC_ID_MPEG2VIDEO) {
230         /* just for testing, we also add B frames */
231         c->max_b_frames = 2;
232     }
233     if (c->codec_id == AV_CODEC_ID_MPEG1VIDEO) {
234         /* Needed to avoid using macroblocks in which some coeffs overflow.
235          * This does not happen with normal video, it just happens here as
236          * the motion of the chroma plane does not match the luma plane. */
237         c->mb_decision = 2;
238     }
239     /* Some formats want stream headers to be separate. */
240     if (oc->oformat->flags & AVFMT_GLOBALHEADER)
241         c->flags |= CODEC_FLAG_GLOBAL_HEADER;
242 }
243
244 static AVFrame *alloc_picture(enum AVPixelFormat pix_fmt, int width, int height)
245 {
246     AVFrame *picture;
247     int ret;
248
249     picture = av_frame_alloc();
250     if (!picture)
251         return NULL;
252
253     picture->format = pix_fmt;
254     picture->width  = width;
255     picture->height = height;
256
257     /* allocate the buffers for the frame data */
258     ret = av_frame_get_buffer(picture, 32);
259     if (ret < 0) {
260         fprintf(stderr, "Could not allocate frame data.\n");
261         exit(1);
262     }
263
264     return picture;
265 }
266
267 static void open_video(AVFormatContext *oc, OutputStream *ost)
268 {
269     AVCodecContext *c;
270
271     c = ost->st->codec;
272
273     /* open the codec */
274     if (avcodec_open2(c, NULL, NULL) < 0) {
275         fprintf(stderr, "could not open codec\n");
276         exit(1);
277     }
278
279     /* Allocate the encoded raw picture. */
280     ost->frame = alloc_picture(c->pix_fmt, c->width, c->height);
281     if (!ost->frame) {
282         fprintf(stderr, "Could not allocate picture\n");
283         exit(1);
284     }
285
286     /* If the output format is not YUV420P, then a temporary YUV420P
287      * picture is needed too. It is then converted to the required
288      * output format. */
289     ost->tmp_frame = NULL;
290     if (c->pix_fmt != AV_PIX_FMT_YUV420P) {
291         ost->tmp_frame = alloc_picture(AV_PIX_FMT_YUV420P, c->width, c->height);
292         if (!ost->tmp_frame) {
293             fprintf(stderr, "Could not allocate temporary picture\n");
294             exit(1);
295         }
296     }
297 }
298
299 /* Prepare a dummy image. */
300 static void fill_yuv_image(AVFrame *pict, int frame_index,
301                            int width, int height)
302 {
303     int x, y, i, ret;
304
305     /* when we pass a frame to the encoder, it may keep a reference to it
306      * internally;
307      * make sure we do not overwrite it here
308      */
309     ret = av_frame_make_writable(pict);
310     if (ret < 0)
311         exit(1);
312
313     i = frame_index;
314
315     /* Y */
316     for (y = 0; y < height; y++)
317         for (x = 0; x < width; x++)
318             pict->data[0][y * pict->linesize[0] + x] = x + y + i * 3;
319
320     /* Cb and Cr */
321     for (y = 0; y < height / 2; y++) {
322         for (x = 0; x < width / 2; x++) {
323             pict->data[1][y * pict->linesize[1] + x] = 128 + y + i * 2;
324             pict->data[2][y * pict->linesize[2] + x] = 64 + x + i * 5;
325         }
326     }
327 }
328
329 static void write_video_frame(AVFormatContext *oc, OutputStream *ost)
330 {
331     int ret;
332     AVCodecContext *c;
333     static struct SwsContext *img_convert_ctx;
334
335     c = ost->st->codec;
336
337     if (frame_count >= STREAM_NB_FRAMES) {
338         /* No more frames to compress. The codec has a latency of a few
339          * frames if using B-frames, so we get the last frames by
340          * passing the same picture again. */
341     } else {
342         if (c->pix_fmt != AV_PIX_FMT_YUV420P) {
343             /* as we only generate a YUV420P picture, we must convert it
344              * to the codec pixel format if needed */
345             if (img_convert_ctx == NULL) {
346                 img_convert_ctx = sws_getContext(c->width, c->height,
347                                                  AV_PIX_FMT_YUV420P,
348                                                  c->width, c->height,
349                                                  c->pix_fmt,
350                                                  sws_flags, NULL, NULL, NULL);
351                 if (img_convert_ctx == NULL) {
352                     fprintf(stderr,
353                             "Cannot initialize the conversion context\n");
354                     exit(1);
355                 }
356             }
357             fill_yuv_image(ost->tmp_frame, frame_count, c->width, c->height);
358             sws_scale(img_convert_ctx, ost->tmp_frame->data, ost->tmp_frame->linesize,
359                       0, c->height, ost->frame->data, ost->frame->linesize);
360         } else {
361             fill_yuv_image(ost->frame, frame_count, c->width, c->height);
362         }
363     }
364
365     if (oc->oformat->flags & AVFMT_RAWPICTURE) {
366         /* Raw video case - the API will change slightly in the near
367          * future for that. */
368         AVPacket pkt;
369         av_init_packet(&pkt);
370
371         pkt.flags        |= AV_PKT_FLAG_KEY;
372         pkt.stream_index  = ost->st->index;
373         pkt.data          = (uint8_t *)ost->frame;
374         pkt.size          = sizeof(AVPicture);
375
376         ret = av_interleaved_write_frame(oc, &pkt);
377     } else {
378         AVPacket pkt = { 0 };
379         int got_packet;
380         av_init_packet(&pkt);
381
382         /* encode the image */
383         ret = avcodec_encode_video2(c, &pkt, ost->frame, &got_packet);
384         /* If size is zero, it means the image was buffered. */
385         if (!ret && got_packet && pkt.size) {
386             av_packet_rescale_ts(&pkt, c->time_base, ost->st->time_base);
387             pkt.stream_index = ost->st->index;
388
389             /* Write the compressed frame to the media file. */
390             ret = av_interleaved_write_frame(oc, &pkt);
391         } else {
392             ret = 0;
393         }
394     }
395     if (ret != 0) {
396         fprintf(stderr, "Error while writing video frame\n");
397         exit(1);
398     }
399     frame_count++;
400 }
401
402 static void close_video(AVFormatContext *oc, OutputStream *ost)
403 {
404     avcodec_close(ost->st->codec);
405     av_frame_free(&ost->frame);
406     av_frame_free(&ost->tmp_frame);
407 }
408
409 /**************************************************************/
410 /* media file output */
411
412 int main(int argc, char **argv)
413 {
414     OutputStream video_st, audio_st;
415     const char *filename;
416     AVOutputFormat *fmt;
417     AVFormatContext *oc;
418     double audio_pts, video_pts;
419     int have_video = 0, have_audio = 0;
420     int i;
421
422     /* Initialize libavcodec, and register all codecs and formats. */
423     av_register_all();
424
425     if (argc != 2) {
426         printf("usage: %s output_file\n"
427                "API example program to output a media file with libavformat.\n"
428                "The output format is automatically guessed according to the file extension.\n"
429                "Raw images can also be output by using '%%d' in the filename\n"
430                "\n", argv[0]);
431         return 1;
432     }
433
434     filename = argv[1];
435
436     /* Autodetect the output format from the name. default is MPEG. */
437     fmt = av_guess_format(NULL, filename, NULL);
438     if (!fmt) {
439         printf("Could not deduce output format from file extension: using MPEG.\n");
440         fmt = av_guess_format("mpeg", NULL, NULL);
441     }
442     if (!fmt) {
443         fprintf(stderr, "Could not find suitable output format\n");
444         return 1;
445     }
446
447     /* Allocate the output media context. */
448     oc = avformat_alloc_context();
449     if (!oc) {
450         fprintf(stderr, "Memory error\n");
451         return 1;
452     }
453     oc->oformat = fmt;
454     snprintf(oc->filename, sizeof(oc->filename), "%s", filename);
455
456     /* Add the audio and video streams using the default format codecs
457      * and initialize the codecs. */
458     if (fmt->video_codec != AV_CODEC_ID_NONE) {
459         add_video_stream(&video_st, oc, fmt->video_codec);
460         have_video = 1;
461     }
462     if (fmt->audio_codec != AV_CODEC_ID_NONE) {
463         add_audio_stream(&audio_st, oc, fmt->audio_codec);
464         have_audio = 1;
465     }
466
467     /* Now that all the parameters are set, we can open the audio and
468      * video codecs and allocate the necessary encode buffers. */
469     if (have_video)
470         open_video(oc, &video_st);
471     if (have_audio)
472         open_audio(oc, &audio_st);
473
474     av_dump_format(oc, 0, filename, 1);
475
476     /* open the output file, if needed */
477     if (!(fmt->flags & AVFMT_NOFILE)) {
478         if (avio_open(&oc->pb, filename, AVIO_FLAG_WRITE) < 0) {
479             fprintf(stderr, "Could not open '%s'\n", filename);
480             return 1;
481         }
482     }
483
484     /* Write the stream header, if any. */
485     avformat_write_header(oc, NULL);
486
487     for (;;) {
488         /* Compute current audio and video time. */
489         if (have_audio)
490             audio_pts = (double)audio_st.st->pts.val * audio_st.st->time_base.num / audio_st.st->time_base.den;
491         else
492             audio_pts = 0.0;
493
494         if (have_video)
495             video_pts = (double)video_st.st->pts.val * video_st.st->time_base.num /
496                         video_st.st->time_base.den;
497         else
498             video_pts = 0.0;
499
500         if ((!have_audio || audio_pts >= STREAM_DURATION) &&
501             (!have_video || video_pts >= STREAM_DURATION))
502             break;
503
504         /* write interleaved audio and video frames */
505         if (!have_video || (have_video && have_audio && audio_pts < video_pts)) {
506             write_audio_frame(oc, &audio_st);
507         } else {
508             write_video_frame(oc, &video_st);
509         }
510     }
511
512     /* Write the trailer, if any. The trailer must be written before you
513      * close the CodecContexts open when you wrote the header; otherwise
514      * av_write_trailer() may try to use memory that was freed on
515      * av_codec_close(). */
516     av_write_trailer(oc);
517
518     /* Close each codec. */
519     if (have_video)
520         close_video(oc, &video_st);
521     if (have_audio)
522         close_audio(oc, &audio_st);
523
524     /* Free the streams. */
525     for (i = 0; i < oc->nb_streams; i++) {
526         av_freep(&oc->streams[i]->codec);
527         av_freep(&oc->streams[i]);
528     }
529
530     if (!(fmt->flags & AVFMT_NOFILE))
531         /* Close the output file. */
532         avio_close(oc->pb);
533
534     /* free the stream */
535     av_free(oc);
536
537     return 0;
538 }