ffmpegpp/ffmpegpp.cpp

#include "ffmpegpp.hpp"

extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>
}

#include <iostream>
#include <fstream>

namespace av {

codec find_encoder(codec_id id) {
    auto ptr = avcodec_find_encoder(static_cast<AVCodecID>(id));
    if(!ptr) throw error("Could not find codec");
    return {ptr};
}


frame::frame()
    : ptr(av_frame_alloc()) {}

void frame::deleter::operator()(AVFrame *p) {
    av_frame_free(&p);
}


packet::packet()
    : ptr(av_packet_alloc()) {}

void packet::deleter::operator()(AVPacket *p) {
    av_packet_free(&p);
}


format_context::format_context() : ptr(avformat_alloc_context()) {}

format_context::format_context(const std::string &url) {
    AVFormatContext * p = nullptr;
    avformat_open_input(&p, url.c_str(), nullptr, nullptr);
    // This might be optional, but AFAIK people always do this after opening input
    avformat_find_stream_info(p, nullptr);
    ptr.reset(p);
}

codec_context format_context::context_from_stream(int i) const {
    (void)ptr->streams[i]->codecpar;
    return {};
}

void format_context::deleter::operator()(AVFormatContext *p) {
    avformat_free_context(p);
}


codec_context::codec_context()
    : ptr(avcodec_alloc_context3(nullptr)) {}

codec_context::codec_context(codec const & c)
    : ptr(avcodec_alloc_context3(c.ptr)) {}

void codec_context::open(const codec &c) {
    if (avcodec_open2(ptr.get(), c.ptr, nullptr) < 0) {
        throw error("Open coded failed");
    }
}

void codec_context::send_frame(const frame & f) {
    if (avcodec_send_frame(ptr.get(), f.ptr.get()) < 0) {
        throw error("Could not send frame");
    }
}

int codec_context::receive_packet(packet & p) {
    auto ret = avcodec_receive_packet(ptr.get(), p.ptr.get());
    if (ret >= 0 || ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
        return ret;
    if (ret < 0)
        throw error("Could not receive packet");
    return ret;
}

void codec_context::send_packet(const packet & p) {
    if (avcodec_send_packet(ptr.get(), p.ptr.get()) < 0)
        throw error("send packet error");
}

int codec_context::receive_frame(frame & f) {
    auto ret = avcodec_receive_frame(ptr.get(), f.ptr.get());
    if (ret >= 0 || ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
        return ret;
    if (ret < 0)
        throw error("could not receive frame");
    return ret;
}

void codec_context::set_parameters(const AVCodecParameters * par) {
    if (avcodec_parameters_to_context(ptr.get(), par) < 0)
        throw error("error with parameters");
}

void codec_context::deleter::operator()(AVCodecContext *p) {
    avcodec_free_context(&p);
}


void encode_as_jpg(const frame & frame, const std::string & filename){
    const auto codec = av::find_encoder(AV_CODEC_ID_MJPEG);

    codec_context codec_ctx(codec);
    codec_ctx.ptr->pix_fmt = static_cast<AVPixelFormat>(frame.ptr->format);
    codec_ctx.ptr->width = frame.ptr->width;
    codec_ctx.ptr->height = frame.ptr->height;
    codec_ctx.ptr->time_base = av_make_q(1, 1);

    codec_ctx.open(codec);
    codec_ctx.send_frame(frame);

    packet p;
    std::ofstream file(filename);

    int ret = 0;
    while (ret >= 0) {
        ret = codec_ctx.receive_packet(p);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
            return;

        file.write(reinterpret_cast<char*>(p.ptr->data), p.ptr->size);
        av_packet_unref(p.ptr.get());
    }
}

frame open_image(std::string const & filename){
    format_context format_ctx(filename);

    AVCodec * ptr;
    const auto idx = av_find_best_stream(format_ctx.ptr.get(), AVMEDIA_TYPE_VIDEO, -1, -1, &ptr, 0);
    if (idx < 0)
        std::clog << "error finding best stream" << std::endl;

    const auto st = format_ctx.ptr->streams[idx];
    const auto dc = avcodec_find_decoder(st->codecpar->codec_id);

    if (!dc)
        std::clog << "error finding decoder" << std::endl;

    codec_context codec_ctx({dc});
    codec_ctx.set_parameters(st->codecpar);
    codec_ctx.open({dc});

    packet p;

    // TODO: parse
    av_parser_parse2();

    codec_ctx.send_packet(p);

    frame f;
    codec_ctx.receive_frame(f);
    return f;
}

}