converter.cpp

#include "Converter.h"
#include "InternetVideo.h"

Converter::Converter()
{
    av_register_all();
    avcodec_register_all();

    m_Converting = false;

    connect(this, SIGNAL(ConvertingFinished(InternetVideo*)), this, SLOT(onConvertingFinished(InternetVideo*)));
    m_FormatCtx = NULL;

    m_DecodedAudioFrame = nullptr;
    m_DecodedVideoFrame = nullptr;
}

void Converter::AddToQueue(InternetVideo* v)
{
    if (!v) return;
    m_Queue.enqueue(v);
}

void Converter::onConvertingFinished(InternetVideo *v)
{
    if (m_Queue.isEmpty())
    {
        emit QueueFinished();
        return;
    }

    Convert(m_Queue.dequeue());
}

void Converter::StartConverting()
{
    if (m_Converting) return;

    m_Converting = true;
    emit ConvertingFinished(nullptr);
}

void Converter::Convert(InternetVideo *v)
{
    if (!v)
    {
        //TODO, output error message
        emit ConvertingFinished(v);
        return;
    }

    QFile* file = v->GetDownloadedFile();
    file->close(); //close file just in case

    //QFileInfo fileInfo(file->fileName());


    QByteArray s = file->fileName().toUtf8();
    char *path = s.data();


    if (!OpenInput(path))
        printf("ERROR"); //TODO

    QString targetFilename = v->GetTitle() + "." + m_FileEnding;
    targetFilename = targetFilename.simplified();
    targetFilename.replace(QRegExp("[" + QRegExp::escape("\\/:*?\"<>|") + "]"), QString("_"));

    if (!OpenOutput(targetFilename))
        printf("ERROR");

    AVFrame* frame = avcodec_alloc_frame();
    float audio_pts;
    while (1)
    {
        AVPacket packet;
        av_init_packet(&packet);
        int rv = decode_frame(&packet, frame);
        if (rv < 0)
        {
            break;
        }
        if (audio_st)
            audio_pts = (double)audio_st->pts.val * audio_st->time_base.num / audio_st->time_base.den;
        else
            audio_pts = 0.0f;

        if (m_AudioCodecCtx)
        {
            if (got_frame)
            {
                write_audio_frame(frame, audio_dst_data, audio_dst_bufsize);
            }
        }
        if (audio_dst_data[0])
        {
            av_freep(&audio_dst_data[0]);
            audio_dst_data[0] = NULL;
        }

        av_free_packet(&packet);
    }
    write_delayed_frames(m_OutputFormatCtx, audio_st);
    av_write_trailer(m_OutputFormatCtx);

    if (audio_st)
        avcodec_close(audio_st->codec);
    if (video_st)
        avcodec_close(video_st->codec);

    avresample_close(m_avresampleCtx);
    avcodec_free_frame(&frame);
    emit ConvertingFinished(v);
}

void Converter::SetFormatSettings(FileFormat format, int audioQuality, int videoQuality)
{
    m_Format = format;
    m_audioOnly = true;
    switch (m_Format)
    {
    case FileFormat::mp4:
        m_VideoQuality = h264Qualities[videoQuality];
        m_AudioQuality = mp3Qualities[audioQuality];
        m_FileEnding = "mp4";
        m_audioOnly = false;
        break;
    case FileFormat::mov:
        m_VideoQuality = h264Qualities[videoQuality];
        m_AudioQuality = mp3Qualities[audioQuality];
        m_FileEnding = "mov";
        m_audioOnly = false;
        break;
    case FileFormat::MP3:
        m_AudioQuality = mp3Qualities[audioQuality];
        m_FileEnding = "mp3";
        break;
    case FileFormat::FLAC:
        m_VideoQuality = h264Qualities[videoQuality];
        m_FileEnding = "flac";
        break;
    case FileFormat::WAV:
        m_AudioQuality = losslessQualities[audioQuality];
        m_FileEnding = "wav";
        break;
    }
}

