/**

* Smoke sample for the HW encoding with current AvCpp API state.

*

* Thanks to https://github.com/mojie126, for more-less complete sample.

*

*/

#include <iostream>

#include <set>

#include <map>

#include <memory>

#include <functional>

#include "avcpp/av.h"

#include "avcpp/ffmpeg.h"

#include "avcpp/codec.h"

#include "avcpp/packet.h"

#include "avcpp/videorescaler.h"

#include "avcpp/audioresampler.h"

#include "avcpp/avutils.h"

// API2

#include "avcpp/format.h"

#include "avcpp/formatcontext.h"

#include "avcpp/codec.h"

#include "avcpp/codeccontext.h"

using namespace std;

using namespace av;

// check CUDA accessibility

bool isCudaAvailable()

{

AVHWDeviceType type = av_hwdevice_find_type_by_name("cuda");

return type != AV_HWDEVICE_TYPE_NONE;

}

// check QSV accessibility

bool isQsvAvailable()

{

AVHWDeviceType type = av_hwdevice_find_type_by_name("qsv");

return type != AV_HWDEVICE_TYPE_NONE;

}

// Get ecnoder

const av::Codec getEncoder(bool use_cuda, bool use_qsv)

{

if (use_cuda) {

auto enc = av::findEncodingCodec("h264_nvenc");

if (!enc.isNull())

return enc;

}

if (use_qsv) {

auto enc = av::findEncodingCodec("h264_qsv");

if (!enc.isNull())

return enc;

}

return av::findEncodingCodec("libx264"); // Fall-back to SW one

}

// HW context

AVHWDeviceType getHwDeviceType(bool use_cuda, bool use_qsv)

{

if (use_cuda)

return av_hwdevice_find_type_by_name("cuda");

if (use_qsv)

return av_hwdevice_find_type_by_name("qsv");

return AV_HWDEVICE_TYPE_NONE;

}

