QtPotato/src/subtitles/subtitleextractor.cpp

#include "subtitleextractor.h"

#include "trackselection.h"

#include <KLocalizedString>

#include <memory>

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

namespace
{
struct FormatContextDeleter {
    void operator()(AVFormatContext *context) const
    {
        avformat_close_input(&context);
    }
};
struct CodecContextDeleter {
    void operator()(AVCodecContext *context) const
    {
        avcodec_free_context(&context);
    }
};
struct PacketDeleter {
    void operator()(AVPacket *packet) const
    {
        av_packet_free(&packet);
    }
};

QString ffmpegError(int code)
{
    char buffer[AV_ERROR_MAX_STRING_SIZE] = {};
    av_strerror(code, buffer, sizeof(buffer));
    return QString::fromUtf8(buffer);
}

qint64 toMs(int64_t value, AVRational timeBase)
{
    return av_rescale_q(value, timeBase, AVRational{1, 1000});
}

// ASS events from the decoder look like "ReadOrder,Layer,Style,Name,MarginL,MarginR,MarginV,Effect,Text".
QString assRectText(const char *ass)
{
    const QString event = QString::fromUtf8(ass);
    qsizetype from = 0;
    for (int field = 0; field < 8; ++field) {
        from = event.indexOf(u',', from);
        if (from < 0) {
            return {};
        }
        ++from;
    }
    return Subtitles::assEventText(QStringView(event).mid(from));
}
}

namespace Subtitles
{
ExtractResult extractEmbedded(const QString &path, int streamIndex, const std::atomic_bool &cancelled)
{
    AVFormatContext *rawFormat = nullptr;
    if (const int ret = avformat_open_input(&rawFormat, path.toUtf8().constData(), nullptr, nullptr); ret < 0) {
        return {{}, i18n("Could not open %1: %2", path, ffmpegError(ret))};
    }
    std::unique_ptr<AVFormatContext, FormatContextDeleter> format(rawFormat);

    if (streamIndex < 0 || streamIndex >= static_cast<int>(format->nb_streams)) {
        return {{}, i18n("Subtitle stream %1 does not exist in this file.", streamIndex)};
    }
    AVStream *stream = format->streams[streamIndex];

    const AVCodecDescriptor *descriptor = avcodec_descriptor_get(stream->codecpar->codec_id);
    if (stream->codecpar->codec_type != AVMEDIA_TYPE_SUBTITLE || !descriptor) {
        return {{}, i18n("Stream %1 is not a subtitle stream.", streamIndex)};
    }
    if (descriptor->props & AV_CODEC_PROP_BITMAP_SUB) {
        return {{}, i18n("This subtitle track is made of images (%1) and can't be translated.", QString::fromUtf8(descriptor->name))};
    }

    const AVCodec *codec = avcodec_find_decoder(stream->codecpar->codec_id);
    if (!codec) {
        return {{}, i18n("No decoder for subtitle format %1.", QString::fromUtf8(descriptor->name))};
    }
    std::unique_ptr<AVCodecContext, CodecContextDeleter> decoder(avcodec_alloc_context3(codec));
    avcodec_parameters_to_context(decoder.get(), stream->codecpar);
    decoder->pkt_timebase = stream->time_base;
    if (const int ret = avcodec_open2(decoder.get(), codec, nullptr); ret < 0) {
        return {{}, i18n("Could not open the subtitle decoder: %1", ffmpegError(ret))};
    }

    // Only the subtitle stream matters; let the demuxer skip the rest.
    for (unsigned i = 0; i < format->nb_streams; ++i) {
        if (static_cast<int>(i) != streamIndex) {
            format->streams[i]->discard = AVDISCARD_ALL;
        }
    }

    ExtractResult result;
    std::unique_ptr<AVPacket, PacketDeleter> packet(av_packet_alloc());
    while (av_read_frame(format.get(), packet.get()) >= 0) {
        if (cancelled) {
            return {};
        }
        if (packet->stream_index != streamIndex || packet->pts == AV_NOPTS_VALUE) {
            av_packet_unref(packet.get());
            continue;
        }

        AVSubtitle subtitle{};
        int gotSubtitle = 0;
        if (avcodec_decode_subtitle2(decoder.get(), &subtitle, &gotSubtitle, packet.get()) >= 0 && gotSubtitle) {
            const qint64 start = toMs(packet->pts, stream->time_base);
            qint64 end = start + toMs(packet->duration, stream->time_base);
            if (packet->duration <= 0) {
                end = subtitle.end_display_time > 0 && subtitle.end_display_time != UINT32_MAX ? start + subtitle.end_display_time : start + 3000;
            }

            for (unsigned i = 0; i < subtitle.num_rects; ++i) {
                const AVSubtitleRect *rect = subtitle.rects[i];
                QString text;
                if (rect->type == SUBTITLE_ASS && rect->ass) {
                    text = assRectText(rect->ass);
                } else if (rect->type == SUBTITLE_TEXT && rect->text) {
                    text = QString::fromUtf8(rect->text).trimmed();
                }
                if (!text.isEmpty()) {
                    result.cues.append({start, end, text});
                }
            }
            avsubtitle_free(&subtitle);
        }
        av_packet_unref(packet.get());
    }
    return result;
}

SourceResult findTranslationSource(const QString &path)
{
    AVFormatContext *rawFormat = nullptr;
    if (const int ret = avformat_open_input(&rawFormat, path.toUtf8().constData(), nullptr, nullptr); ret < 0) {
        return {-1, i18n("Could not open %1: %2", path, ffmpegError(ret))};
    }
    std::unique_ptr<AVFormatContext, FormatContextDeleter> format(rawFormat);

    const auto language = [](const AVStream *stream) {
        const AVDictionaryEntry *tag = av_dict_get(stream->metadata, "language", nullptr, 0);
        return tag ? QString::fromUtf8(tag->value) : QString();
    };
    // Track ids start at 1, since 0 means "none"; id - 1 is the stream index.
    QList<TrackSelection::Track> subtitles;
    QString audioLanguage;
    bool audioIsDefault = false;
    for (unsigned i = 0; i < format->nb_streams; ++i) {
        const AVStream *stream = format->streams[i];
        const bool isDefault = stream->disposition & AV_DISPOSITION_DEFAULT;
        if (stream->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
            if (audioLanguage.isEmpty() || (isDefault && !audioIsDefault)) {
                audioLanguage = language(stream);
                audioIsDefault = isDefault;
            }
        } else if (stream->codecpar->codec_type == AVMEDIA_TYPE_SUBTITLE) {
            subtitles.append({static_cast<int>(i) + 1,
                              language(stream),
                              QString::fromUtf8(avcodec_get_name(stream->codecpar->codec_id)),
                              bool(stream->disposition & AV_DISPOSITION_FORCED),
                              isDefault});
        }
    }
    return {TrackSelection::pickTranslationSource(subtitles, {QStringLiteral("en"), audioLanguage}, 0) - 1, {}};
}
}