efmvs.cpp

  /*
 * Convert from OpenCV image and write movie with FFmpeg
 *
 * Copyright (c) 2018 Jishnu Jaykumar P
 */
#include <iostream>
#include <vector>
#include <cstring> //http://cs.stmarys.ca/~porter/csc/ref/c_cpp_strings.html

// FFmpeg
extern "C" {
  #include <libavutil/motion_vector.h>
// #include "../include/libavutil/motion_vector.h"
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>
#include <libavcodec/avcodec.h>
#include <libavutil/avutil.h>
#include <libavutil/pixdesc.h>
#include <libswscale/swscale.h>
#include <libavutil/imgutils.h>
#include <json/json.h>
#include <libavutil/time.h>
}

// OpenCV
#include <opencv2/opencv.hpp>

using namespace std;
using namespace cv;

class FMV{
private:
  AVFormatContext *pFormatCtx = NULL;
  // static AVCodecContext *video_dec_ctx = NULL;
  AVStream *video_stream = NULL;
  const char *src_filename = NULL;
  AVCodecContext    *pCodecCtxOrig = NULL;
  AVCodecContext    *pCodecCtx = NULL;
  AVCodec           *pCodec = NULL;

  int video_stream_idx = -1;
  AVFrame *pFrame = NULL;
  AVFrame *pFrameRGB = NULL;
  AVPacket *pkt=NULL;
  enum AVMediaType type;
  int video_frame_count = 0;
  int               frameFinished;
  int               numBytes;
  uint8_t           *buffer = NULL;
  struct SwsContext *sws_ctx = NULL;
  int argc;
  bool gotFrame;
  Mat cv_mat;
  json_object *mv_json_array;
  // compatibility with newer API
  #if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(55,28,1)
  #define av_frame_alloc avcodec_alloc_frame
  #define av_frame_free avcodec_free_frame
  #endif

public:
  FMV(int argc, char* videopath){
    // cout<<"line 50: FMV constructor called\n";
    this->argc = argc;
    this->src_filename=videopath;
  };

  const char* getSrcFilename(){
    return this->src_filename;
  }

  void setFrameFlag(bool flag){
    this->gotFrame=flag;
  }

  bool getFrameFlag(){
    return this->gotFrame;
  }

  json_object* getMVJSON(){
    return this->mv_json_array;
  }

  int64_t getPseudoRealTimeStamp(){
      this->pFormatCtx->start_time_realtime = av_gettime();
      return this->pFormatCtx->start_time_realtime;
  }

  int64_t getPseudoTimeStamp(){
      // this->pFormatCtx->start_time_realtime = av_gettime();
      return this->pFormatCtx->start_time;
  }

  void connect(bool dump_av_format){
    // cout<<"line 58: connect() called\n";
    int ret = 0;
    AVPacket pkt = { 0 };
    this->code_meta_data();
    if (argc != 2) {
        if (this->getSrcFilename()==NULL) {
          cout<< "Error: No input file given\n";
          exit(1);
        }else{
        fprintf(stderr, "Usage: %s <video>\n", this->getSrcFilename());
        exit(1);
      }
    }

    // src_filename = videopath;

    if (avformat_open_input(&(this->pFormatCtx), this->getSrcFilename(), NULL, NULL) < 0) {
        fprintf(stderr, "Error:Could not open source file %s\n", this->getSrcFilename());
        exit(1);
    }

    if (avformat_find_stream_info(this->pFormatCtx, NULL) < 0) {
        fprintf(stderr, "Could not find stream information\n");
        exit(1);
    }

    this->type = AVMEDIA_TYPE_VIDEO;

    open_codec_context();
    if(dump_av_format)
      av_dump_format(this->pFormatCtx, 0, this->src_filename, 0);

  }

