feat: translate subtitles with an OpenAI-compatible LLM

src/subtitles/subtitleextractor.cpp

@@ -0,0 +1,141 @@
+#include "subtitleextractor.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;
+}
+}