bool Converter::OpenInput(const char* path)
{
    /* open input file, and allocate format context */
    if (avformat_open_input(&m_FormatCtx, path, NULL, NULL) < 0)
    {
        fprintf(stderr, "Could not open source file %s\n", path);
        return false;
    }

    // Retrieve stream information
    if (avformat_find_stream_info(m_FormatCtx, NULL) < 0)
        return false; // Couldn't find stream information

    // Dump information about file onto standard error
    av_dump_format(m_FormatCtx, 0, path, 0);

    m_VideoStreamIndex = -1;
    m_AudioStreamIndex = -1;

    // Find the first video stream
    for (int i = 0; i < m_FormatCtx->nb_streams; i++)
    {
        if (m_FormatCtx->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO && m_VideoStreamIndex == -1)
        {
            m_VideoStreamIndex = i;
        }
        if (m_FormatCtx->streams[i]->codec->codec_type == AVMEDIA_TYPE_AUDIO && m_AudioStreamIndex == -1)
        {
            m_AudioStreamIndex = i;
        }
    }

    if (m_VideoStreamIndex == -1 || m_AudioStreamIndex == -1)
        return false;


    // Get a pointer to the codec context for the audio stream
    m_VideoCodecCtx = m_FormatCtx->streams[m_VideoStreamIndex]->codec;
    m_AudioCodecCtx = m_FormatCtx->streams[m_AudioStreamIndex]->codec;

    // Find the decoder for the video stream
    m_VideoCodec = avcodec_find_decoder(m_VideoCodecCtx->codec_id);
    if (m_VideoCodec == NULL) {
        fprintf(stderr, "Unsupported video codec!\n");
        return false; // Codec not found
    }
    // Find the decoder for the audio stream
    m_AudioCodec = avcodec_find_decoder(m_AudioCodecCtx->codec_id);
    if (m_AudioCodec == NULL) {
        fprintf(stderr, "Unsupported audio codec!\n");
        return false; // Codec not found
    }

    // Open video codec
    AVDictionary *codecDictOptions = NULL;
    if (avcodec_open2(m_VideoCodecCtx, m_VideoCodec, &codecDictOptions) < 0)
        return -1; // Could not open codec

    // Open audio codec
    if (avcodec_open2(m_AudioCodecCtx, m_AudioCodec, &codecDictOptions) < 0)
        return -1; // Could not open codec

    // Set up SWR context once you've got codec information
    m_avresampleCtx = avresample_alloc_context();
    av_opt_set_int(m_avresampleCtx, "in_channel_layout", m_AudioCodecCtx->channel_layout, 0);
    av_opt_set_int(m_avresampleCtx, "out_channel_layout", av_get_default_channel_layout(m_AudioCodecCtx->channels), 0);
    av_opt_set_int(m_avresampleCtx, "in_sample_rate", m_AudioCodecCtx->sample_rate, 0);
    av_opt_set_int(m_avresampleCtx, "out_sample_rate", m_AudioQuality.samplingRate, 0);
    av_opt_set_int(m_avresampleCtx, "in_sample_fmt", m_AudioCodecCtx->sample_fmt, 0);
    av_opt_set_int(m_avresampleCtx, "out_sample_fmt", AV_SAMPLE_FMT_S16P, 0);

    AV_SAMPLE_FMT_FLTP
    avresample_open(m_avresampleCtx);

    m_Audio_channels = m_AudioCodecCtx->channels;

    // Allocate audio frame
    if (m_DecodedAudioFrame == NULL) m_DecodedAudioFrame = avcodec_alloc_frame();
    int nb_planes = 0;
    AVStream* audio_stream = m_FormatCtx->streams[m_AudioStreamIndex];
    nb_planes = av_sample_fmt_is_planar(m_AudioCodecCtx->sample_fmt) ? m_AudioCodecCtx->channels : 1;
    int tempSize = sizeof(uint8_t *) * nb_planes;
    audio_dst_data = (uint8_t**)av_mallocz(tempSize);
    if (!audio_dst_data)
    {
        fprintf(stderr, "Could not allocate audio data buffers\n");
    }
    else
    {
        for (int i = 0; i < nb_planes; i++)
        {
            audio_dst_data[i] = NULL;
        }
    }

    if (m_DecodedVideoFrame == NULL) m_DecodedVideoFrame = avcodec_alloc_frame();
    // Determine required buffer size and allocate buffer
    int numBytes = avpicture_get_size(PIX_FMT_RGB24, m_VideoCodecCtx->width, m_VideoCodecCtx->height);
    video_dst_data = (uint8_t *)av_mallocz(numBytes*sizeof(uint8_t));

    return true;
}