  int open_codec_context()
  {
      int ret;
      AVStream *st;
      AVCodecContext *dec_ctx = NULL;
      AVCodec *dec = NULL;
      AVDictionary *opts = NULL;

      ret = av_find_best_stream(this->pFormatCtx, this->type, -1, -1, &dec, 0);
      if (ret < 0) {
          fprintf(stderr, "Could not find %s stream in input file '%s'\n",
                  av_get_media_type_string(this->type), this->getSrcFilename());
          return ret;
      } else {
          int stream_idx = ret;
          st = this->pFormatCtx->streams[stream_idx];

          dec_ctx = avcodec_alloc_context3(dec);
          if (!dec_ctx) {
              fprintf(stderr, "Failed to allocate codec\n");
              return AVERROR(EINVAL);
          }

          ret = avcodec_parameters_to_context(dec_ctx, st->codecpar);
          if (ret < 0) {
              fprintf(stderr, "Failed to copy codec parameters to codec context\n");
              return ret;
          }

          /* Init the video decoder */
          av_dict_set(&opts, "flags2", "+export_mvs", 0);
          if ((ret = avcodec_open2(dec_ctx, dec, &opts)) < 0) {
              fprintf(stderr, "Failed to open %s codec\n",
                      av_get_media_type_string(this->type));
              return ret;
          }

          this->video_stream_idx = stream_idx;
          this->video_stream = this->pFormatCtx->streams[this->video_stream_idx];
          this->pCodecCtx = dec_ctx;
      }

      return 0;
  }

