root/vic/branches/mpeg4/codec/encoder-h264.cpp @ 3986

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <errno.h>
#include <assert.h>
#include <iostream>

#include "inet.h"
#include "net.h"
#include "rtp.h"
#include "vic_tcl.h"
#include "transmitter.h"
#include "pktbuf-rtp.h"
#include "module.h"

#include "databuffer.h"
#include "x264encoder.h"
#include "deinterlace.h"
/*extern "C"
{
#include "base64.h"
}*/

static Transmitter *tx;
static RTP_BufferPool *pool;
static u_int32_t ts;
static unsigned char frame_seq;

class H264Encoder:public TransmitterModule
{
  public:
    H264Encoder();
    ~H264Encoder();

    void setq(int q);
    int command(int argc, const char *const *argv);

    void size(int w, int h);
    int consume(const VideoFrame *);

  protected:
    int fps, kbps, gop;
    bool state;
    //UCHAR* bitstream;

    x264Encoder *enc;
    DataBuffer *fOut;
    Deinterlace deinterlacer;

    FILE *fptr;
};

static class H264EncoderMatcher:public Matcher
{
  public:
    H264EncoderMatcher():Matcher("module")
    {
    }
    TclObject *match(const char *fmt)
    {
        if (strcasecmp(fmt, "h264") == 0)
            return (new H264Encoder);
        return (0);
    }
}

encoder_matcher_h264;

H264Encoder::H264Encoder():TransmitterModule(FT_YUV_CIF)
{
    enc = new x264Encoder();
    state = false;
    fptr = NULL;
    frame_seq = 0;
    fps = 20;
    kbps = 512;
    gop = 20;
    fOut=NULL;
}

H264Encoder::~H264Encoder()
{
    delete enc;
    if (fOut) delete fOut;
}

void H264Encoder::size(int w, int h)
{
    debug_msg("H264: WxH %dx%d\n", w, h);
    Module::size(w, h);
    fOut = new DataBuffer(w * h * 3 >> 2);
}

int H264Encoder::command(int argc, const char *const *argv)
{
    if (argc == 3) {
        if (strcmp(argv[1], "q") == 0) {
            gop = atoi(argv[2]);
            return (TCL_OK);
        }
        else if (strcmp(argv[1], "fps") == 0) {
            fps = atoi(argv[2]);
            //std::cout << "H264: fps " << fps << "\n";
            return (TCL_OK);
        }
        else if (strcmp(argv[1], "kbps") == 0) {
            kbps = atoi(argv[2]);
            if (kbps < 64)
                kbps = 64;
            //std::cout << "H264: kbps " << kbps << "\n";
            return (TCL_OK);
        }
        else if (strcmp(argv[1], "hq") == 0) {
            int enable_hq = atoi(argv[2]);
            return (TCL_OK);
        }
    }
    return (TransmitterModule::command(argc, argv));
    return 0;
}

int H264Encoder::consume(const VideoFrame * vf)
{
    pktbuf *pb;
    rtphdr *rh;
    ts = vf->ts_;
    tx = tx_;
    pool = pool_;

    int numNAL, i, sent_size = 0;
    int frame_size = 0;
    unsigned char f_seq = 0;
    //unsigned char f_total_pkt = 0;
    int RTP_HDR_LEN = sizeof(rtphdr);
    int NAL_FRAG_THRESH = tx->mtu() - RTP_HDR_LEN; /* payload max in one packet */
    //debug_msg( "MTU=%d, RTP_HDR_LEN=%d\n", NAL_FRAG_THRESH, RTP_HDR_LEN);


    tx->flush();

    if (!state) {
            state = true;
            size(vf->width_, vf->height_);
            debug_msg("init x264 encoder with kbps:%d, fps:%d", kbps, fps);
            enc->init(vf->width_, vf->height_, kbps, fps);
            enc->setGOP(gop);
            frame_size = vf->width_ * vf->height_;
    }

    frame_size = vf->width_ * vf->height_;
    deinterlacer.render(vf->bp_, vf->width_, vf->height_);
    enc->encodeFrame(vf->bp_);
    numNAL = enc->numNAL();

    //Send out numNAL packets framed according to RFC3984
    for (i = 0; i < numNAL; i++) {

        bool firstFragment = true;
        int FU_HDR_LEN = 1;
        int offset = 0;

        enc->getNALPacket(i, fOut);
        int nalSize1 = fOut->getDataSize();
        int nalSize = nalSize1-5;
        char *data1 = fOut->getData();
        char *data = fOut->getData();
        uint8_t NALhdr = data1[4]; //SV-XXX why does our x.264 provide 4-byte StartSync in the NALU?
        uint8_t NALtype = NALhdr & 0x1f;
        //debug_msg( "Got NALhdr=0x%02x, NALtype=0x%02x from encoded frame.\n", NALhdr, NALtype);
        memcpy(data, &data1[5], nalSize);

        sent_size += nalSize;


        while (nalSize > 0) {
            pb = pool->alloc(vf->ts_, RTP_PT_H264);
            rh = (rtphdr *) pb->data;

            if (nalSize <= NAL_FRAG_THRESH) {
                //==============================================
                //Single NAL or last fragment of FU-A
                //==============================================

                rh->rh_flags |= htons(RTP_M);   // set M bit
                pb->len = nalSize + RTP_HDR_LEN + FU_HDR_LEN;

                debug_msg( "NAL : ");

                if (FU_HDR_LEN==2) {
                        //==============================================
                        //Last fragment of FU-A
                        //==============================================
                        pb->data[12] = 0x00 | (NALhdr & 0x60) | 28;     //FU indicator
                        pb->data[13] = 0x40  | NALtype;                 //FU header
                        
                        debug_msg( "FU_Indicator=0x%02x, FU_Header=0x%02x, ", pb->data[12], pb->data[13]);
                } 
                else {
                        pb->data[12] = NALhdr;                          //NAL Header
                        debug_msg( "-----------------, --------------, ");
                }

                memcpy(&pb->data[RTP_HDR_LEN + FU_HDR_LEN], data + offset, nalSize);

                debug_msg( "i=%d/%d, nalSize=%4d len=%4d firstFrag=%d offset=%4d\n", i, numNAL, nalSize, pb->len, firstFragment, offset);

                nalSize = 0;
                offset = 0;

            } else {
                //==============================================
                //FU-A (not the last fragment though)
                //==============================================

                FU_HDR_LEN = 2;
                pb->len = (NAL_FRAG_THRESH - FU_HDR_LEN) + RTP_HDR_LEN + FU_HDR_LEN;

                pb->data[12] = 0x00 | (NALhdr & 0x60) | 28;                     //FU indicator
                pb->data[13] = ( (firstFragment) ? 0x80 : 0x00 ) | NALtype;     //FU header

                memcpy(&pb->data[RTP_HDR_LEN + FU_HDR_LEN], data + offset, NAL_FRAG_THRESH - FU_HDR_LEN);

                //debug_msg( "FU-A: FU_Indicator=0x%02x, FU_Header=0x%02x, i=%d/%d, nalSize=%4d len=%4d firstFrag=%d offset=%4d\n",  pb->data[12], pb->data[13], i, numNAL, nalSize, pb->len, firstFragment, offset);

                nalSize -= (NAL_FRAG_THRESH-FU_HDR_LEN);
                offset += (NAL_FRAG_THRESH-FU_HDR_LEN);
                firstFragment = false;
            }

            tx->send(pb);
            f_seq++;
        }

    }

    frame_seq++;


    return (kbps*1024) / (fps*8);
}