bool Converter::OpenOutput(QString targetFilename)
{

    QByteArray fName = targetFilename.toLatin1();
    const char* fileName = fName.data();

    QByteArray tPath = m_TargetPath.toLatin1();
    const char* targetPath = tPath.data();

    QByteArray ext = m_FileEnding.toLatin1();

    /* allocate the output media context */
    AVOutputFormat *opfmt = av_guess_format(ext.data(), fileName, NULL);

    switch (m_Format)
    {
    case FileFormat::MP3:
    case FileFormat::FLAC:
    case FileFormat::WAV:
        opfmt->video_codec = AV_CODEC_ID_NONE;
        break;
    case FileFormat::mp4:
    case FileFormat::mov:
        opfmt->audio_codec = AV_CODEC_ID_MP3;
        opfmt->video_codec = AV_CODEC_ID_H264;
        break;
    }

    AVStream *out_stream;
    AVStream *in_stream;
    AVCodecContext *dec_ctx, *enc_ctx;

    m_OutputFormatCtx = avformat_alloc_context();
    if (!m_OutputFormatCtx) {
        fprintf(stderr, "Memory error\n");
        return false;
    }

    sprintf(m_OutputFormatCtx->filename, "%s/%s", targetPath, fileName);
    m_OutputFormatCtx->oformat = opfmt;

    // Add the audio stream using the default format codecs
    // and initialize the codecs.
    audio_st = NULL;
    video_st = NULL;
    if (opfmt->audio_codec != AV_CODEC_ID_NONE)
    {
        audio_st = add_audio_stream(m_OutputFormatCtx, &m_OutputAudioCodec, opfmt->audio_codec);
    }
    if (opfmt->video_codec != AV_CODEC_ID_NONE)
    {
        video_st = add_video_stream(m_OutputFormatCtx, &m_OutputVideoCodec, opfmt->video_codec);
    }

    // Now that all the parameters are set, we can open the audio and
    // video codecs and allocate the necessary encode buffers.
    if (audio_st)
    {
        int rv = open_audio(m_OutputFormatCtx, m_OutputAudioCodec, audio_st);
        if (rv < 0) return false;
    }
    if (video_st)
    {
        int rv = open_video(m_OutputFormatCtx, m_OutputVideoCodec, video_st);
        if (rv < 0) return false;
    }

    // open the output file, if needed
    if (!(opfmt->flags & AVFMT_NOFILE))
    {
        if (avio_open(&m_OutputFormatCtx->pb, m_OutputFormatCtx->filename, AVIO_FLAG_WRITE) < 0) {
            fprintf(stderr, "Could not open '%s'\n", m_OutputFormatCtx->filename);
            return false;
        }
        // Write the stream header, if any.

        int res = avformat_write_header(m_OutputFormatCtx, NULL);
        if (res < 0)
        {
            fprintf(stderr, "Error occurred when opening output file: %d\n", res);
            return false;
        }

    }

    return true;
}

static const char *get_error_text(const int error)
{
    static char error_buffer[255];
    av_strerror(error, error_buffer, sizeof(error_buffer));
    return error_buffer;
}

