ffmpeggrabber.cpp

来自「一个语言识别引擎」· C++ 代码 · 共 531 行 · 第 1/2 页

// -*- mode:C++; tab-width:4; c-basic-offset:4; indent-tabs-mode:nil -*-

/*
 * Copyright (C) 2006 Paul Fitzpatrick, Jonas Ruesch
 * CopyPolicy: Released under the terms of the GNU GPL v2.0.
 *
 */


#include <yarp/FfmpegGrabber.h>
#include <yarp/os/all.h>
#include <yarp/sig/all.h>
#include <yarp/Logger.h>

// based on Martim Boehme's example at:
//   http://www.inb.uni-luebeck.de/~boehme/using_libavcodec.html

#include <ffmpeg/avcodec.h>
#include <ffmpeg/avformat.h>

#include <stdio.h>

using namespace yarp;
using namespace yarp::os;
using namespace yarp::dev;
using namespace yarp::sig;
using namespace yarp::sig::file;

#define DBG if (0)


class DecoderState {
public:
    int      bytesRemaining;
    int      bytesDecoded;
    int      frameFinished;
    int      index;

    AVCodecContext  *pCodecCtx;
    AVCodec         *pCodec;

    // video buffers
    AVFrame         *pFrame; 
    AVFrame         *pFrameRGB;
    uint8_t         *buffer;
    int16_t         *audioBuffer;
    int16_t         *audioBufferAt;
    int audioBufferLen;

    DecoderState() {
        index = -1;
        pCodec = NULL;
        pCodecCtx = NULL;
        pFrame = NULL;
        pFrameRGB = NULL;
        buffer = NULL;
        audioBuffer = NULL;
        audioBufferAt = NULL;
        audioBufferLen = 0;
        frameFinished = 0;
    }

    bool isFinished() {
        return frameFinished!=0;
    }

    int getIndex() {
        return index;
    }

    virtual ~DecoderState() {
        if (pCodecCtx!=NULL) {
            avcodec_close(pCodecCtx);
        }
        if (audioBuffer!=NULL) {
            delete [] audioBuffer;
        }
        if (buffer!=NULL) {
            delete [] buffer;
        }
        if (pFrameRGB!=NULL) {
            av_free(pFrameRGB);
        }
        if (pFrame!=NULL) {
            av_free(pFrame);
        }
    }

    int getStream(AVFormatContext *pFormatCtx, CodecType code, 
                  const char *name) {
        // Find the first stream
        int videoStream=-1;
        for(int i=0; i<pFormatCtx->nb_streams; i++)
            if(pFormatCtx->streams[i]->codec->codec_type==code)
                {
                    videoStream=i;
                    printf("First %s stream is stream #%d\n", name, i);
                    break;
                }
        if(videoStream==-1) {
            printf("Could not find %s stream\n", name);
        }
        index = videoStream;

        return index;
    }

    bool getCodec(AVFormatContext *pFormatCtx) {
        // Get a pointer to the codec context for the video stream
        pCodecCtx=pFormatCtx->streams[index]->codec;

        // Find the decoder for the video stream
        pCodec=avcodec_find_decoder(pCodecCtx->codec_id);
        if(pCodec==NULL) {
            printf("Codec not found\n");
            return false; // Codec not found
        }
        
        // Open codec
        if(avcodec_open(pCodecCtx, pCodec)<0) {
            printf("Could not open codec\n");
            return false; // Could not open codec
        }

        return true;
    }


    bool allocateImage() {
        // Allocate video frame
        pFrame=avcodec_alloc_frame();
        
        // Allocate an AVFrame structure
        pFrameRGB=avcodec_alloc_frame();
        if(pFrameRGB==NULL) {
            printf("Could not allocate a frame\n");
            return false;
        }
        
        // Determine required buffer size and allocate buffer
        int numBytes=avpicture_get_size(PIX_FMT_RGB24, pCodecCtx->width,
                                        pCodecCtx->height);
        buffer=new uint8_t[numBytes];
        
        // Assign appropriate parts of buffer to image planes in pFrameRGB
        avpicture_fill((AVPicture *)pFrameRGB, buffer, PIX_FMT_RGB24,
                       pCodecCtx->width, pCodecCtx->height);
        return true;
    }

    bool allocateSound() {
        audioBuffer = new int16_t[AVCODEC_MAX_AUDIO_FRAME_SIZE];
        audioBufferAt = audioBuffer;
        audioBufferLen = 0;
        printf("channels %d, sample_rate %d, frame_size %d\n",
               pCodecCtx->channels,
               pCodecCtx->sample_rate,
               pCodecCtx->frame_size);
        return true;
    }

    int getWidth() {
        return pCodecCtx->width;
    }

    int getHeight() {
        return pCodecCtx->height;
    }


    int getRate() {
        return pCodecCtx->sample_rate;
    }

    int getChannels() {
        return pCodecCtx->channels;
    }

    bool getAudio(AVPacket& packet,Sound& sound) {
        int ct = 0;
        int bytesRead = 0;
        int bytesWritten = 0;
        while (bytesRead<packet.size) {
            int r = avcodec_decode_audio(pCodecCtx, 
                                         audioBuffer+bytesWritten, 
                                         &ct,
                                         packet.data+bytesRead, 
                                         packet.size-bytesRead);
            if (r<0) {
                printf("error decoding audio\n");
                return false;
            }
            DBG printf("audio bytes %d return %d\n", ct, r);
            int num_channels = getChannels();
            int num_rate = getRate();
            //audioBufferAt += ct;
            //audioBufferLen += ct;
            bytesRead += r;
            bytesWritten += ct;
            if (bytesRead==packet.size) {
                int num_samples = bytesWritten/(sizeof(int16_t)*num_channels);
                sound.resize(num_samples,num_channels);
                sound.setFrequency(num_rate);
                
                int idx = 0;
                for (int i=0; i<num_samples; i++) {
                    for (int j=0; j<num_channels; j++) {
                        sound.set(audioBuffer[idx],i,j);
                        idx++;
                    }
                }
                //audioBufferAt = audioBuffer;
                //audioBufferLen = 0;
            }
        }
        return true;
    }

    bool getVideo(AVPacket& packet) {
        // Decode video frame
        avcodec_decode_video(pCodecCtx, pFrame, &frameFinished, 
                             packet.data, packet.size);
        
        // Did we get a video frame?
        if(frameFinished) {
            // Convert the image from its native format to RGB
            img_convert((AVPicture *)pFrameRGB, PIX_FMT_RGB24, 
                        (AVPicture*)pFrame, pCodecCtx->pix_fmt, 
                        pCodecCtx->width, 
                        pCodecCtx->height);
        }
        return frameFinished;
    }


    bool getVideo(ImageOf<PixelRgb>& image) {
        if (frameFinished) {
            FlexImage flex;
            flex.setPixelCode(VOCAB_PIXEL_RGB);
            flex.setQuantum((pFrameRGB->linesize[0]));
            flex.setExternal(pFrameRGB->data[0],
                             pCodecCtx->width,
                             pCodecCtx->height);
            image.copy(flex); 
        }

        return frameFinished;
    }

    bool haveFrame() {
        return frameFinished;
    }
};

class FfmpegHelper {
public:
    DecoderState videoDecoder;
    DecoderState audioDecoder;
};


#define HELPER(x) (*((FfmpegHelper*)x))



bool FfmpegGrabber::openV4L(yarp::os::Searchable & config,