VideoSink_FFmpeg.cpp

#include "VideoSink_FFmpeg.hh"
#include <Base/Image/ImageIO.hh>
#include <Base/Image/ImageConvert.hh>
#include <Filter/Rescale.hh>
#include <time.h>

using namespace BIAS;
using namespace std;

VideoSink_FFmpeg::VideoSink_FFmpeg()
{
  av_register_all();
  Init_();
  SetDefaults();
}

VideoSink_FFmpeg::~VideoSink_FFmpeg()
{
  Close();
}

void VideoSink_FFmpeg::Init_() {
  pFormatCtx_  = NULL;
  pInputPict_  = NULL;
  pSwsCtx_     = NULL;
  pPict_       = NULL;
  pPictBuffer_ = NULL;
  pBuffer_     = NULL;
  bufferSize_  = 0;
  pts_         = 0;
  errMsg_      = "";
}

void VideoSink_FFmpeg::SetDefaults()
{
  codecId_     = CODEC_ID_MPEG1VIDEO;
  width_       = 352;
  height_      = 288;
  fps_.num     = 25;
  fps_.den     = 1;
  bitrate_     = 1000000;
  gopSize_     = 12;
  maxBFrames_  = 0;
  realtime_    = false;
  override_    = true;
}

void VideoSink_FFmpeg::SetEncoder(enum CodecID codecId)
{
  codecId_ = codecId;
}

void VideoSink_FFmpeg::SetSize(int width, int height)
{
  width_ = width;
  height_ = height;
}

void VideoSink_FFmpeg::SetFPS(float fps)
{
  //fps_ = av_d2q(fps, AV_TIME_BASE);
  char str[32];
  sprintf(str, "%f", fps);
  av_parse_video_frame_rate(&fps_, str);
}

void VideoSink_FFmpeg::SetBitrate(int bitrate)
{
  bitrate_ = bitrate;
}

void VideoSink_FFmpeg::SetGopSize(int gopSize)
{
  gopSize_ = gopSize;
}

void VideoSink_FFmpeg::SetMaxBFrames(int maxBFrames)
{
  maxBFrames_ = maxBFrames;
}

void VideoSink_FFmpeg::SetRealtime(bool realtime)
{
  realtime_ = realtime;
}

void VideoSink_FFmpeg::SetOverride(bool override)
{
  override_ = override;
}

