AvWriter.hpp

Go to the documentation of this file.

#ifndef AVWRITER_HPP_
#define AVWRITER_HPP_

#include <iostream>
#include <sstream>

extern "C" {
#include <libswscale/swscale.h>
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavcodec/version.h>
#include <libavutil/pixfmt.h>
#include <libavutil/mathematics.h>
#include <stdio.h>
}

#include "color_spaces.hpp"
#include "error.hpp"
#include "ContainerWriter.hpp"


template <typename T>
struct av_static_cast_func
{
                template <typename T1> // T1 models type statically convertible to T
                T operator()(const T1& x) const { return static_cast<T>(x); }
};


namespace slip {

template<class ContainerVideo, typename T, std::size_t Nb_components>
class AvWriter: public ContainerWriter<ContainerVideo, T> {
public:
                typedef ContainerVideo container_type;
                typedef T value_type;
                typedef ContainerWriter<ContainerVideo, T> base;


                AvWriter()
                :base(),pFormatCtx(NULL),video_stream_idx_(-1),pictureIN(NULL),pictureOUT(NULL),pSwsCtxOUT(NULL),
                 codec_id_(CODEC_ID_NONE),dimension_optimize_(0),bit_rate_(0),output_video_width_(0),output_video_height_(0),
                 output_video_nbframes_(0),time_base_((AVRational){0,1}),input_pxfmt_(PIX_FMT_GRAY8),
                 output_pxfmt_(PIX_FMT_RGB24),row_stride(0),finished(0),total_nb_slice(0),initialized(false){}

                AvWriter(std::string output_filename, std::size_t output_video_width, std::size_t output_video_height,
                                                std::size_t output_video_nbframes, int bit_rate, AVRational time_base, bool dimension_optimize = 0,
                                                CodecID codec_id = CODEC_ID_NONE)
                :base(output_filename),pFormatCtx(NULL),video_stream_idx_(-1),pictureIN(NULL),pictureOUT(NULL),pSwsCtxOUT(NULL),
                 codec_id_(codec_id),dimension_optimize_(dimension_optimize), bit_rate_(bit_rate),output_video_width_(output_video_width),
                 output_video_height_(output_video_height),output_video_nbframes_(output_video_nbframes),time_base_(time_base),
                 input_pxfmt_(PIX_FMT_GRAY8),output_pxfmt_(PIX_FMT_RGB24),row_stride(0),finished(0),total_nb_slice(0),initialized(false)
                {
                                initialize();
                }

                virtual ~AvWriter(){
                                release();
                }

                void initialize( std::size_t output_video_width, std::size_t output_video_height,
                                                std::size_t output_video_nbframes, int bit_rate, AVRational time_base,
                                                bool dimension_optimize = 0, CodecID codec_id = CODEC_ID_NONE);


                void release();

                int write(const container_type & data);

private:
                AVFormatContext *pFormatCtx;
                int video_stream_idx_;
                AVFrame * pictureIN;
                AVFrame * pictureOUT;
                struct SwsContext * pSwsCtxOUT;

                CodecID codec_id_;
                bool dimension_optimize_;
                int bit_rate_;
                std::size_t output_video_width_;
                std::size_t output_video_height_;
                std::size_t output_video_nbframes_;
                AVRational time_base_;
                PixelFormat input_pxfmt_;
                PixelFormat output_pxfmt_;
                int row_stride;
                bool finished;
                std::size_t total_nb_slice;
                bool initialized;


