VideoSink.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 "VideoSink.hh"
#include <VideoSource/VideoSource_Disk.hh>

using namespace std;

BIAS::VideoSink::VideoSink()
  : _Bitrate(-1), _BPP(4)
{
  static bool firstCall = true;
  if (firstCall)
  {
    avcodec_init();
    av_register_all();
    firstCall = false;
  }
}
BIAS::VideoSink::VideoSink(std::string Video)
  : _Video(Video)
{
  VideoSink();
}
BIAS::VideoSink::VideoSink(std::vector<std::string> Images)
  : _ImgFiles(Images)
{
  VideoSink();
}
BIAS::VideoSink::VideoSink(std::vector<std::string> Images, std::string Video)
  : _ImgFiles(Images), _Video(Video)
{
  VideoSink();
}
BIAS::VideoSink::~VideoSink()
{
}

void BIAS::VideoSink::AddImage(std::string Image)
{
  _ImgFiles.push_back(Image);
}
void BIAS::VideoSink::AddImages(std::vector<std::string> Images)
{
  _ImgFiles.insert(_ImgFiles.end(), Images.begin(), Images.end());
}

void BIAS::VideoSink::SetOutputVideo(std::string Video)
{
  _Video = Video;
}
void BIAS::VideoSink::SetFPS(float FPS)
{
  _FPS = FPS;
}

void BIAS::VideoSink::SetBitrate(int Bitrate)
{
  _Bitrate = Bitrate;
  _BPP = -1;
}
void BIAS::VideoSink::SetBPP(double BPP)
{
  _BPP = BPP;
  _Bitrate = -1;
}

/*
void MakeImagePlanar(BIAS::Image<char> &img)
{
  int size = img.GetWidth() * img.GetHeight();
  int channels = img.GetChannelCount();
  char *cp = (char*)malloc(size * channels);
  char *dat = img.GetImageData();
  memcpy(cp, dat, size * channels);
  for (int i = 0; i < size; i++)
  {
    for (int c = 0; c < channels; c++)
    {
      *(dat + c * size + i) = *(cp + channels * i + c);
    }
  }
  free(cp);
  img.SetInterleaved(false);
}
*/

void BIAS::VideoSink::_RGB888pToYUV420p(uint8_t *rgb, uint8_t *yuv,
                    int width, int height)
{
  /*
    Y  =      (0.257 * R) + (0.504 * G) + (0.098 * B) + 16
    Cr = V =  (0.439 * R) - (0.368 * G) - (0.071 * B) + 128
    Cb = U = -(0.148 * R) - (0.291 * G) + (0.439 * B) + 128
    (www.fourcc.org/fccyvrgb.php)
  */
  int size = width * height;
  int r;
  int g;
  int b;
  int p;
  uint8_t *u;
  
  for (int y = 0; y < height; y++)
  {
    for (int x = 0; x < width; x++)
    {
      p = y*width + x;
      r = *(rgb + p);
      g = *(rgb + size + p);
      b = *(rgb + 2*size + p);
      
      // Y'
      *(yuv + p) = (0.257*r) + (0.504*g) + (0.098*b) + 16;
      // Only 1 U and V value for 2x2 Y' values
      if (y % 2 == 0 && x % 2 == 0)
      {
    // Calculate average color of the 4 pixels
    r = (r
         + *(rgb + p + 1)
         + *(rgb + p + width)
         + *(rgb + p + width + 1)) / 4;
    g = (g
         + *(rgb + size + p + 1)
         + *(rgb + size + p + width)
         + *(rgb + size + p + width + 1)) / 4;
    b = (r
         + *(rgb + 2*size + p + 1)
         + *(rgb + 2*size + p + width)
         + *(rgb + 2*size + p + width + 1)) / 4;
    // U
    u = yuv + size + (y/2) * (width/2) + x/2;
    *u = -(0.148*r) - (0.291*g) + (0.439*b) + 128;
    // V
    *(u + size/4) = (0.439*r) - (0.368*g) - (0.071*b) + 128;
      }
    }
  }
}

