sjero.net Git - qpsnr/blob - src/qav.cpp

   1 /*
   2 *       qpsnr (C) 2010 E. Oriani, ema <AT> fastwebnet <DOT> it
   3 *
   4 *       This file is part of qpsnr.
   5 *
   6 *       qpsnr is free software: you can redistribute it and/or modify
   7 *       it under the terms of the GNU General Public License as published by
   8 *       the Free Software Foundation, either version 3 of the License, or
   9 *       (at your option) any later version.
  10 *
  11 *       qpsnr is distributed in the hope that it will be useful,
  12 *       but WITHOUT ANY WARRANTY; without even the implied warranty of
  13 *       MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  14 *       GNU General Public License for more details.
  15 *
  16 *       You should have received a copy of the GNU General Public License
  17 *       along with qpsnr.  If not, see <http://www.gnu.org/licenses/>.
  18 */
  19
  20 #include "qav.h"
  21 #include "settings.h"
  22 #include <stdexcept>
  23 #include <sstream>
  24
  25 qav::qvideo::qvideo(const char* file, int _out_width, int _out_height) : frnum(0), videoStream(-1), out_width(_out_width), out_height(_out_height) {
  26         const char* pslash = strrchr(file, '/');
  27         if (pslash) fname = pslash+1;
  28         else fname = file;
  29         pFormatCtx=NULL;
  30         if (avformat_open_input(&pFormatCtx, file, NULL, NULL)!=0)
  31                 throw std::runtime_error("Can't open file");
  32         if (av_find_stream_info(pFormatCtx)<0) {
  33                 av_close_input_file(pFormatCtx);
  34                 throw std::runtime_error("Multimedia type not supported");
  35         }
  36         LOG_INFO << "File info for (" << file << ")..." << std::endl;
  37         av_dump_format(pFormatCtx, 0, file, false);
  38         // find video stream (first)
  39         for (int i=0; i<pFormatCtx->nb_streams; i++)
  40                 if (AVMEDIA_TYPE_VIDEO == pFormatCtx->streams[i]->codec->codec_type) {
  41                         videoStream=i;
  42                         break;
  43                 }
  44         if (-1==videoStream) {
  45                 av_close_input_file(pFormatCtx);
  46                 throw std::runtime_error("Can't find video stream");
  47         }
  48         // Get a pointer to the codec context for the video stream
  49         pCodecCtx=pFormatCtx->streams[videoStream]->codec;
  50         pCodec=avcodec_find_decoder(pCodecCtx->codec_id);
  51         if(!pCodec) {
  52                 av_close_input_file(pFormatCtx);
  53                 throw std::runtime_error("Can't find codec for video stream");
  54         }
  55         if(avcodec_open(pCodecCtx, pCodec)<0) {
  56                 av_close_input_file(pFormatCtx);
  57                 throw std::runtime_error("Can't open codec for video stream");
  58         }
  59         // alloacate data to extract frames
  60         pFrame = avcodec_alloc_frame();
  61         if (!pFrame) {
  62                 avcodec_close(pCodecCtx);
  63                 av_close_input_file(pFormatCtx);
  64                 throw std::runtime_error("Can't allocated frame for video stream");
  65         }
  66         // populate the out_width/out_height members
  67         if (out_width > 0 && out_height > 0) {
  68                 LOG_INFO << "Output frame size for (" << file << ") is: " << out_width << 'x' << out_height << std::endl;
  69         } else if (-1 == out_width && -1 == out_height) {
  70                 out_width = pCodecCtx->width;
  71                 out_height = pCodecCtx->height;
  72                 LOG_INFO << "Output frame size for (" << file << ") (default) is: " << out_width << 'x' << out_height << std::endl;
  73         } else {
  74                 avcodec_close(pCodecCtx);
  75                 av_close_input_file(pFormatCtx);
  76                 throw std::runtime_error("Invalid output frame size for video stream");
  77         }
  78         // just report if we're using a different video size
  79         if (out_width!=pCodecCtx->width || out_height!=pCodecCtx->height)
  80                 LOG_WARNING << "Video (" << file <<") will get scaled: " << pCodecCtx->width << 'x' << pCodecCtx->height << " (in), " << out_width << 'x' << out_height << " (out)" << std::endl;
  81         // sw context
  82         img_convert_ctx = sws_getContext(pCodecCtx->width, pCodecCtx->height, pCodecCtx->pix_fmt,
  83                                         out_width, out_height, PIX_FMT_RGB24, SWS_BICUBIC, NULL, NULL, NULL);
  84         if (!img_convert_ctx) {
  85                 av_free(pFrame);
  86                 avcodec_close(pCodecCtx);
  87                 av_close_input_file(pFormatCtx);
  88                 throw std::runtime_error("Can't allocated sw_scale context");
  89         }
  90 }
  91
  92 qav::scr_size qav::qvideo::get_size(void) const {
  93         return scr_size(out_width, out_height);
  94 }
  95
  96 int qav::qvideo::get_fps_k(void) const {
  97         if (pFormatCtx->streams[videoStream]->r_frame_rate.den)
  98                 return 1000*pFormatCtx->streams[videoStream]->r_frame_rate.num/pFormatCtx->streams[videoStream]->r_frame_rate.den;
  99         return 0;
 100 }
 101
 102 bool qav::qvideo::get_frame(std::vector<unsigned char>& out, int *_frnum) {
 103         out.resize(avpicture_get_size(PIX_FMT_RGB24, out_width, out_height));
 104         AVPacket        packet;
 105         bool            is_read = false;
 106         av_init_packet(&packet);
 107         while (av_read_frame(pFormatCtx, &packet)>=0) {
 108                 if (packet.stream_index==videoStream) {
 109                         int frameFinished = 0;
 110                         // Decode video frame
 111                         //avcodec_decode_video(pCodecCtx, pFrame, &frameFinished, packet.data, packet.size);
 112                         avcodec_decode_video2(pCodecCtx, pFrame, &frameFinished, &packet); // packet.data, packet.size);
 113                         if(frameFinished) {
 114                                 AVPicture picRGB;
 115                                 // Assign appropriate parts of buffer to image planes in pFrameRGB
 116                                 avpicture_fill((AVPicture*)&picRGB, (unsigned char*)&out[0], PIX_FMT_RGB24, out_width, out_height);
 117                                 // Convert the image from its native format to RGB
 118                                 //img_convert((AVPicture*)&picRGB, PIX_FMT_RGB24, (AVPicture*)pFrame, pCodecCtx->pix_fmt, pCodecCtx->width, pCodecCtx->height);
 119                                 sws_scale(img_convert_ctx, pFrame->data, pFrame->linesize, 0, pCodecCtx->height, picRGB.data, picRGB.linesize);
 120                                 ++frnum;
 121                                 if (_frnum) *_frnum = frnum;
 122                                 if (settings::SAVE_IMAGES) {
 123                                         if ((settings::SKIP_FRAMES < frnum) && (-1 == settings::MAX_FRAMES || frnum <= settings::MAX_FRAMES))
 124                                                 save_frame(&out[0]);
 125                                 }
 126                                 is_read=true;
 127                         }
 128                 }
 129                 av_free_packet(&packet);
 130                 if (is_read) return true;
 131         }
 132         return false;
 133 }
 134
 135 /*bool SaveTGA(char *name, const unsigned char *data, int sizeX, int sizeY) {
 136         BYTE    TGAheader[12]={0,0,2,0,0,0,0,0,0,0,0,0};
 137         BYTE    header[6];
 138         header[1] = (BYTE)(sizeX/256);
 139         header[0] = (BYTE) sizeX%256;
 140         header[3] = (BYTE)(sizeY/256);
 141         header[2] = (BYTE) sizeY%256;
 142         header[4] = (BYTE) 24;
 143         header[5] = (BYTE) 0x00;
 144         int fh = open(name, _O_CREAT|_O_TRUNC|_O_WRONLY);
 145         if (-1 == fh) return false;
 146         if (12*sizeof(BYTE) != write(fh, TGAheader, 12*sizeof(BYTE))) return false;
 147         if (6*sizeof(BYTE) != write(fh, header, 6*sizeof(BYTE))) return false;
 148         for (int i = 0; i < sizeY; i++) {
 149                 if (3*sizeX != write(fh, data + sizeX*i*3, sizeX*3)) {
 150                         close(fh);
 151                         return false;
 152                 }
 153         }
 154         close(fh);
 155         return true;
 156 }*/
 157
 158 void qav::qvideo::save_frame(const unsigned char *buf, const char* __fname) {
 159         FILE            *pFile;
 160         std::string     s_fname;
 161         char            num_buf[32];
 162
 163         //
 164         sprintf(num_buf, ".%08d.ppm", frnum);
 165         num_buf[31] = '\0';
 166         std::ostringstream oss;
 167         oss << ((__fname) ? __fname : fname.c_str()) << num_buf;
 168         // Open file
 169         pFile=fopen(oss.str().c_str(), "wb");
 170         if(pFile==NULL)
 171                 return;
 172
 173         // Write header
 174         fprintf(pFile, "P6\n%d %d\n255\n", out_width, out_height);
 175
 176         // Write pixel data
 177         for(int y=0; y<out_height; y++)
 178                 fwrite(buf+y*out_width*3, 1, out_width*3, pFile);
 179
 180         // Close file
 181         fclose(pFile);
 182 }
 183
 184 qav::qvideo::~qvideo() {
 185         sws_freeContext(img_convert_ctx);
 186         av_free(pFrame);
 187         avcodec_close(pCodecCtx);
 188         av_close_input_file(pFormatCtx);
 189 }
 190