  int decode_packet()
  {
    // cout<<"line 134: decode_packet() called\n";
    int ret = avcodec_send_packet(this->pCodecCtx, this->pkt);

    /*Creating a json array*/
    json_object *jarray;

    /*Creating a json object*/
    json_object * jobj;

    if (ret < 0) {
        fprintf(stderr, "Error while sending a packet to the decoder\n");
        return ret;
    }

    jarray = json_object_new_array();

    while (ret >= 0)  {
        ret = avcodec_receive_frame(this->pCodecCtx, this->pFrame);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            break;
        } else if (ret < 0) {
            fprintf(stderr, "Error while receiving a frame from the decoder: \n");
            return ret;
        }

        if (ret >= 0) {
            int i;
            AVFrameSideData *sd;
	          // this->video_frame_count++;
            this->increase_current_frame_number();
            sws_scale(this->sws_ctx, (uint8_t const * const *)this->pFrame->data,
                this->pFrame->linesize, 0, this->pCodecCtx->height,
                this->pFrameRGB->data, this->pFrameRGB->linesize);

            //convert frame to Mat
            this->cv_mat=Mat(this->pCodecCtx->height, this->pCodecCtx->width, CV_8UC3, this->pFrameRGB->data[0], this->pFrameRGB->linesize[0]);

            // imshow("frame", cv_mat);
            // waitKey(10);

            // cout<<cv_mat;

            // Save the frame to disk
            // this->SaveFrame(this->pFrameRGB, this->pCodecCtx->width, this->pCodecCtx->height,video_frame_count);


            sd = av_frame_get_side_data(this->pFrame, AV_FRAME_DATA_MOTION_VECTORS);
            if (sd) {
                const AVMotionVector *mvs = (const AVMotionVector *)sd->data;
                for (i = 0; i < sd->size / sizeof(*mvs); i++) {
                    const AVMotionVector *mv = &mvs[i];
                    jobj = json_object_new_object();
                    json_object_object_add(jobj,"source", json_object_new_int (mv->source));
                    json_object_object_add(jobj,"width", json_object_new_int (mv->w));
                    json_object_object_add(jobj,"height", json_object_new_int (mv->h));

                    if(mv->source<0){
                      json_object_object_add(jobj,"src_x", json_object_new_int ((mv->src_x)/abs(mv->source)));
                      json_object_object_add(jobj,"src_y", json_object_new_int ((mv->src_y)/abs(mv->source)));
                      json_object_object_add(jobj,"dst_x", json_object_new_int ((mv->dst_x/abs(mv->source))));
                      json_object_object_add(jobj,"dst_y", json_object_new_int ((mv->dst_y/abs(mv->source))));
                      json_object_object_add(jobj,"dx", json_object_new_int (((mv->dst_x - mv->src_x)/abs(mv->source))));
                      json_object_object_add(jobj,"dy", json_object_new_int (((mv->dst_y - mv->src_y)/abs(mv->source))));
                    }else{
                      json_object_object_add(jobj,"src_x", json_object_new_int ((mv->dst_x/abs(mv->source))));
                      json_object_object_add(jobj,"src_y", json_object_new_int ((mv->dst_y/abs(mv->source))));
                      json_object_object_add(jobj,"dst_x", json_object_new_int ((mv->src_x/abs(mv->source))));
                      json_object_object_add(jobj,"dst_y", json_object_new_int ((mv->src_y/abs(mv->source))));
                      json_object_object_add(jobj,"dx", json_object_new_int (((mv->src_x - mv->dst_x)/abs(mv->source))));
                      json_object_object_add(jobj,"dy", json_object_new_int (((mv->src_y - mv->dst_y)/abs(mv->source))));
                    }
                    /*Adding the above created json object to the array*/
                    json_object_array_add(jarray,jobj);
                    /*Now printing the json object*/
                    // printf ("\n\nThe json object created: %s\n",json_object_to_json_string(jarray));
                    jobj = NULL;
                }
                this->mv_json_array=jarray;
                jarray=NULL;
                // printf("%s\n", strcat( "./output/", strcat(video_frame_count, ".json")));
            }

	    printf("\rTotal Processed Frames:%d", this->video_frame_count);
	    fflush(stdout);
      // cout<<"Total Processed Frames:"<<this->video_frame_count<<endl;
            //Print frame data
            // print_frame_data(frame);
            av_frame_unref(this->pFrame);
        }
    }
    this->setFrameFlag(true);
    return this->video_frame_count;
    // return 0;
  }

  void code_meta_data(){
    printf("\n");
    printf("*******************************************************************************************\n");
    printf("*       Tool : MV-Tractus                                                                 *\n");
    printf("*     Author : Jishnu Jaykumar Padalunkal (https://jishnujayakumar.github.io)             *\n");
    printf("*  Used Libs : FFmpeg, OpenCV                                                             *\n");
    printf("*Description : A simple tool to extract motion vectors and frames from H264 video streams *\n");
    printf("*******************************************************************************************\n");
    printf("\n");
    printf("--------------------------------------------------------------------------------------\n");
    printf("framenum,source,blockw,blockh,srcx,srcy,dstx,dsty,motion_x,motion_y,motion_scale,flags\n");
    printf("--------------------------------------------------------------------------------------\n");
  }

  int extract_motion_vectors(){
      // cout<<"line 139: extract_motion_vectors() called\n";
      int ret = 0;
      AVPacket pkt = { 0 };
      // cout<<"1"<<endl;
      if (!this->video_stream) {
          fprintf(stderr, "Could not find video stream in the input, aborting\n");
          ret = 1;
          // cout<<"2"<<endl;
          goto end;
      }

      // cout<<"3"<<endl;
      this->pFrame = av_frame_alloc();
      if (!this->pFrame) {
          fprintf(stderr, "Could not allocate frame\n");
          ret = AVERROR(ENOMEM);
          // cout<<"4"<<endl;
          goto end;
      }

      // cout<<"5"<<endl;

      // Allocate an AVFrame structure
      this->pFrameRGB=av_frame_alloc();
      if(this->pFrameRGB==NULL)
        return -1;

        // Determine required buffer size and allocate buffer
        this->numBytes=avpicture_get_size(AV_PIX_FMT_BGR24, this->pCodecCtx->width,
      			      this->pCodecCtx->height);

        this->buffer=(uint8_t *)av_malloc(this->numBytes*sizeof(uint8_t));

      // Assign appropriate parts of buffer to image planes in pFrameRGB
      // Note that pFrameRGB is an AVFrame, but AVFrame is a superset
      // of AVPicture
      avpicture_fill((AVPicture *)this->pFrameRGB, this->buffer, AV_PIX_FMT_RGB24,
    		 this->pCodecCtx->width, this->pCodecCtx->height);
      // initialize SWS context for software scaling
      this->sws_ctx = sws_getContext(
             this->pCodecCtx->width,
    			   this->pCodecCtx->height,
    			   this->pCodecCtx->pix_fmt,
    			   this->pCodecCtx->width,
    			   this->pCodecCtx->height,
    			   AV_PIX_FMT_BGR24,
    			   SWS_BILINEAR,
    			   NULL,
    			   NULL,
    			   NULL
    			   );



      /* read frames from the file */
      while (av_read_frame(this->pFormatCtx, &pkt) >= 0) {

        // cout<<"6"<<endl;
          this->pkt = &pkt;
          if (this->pkt->stream_index == this->video_stream_idx){
                ret=this->decode_packet();
          }

          av_packet_unref(&pkt);
          av_packet_unref(this->pkt);
          if (ret < 0 || this->getFrameFlag()){
              av_freep(&(this->buffer));
              this->setFrameFlag(false);
              return this->video_frame_count;
              break;
            }
      }

      /* flush cached frames */
      // this->decode_packet(NULL);
      printf("\n--------------------------------------------------------------------------------------\n");

  end:
      // cout<<"RET "<<endl;
      av_frame_free(&(this->pFrame));
      return ret < 0;
  }

  /* flush cached frames */
  void flush_cached_frames(){
    this->pkt = NULL;
    this->decode_packet();
  }

  void freeStuff(){
    av_frame_free(&(this->pFrame));
    avcodec_free_context(&(this->pCodecCtx));
    avformat_close_input(&(this->pFormatCtx));
  }

  int get_current_frame_number(){
    return this->video_frame_count;
  }

  void increase_current_frame_number(){
    this->video_frame_count = this->video_frame_count%100;
    this->video_frame_count++;
  }

};

