VideoSource_FFmpeg.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_FFmpeg.hh"
#include <sstream>

using namespace BIAS;
using namespace std;

VideoSource_FFmpeg::VideoSource_FFmpeg() : VideoSource()
{
  av_register_all();
  Init_();
}

VideoSource_FFmpeg::~VideoSource_FFmpeg()
{
  CloseDevice();
}

void VideoSource_FFmpeg::Init_()
{
  pFormatCtx_    = NULL;
  streamIndex_   = -1;
  pFrame_        = NULL;
  pOutputFrame_  = NULL;
  pBuffer_       = NULL;
  outputPixFmt_  = PIX_FMT_RGB24;
  pSwsCtx_       = NULL;
  numFrames_     = 0;
  frameNumber_   = 0;
  timeBaseNum_   = 1;
  timeBaseDen_   = 1;
  timestamp_     = -1;
  dateStr_       = NULL;
  duration_      = -1;
  errMsg_        = "";
  pFormatCtx_    = NULL;
  errMsg_        = "";
  frameNumber_   = 0;
  pts_           = 0;

  ColorMode_     = ImageBase::CM_RGB;
  BytesPerPixel_ = 1.0;
  ColorChannels_ = 3;
  DeviceChannel_ = streamIndex_;
}

int VideoSource_FFmpeg::OpenDevice(const char *filename)
{
  if (CompleteInitialized_) {
    errMsg_ = "VideoSource_FFmpeg already initialized. Call CloseDevice() first.";
    return -1;
  }

  AVFormatContext *ic;
  AVStream *st;
  AVCodecContext *c;
  AVCodec *codec;

  // Open video file
  if (av_open_input_file(&ic, filename, NULL, 0, NULL) < 0) {
    errMsg_ = "Could not open video file";
    return -1;
  }

  // Retrieve stream information
  if (av_find_stream_info(ic) < 0) {
    errMsg_ = "Could not find stream information";
    return -1;
  }

  // Dump format information
  dump_format(ic, 0, filename, false);

  // Select video stream
  streamIndex_ = -1;
  if (DeviceChannel_ >= 0) {
    // Select stream specified by user
    if (ic->streams[DeviceChannel_]->codec->codec_type == CODEC_TYPE_VIDEO)
      streamIndex_ = DeviceChannel_;
    else
      cout << "Video stream not found. Falling back to first video stream" << endl;
  }
  if (streamIndex_ == -1) {
    // Select first video stream
    for (unsigned int i = 0; i < ic->nb_streams; i++) {
      if (ic->streams[i]->codec->codec_type == CODEC_TYPE_VIDEO) {
        streamIndex_ = DeviceChannel_ = i;
        break;
      }
    }
  }
  if (streamIndex_ == -1) {
    errMsg_ = "Could not find a video stream";
    return -1;
  }

  st = ic->streams[streamIndex_];
  c = st->codec;

  // Find decoder
  codec = avcodec_find_decoder(c->codec_id);
  if (!codec) {
    errMsg_ = "Could not find decoder";
    return -1;
  }

  // Open decoder
  if (avcodec_open(c, codec) < 0) {
    errMsg_ = "Could not open decoder";
    return -1;
  }

  // Allocate video frame
  pFrame_ = avcodec_alloc_frame();
  if (!pFrame_) {
    errMsg_ = "Could not allocate video frame";
    return -1;
  }

  // Allocate output frame
  pOutputFrame_ = avcodec_alloc_frame();
  if (!pOutputFrame_) {
    errMsg_ = "Could not allocate output frame";
    return -1;
  }

  // Get some video stream parameters
  Width_ = c->width;
  Height_ = c->height;
  FramesPerSecond_ = st->r_frame_rate.num / (float) st->r_frame_rate.den;
  timeBaseNum_ = st->time_base.num;
  timeBaseDen_ = st->time_base.den;
  timestamp_ = ic->timestamp;
  duration_ = ic->duration;
  numFrames_ = st->nb_frames;

  // Try to get date string from metadata or filename
  AVMetadataTag *mtag;
  mtag = av_metadata_get(ic->metadata, "date", NULL, 0);
  if (mtag)
    dateStr_ = mtag->value;
  mtag = av_metadata_get(st->metadata, "date", NULL, 0);
  if (mtag)
    dateStr_ = mtag->value;
  if (!dateStr_)
    dateStr_ = ExtractDateString_(filename);

  // Get video frame size
  ImageSize_ = avpicture_get_size(outputPixFmt_, Width_, Height_);
  if (ImageSize_ < 0) {
    errMsg_ = "Could not calculate video frame size";
    return -1;
  }

  // Allocate video frame buffer
  pBuffer_ = (uint8_t *) av_malloc(ImageSize_ * sizeof(uint8_t));
  if (!pBuffer_) {
    errMsg_ = "Could not allocate video frame buffer";
    return -1;
  }

  // Fill video frame buffer
  if (avpicture_fill((AVPicture *) pOutputFrame_, pBuffer_, outputPixFmt_, Width_, Height_) < 0) {
    errMsg_ = "Could not fill video frame buffer";
    return -1;
  }

  pFormatCtx_ = ic;
  CompleteInitialized_ = true;

  return 0;
}

