source: orxonox.OLD/trunk/src/world_entities/recorder.cc @ 7897

/*
   orxonox - the future of 3D-vertical-scrollers

   Copyright (C) 2006 orx

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

### File Specific:
   main-programmer: David Hasenfratz
   co-programmer:
*/

#include "recorder.h"

#include "util/loading/load_param.h"
#include "util/loading/factory.h"
#include "util/loading/resource_manager.h"
#include "state.h"

using namespace std;

CREATE_FACTORY(Recorder, CL_RECORDER);


Recorder::Recorder (const TiXmlElement* root)
{
  this->setClassID(CL_RECORDER, "Recorder");

  // initialize libavcodec, and register all codecs and formats
  av_register_all();

  // Default values
  stream_duration = 10;
  stream_frame_rate = 20;

  this->loadParams(root);

  this->toList(OM_COMMON);
}


Recorder::~Recorder ()
{

}


void Recorder::loadParams(const TiXmlElement* root)
{
  WorldEntity::loadParams(root);

  LoadParam(root, "duration", this, Recorder, setStreamDuration);

  LoadParam(root, "fps", this, Recorder, setFPS);

  LoadParam(root, "name", this, Recorder, initVideo);
}


void Recorder::setStreamDuration(float duration)
{
  this->stream_duration = duration;
}


void Recorder::setFPS(float fps)
{
  this->stream_frame_rate = fps;
}


void Recorder::initVideo(const std::string& filename)
{
  frame_count = 0;
  time = 0;
  stream_nb_frames = (int)(stream_duration * stream_frame_rate);

  // auto detect the output format from the name, default is mpeg
  output_format = guess_format(NULL, filename.c_str(), NULL);
  if (!output_format)
  {
      PRINTF(0)("Could not deduce output format from file extension: using MPEG.\n");
      output_format = guess_format("mpeg", NULL, NULL);
  }
  if (!output_format)
    PRINTF(1)("Could not find suitable output format\n");

  // allocate the output media context
  format_context = av_alloc_format_context();
  if (!format_context)
    PRINTF(1)("Memory error\n");

  format_context->oformat = output_format;
  snprintf(format_context->filename, sizeof(format_context->filename), "%s", filename.c_str());

  // add video stream using the default format codec and initialize the codec
  if (output_format->video_codec != CODEC_ID_NONE)
      this->addVideoStream();

  // set the output parameters (must be done even if no parameters)
  if (av_set_parameters(format_context, NULL) < 0)
    PRINTF(1)("Invalid output format parameters\n");

  // print some information
  dump_format(format_context, 0, filename.c_str(), 1);

  // now that all the parameters are set, we can open the
  // video codec and allocate the necessary encode buffer
  if(video_stream)
    this->openVideo();

  // open the output file, if needed
  if(!(output_format->flags & AVFMT_NOFILE))
  {
    if(url_fopen(&format_context->pb, filename.c_str(), URL_WRONLY) < 0)
      PRINTF(1)("Could not open %s\n", filename.c_str());
  }

  // write the stream header, if any
  av_write_header(format_context);
}

void Recorder::closeVideo()
{
  avcodec_close(video_stream->codec);
  av_free(picture->data[0]);
  av_free(picture);
  av_free(buffer);

  // write the trailer, if any
  av_write_trailer(format_context);

  // free the streams
  for(int i = 0; i < format_context->nb_streams; i++)
    av_freep(&format_context->streams[i]);

  // close the output file
  if (!(output_format->flags & AVFMT_NOFILE))
    url_fclose(&format_context->pb);

  // free the stream
  av_free(format_context);
}


void Recorder::openVideo()
{
  codec_context = video_stream->codec;

  // find the video encoder
  codec = avcodec_find_encoder(codec_context->codec_id);
  if(!codec)
    PRINTF(1)("codec not found\n");

  // open the codec
  if(avcodec_open(codec_context, codec) < 0)
    PRINTF(1)("could not open codec\n");

  buffer = NULL;
  if(!(format_context->oformat->flags & AVFMT_RAWPICTURE))
  {
      // allocate output buffer
      // XXX: API change will be done
      buffer_size = 200000;
      buffer = new uint8_t[buffer_size];
  }

  // allocate the encoded raw picture
  this->allocPicture();
  if(!picture)
    PRINTF(0)("Could not allocate picture\n");
}


void Recorder::allocPicture()
{
  picture = avcodec_alloc_frame();
  if(!picture)
  {
    picture = NULL;
    return;
  }
  size = avpicture_get_size(codec_context->pix_fmt, width, height);
  picture_buf = new uint8_t[size];
  if(!picture_buf)
  {
      av_free(picture);
      return;
  }
  avpicture_fill((AVPicture *)picture, picture_buf,
                  codec_context->pix_fmt, width, height);



  RGB_frame = avcodec_alloc_frame();

  // Determine required buffer size and allocate buffer
  size = avpicture_get_size(PIX_FMT_RGB24, width, height);
  picture_buf = new uint8_t[size];

  // Assign appropriate parts of buffer to image planes in RGB_frame
  avpicture_fill((AVPicture *)RGB_frame, picture_buf, PIX_FMT_RGB24, width, height);
}