bool BIAS::VideoSink::Encode()
{
  int outSize;
  int size;
  int outbufSize;
  FILE *f;
  uint8_t *outbuf;
  uint8_t *pictBuf;
  uint8_t *yuvBuf;
  AVFrame *pict;
  AVFrame *yuvPict;
  int w;
  int h;
  Image<char> img;

  if (!_SetupVideoSource())
    return false;

  w = _VSD.GetWidth();
  h = _VSD.GetHeight();
  size = w * h;

  _Codec = avcodec_find_encoder(CODEC_ID_MPEG1VIDEO);
  if (!_Codec)
  {
    _ErrMsg = "Codec not found";
    return false;
  }
  cout << "Privdatasize: " << _Codec->priv_data_size << endl;

  _CodecCtx = avcodec_alloc_context();
  pict = avcodec_alloc_frame();
  yuvPict = avcodec_alloc_frame();

  int bitrate = _Bitrate;
  if (bitrate < 0)
    bitrate = _BPP * size;
  _CodecCtx->bit_rate = bitrate;
  _CodecCtx->width = w;
  _CodecCtx->height = h;
  _CodecCtx->time_base.num = 1;
  _CodecCtx->time_base.den = _FPS;
  _CodecCtx->gop_size = 10;
  _CodecCtx->pix_fmt = PIX_FMT_YUV420P;
  _CodecCtx->priv_data = (void*)"title=BIAS FFMPEG Encoder";
  
  if (avcodec_open(_CodecCtx, _Codec) < 0)
  {
    _ErrMsg = "Could not open codec";
    return false;
  }

  f = fopen(_Video.c_str(), "wb");
  if (!f)
  {
    _ErrMsg = string("Could not open output file: ") + _Video;
    return false;
  }

  outbufSize = 100000;
  outbuf = (uint8_t*)malloc(outbufSize);
  if (!_IsYUV)
    pictBuf = (uint8_t*)malloc(size * 3); // Size for RGB
  else
    pictBuf = (uint8_t*)malloc((size * 3) / 2); // Size for YUV420

  pict->linesize[0] = w;
  pict->linesize[1] = w;
  pict->linesize[2] = w;
  pict->data[0] = pictBuf;
  pict->data[1] = pictBuf + size;
  pict->data[2] = pictBuf + size + size;

  if (!_IsYUV)
  {
    yuvBuf = (uint8_t*)malloc((size * 3) / 2); // Size for YUV420

    yuvPict->linesize[0] = w;
    yuvPict->linesize[1] = w / 2;
    yuvPict->linesize[2] = w / 2;
    yuvPict->data[0] = yuvBuf;
    yuvPict->data[1] = yuvBuf + size;
    yuvPict->data[2] = yuvBuf + size + size / 4;
  }
  else
  {
    yuvPict = pict;
    yuvBuf = pictBuf;
  }

  _VSD.PreGrab();
  int i = 0;
  while (_VSD.GetNumFramesRemaining() != 0)
  {
    fflush(stdout);
    
    if (_VSD.GrabSingle(img) != 0)
    {
      _ErrMsg = "Error grabbing image";
      return false;
    }

    // Need to convert colormodel?
    if (!_IsYUV)
    {
      ImageConvert::ToPlanar(img, img);
      //memcpy(pictBuf, img.GetImageData(), size * 3);
      //sws_scale(_SwsCtx, pict->data, pict->linesize, 0,
      //          h, yuvPict->data, yuvPict->linesize);
      _RGB888pToYUV420p((uint8_t*)img.GetImageData(), yuvBuf, w, h);
    }
    else
    {
      memcpy(yuvBuf, img.GetImageData(), (size * 3) / 2);
    }

    outSize = avcodec_encode_video(_CodecCtx, outbuf, outbufSize, yuvPict);
    fwrite(outbuf, 1, outSize, f);

    cout << "Frame " << i << ", remaining: "
         << _VSD.GetNumFramesRemaining() << endl;
    i++;
  }
  _VSD.PostGrab();

  outbuf[0] = 0x00;
  outbuf[1] = 0x00;
  outbuf[2] = 0x01;
  outbuf[3] = 0xb7;
  fwrite(outbuf, sizeof(uint8_t), 4, f);
  fclose(f);
  free(pictBuf);
  if (yuvBuf)
    free(yuvBuf);
  free(outbuf);
  
  avcodec_close(_CodecCtx);
  av_free(_CodecCtx);
  av_free(pict);
  if (yuvPict)
    av_free(yuvPict);

  return true;
}

std::string BIAS::VideoSink::GetErrMsg()
{
  return _ErrMsg;
}


bool BIAS::VideoSink::_SetupVideoSource()
{
  _VSD.SetLoopMode(false);
  _VSD.SetColorModel(ImageBase::CM_YUV420P);
  if (_VSD.OpenDevice(_ImgFiles) != 0)
  {
    _ErrMsg = "Loading images failed";
    return false;
  }
  ImageBase::EColorModel cm = _VSD.GetColorModel();
  if (cm == ImageBase::CM_RGB)
  {
    //_SetupSwsContext(PIX_FMT_RGB24);
    _IsYUV = false;
  }
  else if (cm != ImageBase::CM_YUV420P)
  {
    _ErrMsg = "VideoSource returns unsupported colormodel";
    return false;
  }
  else
    _IsYUV = true;

  return true;
}

/*
bool BIAS::VideoSink::_SetupSwsContext(PixelFormat fromFmt, int width,
                                       int height)
{
  int w = _VSD.GetWidth();
  int h = _VSD.GetHeight();
  
  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, fromFmt,
                           width, height, PIX_FMT_YUV420P,
                           SWS_BICUBIC, NULL, NULL, NULL);
  if (_SwsCtx == NULL)
  {
    _ErrMsg = "Cannot initialize the conversion context!";
    return false;
  }
  return true;
}
*/