AVStream * Converter::add_audio_stream(AVFormatContext *oc, AVCodec **codec, enum AVCodecID codec_id)
{
    AVCodecContext *c;
    AVStream *st;

    /* find the audio encoder */
    *codec = avcodec_find_encoder(codec_id);
    if (!(*codec)) {
        fprintf(stderr, "Could not find codec\n");
        return nullptr;
    }

    st = avformat_new_stream(oc, *codec);
    if (!st) {
        fprintf(stderr, "Could not allocate stream\n");
        return nullptr;
    }

    //st->id = 0;

    c = st->codec;

    /* put sample parameters */
    c->sample_fmt = AV_SAMPLE_FMT_S16P;
    c->bit_rate = m_AudioQuality.bitrate;
    c->sample_rate = m_AudioQuality.samplingRate;
    c->channel_layout = av_get_default_channel_layout(m_AudioCodecCtx->channels);
    c->channels = m_AudioCodecCtx->channels;

    // some formats want stream headers to be separate
    if (oc->oformat->flags & AVFMT_GLOBALHEADER)
        c->flags |= CODEC_FLAG_GLOBAL_HEADER;

    return st;
}

AVStream * Converter::add_video_stream(AVFormatContext *oc, AVCodec **codec, enum AVCodecID codec_id)
{
    AVCodecContext *c;
    AVStream *st;

    /* find the audio encoder */
    *codec = avcodec_find_encoder(codec_id);
    if (!(*codec)) {
        fprintf(stderr, "Could not find codec\n");
        return nullptr;
    }

    st = avformat_new_stream(oc, *codec);
    if (!st) {
        fprintf(stderr, "Could not allocate stream\n");
        return nullptr;
    }

    c = st->codec;

    /* put sample parameters */
    c->width = m_VideoCodecCtx->width;
    c->height = m_VideoCodecCtx->height;
    c->sample_aspect_ratio = m_VideoCodecCtx->sample_aspect_ratio;
    c->pix_fmt = (*codec)->pix_fmts[0];
    c->time_base = m_VideoCodecCtx->time_base;
    c->bit_rate = m_VideoQuality.bitrate;
    c->sample_rate = m_VideoQuality.bitrate;

    // some formats want stream headers to be separate
    if (oc->oformat->flags & AVFMT_GLOBALHEADER)
        c->flags |= CODEC_FLAG_GLOBAL_HEADER;

    return st;
}

int Converter::open_audio(AVFormatContext *oc, AVCodec *codec, AVStream *st)
{
    int ret = 0;
    AVCodecContext *c;

    st->duration = m_FormatCtx->duration;
    c = st->codec;

    /* open it */
    ret = avcodec_open2(c, codec, NULL);
    if (ret < 0)
    {
        fprintf(stderr, "could not open codec\n");
        return false;
    }
    if (c->codec->capabilities & CODEC_CAP_VARIABLE_FRAME_SIZE)
        audio_input_frame_size = 10000;
    else
        audio_input_frame_size = c->frame_size;
    int tempSize = audio_input_frame_size * av_get_bytes_per_sample(c->sample_fmt) * c->channels;
    return ret;

}

int Converter::open_video(AVFormatContext *oc, AVCodec *codec, AVStream *st)
{
    int ret = 0;
    AVCodecContext *c;

    st->duration = m_FormatCtx->duration;
    c = st->codec;

    /* open it */
    ret = avcodec_open2(c, codec, NULL);
    if (ret < 0)
    {
        fprintf(stderr, "could not open codec\n");
        return false;
    }

    return ret;
}

int Converter::decode_frame(AVPacket* packet, AVFrame* outputFrame)
{
    int rv = 0;
    got_frame = 0;
    if (m_FormatCtx == NULL)
    {
        return rv;
    }

    int ret = 0;
    audiobufsize = 0;
    rv = av_read_frame(m_FormatCtx, packet);
    if (rv < 0)
    {
        return rv;
    }

    rv = decode_packet(packet, outputFrame);
    return rv;
}