int main(int argc, char **argv) {

if (argc < 3)

return 1;

av::init();

av::setFFmpegLoggingLevel(AV_LOG_DEBUG);

string uri{argv[1]};

string out{argv[2]};

error_code ec;

//

// INPUT

//

FormatContext ictx;

ssize_t videoStream = -1;

// VideoDecoderContext vdec;

Stream vst;

// int count = 0;

ictx.openInput(uri, ec);

if (ec) {

cerr << "Can't open input\n";

return 1;

}

ictx.findStreamInfo();

for (size_t i = 0; i < ictx.streamsCount(); ++i) {

auto st = ictx.stream(i);

if (st.mediaType() == AVMEDIA_TYPE_VIDEO) {

videoStream = i;

vst = st;

break;

}

}

bool use_cuda = isCudaAvailable();

bool use_qsv = isQsvAvailable();

// Get encoder and hardware context

auto _encoder = getEncoder(use_cuda, use_qsv);

AVHWDeviceType hw_type = getHwDeviceType(use_cuda, use_qsv); // --htrd: never use

clog << "test hw: " << av_hwdevice_get_type_name(hw_type) << endl;

//av::VideoDecoderContext vdec(vst, _encoder);

av::VideoDecoderContext vdec;

if (vst.isNull()) {

cerr << "Video stream not found\n";

return 1;

}

if (vst.isValid()) {

vdec = VideoDecoderContext(vst);

vdec.setRefCountedFrames(true);

vdec.open(Codec(), ec);

if (ec) {

cerr << "Can't open codec\n";

return 1;

}

}

//

// OUTPUT

//

OutputFormat ofrmt;

FormatContext octx;

ofrmt.setFormat(string(), out);

octx.setFormat(ofrmt);

//Codec ocodec = findEncodingCodec(ofrmt);

VideoEncoderContext encoder{_encoder};

// Settings

encoder.setWidth(vdec.width());

encoder.setHeight(vdec.height());

if (vdec.pixelFormat() > -1)

encoder.setPixelFormat(vdec.pixelFormat());

encoder.setTimeBase(Rational{1, 1000});

encoder.setBitRate(vdec.bitRate());

// MKV format wants global header, so

// this fixes the issue with MKV container

if (ofrmt.flags() & AVFMT_GLOBALHEADER)

encoder.addFlags(AV_CODEC_FLAG_GLOBAL_HEADER);

encoder.open(Codec(), ec);

if (ec) {

cerr << "Can't opent encodec\n";

return 1;

}

Stream ost = octx.addStream(encoder);

ost.setFrameRate(vst.frameRate());

ost.setTimeBase(vst.timeBase());

ost.setAverageFrameRate(vst.averageFrameRate());

octx.openOutput(out, ec);

if (ec) {

cerr << "Can't open output\n";

return 1;

}

octx.dump();

octx.writeHeader();

octx.flush();

//

// PROCESS

//

for (int64_t opkt_index = 0;;) {

// READING

Packet pkt = ictx.readPacket(ec);

if (ec) {

clog << "Packet reading error: " << ec << ", " << ec.message() << endl;

break;

}

bool flushDecoder = false;

// !EOF

if (pkt) {

if (pkt.streamIndex() != videoStream) {

continue;

}

clog << "Read packet: pts=" << pkt.pts() << ", dts=" << pkt.dts() << " / " << pkt.pts().seconds() << " / " << pkt.timeBase() << " / st: " << pkt.streamIndex() << endl;

} else {

flushDecoder = true;

}

do {

// DECODING

auto frame = vdec.decode(pkt, ec);

//count++;

//if (count > 200)

// break;

bool flushEncoder = false;

if (ec) {

cerr << "Decoding error: " << ec << endl;

return 1;

} else if (!frame) {

//cerr << "Empty frame\n";

//flushDecoder = false;

//continue;

if (flushDecoder) {

flushEncoder = true;

}

}

if (frame) {

clog << "Frame: pts=" << frame.pts() << " / " << frame.pts().seconds() << " / " << frame.timeBase() << ", " << frame.width() << "x" << frame.height() << ", size=" << frame.size() << ", ref=" << frame.isReferenced() << ":" << frame.refCount() << " / type: " << frame.pictureType() << endl;

// Change timebase

frame.setTimeBase(encoder.timeBase());

frame.setStreamIndex(0);

frame.setPictureType();

clog << "Frame: pts=" << frame.pts() << " / " << frame.pts().seconds() << " / " << frame.timeBase() << ", " << frame.width() << "x" << frame.height() << ", size=" << frame.size() << ", ref=" << frame.isReferenced() << ":" << frame.refCount() << " / type: " << frame.pictureType() << endl;

}

if (frame || flushEncoder) {

do {

// Encode

Packet opkt = frame ? encoder.encode(frame, ec) : encoder.encode(ec);

if (ec) {

cerr << "Encoding error: " << ec << endl;

return 1;

} else if (!opkt) {

//cerr << "Empty packet\n";

//continue;

break;

}

// Only one output stream

opkt.setStreamIndex(0);

opkt.setTimeBase(1 / ost.frameRate());

// Trick: when timeBase is a 1/FPS, PTS value - just an index of the packet

// CUDA encoder provides wrong PTS

opkt.setPts({opkt_index++, 1 / ost.frameRate()});

opkt.setDts(opkt.pts());

clog << "Write packet: pts=" << opkt.pts() << ", dts=" << opkt.dts() << " / " << opkt.pts().seconds() << " / " << opkt.timeBase() << " / st: " << opkt.streamIndex() << endl;

octx.writePacket(opkt, ec);

if (ec) {

cerr << "Error write packet: " << ec << ", " << ec.message() << endl;

return 1;

}

} while (flushEncoder);

}

if (flushEncoder)

break;

} while (flushDecoder);

if (flushDecoder)

break;

}

octx.writeTrailer();

}