int VideoSink_FFmpeg::Open(string filename)
{
  AVFormatContext *oc;
  AVOutputFormat *oformat;
  AVStream *st;
  AVCodecContext *c;
  AVCodec *codec;

  // Close video file if already opened
  Close();

  // Check parameters

  // Check if output file already exists
  if (!override_ && url_exist(oc->filename)) {
    errMsg_ = "Output file already exists.";
    return -1;
  }

  // Get output format from filename extension
  oformat = av_guess_format(NULL, filename.c_str(), NULL);
  if (!oformat) {
    errMsg_ = "Unknown file extension";
    return -1;
  }

  // Check output format for realtime support
  if (realtime_ && (strcmp(oformat->name, "matroska") && strcmp(oformat->name, "flv"))) {
    errMsg_ = "Realtime encoding is only available for .mkv and .flv files";
    return -1;
  }

  // Select video encoder
  codec = avcodec_find_encoder(codecId_);
  if (!codec || codec->type != CODEC_TYPE_VIDEO) {
    errMsg_ = "Unknown video encoder";
    return -1;
  }

  // Width and hight must be a multiple of two
  if (width_ % 2 || height_ % 2) {
    errMsg_ = "Resolution must be a multiple of two";
    return -1;
  }

  // Now set up the real stuff

  // Allocate format context
  oc = avformat_alloc_context();
  if (!oc) {
    errMsg_ = "Not enough memory to allocate format context";
    return -1;
  }

  // Set format parameters
  oc->oformat = oformat;
  av_strlcpy(oc->filename, filename.c_str(), sizeof(oc->filename));
  oc->preload = (int)(0.5 * AV_TIME_BASE);
  oc->max_delay = (int)(0.7 * AV_TIME_BASE);
  oc->flags |= AVFMT_FLAG_NONBLOCK;

  // FIXME: Timestamp is not set in encoded file at all
  oc->timestamp = av_gettime();

  // Alternatively set date in metadata
  time_t rawtime;
  struct tm *timeinfo;
  char datestr[20];
  time(&rawtime);
  timeinfo = localtime(&rawtime);
  strftime(datestr, 20, "%Y-%m-%d %H:%M:%S", timeinfo);
  av_metadata_set2(&oc->metadata, "date", datestr, 0);

  // Add a new stream to video file
  st = av_new_stream(oc, codec->id);
  if (!st) {
    errMsg_ = "Could not create video stream";
    return -1;
  }

  // Set default codec context values
  avcodec_get_context_defaults2(st->codec, AVMEDIA_TYPE_VIDEO);

  // Set codec context parameters
  c = st->codec;
  c->codec_id = codec->id;
  c->codec_type = CODEC_TYPE_VIDEO;
  c->codec_tag = av_codec_get_tag(oformat->codec_tag, c->codec_id);
  c->width = width_;
  c->height = height_;
  c->gop_size = gopSize_;
  c->bit_rate = bitrate_;
  c->max_b_frames = maxBFrames_;

  // TODO: Unused paramaters
  //c->sample_aspect_ratio;
  //c->bits_per_raw_sample;
  //c->chroma_sample_location;

  // Set fps and timebase
  if (codec && codec->supported_framerates)
    fps_ = codec->supported_framerates[av_find_nearest_q_idx(fps_, codec->supported_framerates)];
  c->time_base.den = st->r_frame_rate.num = fps_.num;
  c->time_base.num = st->r_frame_rate.den = fps_.den;
  st->time_base.den = realtime_ ? 1000 : fps_.num;
  st->time_base.num = realtime_ ? 1: fps_.den;
  av_metadata_set2(&st->metadata, "date", datestr, 0);

  // Select pixel format
  switch (codec->id) {
  case CODEC_ID_HUFFYUV:
    c->pix_fmt = PIX_FMT_YUV422P;
    break;
  case CODEC_ID_MJPEG:
  case CODEC_ID_LJPEG:
    c->pix_fmt = PIX_FMT_YUVJ420P;
    break;
  case CODEC_ID_FFV1:
  default:
    c->pix_fmt = PIX_FMT_YUV420P;
    break;
  }

  // Set stream parameters
  st->sample_aspect_ratio = c->sample_aspect_ratio;

  // Set global header if needed
  if (oc->oformat->flags & AVFMT_GLOBALHEADER)
    c->flags |= CODEC_FLAG_GLOBAL_HEADER;

  // Set video file parameters
  if (av_set_parameters(oc, NULL) < 0) {
    errMsg_ = "invalid encoding parameters";
    return -1;
  }

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

  // Open output file
  if (url_fopen(&oc->pb, oc->filename, URL_WRONLY) < 0) {
    errMsg_ = "Could not open output file";
    return -1;
  }

  // Allocate output buffer
  bufferSize_ = c->width * c->height * 3;
  pBuffer_ = (uint8_t *) av_malloc(bufferSize_);
  if (!pBuffer_) {
    errMsg_ = "Could not allocate output buffer";
    return -1;
  }

  // Allocate encoded raw picture
  pPict_ = avcodec_alloc_frame();
  if (!pPict_) {
    errMsg_ = "Could not allocate raw picture";
    return -1;
  }

  // Allocate encoded raw picture buffer
  int size = avpicture_get_size(c->pix_fmt, c->width, c->height);
  pPictBuffer_ = (uint8_t *) av_malloc(size);
  if (!pPictBuffer_) {
    errMsg_ = "Could not allocate raw picture buffer";
    return -1;
  }

  // Setup the planes
  if (!avpicture_fill((AVPicture *) pPict_, pPictBuffer_, c->pix_fmt, c->width, c->height)) {
    errMsg_ = "Could not fill raw picture buffer";
    return -1;
  }

  // Allocate input picture
  pInputPict_ = avcodec_alloc_frame();
  if (!pInputPict_) {
    errMsg_ = "Could not allocate input picture";
    return -1;
  }

  // Write header
  if (av_write_header(oc) < 0) {
    errMsg_ = "Could not write header";
    return -1;
  }

  // Dump format information
  dump_format(oc, 0, oc->filename, 1);

  pFormatCtx_ = oc;

  return 0;
}

