VideoSource_FFMPEGVideo.cpp

/* 
This file is part of the BIAS library (Basic ImageAlgorithmS).

Copyright (C) 2003-2009    (see file CONTACT for details)
  Multimediale Systeme der Informationsverarbeitung
  Institut fuer Informatik
  Christian-Albrechts-Universitaet Kiel


BIAS is free software; you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation; either version 2.1 of the License, or
(at your option) any later version.

BIAS 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 Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with BIAS; if not, write to the Free Software
Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
*/

#include "VideoSource_FFMPEGVideo.hh"

using namespace std;
using namespace BIAS;

BIAS::VideoSource_FFMPEGVideo::VideoSource_FFMPEGVideo()
  : _vidStreams(0),
    _streamIndex(0),
    _fmtCtx(NULL),
    _codecCtx(NULL),
    _swsCtx(NULL),
    _codec(NULL),
    _frame(NULL),
    _frameRGB(NULL),
    _buffer(NULL)
{
  static bool firstCall = true;
  if (firstCall)
  {
    av_register_all();
    firstCall = false;
  }
}
BIAS::VideoSource_FFMPEGVideo::~VideoSource_FFMPEGVideo()
{
}

bool BIAS::VideoSource_FFMPEGVideo::Open(const char *file)
{
  // Open file and retreive info
  if (av_open_input_file(&_fmtCtx, file, NULL, 0, NULL) != 0)
  {
    _errMsg = string("Could not open file ") + file;
    return false;
  }
  if (av_find_stream_info(_fmtCtx) < 0)
  {
    _errMsg = "Could not find stream info";
    return false;
  }

  // Find video streams
  for (unsigned int i = 0; i < _fmtCtx->nb_streams; i++)
  {
    if (_fmtCtx->streams[i]->codec->codec_type == CODEC_TYPE_VIDEO)
    {
      _vidStreams.push_back(i);
    }
  }
  if (CountVideoStreams() == 0)
  {
    _errMsg = "No video streams found";
    return false;
  }
  
  return true;
}

bool BIAS::VideoSource_FFMPEGVideo::Init(int width, int height, PixelFormat pixFmt,
                 unsigned int streamIndex)
{
  // Get codec context and decoder
  _streamIndex = streamIndex;
  if (_streamIndex >= CountVideoStreams())
    _streamIndex = 0;

  _codecCtx = _fmtCtx->streams[_streamIndex]->codec;
  _codec = avcodec_find_decoder(_codecCtx->codec_id);
  if (_codec == NULL)
  {
    _errMsg = "Codec not found";
    return false;
  }
  if (avcodec_open(_codecCtx, _codec) < 0)
  {
    _errMsg = "Could not open codec";
    return false;
  }

  // Get Sws context
  int w = _codecCtx->width;
  int h = _codecCtx->height;

  if (width == 0 && height == 0)
  {
    width = w;
    height = h;
  }
  // Resize but maintain aspect ratio
  else if (width != 0 && height == 0)
  {
    height = (float)width / w * h;
  }
  else if (width == 0 && height != 0)
  {
    width = (float)height / h * w;
  }

  _swsCtx = sws_getContext(w, h, _codecCtx->pix_fmt,
               width, height, pixFmt,
                           SWS_BICUBIC, NULL, NULL, NULL);
  if (_swsCtx == NULL)
  {
    _errMsg = "Cannot initialize the conversion context!";
    return false;
  }

  _frame = avcodec_alloc_frame();
  _frameRGB = avcodec_alloc_frame();
  int nBytes = avpicture_get_size(pixFmt, _codecCtx->width,
                                  _codecCtx->height);
  _buffer = (uint8_t*)av_malloc((nBytes + 8) * sizeof(uint8_t));
  avpicture_fill((AVPicture*)_frameRGB, _buffer, pixFmt, w, h);

  return true;
}

void BIAS::VideoSource_FFMPEGVideo::Close()
{
  if (_fmtCtx != NULL)
    av_close_input_file(_fmtCtx);
  if (_codecCtx != NULL)
    avcodec_close(_codecCtx);
  if (_frame != NULL)
    av_free(_frame);
  if (_frameRGB != NULL)
    av_free(_frameRGB);
  if (_buffer != NULL)
    av_free(_buffer);

  _fmtCtx = NULL;
  _codecCtx = NULL;
  _frame = NULL;
  _frameRGB = NULL;
  _buffer = NULL;
}


unsigned int BIAS::VideoSource_FFMPEGVideo::CountVideoStreams()
{
  return _vidStreams.size();
}


int BIAS::VideoSource_FFMPEGVideo::Width()
{
  if (_codecCtx != NULL)
    return _codecCtx->width;
  else
    return _fmtCtx->streams[_streamIndex]->codec->width;
}

int BIAS::VideoSource_FFMPEGVideo::Height()
{
  if (_codecCtx != NULL)
    return _codecCtx->height;
  else
    return _fmtCtx->streams[_streamIndex]->codec->height;
}


bool BIAS::VideoSource_FFMPEGVideo::GetFrame(BIAS::ImageBase &img)
{
  int frameFinished;
  AVPacket packet;
  bool ret = false;

  while (av_read_frame(_fmtCtx, &packet) >= 0 && !ret)
  {
    if (packet.stream_index == _vidStreams[_streamIndex])
    {
      avcodec_decode_video2(_codecCtx, _frame, &frameFinished,&packet);      
      if (frameFinished)
      {
    void *data = img.GetImageData();
    sws_scale(_swsCtx, _frame->data, _frame->linesize, 0, _codecCtx->height,
          (uint8_t**)&data, _frameRGB->linesize);
    ret = true;
      }
    }
  }

  av_free_packet(&packet);
  return ret;
}

void BIAS::VideoSource_FFMPEGVideo::Seek(int64_t time)
{
  av_seek_frame(_fmtCtx, _vidStreams[_streamIndex], time, 0);
}

AVRational BIAS::VideoSource_FFMPEGVideo::TimeBase()
{
  return _codecCtx->time_base;
}

unsigned int BIAS::VideoSource_FFMPEGVideo::GetCurrentDTS() const
{
  return _fmtCtx->streams[_vidStreams[_streamIndex]]->cur_dts;
}

unsigned int BIAS::VideoSource_FFMPEGVideo::GetLength() const
{
  AVStream *s = _fmtCtx->streams[_vidStreams[_streamIndex]];
  return s->duration;
}


string BIAS::VideoSource_FFMPEGVideo::ErrMsg()
{
  return _errMsg;
}


PixelFormat BIAS::VideoSource_FFMPEGVideo::Colormodel2PixFmt(ImageBase::EColorModel cm)
{
  switch (cm)
  {
  case ImageBase::CM_RGB:
    return PIX_FMT_RGB24;
  case ImageBase::CM_BGR:
    return PIX_FMT_BGR24;
  case ImageBase::CM_Grey:
    return PIX_FMT_GRAY8;
  default:
    return PIX_FMT_NONE;
  }
}