libpipi/pipi/sequence.c

/*
 *  libpipi       Pathetic image processing interface library
 *  Copyright (c) 2004-2009 Sam Hocevar <sam@hocevar.net>
 *                All Rights Reserved
 *
 *  $Id$
 *
 *  This library is free software. It comes without any warranty, to
 *  the extent permitted by applicable law. You can redistribute it
 *  and/or modify it under the terms of the Do What The Fuck You Want
 *  To Public License, Version 2, as published by Sam Hocevar. See
 *  http://sam.zoy.org/wtfpl/COPYING for more details.
 */

/*
 * codec.c: image I/O functions
 */

#include "config.h"

#if defined _WIN32
#   undef _CRT_SECURE_NO_WARNINGS
#   define _CRT_SECURE_NO_WARNINGS /* I know how to use snprintf, thank you */
#   define snprintf _snprintf
#endif

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#if defined USE_FFMPEG
#   include <libavformat/avformat.h>
#   include <libswscale/swscale.h>
#endif

#include "pipi.h"
#include "pipi_internals.h"

#if defined USE_FFMPEG
typedef struct
{
    uint8_t *buf;
    size_t buf_len;

    AVFormatContext *fmt_ctx;
    AVStream *stream;
    AVCodecContext *cod_ctx;
    AVCodec *codec;
    AVFrame *frame;

    struct SwsContext *sws_ctx;
    int src_width, src_height;
}
ffmpeg_codec_t;
#endif

#define PAR_NUM 1
#define PAR_DEN 1
#define BITRATE (16 * 1024 * 1024)