// Define C functions for the C++ class - as ctypes can only talk to C...

extern "C"
{
    FMV* FMV_new(int argc, char **argv) {return new FMV(argc, argv[1]);}

    const char* FMV_getSrcFilename(FMV* fmv) {return fmv->getSrcFilename();}

    void FMV_setFrameFlag(FMV* fmv, bool flag){fmv->setFrameFlag(flag);}

    bool FMV_getFrameFlag(FMV* fmv){return fmv->getFrameFlag();}

    void FMV_connect(FMV* fmv, bool dump_av_format){fmv->connect(dump_av_format);}

    int FMV_open_codec_context(FMV* fmv){return fmv->open_codec_context();}

    int FMV_decode_packet(FMV* fmv){return fmv->decode_packet();}

    int FMV_extract_motion_vectors(FMV* fmv){return fmv->extract_motion_vectors();}

    /* flush cached frames */
    void FMV_flush_cached_frames(FMV* fmv){fmv->flush_cached_frames();}

    void FMV_freeStuff(FMV* fmv){fmv->freeStuff();}

    int FMV_get_current_frame_number(FMV* fmv){return fmv->get_current_frame_number();}

    void FMV_increase_current_frame_number(FMV* fmv){fmv->increase_current_frame_number();}

}

int main(int argc, char **argv){
  // cout << "line 21: main() called\n";
  FMV* fmv = new FMV(argc, argv[1]);
  fmv->connect(true);// Give true to dump_av_format
  int i = fmv->get_current_frame_number();
  int frame_num=0;
  // while(i<10){
  while(frame_num==fmv->get_current_frame_number()){
    // if (i==408) {
      // cout<<endl<<"PRE-i:"<<i<<endl;
    //   break;
    // }

    frame_num=fmv->extract_motion_vectors();
    // cout<<"FN:"<<frame_num<<endl;
    // cout<<"HHHH"<<endl;
    // printf ("\n\nThe json object created: %s\n",json_object_to_json_string(fmv->getMVJSON()));
    i=fmv->get_current_frame_number();
    cout<<"\t|\t"<<fmv->getPseudoRealTimeStamp()<<endl;
    // cout<<"IIII\n";
    // cout<<endl<<"POST-i:"<<fmv->get_current_frame_number()<<endl;
    // break;
  }
  // fmv->flush_cached_frames();
  fmv->freeStuff();
  return 0;
}