int Converter::decode_packet(AVPacket* packet, AVFrame* outputFrame)
{
    int rv = 0;

    //audio stream?
    if (packet->stream_index == m_AudioStreamIndex)
    {
        /* decode audio frame */
        rv = avcodec_decode_audio4(m_AudioCodecCtx, m_DecodedAudioFrame, &got_frame, packet);
        if (rv < 0)
        {
            fprintf(stderr, "Error decoding audio frame\n");
            //return ret;
        }
        else
        {
            if (got_frame)
            {
                if (audio_dst_data[0] == NULL)
                {
                    rv = av_samples_alloc(audio_dst_data, &audio_dst_linesize, m_AudioCodecCtx->channels, m_DecodedAudioFrame->nb_samples, (AVSampleFormat)m_DecodedAudioFrame->format, 1);
                    if (rv < 0)
                    {
                        fprintf(stderr, "Could not allocate audio buffer\n");
                        return AVERROR(ENOMEM);
                    }

                    /* TODO: extend return code of the av_samples_* functions so that this call is not needed */
                    audio_dst_bufsize = av_samples_get_buffer_size(NULL, audio_st->codec->channels,
                        m_DecodedAudioFrame->nb_samples, (AVSampleFormat)m_DecodedAudioFrame->format, 1);

                    //int16_t* outputBuffer = ...;
                    avresample_convert(m_avresampleCtx, audio_dst_data, 0, outputFrame->nb_samples, m_DecodedAudioFrame->extended_data, 0, m_DecodedAudioFrame->nb_samples);
                }
            }
        }
    }

    return rv;
}

void Converter::write_audio_frame(AVFrame *frame, uint8_t ** audio_src_data, int audio_src_bufsize)
{
    AVFormatContext *oc = m_OutputFormatCtx;
    AVStream *st = audio_st;
    if (oc == NULL || st == NULL) return;
    AVCodecContext *c;
    AVPacket pkt = { 0 }; // data and size must be 0;
    int got_packet;

    av_init_packet(&pkt);
    c = st->codec;

    frame->nb_samples = audio_input_frame_size;
    int buf_size = audio_src_bufsize * av_get_bytes_per_sample(c->sample_fmt) * c->channels;
    avcodec_fill_audio_frame(frame, c->channels, c->sample_fmt, *audio_src_data, buf_size, 1);
    avcodec_encode_audio2(c, &pkt, frame, &got_packet);

    if (got_packet)
    {
        if (pkt.pts != AV_NOPTS_VALUE)
            pkt.pts = av_rescale_q(pkt.pts, st->codec->time_base, st->time_base);
        if (pkt.dts != AV_NOPTS_VALUE)
            pkt.dts = av_rescale_q(pkt.dts, st->codec->time_base, st->time_base);
        if (c && c->coded_frame && c->coded_frame->key_frame)
            pkt.flags |= AV_PKT_FLAG_KEY;

        pkt.stream_index = st->index;
        pkt.flags |= AV_PKT_FLAG_KEY;

        /* Write the compressed frame to the media file. */
        if (av_interleaved_write_frame(oc, &pkt) != 0)
        {
            fprintf(stderr, "Error while writing audio frame\n");
            return;
        }
    }

    av_free_packet(&pkt);
}

void Converter::write_delayed_frames(AVFormatContext *oc, AVStream *st)
{
    AVCodecContext *c = st->codec;
    int got_output = 0;
    int ret = 0;
    AVPacket pkt;
    pkt.data = NULL;
    pkt.size = 0;
    av_init_packet(&pkt);
    int i = 0;

    for (got_output = 1; got_output; i++)
    {
        ret = avcodec_encode_audio2(c, &pkt, NULL, &got_output);
        if (ret < 0)
        {
            fprintf(stderr, "error encoding frame\n");
            exit(1);
        }
        static int64_t tempPts = 0;
        static int64_t tempDts = 0;
        /* If size is zero, it means the image was buffered. */
        if (got_output)
        {
            if (pkt.pts != AV_NOPTS_VALUE)
                pkt.pts = av_rescale_q(pkt.pts, st->codec->time_base, st->time_base);
            if (pkt.dts != AV_NOPTS_VALUE)
                pkt.dts = av_rescale_q(pkt.dts, st->codec->time_base, st->time_base);
            if (c && c->coded_frame && c->coded_frame->key_frame)
                pkt.flags |= AV_PKT_FLAG_KEY;

            pkt.stream_index = st->index;
            /* Write the compressed frame to the media file. */
            ret = av_interleaved_write_frame(oc, &pkt);
        }
        else
        {
            ret = 0;
        }
        av_free_packet(&pkt);
    }
}