int VideoSource_FFmpeg::CloseDevice()
{
  AVFormatContext *ic = pFormatCtx_;
  AVStream *st = NULL;
  AVCodecContext *c = NULL;

  if (ic) {
    st = ic->streams[streamIndex_];
    if (st) {
      c = st->codec;
    }
  }

  if (Active_)
    PostGrab();

  Grabbing_ = false;
  CompleteInitialized_ = false;

  // Free buffers and frames
  av_free(pBuffer_);
  av_free(pOutputFrame_);
  av_free(pFrame_);

  // Close the codec
  if (c)
    avcodec_close(c);

  // Close the video file
  if (ic)
    av_close_input_file(ic);

  Init_();

  return 0;
}

int VideoSource_FFmpeg::GrabSingle(Camera<unsigned char> &image)
{
    assert(!image.IsEmpty());

    if (!Active_)
        BEXCEPTION("VideoSource Object not active, use PreGrab()");

    if (Grabbing_)
        BEXCEPTION("VideoSource Object is already grabbing()");

    Grabbing_ = true;

  AVFormatContext *ic = pFormatCtx_;
  AVStream *st = ic->streams[streamIndex_];
  AVCodecContext *c = st->codec;
  AVPacket pkt;
  int ret, gotPicture;
  MetaData *m;
  ostringstream ascii;

  // Read a frame
  bool done = false;
  while (!done) {
    ret = av_read_frame(ic, &pkt);
    if (ret < 0) {
      errMsg_ = ret == EOF ? "End of video file reached" : "Could not read video frame";
      Grabbing_ = false;
      return ret;
    }

    if (pkt.stream_index == streamIndex_) {
      ret = avcodec_decode_video2(c, pFrame_, &gotPicture, &pkt);
      if (ret < 0) {
        errMsg_ = "Could not decode video frame";
        Grabbing_ = false;
        return -1;
      }

      if (gotPicture) {
        // Color conversion
        pSwsCtx_ = sws_getCachedContext(pSwsCtx_, c->width, c->height, c->pix_fmt, c->width, c->height, outputPixFmt_, SWS_BICUBIC, NULL, NULL, NULL);
        if (!pSwsCtx_) {
          errMsg_ = "Could not initialize a conversion context";
          Grabbing_ = false;
          return -1;
        }
        void *data = image.GetImageData();
        sws_scale(pSwsCtx_, pFrame_->data, pFrame_->linesize, 0, c->height, (uint8_t **) &data, pOutputFrame_->linesize);

        // Generate new uid of image
        image.InvalidateUID();
        image.SetUID(UUID::GenerateUUID());

        // Calculate pts in microseconds
        pts_ = pkt.pts < 0 ? pkt.dts : pkt.pts;
        pts_ = pts_ * 1000000 * st->time_base.num / st->time_base.den;

        // Set metadata of image
        m = image.GetMetaData();
        ascii.str("");
        ascii.clear();
        ascii << frameNumber_;
        m->Add("#[num]", ascii.str());
        ascii.str("");
        ascii.clear();
        ascii << pts_;
        m->Add("#[pts]", ascii.str());
        image.SetMetaData(*m);

        frameNumber_++;
        done = true;
      }
    }
  }

  Grabbing_ = false;

  return 0;
}