int VideoSink_FFmpeg::Close()
{
  AVFormatContext *oc = pFormatCtx_;
  AVStream *st = NULL;
  AVCodecContext *c = NULL;

  if (oc) {
    st = oc->streams[0];
    if (st) {
      c = st->codec;
    }
  }

  // Write trailer
  if (oc)
    av_write_trailer(oc);

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

  // Close video file
  if (oc)
    url_fclose(oc->pb);

  // Free pointers
  av_free(pPict_);
  av_free(pInputPict_);
  av_free(pBuffer_);
  av_free(c);
  av_free(st);
  av_free(oc);

  Init_();

  return 0;
}

int VideoSink_FFmpeg::AddFrame(const ImageBase& img)
{
  int ret = 0;
  if (!pFormatCtx_) {
    errMsg_ = "Video output file is not opened";
    return -1;
  }

  AVFormatContext *oc = pFormatCtx_;
  AVStream *st = oc->streams[0];
  AVCodecContext *c = st->codec;
  AVPacket pkt;
  int64_t delta;
  enum PixelFormat inputPixFmt;

  // Get realtime presentation timestamp
  if (realtime_) {
    if (!oc->start_time_realtime) {
      oc->start_time_realtime = av_gettime();
      pts_ = 0;
    }
    else {
      delta = av_gettime() - oc->start_time_realtime;
      pts_ = delta * (float) st->time_base.den / (float) st->time_base.num / (float) 1000000;
    }
  }


  //make working copy
  Image<unsigned char> workImg,tmpImg;
  if(img.GetStorageType() != ImageBase::ST_unsignedchar){
    Image<float> tmpFL(img);
    ret =tmpFL.ScaleShiftBetween(0,255);
    if(ret != 0) { BIASERR("Could not ScaleShift float images!");  return -1;}

    ret = ImageConvert::ConvertST(tmpFL,workImg,ImageBase::ST_unsignedchar);
    if(ret != 0){ BIASERR("Could not convert Storagetype of image!"); return -1;}
  }
  else
    workImg=img;

  // if invalid color model set to grey or RGB
  if(workImg.GetColorModel() == ImageBase::CM_invalid){
    if(workImg.GetChannelCount() == 3)
      workImg.SetColorModel(ImageBase::CM_RGB);
    else if(workImg.GetChannelCount() == 2)
      workImg.SetColorModel(ImageBase::CM_UYVY422);
    else
      workImg.SetColorModel(ImageBase::CM_Grey);
  }

  // Get pixel format of image and convert, else copy to workImg
  inputPixFmt = GetPixelFormat_(workImg.GetColorModel());
  if (inputPixFmt == PIX_FMT_NONE) {
    tmpImg  = workImg;
    ret = ImageConvert::ToRGB(tmpImg,workImg);
    if(ret != 0) BIASERR("Could not convert image to rgb!");
  }

  inputPixFmt = GetPixelFormat_(workImg.GetColorModel());
  if (inputPixFmt == PIX_FMT_NONE) {
    errMsg_ = "Input color model not supported";
    return -1;
  }

  // Check size of image
  if (workImg.GetSize() != (unsigned) avpicture_get_size(inputPixFmt, c->width, c->height)) {
    Rescale<unsigned char, unsigned char> rescaler;
    if(workImg.GetSize() >(unsigned) avpicture_get_size(inputPixFmt, c->width, c->height)){
      tmpImg  = workImg;
      ret = rescaler.Downsample(tmpImg,workImg,c->width,c->height);
      if(ret != 0) BIASERR("Could not downsample image to requested size!");
    }
    else{
      tmpImg  = workImg;
      ret = rescaler.Upsample(tmpImg,workImg,c->width,c->height);
      if(ret != 0) BIASERR("Could not upsample image to requested size!");
    }
  }
  if (workImg.GetSize() != (unsigned) avpicture_get_size(inputPixFmt, c->width, c->height)) {
    stringstream m;
    m<<"Image size:"<<workImg.GetWidth()<<"x"<<workImg.GetHeight()
        <<"does not match video output size:"<<c->width<<"x"<<c->height;
    errMsg_ = m.str();
    return -1;
  }

  // Copy image to input picture
  avpicture_fill((AVPicture *) pInputPict_, (uint8_t *) workImg.GetImageData(), inputPixFmt, c->width, c->height);

  // Get color convert context
  pSwsCtx_ = sws_getCachedContext(pSwsCtx_, workImg.GetWidth(), workImg.GetHeight(), inputPixFmt, c->width, c->height, c->pix_fmt, SWS_BICUBIC, NULL, NULL, NULL);
  if (!pSwsCtx_) {
    errMsg_ = "Could not find color convert context";
    return -1;
  }

  // Convert color of input picture
  if (sws_scale(pSwsCtx_, pInputPict_->data, pInputPict_->linesize, 0, c->height, pPict_->data, pPict_->linesize) < 0) {
    errMsg_ = "Could not convert color of input picture";
    return -1;
  }

  // Init packet
  av_init_packet(&pkt);
  pkt.stream_index = st->index;

  if (oc->oformat->flags & AVFMT_RAWPICTURE) {
    // Write raw picture
    pkt.data = (uint8_t *) pPict_;
    pkt.size = sizeof(AVPicture);
    pkt.pts = av_rescale_q(pts_++, c->time_base, st->time_base);
    pkt.flags |= AV_PKT_FLAG_KEY;

    // Write the frame to the video file
    ret = av_interleaved_write_frame(oc, &pkt);
    if (ret < 0) {
      errMsg_ = "Could not write raw frame";
      return -1;
    }
  }
  else {
    // Standard encoding
    pPict_->quality = st->quality;
    if (realtime_)
      pPict_->pts = pts_;

    // Encode video picture
    ret = avcodec_encode_video(c, pBuffer_, bufferSize_, pPict_);
    if (ret < 0) {
      errMsg_ = "Video encoding failed";
      return -1;
    }

    // Dump some info about the encoded frame
    //printf("frame=%5d dts=%5d pts=%5ld\n", c->coded_frame->coded_picture_number, c->coded_frame->display_picture_number, c->coded_frame->pts);

    if (ret > 0) {
      pkt.data = pBuffer_;
      pkt.size = ret;
      if (!realtime_ && c->coded_frame->pts != (unsigned) AV_NOPTS_VALUE)
        pkt.pts = av_rescale_q(c->coded_frame->pts, c->time_base, st->time_base);
      else
        pkt.pts = c->coded_frame->pts;
      if (c->coded_frame->key_frame)
        pkt.flags |= AV_PKT_FLAG_KEY;

      // Write the frame to the video file
      ret = av_interleaved_write_frame(oc, &pkt);
      if (ret < 0) {
        errMsg_ = "Could not write frame";
        return -1;
      }
    }
  }

  return 0;
}

