avc2avi.c

H.264 source codes

/*****************************************************************************
 * avc2avi.c: raw h264 -> AVI
 *****************************************************************************
 * Copyright (C) 2004 Laurent Aimar
 * $Id: avc2avi.c,v 1.1 2004/06/03 19:27:08 fenrir Exp $
 *
 * Authors: Laurent Aimar <[e-mail]>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111, USA.
 *****************************************************************************/

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

#include <signal.h>
#define _GNU_SOURCE
#include <getopt.h>

#ifdef _MSC_VER
#include <io.h>     /* _setmode() */
#include <fcntl.h>  /* _O_BINARY */
#endif

#include "../common/bs.h"

#define DATA_MAX 3000000
uint8_t data[DATA_MAX];

/* Ctrl-C handler */
static int     i_ctrl_c = 0;
static void    SigIntHandler( int a )
{
    i_ctrl_c = 1;
}

typedef struct
{
    char *psz_fin;
    char *psz_fout;

    float f_fps;
    char  fcc[4];
} cfg_t;

typedef struct
{
    int i_data;
    int i_data_max;
    uint8_t *p_data;
} vbuf_t;

void vbuf_init( vbuf_t * );
void vbuf_add( vbuf_t *, int i_data, void *p_data );
void vbuf_reset( vbuf_t * );

typedef struct
{
    FILE *f;

    float f_fps;
    char  fcc[4];

    int   i_width;
    int   i_height;

    int64_t i_movi;
    int64_t i_movi_end;
    int64_t i_riff;

    int      i_frame;
    int      i_idx_max;
    uint32_t *idx;
} avi_t;

void avi_init( avi_t *, FILE *, float, char fcc[4] );
void avi_write( avi_t *, vbuf_t *, int  );
void avi_end( avi_t * );

enum nal_unit_type_e
{
    NAL_UNKNOWN = 0,
    NAL_SLICE   = 1,
    NAL_SLICE_DPA   = 2,
    NAL_SLICE_DPB   = 3,
    NAL_SLICE_DPC   = 4,
    NAL_SLICE_IDR   = 5,    /* ref_idc != 0 */
    NAL_SEI         = 6,    /* ref_idc == 0 */
    NAL_SPS         = 7,
    NAL_PPS         = 8
    /* ref_idc == 0 for 6,9,10,11,12 */
};
enum nal_priority_e
{
    NAL_PRIORITY_DISPOSABLE = 0,
    NAL_PRIORITY_LOW        = 1,
    NAL_PRIORITY_HIGH       = 2,
    NAL_PRIORITY_HIGHEST    = 3,
};

typedef struct
{
    int i_ref_idc;  /* nal_priority_e */
    int i_type;     /* nal_unit_type_e */

    /* This data are raw payload */
    int     i_payload;
    uint8_t *p_payload;
} nal_t;

typedef struct
{
    int i_width;
    int i_height;

    int i_nal_type;
    int i_ref_idc;
    int i_idr_pic_id;
    int i_frame_num;
    int i_poc;

    int b_key;
    int i_log2_max_frame_num;
    int i_poc_type;
    int i_log2_max_poc_lsb;
} h264_t;

void h264_parser_init( h264_t * );
void h264_parser_parse( h264_t *h, nal_t *n, int *pb_nal_start );


static int nal_decode( nal_t *nal, void *p_data, int i_data );

static void Help( void );
static int  Parse( int argc, char **argv, cfg_t * );
static int  ParseNAL( nal_t *nal, avi_t *a, h264_t *h, int *pb_slice );

/****************************************************************************
 * main:
 ****************************************************************************/