int VideoSource_FFmpeg::SeekFrame(int64_t frame)
{
  AVFormatContext *ic = pFormatCtx_;
  AVStream *st = ic->streams[streamIndex_];
  AVCodecContext *c = st->codec;

  /*if (avformat_seek_file(ic, streamIndex_, 0, frame, frame, 
AVSEEK_FLAG_FRAME) < 0) {
    errMsg_ ="Could not seek frame";
    return -1;
  }*/

  int64_t timestamp = frame * st->time_base.den / st->time_base.num / 
FramesPerSecond_;
  cout << timestamp << endl;
  if (av_seek_frame(ic, streamIndex_, timestamp, 0) < 0) {
    errMsg_ ="Could not seek frame";
    return -1;
  }
  avcodec_flush_buffers(c);

  return 0;
}

int VideoSource_FFmpeg::SeekMs(int ms)
{
  AVFormatContext *ic = pFormatCtx_;
  AVStream *st = ic->streams[streamIndex_];
  AVCodecContext *c = st->codec;
  int64_t ts = ms * st->time_base.den / st->time_base.num / 1000;

  /*if (avformat_seek_file(ic, streamIndex_, 0, ts, ts, 0) < 0) {
    errMsg_ = "Could not seek time";
    return -1;
  }*/
  if (av_seek_frame(ic, streamIndex_, ts, 0) < 0) {
    errMsg_ ="Could not seek frame";
    return -1;
  }
  avcodec_flush_buffers(c);

  return 0;
}

int VideoSource_FFmpeg::SetColorModel(ImageBase::EColorModel mode)
{
  switch (mode) {
  case ImageBase::CM_Grey:
    outputPixFmt_ = PIX_FMT_GRAY8;
    ColorMode_ = mode;
    ColorChannels_ = 1;
    break;
  case ImageBase::CM_RGB:
    outputPixFmt_ = PIX_FMT_RGB24;
    ColorMode_ = mode;
    ColorChannels_ = 3;
    break;
  case ImageBase::CM_BGR:
    outputPixFmt_ = PIX_FMT_BGR24;
    ColorMode_ = mode;
    ColorChannels_ = 3;
    break;
  case ImageBase::CM_RGBA:
    outputPixFmt_ = PIX_FMT_RGBA;
    ColorMode_ = mode;
    ColorChannels_ = 4;
    break;
  case ImageBase::CM_BGRA:
    outputPixFmt_ = PIX_FMT_BGRA;
    ColorMode_ = mode;
    ColorChannels_ = 4;
    break;
  case ImageBase::CM_YUYV422:
    outputPixFmt_ = PIX_FMT_YUYV422;
    ColorMode_ = mode;
    ColorChannels_ = 2;
    break;
  case ImageBase::CM_UYVY422:
    outputPixFmt_ = PIX_FMT_UYVY422;
    ColorMode_ = mode;
    ColorChannels_ = 2;
    break;
  case ImageBase::CM_YUV420P:
    outputPixFmt_ = PIX_FMT_YUV420P;
    ColorMode_ = mode;
    ColorChannels_ = 2;
    break;
  default:
    errMsg_ = "Unsupported color model";
    return -1;
    break;
  }

  return 0;
}

char * VideoSource_FFmpeg::ExtractDateString_(string str)
{
  bool done = false;
  int pos = 0;
  int size = str.size();
  char *dateStr;
  int dateStrSize = 0;

  dateStr = (char *) malloc(20);
  dateStr[0] = '\0';

  while (!done && pos < size) {
    const char c = str[pos++];
    switch (c) {
    case '0':
    case '1':
    case '2':
    case '3':
    case '4':
    case '5':
    case '6':
    case '7':
    case '8':
    case '9':
      {
        switch (dateStrSize) {
        case 4:
        case 7:
          dateStr[dateStrSize] = '-';
          dateStrSize++;
          break;
        case 10:
          dateStr[dateStrSize] = ' ';
          dateStrSize++;
          break;
        case 13:
        case 16:
          dateStr[dateStrSize] = ':';
          dateStrSize++;
          break;
        }
        dateStr[dateStrSize] = c;
        dateStrSize++;
        if (dateStrSize >= 19)
          done = true;
        break;
      }
    case '_':
      {
        if (dateStrSize == 10) {
          dateStr[dateStrSize] = ' ';
          dateStrSize++;
        }
        else {
          dateStr[0] = '\0';
          dateStrSize = 0;
        }
        break;
      }
    default:
      dateStr[0] = '\0';
      dateStrSize = 0;
      break;
    }
  }

  dateStr[dateStrSize] = '\0';
  if (!done || parse_date(dateStr, 0) < 0)
    return NULL;
  else
    return dateStr;
}