ffconv/ffcpp/MediaFile.cpp

#include "MediaFile.h"
#include "ffcpp.h"
#include <stdexcept>

namespace ffcpp {

	MediaFile::MediaFile(const std::string& src, Mode mode) {
		_formatCtx = nullptr;
		_mode = mode;
		if(mode == Mode::Read) {
			int res = avformat_open_input(&_formatCtx, src.c_str(), nullptr, nullptr);
			throwIfError(res, "cannot open input file");

			res = avformat_find_stream_info(_formatCtx, nullptr);
			throwIfError(res, "cannot find stream info");

			_streams.reserve(_formatCtx->nb_streams);
			for(size_t i = 0; i < _formatCtx->nb_streams; ++i) {
				_streams.emplace_back(_formatCtx->streams[i]);
			}
		} else if(mode == Mode::Write) {
			int res = avformat_alloc_output_context2(&_formatCtx, nullptr, nullptr, src.c_str());
			throwIfError(res, "cannot allocate format context");

			if (!(_formatCtx->oformat->flags & AVFMT_NOFILE)) {
				res = avio_open(&_formatCtx->pb, src.c_str(), AVIO_FLAG_WRITE);
				throwIfError(res, "cannot init avio context");
			}

			if(_formatCtx->oformat->flags & AVFMT_GLOBALHEADER)
				_formatCtx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

			// Allocate enough space, so streams will never be destroyed
			// due to reallocation of internal buffer of std::vector
			_streams.reserve(100);

			av_dump_format(_formatCtx, 0, src.c_str(), 1);
		}
	}

	bool MediaFile::hasStream(AVMediaType type) const {
		for(size_t i = 0; i < _formatCtx->nb_streams; ++i) {
			if(_formatCtx->streams[i]->codec->codec_type == type) {
				return true;
			}
		}

		return false;
	}

	bool MediaFile::hasVideo() const {
		return hasStream(AVMEDIA_TYPE_VIDEO);
	}

	bool MediaFile::hasAudio() const {
		return hasStream(AVMEDIA_TYPE_AUDIO);
	}

	Stream& MediaFile::getStream(AVMediaType type, size_t index) {
		for(size_t i = 0, curIndex = 0; i < _formatCtx->nb_streams; ++i) {
			if(_formatCtx->streams[i]->codec->codec_type == type) {
				if(curIndex == index) {
					return _streams[i];
				} else {
					curIndex++;
				}
			}
		}

		throw std::runtime_error("cannot find stream");
	}

	Stream& MediaFile::videoStream(size_t index /* = 0 */) {
		return getStream(AVMEDIA_TYPE_VIDEO, index);
	}

	Stream& MediaFile::audioStream(size_t index /* = 0 */) {
		return getStream(AVMEDIA_TYPE_AUDIO, index);
	}

	MediaFile::operator AVFormatContext*() const {
		return _formatCtx;
	}

	MediaFile::~MediaFile() {
		_streams.clear();

		if(_mode == Mode::Write) {
			avformat_free_context(_formatCtx);
		} else {
			avformat_close_input(&_formatCtx);
		}
	}

	Stream& MediaFile::addVideoStream(AVCodecID codecID, int width, int height, AVPixelFormat pixelFormat) {
		AVCodec* codec = avcodec_find_encoder(codecID);
		if(!codec) throw std::runtime_error("cannot find codec");

		AVStream* stream = avformat_new_stream(_formatCtx, codec);
		if(!stream) throw  std::runtime_error("cannot create stream");

		AVCodecContext* ctx = stream->codec;
		ctx->width = width;
		ctx->height = height;

		if(pixelFormat == AV_PIX_FMT_NONE) {
			ctx->pix_fmt = codec->pix_fmts[0];
		} else {
			ctx->pix_fmt = pixelFormat;
		}

		_streams.emplace_back(stream, codec);
		return _streams.back();
	}

	Stream& MediaFile::addAudioStream(AVCodecID codecID, int channels, int sampleRate, AVSampleFormat sampleFormat) {
		AVCodec* codec = avcodec_find_encoder(codecID);
		if(!codec) throw std::runtime_error("cannot find codec");

		AVStream* stream = avformat_new_stream(_formatCtx, codec);
		if(!stream) throw  std::runtime_error("cannot create stream");

		// TODO: Here we need adjust encoder parameters
		AVCodecContext* ctx = stream->codec;
		if(sampleFormat == AV_SAMPLE_FMT_NONE) {
			ctx->sample_fmt = codec->sample_fmts[0];
		} else {
			ctx->sample_fmt = sampleFormat;
		}
		ctx->global_quality = 10;
		ctx->channels = channels;
		ctx->channel_layout = (uint64_t)av_get_default_channel_layout(channels);
		ctx->sample_rate = sampleRate;

		_streams.emplace_back(stream, codec);
		return _streams.back();
	}

	Packet MediaFile::readPacket() {
		AVPacket packet;
		packet.data = nullptr;
		packet.size = 0;
		int res = av_read_frame(_formatCtx, &packet);
		return Packet(packet);
	}

	AVMediaType MediaFile::packetType(const Packet &packet) {
		return _formatCtx->streams[packet.streamIndex()]->codec->codec_type;
	}

	void MediaFile::writeHeader() {
		int res = avformat_write_header(_formatCtx, nullptr);
		throwIfError(res, "error writing header");
	}

	void MediaFile::writeTrailer() {
		int res = av_write_trailer(_formatCtx);
		throwIfError(res, "error writing trailer");
	}

	void MediaFile::writePacket(Packet &packet) {
		int res = av_interleaved_write_frame(_formatCtx, packet);
		throwIfError(res, "cannot write frame to output file");
	}

}