int main( int argc, char **argv )
{
    cfg_t cfg;

    FILE    *fout;
    FILE    *fin;

    vbuf_t  vb;
    avi_t   avi;
    h264_t  h264;

    nal_t nal;
    int i_frame;
    int i_data;
    int b_eof;
    int b_key;
    int b_slice;

#ifdef _MSC_VER
    _setmode(_fileno(stdin), _O_BINARY);    /* thanks to Marcos Morais <morais at dee.ufcg.edu.br> */
    _setmode(_fileno(stdout), _O_BINARY);
#endif

    /* Parse command line */
    if( Parse( argc, argv, &cfg ) < 0 )
    {
        return -1;
    }

    /* Open input */
    if( cfg.psz_fin == NULL || *cfg.psz_fin == '\0' || !strcmp( cfg.psz_fin, "-" ) )
        fin = stdin;
    else
        fin = fopen( cfg.psz_fin, "rb" );
    if( fin == NULL )
    {
        fprintf( stderr, "cannot open input file\n" );
        return -1;
    }

    /* Open output */
    if( cfg.psz_fout == NULL || *cfg.psz_fout == '\0' || !strcmp( cfg.psz_fout, "-" ) )
        fout = stdin;
    else
        fout = fopen( cfg.psz_fout, "wb" );
    if( fout == NULL )
    {
        fprintf( stderr, "cannot open output file\n" );
        return -1;
    }

    /* Init avi */
    avi_init( &avi, fout, cfg.f_fps, cfg.fcc );

    /* Init parser */
    h264_parser_init( &h264 );

    /* Control-C handler */
    signal( SIGINT, SigIntHandler );

    /* Init data */
    b_eof = 0;
    b_key = 0;
    b_slice = 0;
    i_frame = 0;
    i_data  = 0;

    /* Alloc space for a nal, used for decoding pps/sps/slice header */
    nal.p_payload = malloc( DATA_MAX );

    vbuf_init( &vb );

    /* split frame */
    while( !i_ctrl_c )
    {
        uint8_t *p, *p_next, *end;
        int i_size;

        /* fill buffer */
        if( i_data < DATA_MAX && !b_eof )
        {
            int i_read = fread( &data[i_data], 1, DATA_MAX - i_data, fin );
            if( i_read <= 0 )
                b_eof = 1;
            else
                i_data += i_read;
        }
        if( i_data < 3 )
            break;

        end = &data[i_data];

        /* Search begin of a NAL */
        p = &data[0];
        while( p < end - 3 )
        {
            if( p[0] == 0x00 && p[1] == 0x00 && p[2] == 0x01 )
            {
                break;
            }
            p++;
        }

        if( p >= end - 3 )
        {
            fprintf( stderr, "garbage (i_data = %d)\n", i_data );
            i_data = 0;
            continue;
        }

        /* Search end of NAL */
        p_next = p + 3;
        while( p_next < end - 3 )
        {
            if( p_next[0] == 0x00 && p_next[1] == 0x00 && p_next[2] == 0x01 )
            {
                break;
            }
            p_next++;
        }

        if( p_next == end - 3 && i_data < DATA_MAX )
            p_next = end;

        /* Compute NAL size */
        i_size = p_next - p - 3;
        if( i_size <= 0 )
        {
            if( b_eof )
                break;

            fprintf( stderr, "nal too large (FIXME) ?\n" );
            i_data = 0;
            continue;
        }

        /* Nal start at p+3 with i_size length */
        nal_decode( &nal, p +3, i_size < 2048 ? i_size : 2048 );

        b_key = h264.b_key;

        if( b_slice && vb.i_data && ( nal.i_type == NAL_SPS || nal.i_type == NAL_PPS ) )
        {
            avi_write( &avi, &vb, b_key );
            vbuf_reset( &vb );
            b_slice = 0;
        }

        /* Parse SPS/PPS/Slice */
        if( ParseNAL( &nal, &avi, &h264, &b_slice ) && vb.i_data > 0 )
        {
            avi_write( &avi, &vb, b_key );
            vbuf_reset( &vb );
        }

        /* fprintf( stderr, "nal:%d ref:%d\n", nal.i_type, nal.i_ref_idc ); */

        /* Append NAL to buffer */
        vbuf_add( &vb, i_size + 3, p );

        /* Remove this nal */
        memmove( &data[0], p_next, end - p_next );
        i_data -= p_next - &data[0];
    }

    if( vb.i_data > 0 )
    {
        avi_write( &avi, &vb, h264.b_key );
    }

    avi.i_width  = h264.i_width;
    avi.i_height = h264.i_height;

    avi_end( &avi );

    /* free mem */
    free( nal.p_payload );

    fclose( fin );
    fclose( fout );

    return 0;
}

/*****************************************************************************
 * Help:
 *****************************************************************************/
static void Help( void )
{
    fprintf( stderr,
             "avc2avi\n"
             "Syntax: avc2avi [options] [ -i input.h264 ] [ -o output.avi ]\n"
             "\n"
             "  -h, --help                  Print this help\n"
             "\n"
             "  -i, --input                 Specify input file (default: stdin)\n"
             "  -o, --output                Specify output file (default: stdout)\n"
             "\n"
             "  -f, --fps <float>           Set FPS (default: 25.0)\n"
             "  -c, --codec <string>        Set the codec fourcc (default: 'h264')\n"
             "\n" );
}

/*****************************************************************************
 * Parse:
 *****************************************************************************/
static int  Parse( int argc, char **argv, cfg_t *cfg )
{
    /* Set default values */
    cfg->psz_fin = NULL;
    cfg->psz_fout = NULL;
    cfg->f_fps = 25.0;
    memcpy( cfg->fcc, "h264", 4 );

    /* Parse command line options */
    opterr = 0; // no error message
    for( ;; )
    {
        int long_options_index;
        static struct option long_options[] =
        {
            { "help",   no_argument,       NULL, 'h' },
            { "input",  required_argument, NULL, 'i' },
            { "output", required_argument, NULL, 'o' },
            { "fps",    required_argument, NULL, 'f' },
            { "codec",  required_argument, NULL, 'c' },
            {0, 0, 0, 0}
        };

        int c;

        c = getopt_long( argc, argv, "hi:o:f:c:",
                         long_options, &long_options_index);

        if( c == -1 )
        {
            break;
        }

        switch( c )
        {
            case 'h':
                Help();
                return -1;

            case 0:
                break;
            case 'i':
                cfg->psz_fin = strdup( optarg );
                break;
            case 'o':
                cfg->psz_fout = strdup( optarg );
                break;
            case 'f':
                cfg->f_fps = atof( optarg );
                break;
            case 'c':
                memset( cfg->fcc, ' ', 4 );
                memcpy( cfg->fcc, optarg, strlen( optarg ) < 4 ? strlen( optarg ) : 4 );
                break;

            default:
                fprintf( stderr, "unknown option (%c)\n", optopt );