                void initialize();
                void alloc_picture(AVFrame *&picture, PixelFormat pix_fmt, int width, int height);
                void add_video_stream();
                void write_video_frame(const container_type & data, std::size_t nb_slice);
};

} 
namespace slip {

template<class ContainerVideo, typename T, std::size_t Nb_components>
inline
void AvWriter<ContainerVideo,T,Nb_components>::initialize(std::size_t output_video_width, std::size_t output_video_height, std::size_t output_video_nbframes, int bit_rate, AVRational time_base, bool dimension_optimize, CodecID codec_id){
                bit_rate_ = bit_rate;
                output_video_width_ = output_video_width;
                output_video_height_ = output_video_height;
                output_video_nbframes_ = output_video_nbframes;
                time_base_ = time_base;
                codec_id_ = codec_id;
                dimension_optimize_ = dimension_optimize;
                initialize();
}
template<class ContainerVideo, typename T, std::size_t Nb_components>
inline
void AvWriter<ContainerVideo,T,Nb_components>::alloc_picture(AVFrame * &picture, PixelFormat pix_fmt, int width, int height)
{
                //AVFrame *picture;
                uint8_t *picture_buf;
                int size;

                picture = avcodec_alloc_frame();
                if (picture){
                                size = avpicture_get_size(pix_fmt, width, height);
                                picture_buf = reinterpret_cast<uint8_t*>(av_malloc(size));
                                if (!picture_buf)
                                                av_free(picture);
                                else
                                                avpicture_fill((AVPicture *)picture, picture_buf, pix_fmt, width, height);
                }
}



template<class ContainerVideo, typename T, std::size_t Nb_components>
inline
void AvWriter<ContainerVideo,T,Nb_components>::add_video_stream()
{
                AVStream * video_st = NULL;
                /* find the video encoder */
                AVCodec * codec = avcodec_find_encoder(codec_id_);
                if (!codec) {
                                std::ostringstream err;
                                err << FILE_WRITE_ERROR << this->output_filename_ << " | Codec not found.";
                                slip::slip_exception exc(std::string("slip"), std::string("AvWriter::add_video_stream()"), err.str());
                                throw (exc);
                }

                const PixelFormat * possible_pix_fmt = codec->pix_fmts;
                if (possible_pix_fmt){
                                int loss;
#ifndef FF_API_FIND_BEST_PIX_FMT
                                int64_t pix_fmt_mask = (1 << *possible_pix_fmt);
                                possible_pix_fmt++;
                                while(*possible_pix_fmt != -1){
                                                pix_fmt_mask = pix_fmt_mask || (1 << *possible_pix_fmt);
                                                possible_pix_fmt++;
                                }
                                output_pxfmt_ = avcodec_find_best_pix_fmt(pix_fmt_mask,input_pxfmt_,0,&loss);
#else
                output_pxfmt_ = avcodec_find_best_pix_fmt2(possible_pix_fmt,input_pxfmt_,0,&loss);
#endif

                }else{
                                std::cerr << "Warning -> pixel format unknown for this codec, "
                                                                << "PIX_FMT_RGB24 has been chosen by default." << std::endl;
                }

                video_st = avformat_new_stream(pFormatCtx, codec);
                if (!video_st) {
                                std::ostringstream err;
                                err << FILE_WRITE_ERROR << this->output_filename_ << " | Could not alloc stream.";
                                slip::slip_exception exc(std::string("slip"), std::string("AvWriter::add_video_stream()"), err.str());
                                throw (exc);
                }

                /* frames per second */
                video_st->codec->time_base= time_base_;
                video_st->time_base = time_base_;
                video_st->codec->gop_size = time_base_.den; /* =frame rate */
                video_st->codec->pix_fmt = output_pxfmt_;
                if (video_st->codec->codec_id == CODEC_ID_MPEG2VIDEO) {
                                /* just for testing, we also add B frames */
                                video_st->codec->max_b_frames = 2;
                }
                if (video_st->codec->codec_id == CODEC_ID_MPEG1VIDEO) {
                                /* Needed to avoid using macroblocks in which some coeffs overflow.
                                 * This does not happen with normal video, it just happens here as
                                 * the motion of the chroma plane does not match the luma plane. */
                                video_st->codec->mb_decision = 2;
                }

                int width = static_cast<int>(output_video_width_);
                int height = static_cast<int>(output_video_height_);
                //Resize if needed
                if(dimension_optimize_){
                                std::cerr << "Old dimensions = (" << width << ", " << height << ")\n";
                                int linesize_align[AV_NUM_DATA_POINTERS];
                                avcodec_align_dimensions2(video_st->codec,&width,&height,linesize_align);
                                std::cerr << "New dimensions = (" << width << ", " << height << ")\n";
                }
                video_st->codec->codec_type = AVMEDIA_TYPE_VIDEO;
                video_st->codec->bit_rate = bit_rate_;
                video_st->codec->width = width;
                video_st->codec->height = height;

                /* Some formats want stream headers to be separate. */
                if (pFormatCtx->oformat->flags & AVFMT_GLOBALHEADER)
                                video_st->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;

                video_stream_idx_ = video_st->index;

                // open the video codec
                if (avcodec_open2(video_st->codec, codec, NULL) < 0) {
                                std::ostringstream err;
                                err << FILE_WRITE_ERROR << this->output_filename_ << " | Can not open codec.";
                                slip::slip_exception exc(std::string("slip"), std::string("AvWritter::add_video_stream()"), err.str());
                                throw (exc);
                }

                /* Allocate all pictures. */
                //pictureIN
                alloc_picture(pictureIN, input_pxfmt_, output_video_width_, output_video_height_);
                if (!pictureIN) {
                                std::ostringstream err;
                                err << FILE_WRITE_ERROR << this->output_filename_ << " | Could not allocate input picture buffer.";
                                slip::slip_exception exc(std::string("slip"), std::string("AvWritter::add_video_stream()"), err.str());
                                throw (exc);
                }
                //pictureOUT =
                alloc_picture(pictureOUT, video_st->codec->pix_fmt, video_st->codec->width, video_st->codec->height);
                if (!pictureOUT) {
                                std::ostringstream err;
                                err << FILE_WRITE_ERROR << this->output_filename_ << " | Could not allocate output picture buffer.";
                                slip::slip_exception exc(std::string("slip"), std::string("AvWritter::add_video_stream()"), err.str());
                                throw (exc);
                }
}

template<class ContainerVideo, typename T, std::size_t Nb_components>
inline
void AvWriter<ContainerVideo,T,Nb_components>::initialize(){
                if(!initialized){
                                if (Nb_components != 1 && Nb_components != 3){
                                                std::ostringstream err;
                                                err << FILE_WRITE_ERROR << this->output_filename_ << " | Nb_components should be 1 or 3.";
                                                slip::slip_exception exc(std::string("slip"), std::string("AvWriter::initialize()"), err.str());
                                                throw (exc);
                                }
                                if (Nb_components == 3){
                                                input_pxfmt_ = PIX_FMT_RGB24;
                                }

                                avcodec_register_all();
                                av_register_all();
                                //avformat_network_init();

                                /* Autodetect the output format from the name. default is MPEG. */
                                AVOutputFormat *fmt = av_guess_format(NULL, this->output_filename_.c_str(), NULL);
                                if (!fmt) {
                                                std::cerr << "Could not deduce output format from file extension: using MPEG." << std::endl;
                                                fmt = av_guess_format("mpeg", NULL, NULL);
                                }
                                if (!fmt) {
                                                std::ostringstream err;
                                                err << FILE_WRITE_ERROR << this->output_filename_ << " | Could not find suitable output format.";
                                                slip::slip_exception exc(std::string("slip"), std::string("AvWriter::initialize()"), err.str());
                                                throw (exc);
                                }
                                if(codec_id_ != fmt->video_codec && codec_id_ != CODEC_ID_NONE){
                                                std::cerr << "The specified codec id does not match the file extension." << std::endl;
                                }
                                codec_id_ = fmt->video_codec;
                                /* Allocate the output media context. */
                                pFormatCtx = avformat_alloc_context();
                                if (!pFormatCtx) {
                                                std::ostringstream err;
                                                err << FILE_WRITE_ERROR << this->output_filename_ << " | Memory error.";
                                                slip::slip_exception exc(std::string("slip"), std::string("AvWriter::initialize()"), err.str());
                                                throw (exc);
                                }
                                pFormatCtx->oformat = fmt;
                                strcpy(pFormatCtx->filename, this->output_filename_.c_str());

                                /* Add the video stream using the default format codecs
                                 * and initialize the codecs. */
                                //CHECK -> is this if condition useful?
                                if (pFormatCtx->oformat->video_codec != CODEC_ID_NONE) {
                                                add_video_stream();
                                }

                                av_dump_format(pFormatCtx, 0, this->output_filename_.c_str(), 1);
                                /* open the output file*/
                                if (!(pFormatCtx->oformat->flags & AVFMT_NOFILE)) {
                                                //CHECK -> URL_WRONLY instead of AVIO_FLAG_WRITE
                                                if (avio_open(&pFormatCtx->pb, this->output_filename_.c_str(), AVIO_FLAG_WRITE) < 0) {
                                                                std::ostringstream err;
                                                                err << FILE_OPEN_ERROR << this->output_filename_;
                                                                slip::slip_exception exc(std::string("slip"), std::string("AvWriter::initialize()"), err.str());
                                                                throw (exc);
                                                }
                                }

                                /* Write the stream header, if any. */
                                if (avformat_write_header(pFormatCtx, NULL) < 0){
                                                std::ostringstream err;
                                                err << FILE_WRITE_ERROR << this->output_filename_ << " | Cannot write a file header.";
                                                slip::slip_exception exc(std::string("slip"), std::string("AvWriter::initialize()"), err.str());
                                                throw (exc);
                                }
                                row_stride = output_video_width_ * Nb_components;
                                initialized = true;
                }
}

template<class ContainerVideo, typename T, std::size_t Nb_components>
inline
void AvWriter<ContainerVideo,T,Nb_components>::write_video_frame(const container_type & data, std::size_t nb_slice)
{
                int ret = 1;
                AVStream * video_st = pFormatCtx->streams[video_stream_idx_];
                AVCodecContext * pCodecCtx = video_st->codec;
                //CHECK-> replace SWS_FAST_BILINEAR by SWS_BICUBIC
                if (Nb_components == 1){
                                pSwsCtxOUT = sws_getContext(data.cols(),data.rows(),PIX_FMT_GRAY8,pCodecCtx->width,pCodecCtx->height,
                                                                pCodecCtx->pix_fmt,SWS_FAST_BILINEAR, NULL, NULL, NULL);
                }else{
                                pSwsCtxOUT = sws_getContext(data.cols(),data.rows(),PIX_FMT_RGB24,pCodecCtx->width,pCodecCtx->height,
                                                                pCodecCtx->pix_fmt,SWS_FAST_BILINEAR, NULL, NULL, NULL);
                }
                if (pSwsCtxOUT == NULL) {
                                std::ostringstream err;
                                err << FILE_WRITE_ERROR << this->output_filename_ << " | Cannot initialize the YUV conversion context.";
                                slip::slip_exception exc(std::string("slip"), std::string("AvWriter::write_video_frame()"), err.str());
                                throw (exc);
                }

                if (total_nb_slice >= output_video_nbframes_) {
                                /* No more frames 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 {
                                for(std::size_t y=0; y < data.rows(); y++) {
                                                const T * data_ptr_i = reinterpret_cast<const T *>(data[nb_slice][y]);
                                                std::transform(data_ptr_i, data_ptr_i + row_stride,
                                                                                pictureIN->data[0]+y*pictureIN->linesize[0], av_static_cast_func<unsigned char>());
                                }
                                //Scale into pictureOUT
                                sws_scale(pSwsCtxOUT, pictureIN->data, pictureIN->linesize,0, data.rows(), pictureOUT->data,
                                                                pictureOUT->linesize);
                }

                pictureOUT->pts = //static_cast<int64_t>((1000.0 * av_q2d(time_base_)) * 90.0 *
                                                static_cast<double>(total_nb_slice);//);

                if (pFormatCtx->oformat->flags & AVFMT_RAWPICTURE) {
                                /* Raw video case - the API will change slightly in the near
                                 * future for that. */
                                AVPacket pkt;
                                av_init_packet(&pkt);

                                pkt.flags        |= AV_PKT_FLAG_KEY;
                                pkt.stream_index  = video_stream_idx_;
                                pkt.data          = reinterpret_cast<uint8_t *>(pictureOUT);
                                pkt.size          = sizeof(AVPicture);

                                ret = av_write_frame(pFormatCtx, &pkt);
                                av_free_packet(&pkt);
                } else {
                                AVPacket pkt = { 0 };
                                av_init_packet(&pkt);
                                /* encode the image */
#ifndef FF_API_OLD_ENCODE_VIDEO
                                int buf_size = static_cast<int>( 9 * video_st->codec->width * video_st->codec->height * Nb_components + 10000);
                                //std::cerr << "buf_size = " << buf_size << "\n";
                                uint8_t * buf = new uint8_t[buf_size];
                                int sizeout = avcodec_encode_video(pCodecCtx, buf, buf_size, pictureOUT);
                                if(sizeout > 0){
                                                pkt.data = buf;
                                                pkt.size = sizeout;
                                                if (pCodecCtx->coded_frame->pts != static_cast<int64_t>(AV_NOPTS_VALUE))
                                                                pkt.pts= av_rescale_q(pCodecCtx->coded_frame->pts, pCodecCtx->time_base, video_st->time_base);
                                                if(pCodecCtx->coded_frame->key_frame)
                                                                pkt.flags |= AV_PKT_FLAG_KEY;
                                                pkt.stream_index = video_st->index;
                                                /* Write the compressed frame to the media file. */
                                                ret = av_write_frame(pFormatCtx, &pkt);
                                }else{
                                                ret = 0;
                                }
#else
                                int got_packet;
                                ret = avcodec_encode_video2(pCodecCtx, &pkt, pictureOUT, &got_packet);
                                if(!ret && got_packet && pkt.size){
                                                if (pkt.pts != static_cast<int64_t>(AV_NOPTS_VALUE)) {
                                                                pkt.pts = av_rescale_q(pkt.pts,
                                                                                                pCodecCtx->time_base, video_st->time_base);
                                                }
                                                if (pkt.dts != static_cast<int64_t>(AV_NOPTS_VALUE)) {
                                                                pkt.dts = av_rescale_q(pkt.dts,
                                                                                                pCodecCtx->time_base, video_st->time_base);
                                                }

                                                pkt.stream_index = video_st->index;
                                                /* Write the compressed frame to the media file. */
                                                ret = av_write_frame(pFormatCtx, &pkt);
                                }else {
                                                ret = 0;
                                }
#endif
                                av_free_packet(&pkt);
                                if(buf)
                                                delete[] buf;
                }
                if (ret != 0) {
                                std::ostringstream err;
                                err << FILE_WRITE_ERROR << this->output_filename_ << " | Error while writing frame.";
                                slip::slip_exception exc(std::string("slip"), std::string("AvWriter::write()"), err.str());
                                throw (exc);
                                exit(1);
                }
                total_nb_slice++;
}


template<class ContainerVideo, typename T, std::size_t Nb_components>
inline
int AvWriter<ContainerVideo,T,Nb_components>::write(const container_type & data){
                if(!initialized){
                                std::ostringstream err;
                                err << FILE_WRITE_ERROR << this->output_filename_ << " | The writer to be initialized before writing.";
                                slip::slip_exception exc(std::string("slip"), std::string("AvWriter::write()"), err.str());
                                throw (exc);
                }
                if (finished){
                                return 0;
                }

                std::size_t nb_slices = data.slices();
                if(output_video_nbframes_ < (total_nb_slice + nb_slices)){
                                std::cerr << "output_video_nbframes_  = " <<  output_video_nbframes_ << std::endl;
                                std::cerr << "total_nb_slice  = " << total_nb_slice << std::endl;
                                std::cerr << "nb_slices  = " << nb_slices << std::endl;
                                std::ostringstream err;
                                err << FILE_WRITE_ERROR << this->output_filename_ << " | Allocation problem during writing.";
                                slip::slip_exception exc(std::string("slip"), std::string("AvWriter::write()"), err.str());
                                throw (exc);
                }

                /* encode the video */
                for(std::size_t i = 0; i < nb_slices; i++) {
                                write_video_frame(data,i);
                }

                if (total_nb_slice == output_video_nbframes_){
                                if(pFormatCtx->streams[video_stream_idx_]){
                                                double time_base_dble = static_cast<double>(pFormatCtx->streams[video_stream_idx_]->time_base.num) /
                                                                                static_cast<double>(pFormatCtx->streams[video_stream_idx_]->time_base.den);
                                                double video_pts = static_cast<double>(pFormatCtx->streams[video_stream_idx_]->pts.val) * time_base_dble;
                                                double duration = static_cast<double>(output_video_nbframes_) * time_base_dble;
                                                while(video_pts < duration){
                                                                write_video_frame(data,0);
                                                                video_pts = static_cast<double>(pFormatCtx->streams[video_stream_idx_]->pts.val) * time_base_dble ;
                                                }
                                }
                                av_write_trailer(pFormatCtx);
                                finished = true;
                }

                return (!finished);
}

template<class ContainerVideo, typename T, std::size_t Nb_components>
inline
void AvWriter<ContainerVideo,T,Nb_components>::release(){
                if(initialized){


                                av_free(pictureIN->data[0]);
                                av_free(pictureIN);
                                av_free(pictureOUT->data[0]);
                                av_free(pictureOUT);
                                if (!(pFormatCtx->oformat->flags & AVFMT_NOFILE))
                                                /* Close the output file. */
                                                avio_close(pFormatCtx->pb);

                                for (std::size_t i = 0; i < pFormatCtx->nb_streams; i++) {
                                                av_freep(&pFormatCtx->streams[i]->codec);
                                                av_freep(&pFormatCtx->streams[i]);
                                }

                                /* free the stream */
                                av_free(pFormatCtx);

                                video_stream_idx_ = -1;
                                codec_id_ = CODEC_ID_NONE;
                                dimension_optimize_ = false;
                                bit_rate_ = 0;
                                output_video_width_ = 0;
                                output_video_height_ = 0;
                                output_video_nbframes_ = 0;
                                time_base_ = (AVRational){0,1};
                                input_pxfmt_ = PIX_FMT_GRAY8;
                                output_pxfmt_ = PIX_FMT_RGB24;
                                row_stride = 0;
                                finished = false;
                                total_nb_slice = 0;
                                initialized = false;
                }
}




} 
#endif /* AVWRITER_HPP_ */