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