pipi_sequence_t *pipi_open_sequence(char *file, int width, int height, int fps)
{
#if defined USE_FFMPEG
    static int initialised = 0;

    pipi_sequence_t *seq;
    ffmpeg_codec_t *ff;
    uint8_t *tmp;

    seq = malloc(sizeof(pipi_sequence_t));
    seq->w = width;
    seq->h = height;
    seq->fps = fps;

    ff = malloc(sizeof(ffmpeg_codec_t));
    memset(ff, 0, sizeof(*ff));

    seq->codec_priv = ff;

    if (!initialised)
    {
        av_register_all();
        initialised = 1;
    }

    ff->fmt_ctx = avformat_alloc_context();
    if (!ff->fmt_ctx)
        goto error;

    /* Careful here: the Win32 snprintf doesn't seem to add a trailing
     * zero to the truncated output. */
    snprintf(ff->fmt_ctx->filename, sizeof(ff->fmt_ctx->filename),
             file);
    ff->fmt_ctx->filename[sizeof(ff->fmt_ctx->filename) - 1] = '\0';

    ff->fmt_ctx->oformat = guess_format(NULL, file, NULL);
    if (!ff->fmt_ctx->oformat)
        ff->fmt_ctx->oformat = guess_format("mpeg", NULL, NULL);
    if (!ff->fmt_ctx->oformat)
        goto error;

    ff->stream = av_new_stream(ff->fmt_ctx, 0);
    if (!ff->stream)
        goto error;

    ff->stream->sample_aspect_ratio.num = PAR_NUM;
    ff->stream->sample_aspect_ratio.den = PAR_DEN;

    ff->cod_ctx = ff->stream->codec;

    ff->cod_ctx->width = width;
    ff->cod_ctx->height = height;
    ff->cod_ctx->sample_aspect_ratio.num = PAR_NUM;
    ff->cod_ctx->sample_aspect_ratio.den = PAR_DEN;
    ff->cod_ctx->codec_id = ff->fmt_ctx->oformat->video_codec;
    ff->cod_ctx->codec_type = CODEC_TYPE_VIDEO;
    ff->cod_ctx->bit_rate = BITRATE;
    ff->cod_ctx->time_base.num = 1;
    ff->cod_ctx->time_base.den = fps;
    ff->cod_ctx->gop_size = fps * 3 / 4; /* empirical */
    ff->cod_ctx->pix_fmt = PIX_FMT_YUV420P; /* send YUV 420 */
    if (ff->cod_ctx->codec_id == CODEC_ID_MPEG2VIDEO)
        ff->cod_ctx->max_b_frames = 2;
    if (ff->cod_ctx->codec_id == CODEC_ID_MPEG1VIDEO)
        ff->cod_ctx->mb_decision = 2;
    if (ff->fmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
        ff->cod_ctx->flags |= CODEC_FLAG_GLOBAL_HEADER;

    if (av_set_parameters(ff->fmt_ctx, NULL) < 0)
        goto error;

    ff->codec = avcodec_find_encoder(ff->cod_ctx->codec_id);
    if (!ff->codec)
        goto error;
    if (avcodec_open(ff->cod_ctx, ff->codec) < 0)
        goto error;

    ff->frame = avcodec_alloc_frame();
    if (!ff->frame)
        goto error;
    tmp = (uint8_t *)av_malloc(avpicture_get_size(ff->cod_ctx->pix_fmt,
                                                  ff->cod_ctx->width,
                                                  ff->cod_ctx->height));
    if (!tmp)
        goto error;
    avpicture_fill((AVPicture *)ff->frame, tmp, ff->cod_ctx->pix_fmt,
                   ff->cod_ctx->width, ff->cod_ctx->height);

    if (!(ff->fmt_ctx->flags & AVFMT_NOFILE))
        if (url_fopen(&ff->fmt_ctx->pb, file, URL_WRONLY) < 0)
            goto error;

    ff->buf_len = 64 * 1024 * 1024;
    ff->buf = (uint8_t *)av_malloc(ff->buf_len);

    av_write_header(ff->fmt_ctx);

    return seq;

error:
    pipi_close_sequence(seq);
    return NULL;

#else
    return NULL;

#endif
}

int pipi_feed_sequence(pipi_sequence_t *seq, uint8_t *buffer, int width, int height)
{
#if defined USE_FFMPEG
    AVPacket packet;
    size_t bytes;
    int pitch;

    ffmpeg_codec_t *ff = (ffmpeg_codec_t *)seq->codec_priv;

    if (ff->src_width != width || ff->src_height != height)
    {
        ff->src_width = width;
        ff->src_height = height;
        if (ff->sws_ctx)
            sws_freeContext(ff->sws_ctx);
        ff->sws_ctx = NULL;
    }

    if (!ff->sws_ctx)
        ff->sws_ctx = sws_getContext(width, height, PIX_FMT_RGB32,
                                       ff->cod_ctx->width,
                                       ff->cod_ctx->height,
                                       ff->cod_ctx->pix_fmt, SWS_BICUBIC,
                                       NULL, NULL, NULL);
    if (!ff->sws_ctx)
        return -1;

    pitch = width * 4;
    sws_scale(ff->sws_ctx, &buffer, &pitch, 0, height,
              ff->frame->data, ff->frame->linesize);

    bytes = avcodec_encode_video(ff->cod_ctx, ff->buf,
                                 ff->buf_len, ff->frame);
    if (bytes <= 0)
        return 0;

    av_init_packet(&packet);
    if (ff->cod_ctx->coded_frame->pts != 0x8000000000000000LL)
        packet.pts = av_rescale_q(ff->cod_ctx->coded_frame->pts,
                                  ff->cod_ctx->time_base, ff->stream->time_base);
    if (ff->cod_ctx->coded_frame->key_frame)
        packet.flags |= PKT_FLAG_KEY;
    packet.stream_index = 0;
    packet.data = ff->buf;
    packet.size = bytes;

    if (av_interleaved_write_frame(ff->fmt_ctx, &packet) < 0)
        return -1;
#endif

    return 0;
}

int pipi_close_sequence(pipi_sequence_t *seq)
{
#if defined USE_FFMPEG
    ffmpeg_codec_t *ff = (ffmpeg_codec_t *)seq->codec_priv;

    if (ff->fmt_ctx)
    {
        av_write_trailer(ff->fmt_ctx);
    }

    if (ff->buf)
    {
        av_free(ff->buf);
        ff->buf = NULL;
    }

    if (ff->cod_ctx)
    {
        avcodec_close(ff->cod_ctx);
        ff->cod_ctx = NULL;
    }

    if (ff->frame)
    {
        av_free(ff->frame->data[0]);
        av_free(ff->frame);
        ff->frame = NULL;
    }

    if (ff->fmt_ctx)
    {
        av_freep(&ff->fmt_ctx->streams[0]->codec);
        ff->codec = NULL;

        av_freep(&ff->fmt_ctx->streams[0]);
        ff->stream = NULL;

        if (!(ff->fmt_ctx->flags & AVFMT_NOFILE))
            url_fclose(ff->fmt_ctx->pb);

        av_free(ff->fmt_ctx);
        ff->fmt_ctx = NULL;
    }

    if (ff->sws_ctx)
    {
        sws_freeContext(ff->sws_ctx);
        ff->sws_ctx = NULL;
        ff->src_width = ff->src_height = 0;
    }

    free(ff);
    free(seq);
#endif

    return 0;
}