int VideoSink_FFmpeg::AddFrame(string filename)
{
  ImageBase image;
  ImageIO::Load(filename, image);  
  return AddFrame(image);
}

int VideoSink_FFmpeg::AddFrames(vector<string> filenames)
{
  int ret = 0;
  vector<string>::iterator f;
  for (f = filenames.begin(); f != filenames.end() && ret == 0; ++f)
    ret = AddFrame(*f);
  return ret;
}

enum PixelFormat VideoSink_FFmpeg::GetPixelFormat_(enum ImageBase::EColorModel colorModel)
{
  switch (colorModel) {
  case ImageBase::CM_Grey:    return PIX_FMT_GRAY8;
  case ImageBase::CM_GreyA:   return PIX_FMT_GRAY16;
  case ImageBase::CM_RGB:     return PIX_FMT_RGB24;
  case ImageBase::CM_BGR:     return PIX_FMT_BGR24;
  case ImageBase::CM_RGBA:    return PIX_FMT_RGBA;
  case ImageBase::CM_BGRA:    return PIX_FMT_BGRA;
  case ImageBase::CM_YUYV422: return PIX_FMT_YUYV422;
  case ImageBase::CM_UYVY422: return PIX_FMT_UYVY422;
  case ImageBase::CM_YUV420P: return PIX_FMT_YUV420P;
  case ImageBase::CM_YUV444:  return PIX_FMT_YUV444P;
  case ImageBase::CM_YUV411:  return PIX_FMT_YUV411P;
  default:                    return PIX_FMT_NONE;
  }
}