source: libpipi/trunk/pipi/sequence.c @ 4765

Last change on this file since 4765 was 4765, checked in by Sam Hocevar, 9 years ago

Allow to choose between RGB and YUV in pipi_open_sequence().

File size: 9.3 KB
Line 
1/*
2 *  libpipi       Pathetic image processing interface library
3 *  Copyright (c) 2004-2009 Sam Hocevar <sam@hocevar.net>
4 *                All Rights Reserved
5 *
6 *  $Id$
7 *
8 *  This library is free software. It comes without any warranty, to
9 *  the extent permitted by applicable law. You can redistribute it
10 *  and/or modify it under the terms of the Do What The Fuck You Want
11 *  To Public License, Version 2, as published by Sam Hocevar. See
12 *  http://sam.zoy.org/wtfpl/COPYING for more details.
13 */
14
15/*
16 * codec.c: image I/O functions
17 */
18
19#include "config.h"
20
21#if defined _WIN32
22#   undef _CRT_SECURE_NO_WARNINGS
23#   define _CRT_SECURE_NO_WARNINGS /* I know how to use snprintf, thank you */
24#   define snprintf _snprintf
25#endif
26
27#include <stdio.h>
28#include <stdlib.h>
29#include <string.h>
30
31#if defined USE_FFMPEG
32#   include <libavformat/avformat.h>
33#   include <libswscale/swscale.h>
34#endif
35
36#include "pipi.h"
37#include "pipi_internals.h"
38
39#if defined USE_FFMPEG
40typedef struct
41{
42    uint8_t *buf;
43    size_t buf_len;
44
45    AVFormatContext *fmt_ctx;
46    AVStream *stream;
47    AVCodecContext *cod_ctx;
48    AVCodec *codec;
49    AVFrame *frame;
50
51    struct SwsContext *sws_ctx;
52    int src_width, src_height, src_fmt;
53}
54ffmpeg_codec_t;
55#endif
56
57pipi_sequence_t *pipi_open_sequence(char const *file,
58                                    int width, int height, int rgb, int fps,
59                                    int par_num, int par_den, int bitrate)
60{
61#if defined USE_FFMPEG
62    static int initialised = 0;
63
64    pipi_sequence_t *seq;
65    ffmpeg_codec_t *ff;
66    uint8_t *tmp;
67
68    seq = malloc(sizeof(pipi_sequence_t));
69    seq->w = width;
70    seq->h = height;
71    seq->fps = fps;
72    seq->convert_buf = NULL;
73
74    ff = malloc(sizeof(ffmpeg_codec_t));
75    memset(ff, 0, sizeof(*ff));
76
77    seq->codec_priv = ff;
78
79    if (!initialised)
80    {
81        av_register_all();
82        initialised = 1;
83    }
84
85    ff->fmt_ctx = avformat_alloc_context();
86    if (!ff->fmt_ctx)
87        goto error;
88
89    /* Careful here: the Win32 snprintf doesn't seem to add a trailing
90     * zero to the truncated output. */
91    snprintf(ff->fmt_ctx->filename, sizeof(ff->fmt_ctx->filename),
92             file);
93    ff->fmt_ctx->filename[sizeof(ff->fmt_ctx->filename) - 1] = '\0';
94
95    ff->fmt_ctx->oformat = av_guess_format(NULL, file, NULL);
96    if (!ff->fmt_ctx->oformat)
97        ff->fmt_ctx->oformat = av_guess_format("mpeg", NULL, NULL);
98    if (!ff->fmt_ctx->oformat)
99        goto error;
100
101    ff->stream = av_new_stream(ff->fmt_ctx, 0);
102    if (!ff->stream)
103        goto error;
104
105    ff->stream->sample_aspect_ratio.num = par_num;
106    ff->stream->sample_aspect_ratio.den = par_den;
107
108    ff->cod_ctx = ff->stream->codec;
109
110    ff->cod_ctx->width = width;
111    ff->cod_ctx->height = height;
112    ff->cod_ctx->sample_aspect_ratio.num = par_num;
113    ff->cod_ctx->sample_aspect_ratio.den = par_den;
114    ff->cod_ctx->codec_id = ff->fmt_ctx->oformat->video_codec;
115    ff->cod_ctx->codec_type = CODEC_TYPE_VIDEO;
116    ff->cod_ctx->bit_rate = bitrate;
117    ff->cod_ctx->time_base.num = 1;
118    ff->cod_ctx->time_base.den = fps;
119
120    ff->cod_ctx->pix_fmt = PIX_FMT_YUV420P; /* send YUV 420 */
121    if (ff->cod_ctx->codec_id == CODEC_ID_MPEG2VIDEO)
122        ff->cod_ctx->max_b_frames = 2;
123    if (ff->cod_ctx->codec_id == CODEC_ID_MPEG1VIDEO)
124        ff->cod_ctx->mb_decision = 2;
125    if (ff->cod_ctx->codec_id == CODEC_ID_H264)
126    {
127        /* Import x264 slow presets */
128        ff->cod_ctx->coder_type = 1;
129        ff->cod_ctx->flags |= CODEC_FLAG_LOOP_FILTER;
130        ff->cod_ctx->me_cmp |= FF_CMP_CHROMA;
131        ff->cod_ctx->partitions |= X264_PART_I4X4 | X264_PART_I8X8
132                                 | X264_PART_P4X4 | X264_PART_P8X8;
133        ff->cod_ctx->me_method = ME_UMH;
134        ff->cod_ctx->me_subpel_quality = 8;
135        ff->cod_ctx->me_range = 16;
136        ff->cod_ctx->gop_size = 250;
137        ff->cod_ctx->keyint_min = 25;
138        ff->cod_ctx->scenechange_threshold = 40;
139        ff->cod_ctx->i_quant_factor = 0.71f;
140        ff->cod_ctx->b_frame_strategy = 2;
141        ff->cod_ctx->qcompress = 0.6f;
142        ff->cod_ctx->qmin = 10;
143        ff->cod_ctx->qmax = 51;
144        ff->cod_ctx->max_qdiff = 4;
145        ff->cod_ctx->max_b_frames = 3;
146        ff->cod_ctx->refs = 5;
147        ff->cod_ctx->directpred = 3;
148        ff->cod_ctx->trellis = 1;
149        ff->cod_ctx->flags2 |= CODEC_FLAG2_BPYRAMID | CODEC_FLAG2_MIXED_REFS
150                             | CODEC_FLAG2_WPRED | CODEC_FLAG2_8X8DCT
151                             | CODEC_FLAG2_FASTPSKIP;
152        ff->cod_ctx->weighted_p_pred = 2;
153        ff->cod_ctx->rc_lookahead = 50;
154    }
155    if (ff->fmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
156        ff->cod_ctx->flags |= CODEC_FLAG_GLOBAL_HEADER;
157
158    if (av_set_parameters(ff->fmt_ctx, NULL) < 0)
159        goto error;
160
161    ff->codec = avcodec_find_encoder(ff->cod_ctx->codec_id);
162    if (!ff->codec)
163        goto error;
164    if (avcodec_open(ff->cod_ctx, ff->codec) < 0)
165        goto error;
166
167    ff->frame = avcodec_alloc_frame();
168    if (!ff->frame)
169        goto error;
170    tmp = (uint8_t *)av_malloc(avpicture_get_size(ff->cod_ctx->pix_fmt,
171                                                  ff->cod_ctx->width,
172                                                  ff->cod_ctx->height));
173    if (!tmp)
174        goto error;
175    avpicture_fill((AVPicture *)ff->frame, tmp, ff->cod_ctx->pix_fmt,
176                   ff->cod_ctx->width, ff->cod_ctx->height);
177
178    if (!(ff->fmt_ctx->flags & AVFMT_NOFILE))
179        if (url_fopen(&ff->fmt_ctx->pb, file, URL_WRONLY) < 0)
180            goto error;
181
182    ff->buf_len = 64 * 1024 * 1024;
183    ff->buf = (uint8_t *)av_malloc(ff->buf_len);
184
185    ff->src_fmt = rgb ? PIX_FMT_RGB32 : PIX_FMT_YUV444P;
186
187    av_write_header(ff->fmt_ctx);
188
189    return seq;
190
191error:
192    pipi_close_sequence(seq);
193    return NULL;
194
195#else
196    return NULL;
197
198#endif
199}
200
201int pipi_feed_sequence(pipi_sequence_t *seq, uint8_t const *buffer,
202                       int width, int height)
203{
204#if defined USE_FFMPEG
205    AVPacket packet;
206    uint8_t const *buflist[3];
207    int pitchlist[3];
208    size_t bytes;
209    int n;
210
211    ffmpeg_codec_t *ff = (ffmpeg_codec_t *)seq->codec_priv;
212
213    if (ff->src_width != width || ff->src_height != height)
214    {
215        ff->src_width = width;
216        ff->src_height = height;
217        if (ff->sws_ctx)
218            sws_freeContext(ff->sws_ctx);
219        ff->sws_ctx = NULL;
220    }
221
222    if (!ff->sws_ctx)
223    {
224        ff->sws_ctx = sws_getContext(width, height, ff->src_fmt,
225                                     ff->cod_ctx->width,
226                                     ff->cod_ctx->height,
227                                     ff->cod_ctx->pix_fmt, SWS_BICUBIC,
228                                     NULL, NULL, NULL);
229        if (seq->convert_buf)
230        {
231           free(seq->convert_buf);
232           seq->convert_buf = NULL;
233        }
234    }
235    if (!ff->sws_ctx)
236        return -1;
237
238    /* Convert interleaved YUV to planar YUV */
239    if (ff->src_fmt == PIX_FMT_YUV444P)
240    {
241        if (!seq->convert_buf)
242            seq->convert_buf = malloc(width * height * 3);
243
244        for (n = 0; n < width * height; n++)
245        {
246            seq->convert_buf[n] = buffer[4 * n];
247            seq->convert_buf[n + width * height] = buffer[4 * n + 1];
248            seq->convert_buf[n + 2 * width * height] = buffer[4 * n + 2];
249        }
250
251        /* Feed the buffers to FFmpeg */
252        buflist[0] = seq->convert_buf;
253        buflist[1] = seq->convert_buf + 2 * width * height;
254        buflist[2] = seq->convert_buf + width * height;
255        pitchlist[0] = pitchlist[1] = pitchlist[2] = width;
256    }
257    else
258    {
259        buflist[0] = buffer;
260        pitchlist[0] = 4 * width;
261    }
262
263    sws_scale(ff->sws_ctx, buflist, pitchlist, 0, height,
264              ff->frame->data, ff->frame->linesize);
265
266    bytes = avcodec_encode_video(ff->cod_ctx, ff->buf,
267                                 ff->buf_len, ff->frame);
268    if (bytes <= 0)
269        return 0;
270
271    av_init_packet(&packet);
272    if (ff->cod_ctx->coded_frame->pts != 0x8000000000000000LL)
273        packet.pts = av_rescale_q(ff->cod_ctx->coded_frame->pts,
274                                  ff->cod_ctx->time_base, ff->stream->time_base);
275    if (ff->cod_ctx->coded_frame->key_frame)
276        packet.flags |= PKT_FLAG_KEY;
277    packet.stream_index = 0;
278    packet.data = ff->buf;
279    packet.size = bytes;
280
281    if (av_interleaved_write_frame(ff->fmt_ctx, &packet) < 0)
282        return -1;
283#endif
284
285    return 0;
286}
287
288int pipi_close_sequence(pipi_sequence_t *seq)
289{
290#if defined USE_FFMPEG
291    ffmpeg_codec_t *ff = (ffmpeg_codec_t *)seq->codec_priv;
292
293    /* Finish the sequence */
294    if (ff->buf)
295    {
296        av_write_trailer(ff->fmt_ctx);
297    }
298
299    /* Close everything */
300    if (ff->buf)
301    {
302        av_free(ff->buf);
303        ff->buf = NULL;
304    }
305
306    if (ff->cod_ctx)
307    {
308        avcodec_close(ff->cod_ctx);
309        ff->cod_ctx = NULL;
310    }
311
312    if (ff->frame)
313    {
314        av_free(ff->frame->data[0]);
315        av_free(ff->frame);
316        ff->frame = NULL;
317    }
318
319    if (ff->fmt_ctx)
320    {
321        av_freep(&ff->fmt_ctx->streams[0]->codec);
322        ff->codec = NULL;
323
324        av_freep(&ff->fmt_ctx->streams[0]);
325        ff->stream = NULL;
326
327        if (!(ff->fmt_ctx->flags & AVFMT_NOFILE) && ff->fmt_ctx->pb)
328            url_fclose(ff->fmt_ctx->pb);
329
330        av_free(ff->fmt_ctx);
331        ff->fmt_ctx = NULL;
332    }
333
334    if (ff->sws_ctx)
335    {
336        sws_freeContext(ff->sws_ctx);
337        ff->sws_ctx = NULL;
338        ff->src_width = ff->src_height = 0;
339    }
340
341    free(ff);
342    free(seq);
343#endif
344
345    return 0;
346}
Note: See TracBrowser for help on using the repository browser.