// add a video output stream
void Recorder::addVideoStream()
{
  video_stream = av_new_stream(format_context, 0);
  if (!video_stream)
    PRINTF(1)("Could not alloc stream\n");

  codec_context = video_stream->codec;
  codec_context->codec_id = output_format->video_codec;
  codec_context->codec_type = CODEC_TYPE_VIDEO;

  // put sample parameters
  codec_context->bit_rate = 400000;
  // resolution must be a multiple of two
  codec_context->width = State::getResX();
  codec_context->height = State::getResY();

  this->width = codec_context->width;
  this->height = codec_context->height;

  // time base: this is the fundamental unit of time (in seconds) in terms
  // of which frame timestamps are represented. for fixed-fps content,
  // timebase should be 1/framerate and timestamp increments should be
  // identically 1
  codec_context->time_base.den = (int)stream_frame_rate;
  codec_context->time_base.num = 1;
  codec_context->gop_size = 12;  // emit one intra frame every twelve frames at most
  codec_context->pix_fmt = PIX_FMT_YUV420P;

  if (codec_context->codec_id == CODEC_ID_MPEG1VIDEO)
      // needed to avoid using macroblocks in which some coeffs overflow
      // this doesnt happen with normal video, it just happens here as the
      // motion of the chroma plane doesnt match the luma plane
      codec_context->mb_decision=2;

  // some formats want stream headers to be seperate
  if(!strcmp(format_context->oformat->name, "mp4") ||
     !strcmp(format_context->oformat->name, "mov") ||
     !strcmp(format_context->oformat->name, "3gp"))
    format_context->flags |= CODEC_FLAG_GLOBAL_HEADER;
}


void Recorder::tick(float dt)
{
  time += dt;
  if(time < 1/stream_frame_rate)
    return;
  else
    time = 0;

  // compute current video time
  if (video_stream)
      video_pts = (double)video_stream->pts.val * video_stream->time_base.num / video_stream->time_base.den;
  else
      video_pts = 0.0;

  if (!video_stream || video_pts >= stream_duration)
  {
    this->toList(OM_DEAD);
    this->closeVideo();
  }
  else
    // write video frame
    this->writeVideoFrame();
}

void Recorder::writeVideoFrame()
{
  int out_size, err;

  codec_context = video_stream->codec;

  if(frame_count >= stream_nb_frames)
  {
    /* no more frame to compress. The codec has a latency of a few
        frames if using B frames, so we get the last frames by
        passing the same picture again */
  }
  else
    this->fillYuvImage();


  if(format_context->oformat->flags & AVFMT_RAWPICTURE)
  {
    // raw video case. The API will change slightly in the near
    // futur for that
    av_init_packet(&packet);

    packet.flags |= PKT_FLAG_KEY;
    packet.stream_index= video_stream->index;
    packet.data= (uint8_t *)picture;
    packet.size= sizeof(AVPicture);

    err = av_write_frame(format_context, &packet);
  }
  else
  {
    // encode the image
    out_size = avcodec_encode_video(codec_context, buffer, buffer_size, picture);
    // if zero size, it means the image was buffered
    if (out_size > 0)
    {
      av_init_packet(&packet);

      packet.pts= av_rescale_q(codec_context->coded_frame->pts, codec_context->time_base, video_stream->time_base);
      if(codec_context->coded_frame->key_frame)
        packet.flags |= PKT_FLAG_KEY;
      packet.stream_index = video_stream->index;
      packet.data= buffer;
      packet.size= out_size;

      // write the compressed frame in the media file
      err = av_write_frame(format_context, &packet);
    }
    else
      err = 0;
  }

  if(err != 0)
    PRINTF(1)("Error while writing video frame\n");

  frame_count++;
}


void Recorder::fillYuvImage()
{
  unsigned char *outputImage = 0;

  // Allocate the neccessary memory.
  outputImage = (unsigned char*)malloc(width * height * 3);

  // Clear the variable.
  memset(outputImage, 0, width * height * 3);

  // You use the glReadPixels() to read every pixel on the screen
  // that you specify.  You must use one less than each size.
  glReadPixels(0, 0, width, height, GL_RGB, GL_UNSIGNED_BYTE, outputImage);

  int i = 0;
  for(int y=height-1;y>=0;y--)
  {
    for(int x=0;x<width*3;x++)
    {
      RGB_frame->data[0][y*width*3+x] = outputImage[i];
      i++;
    }
  }


  img_convert((AVPicture*)picture, PIX_FMT_YUV420P, (AVPicture*)RGB_frame,
              PIX_FMT_RGB24, width, height);

  // Clear the allocated memory.
  free(outputImage);
}


void Recorder::draw() const
{
}